Batch progression task (#3634)

Batch progression task so that preprocess is run, then
a group of preprocess tasks will be put on batch.

Author: jonathanmetzman

src/clusterfuzz/_internal/bot/tasks/task_creation.py

@@ -299,27 +299,31 @@ def _preprocess(task: Task) -> None:
 
 def start_utask_mains(tasks: List[Task]) -> None:
   """Start utask_main of multiple tasks as batch tasks on batch."""
-  batch_tasks = [
-      batch.BatchTask(task.name, task.job, task.uworker_input) for task in tasks
-  ]
-  batch.create_uworker_main_batch_jobs(batch_tasks)
+  batch_tasks = (
+      batch.BatchTask(task.name, task.job, task.uworker_input)
+      for task in tasks)
+  return batch.create_uworker_main_batch_jobs_bunched(batch_tasks)
 
 
-def schedule_tasks(tasks: List[Task]):
-  """Starts tasks as defined by task objects. If the tasks are not executed
-  remotely, then they are put on the queue. If they are executed remotely, then
-  the utask_mains are scheduled on batch, since preprocess has already been done
-  in this module on this bot."""
-  uworker_tasks = []
-  tasks = [task for task in tasks if task is not None]
+def preprocess_utasks_and_queue_ttasks(tasks: List[Optional[Task]]):
+  """If a task is a utask, then it is preprocessed yielded, otherwise it is
+  added to the queue (when it is a ttask)."""
   for task in tasks:
+    if task is None:
+      continue
     if not task_types.is_remote_utask(task.name, task.job):
-      taskslib.add_task(task.name, task.argument, task.job,
-                        task.queue_for_platform)
+      taskslib.add_task(
+          task.name, task.argument, task.job, queue=task.queue_for_platform)
       logs.log(f'UTask {task.name} not remote.')
       continue
     _preprocess(task)
-    uworker_tasks.append(task)
-  logs.log(f'Starting utask_mains: {len(uworker_tasks)}.')
-  if uworker_tasks:
-    start_utask_mains(uworker_tasks)
+    yield task
+
+
+def schedule_tasks(tasks: List[Task]):
+  """Starts tasks as defined by task objects. If the tasks are not executed
+  remotely, then they are put on the queue. If they are executed remotely, then
+  the utask_mains are scheduled on batch, since preprocess has already been done
+  in this module on this bot."""
+  uworker_tasks = preprocess_utasks_and_queue_ttasks(tasks)
+  return start_utask_mains(uworker_tasks)

src/clusterfuzz/_internal/cron/helpers/tasks_scheduler.py

@@ -13,29 +13,30 @@
 # limitations under the License.
 """Task scheduler used to recreate recurring tasks."""
 
-from clusterfuzz._internal.base import tasks
+from clusterfuzz._internal.base import tasks as taskslib
+from clusterfuzz._internal.bot.tasks import task_creation
 from clusterfuzz._internal.datastore import data_types
 from clusterfuzz._internal.datastore import ndb_utils
-from clusterfuzz._internal.metrics import logs
 
 
 def schedule(task):
   """Creates tasks for open reproducible testcases."""
 
+  testcases = []
   for status in ['Processed', 'Duplicate']:
-    testcases = data_types.Testcase.query(
-        ndb_utils.is_true(data_types.Testcase.open),
-        ndb_utils.is_false(data_types.Testcase.one_time_crasher_flag),
-        data_types.Testcase.status == status)
+    testcases.extend(
+        data_types.Testcase.query(
+            ndb_utils.is_true(data_types.Testcase.open),
+            ndb_utils.is_false(data_types.Testcase.one_time_crasher_flag),
+            data_types.Testcase.status == status))
 
-    for testcase in testcases:
-      try:
-        tasks.add_task(
-            task,
-            testcase.key.id(),
-            testcase.job_type,
-            queue=tasks.queue_for_testcase(testcase))
-      except Exception:
-        logs.log_error('Failed to add task.')
-        continue
-  logs.log(f'Task {task} added successfully')
+  tasks = [
+      task_creation.Task(
+          task,
+          testcase.key.id(),
+          testcase.job_type,
+          queue_for_platform=taskslib.queue_for_testcase(testcase))
+      for testcase in testcases
+  ]
+
+  task_creation.schedule_tasks(tasks)

src/clusterfuzz/_internal/google_cloud_utils/batch.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Cloud Batch helpers."""
 import collections
+import itertools
 import threading
 import uuid
 
@@ -34,6 +35,8 @@
 MAX_DURATION = f'{int(60 * 60 * 2.5)}s'
 RETRY_COUNT = 0
 
+TASK_BUNCH_SIZE = 20
+
 # Controls how many containers (ClusterFuzz tasks) can run on a single VM.
 # THIS SHOULD BE 1 OR THERE WILL BE SECURITY PROBLEMS.
 TASK_COUNT_PER_NODE = 1
@@ -96,17 +99,52 @@ def create_uworker_main_batch_job(module, job_type, input_download_url):
   return result[0]
 
 
+def _bunched(iterator, bunch_size):
+  """Implementation of itertools.py's batched that was added after Python3.7."""
+  # TODO(metzman): Replace this with itertools.batched.
+  assert bunch_size > -1
+  idx = 0
+  bunch = []
+  for item in iterator:
+    idx += 1
+    bunch.append(item)
+    if idx == bunch_size:
+      idx = 0
+      yield bunch
+      bunch = []
+
+  if bunch:
+    yield bunch
+
+
 def create_uworker_main_batch_jobs(batch_tasks):
-  # Define what will be done as part of the job.
+  """Creates batch jobs."""
   job_specs = collections.defaultdict(list)
   for batch_task in batch_tasks:
     spec = _get_spec_from_config(batch_task.command, batch_task.job_type)
     job_specs[spec].append(batch_task.input_download_url)
 
   logs.log('Creating batch jobs.')
-  return [
-      _create_job(spec, input_urls) for spec, input_urls in job_specs.items()
+  jobs = []
+
+  logs.log(f'Starting utask_mains: {job_specs}.')
+  for spec, input_urls in job_specs.items():
+    for input_urls_portion in _bunched(input_urls, MAX_CONCURRENT_VMS_PER_JOB):
+      jobs.append(_create_job(spec, input_urls_portion))
+
+  return jobs
+
+
+def create_uworker_main_batch_jobs_bunched(batch_tasks):
+  """Creates batch jobs 20 tasks at a time, lazily. This is helpful to use when
+  batch_tasks takes a very long time to create."""
+  # Use term bunch instead of "batch" since "batch" has nothing to do with the
+  # cloud service and is thus very confusing in this context.
+  jobs = [
+      create_uworker_main_batch_jobs(bunch)
+      for bunch in _bunched(batch_tasks, TASK_BUNCH_SIZE)
   ]
+  return list(itertools.chain(jobs))
 
 
 def _get_task_spec(batch_workload_spec):

src/clusterfuzz/_internal/tests/appengine/handlers/cron/recurring_tasks_test.py

@@ -91,8 +91,8 @@ def test_execute(self):
     """Tests scheduling of progression tasks."""
     schedule_progression_tasks.main()
     self.mock.add_task.assert_has_calls([
-        mock.call('progression', 1, 'job', queue='jobs-linux'),
-        mock.call('progression', 5, 'job_windows', queue='jobs-windows')
+        mock.call('progression', '1', 'job', queue='jobs-linux'),
+        mock.call('progression', '5', 'job_windows', queue='jobs-windows'),
     ])
 
 
@@ -109,6 +109,6 @@ def test_execute(self):
     """Tests scheduling of progression tasks."""
     schedule_impact_tasks.main()
     self.mock.add_task.assert_has_calls([
-        mock.call('impact', 1, 'job', queue='jobs-linux'),
-        mock.call('impact', 5, 'job_windows', queue='jobs-windows'),
+        mock.call('impact', '1', 'job', queue='jobs-linux'),
+        mock.call('impact', '5', 'job_windows', queue='jobs-windows'),
     ])