| # Copyright (c) 2015 The Chromium Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| import datetime |
| import json |
| import logging |
| import math |
| import os |
| import urllib |
| import uuid |
| import webapp2 |
| |
| from google.appengine.api import taskqueue |
| from google.appengine.api import urlfetch |
| from perf_insights.endpoints.cloud_mapper import cloud_helper |
| from perf_insights.endpoints.cloud_mapper import job_info |
| from perf_insights import cloud_config |
| |
| # If you modify this, you need to change the max_concurrent_requests in |
| # queue.yaml. |
| DEFAULT_TRACES_PER_INSTANCE = 4 |
| |
| class TaskPage(webapp2.RequestHandler): |
| |
| def _QueryForTraces(self, corpus, query): |
| payload = urllib.urlencode({'q': query}) |
| query_url = '%s/query?%s' % (corpus, payload) |
| |
| headers = { |
| 'X-URLFetch-Service-Id': cloud_config.Get().urlfetch_service_id |
| } |
| |
| result = urlfetch.fetch(url=query_url, |
| payload=payload, |
| method=urlfetch.GET, |
| headers=headers, |
| follow_redirects=False, |
| deadline=10) |
| return json.loads(result.content) |
| |
| def _DispatchTracesAndWaitForResult(self, job, traces, num_instances): |
| def _slice_it(li, cols=2): |
| start = 0 |
| for i in xrange(cols): |
| stop = start + len(li[i::cols]) |
| yield li[start:stop] |
| start = stop |
| |
| # TODO(simonhatch): In the future it might be possibly to only specify a |
| # reducer and no mapper. Revisit this. |
| bucket_path = cloud_config.Get().control_bucket_path + "/jobs/" |
| |
| mapper_url = '' |
| reducer_url = '' |
| |
| if job.reducer: |
| reducer_url = '%s%s.reducer' % (bucket_path, job.key.id()) |
| reducer_text = job.reducer.encode('ascii', 'ignore') |
| cloud_helper.WriteGCS(reducer_url, reducer_text) |
| |
| if job.mapper: |
| mapper_url = '%s%s.mapper' % (bucket_path, job.key.id()) |
| mapper_text = job.mapper.encode('ascii', 'ignore') |
| cloud_helper.WriteGCS(mapper_url, mapper_text) |
| |
| version = self._GetVersion() |
| |
| tasks = {} |
| |
| # Split the traces up into N buckets. |
| logging.info('Splitting traces across %d instances.' % num_instances) |
| for current_traces in _slice_it(traces, num_instances): |
| logging.info('Submitting %d traces job.' % len(current_traces)) |
| task_id = str(uuid.uuid4()) |
| |
| payload = { |
| 'revision': job.revision, |
| 'traces': json.dumps(current_traces), |
| 'result': '%s%s.result' % (bucket_path, task_id), |
| 'mapper': mapper_url, |
| 'mapper_function': job.mapper_function, |
| 'timeout': job.function_timeout, |
| } |
| taskqueue.add( |
| queue_name='mapper-queue', |
| url='/cloud_worker/task', |
| target=version, |
| name=task_id, |
| params=payload) |
| tasks[task_id] = {'status': 'IN_PROGRESS'} |
| |
| # On production servers, we could just sit and wait for the results, but |
| # dev_server is single threaded and won't run any other tasks until the |
| # current one is finished. We'll just do the easy thing for now and |
| # queue a task to check for the result. |
| mapper_timeout = int(job.timeout - job.function_timeout) |
| timeout = ( |
| datetime.datetime.now() + datetime.timedelta( |
| seconds=mapper_timeout)).strftime( |
| '%Y-%m-%d %H:%M:%S') |
| taskqueue.add( |
| queue_name='default', |
| url='/cloud_mapper/task', |
| target=version, |
| countdown=1, |
| params={'jobid': job.key.id(), |
| 'type': 'check_map_results', |
| 'reducer': reducer_url, |
| 'tasks': json.dumps(tasks), |
| 'timeout': timeout}) |
| |
| def _GetVersion(self): |
| version = os.environ['CURRENT_VERSION_ID'].split('.')[0] |
| if cloud_config._is_devserver(): |
| version = taskqueue.DEFAULT_APP_VERSION |
| return version |
| |
| def _CheckOnMapResults(self, job): |
| if job.status != 'IN_PROGRESS': |
| return |
| |
| tasks = json.loads(self.request.get('tasks')) |
| reducer_url = self.request.get('reducer') |
| reducer_function = job.reducer_function |
| revision = job.revision |
| timeout = datetime.datetime.strptime( |
| self.request.get('timeout'), '%Y-%m-%d %H:%M:%S') |
| |
| # TODO: There's no reducer yet, so we can't actually collapse multiple |
| # results into one results file. |
| mappers_done = True |
| for task_id, task_values in tasks.iteritems(): |
| if task_values['status'] == 'DONE': |
| continue |
| task_results_path = '%s/jobs/%s.result' % ( |
| cloud_config.Get().control_bucket_path, task_id) |
| stat_result = cloud_helper.StatGCS(task_results_path) |
| if stat_result is not None: |
| logging.info(str(stat_result)) |
| tasks[task_id]['status'] = 'DONE' |
| else: |
| mappers_done = False |
| |
| logging.info("Tasks: %s" % str(tasks)) |
| |
| if not mappers_done and datetime.datetime.now() < timeout: |
| taskqueue.add( |
| url='/cloud_mapper/task', |
| target=self._GetVersion(), |
| countdown=1, |
| params={'jobid': job.key.id(), |
| 'type': 'check_map_results', |
| 'reducer': reducer_url, |
| 'tasks': json.dumps(tasks), |
| 'timeout': self.request.get('timeout')}) |
| return |
| |
| # Clear out any leftover tasks in case we just hit the timeout. |
| self._CancelTasks(tasks) |
| |
| map_results = [] |
| for task_id, _ in tasks.iteritems(): |
| if tasks[task_id]['status'] != 'DONE': |
| continue |
| task_results_path = '%s/jobs/%s.result' % ( |
| cloud_config.Get().control_bucket_path, task_id) |
| map_results.append(task_results_path) |
| |
| # We'll only do 1 reduce job for now, maybe shard it better later |
| logging.info("Kicking off reduce.") |
| task_id = str(uuid.uuid4()) |
| payload = { |
| 'revision': revision, |
| 'traces': json.dumps(map_results), |
| 'result': '%s/jobs/%s.result' % ( |
| cloud_config.Get().control_bucket_path, task_id), |
| 'reducer': reducer_url, |
| 'reducer_function': reducer_function, |
| 'timeout': job.function_timeout, |
| } |
| taskqueue.add( |
| queue_name='mapper-queue', |
| url='/cloud_worker/task', |
| target=self._GetVersion(), |
| name=task_id, |
| params=payload) |
| |
| tasks = {} |
| tasks[task_id] = {'status': 'IN_PROGRESS'} |
| |
| job.running_tasks = [task_id for task_id, _ in tasks.iteritems()] |
| job.put() |
| |
| reduce_tasks = {} |
| reduce_tasks[task_id] = {'status': 'IN_PROGRESS'} |
| |
| # On production servers, we could just sit and wait for the results, but |
| # dev_server is single threaded and won't run any other tasks until the |
| # current one is finished. We'll just do the easy thing for now and |
| # queue a task to check for the result. |
| reducer_timeout = int(job.function_timeout) |
| timeout = ( |
| datetime.datetime.now() + datetime.timedelta( |
| seconds=reducer_timeout)).strftime( |
| '%Y-%m-%d %H:%M:%S') |
| taskqueue.add( |
| queue_name='default', |
| url='/cloud_mapper/task', |
| target=self._GetVersion(), |
| countdown=1, |
| params={'jobid': job.key.id(), |
| 'type': 'check_reduce_results', |
| 'tasks': json.dumps(reduce_tasks), |
| 'timeout': timeout}) |
| |
| def _CancelTasks(self, tasks): |
| task_names = [task_id for task_id, _ in tasks.iteritems()] |
| taskqueue.Queue('mapper-queue').delete_tasks_by_name(task_names) |
| |
| def _CheckOnReduceResults(self, job): |
| if job.status != 'IN_PROGRESS': |
| return |
| |
| tasks = json.loads(self.request.get('tasks')) |
| |
| # TODO: There's really only one reducer job at the moment |
| results = None |
| for task_id, _ in tasks.iteritems(): |
| task_results_path = '%s/jobs/%s.result' % ( |
| cloud_config.Get().control_bucket_path, task_id) |
| stat_result = cloud_helper.StatGCS(task_results_path) |
| if stat_result is not None: |
| tasks[task_id]['status'] = 'DONE' |
| results = task_results_path |
| |
| logging.info("Reduce results: %s" % str(tasks)) |
| |
| if not results: |
| timeout = datetime.datetime.strptime( |
| self.request.get('timeout'), '%Y-%m-%d %H:%M:%S') |
| if datetime.datetime.now() > timeout: |
| self._CancelTasks(tasks) |
| job.status = 'ERROR' |
| job.put() |
| logging.error('Task timed out waiting for results.') |
| return |
| taskqueue.add( |
| url='/cloud_mapper/task', |
| target=self._GetVersion(), |
| countdown=1, |
| params={'jobid': job.key.id(), |
| 'type': 'check_reduce_results', |
| 'tasks': json.dumps(tasks), |
| 'timeout': self.request.get('timeout')}) |
| return |
| |
| logging.info("Finished all tasks.") |
| |
| job.status = 'COMPLETE' |
| job.results = results |
| job.put() |
| |
| def _CalculateNumInstancesNeeded(self, num_traces): |
| return int(math.ceil(float(num_traces) / DEFAULT_TRACES_PER_INSTANCE)) |
| |
| def _RunMappers(self, job): |
| # Get all the traces to process |
| traces = self._QueryForTraces(job.corpus, job.query) |
| |
| # We can probably be smarter about this down the road, maybe breaking |
| # this into many smaller tasks and allowing each instance to run |
| # several tasks at once. For now we'll just break it into a few big ones. |
| num_instances = self._CalculateNumInstancesNeeded(len(traces)) |
| |
| return self._DispatchTracesAndWaitForResult(job, traces, num_instances) |
| |
| def _CreateMapperJob(self, job): |
| if job.status != 'QUEUED': |
| return |
| |
| job.status = 'IN_PROGRESS' |
| job.put() |
| |
| self._RunMappers(job) |
| |
| def post(self): |
| self.response.headers['Content-Type'] = 'text/plain' |
| |
| jobid = self.request.get('jobid') |
| job = job_info.JobInfo.get_by_id(jobid) |
| if not job: |
| return |
| |
| try: |
| if self.request.get('type') == 'create': |
| self._CreateMapperJob(job) |
| elif self.request.get('type') == 'check_map_results': |
| self._CheckOnMapResults(job) |
| elif self.request.get('type') == 'check_reduce_results': |
| self._CheckOnReduceResults(job) |
| except Exception as e: |
| job.status = 'ERROR' |
| job.put() |
| logging.exception('Failed job: %s' % e.message) |
| |
| |
| app = webapp2.WSGIApplication([('/cloud_mapper/task', TaskPage)]) |
