Start to schedule fuzz tasks on batch in OSS-Fuzz (#4397)

Start to schedule fuzz tasks on batch in OSS-Fuzz

The scheduler will work differenly in OSS-Fuzz and Chrome. This only
implements and OSS-Fuzz version.
This version will use job and project weights to decide which fuzzing
jobs to schedule. It then adds these tasks
to the queue for other bots to preprocess and then for the
utask_main_scheduler to actually schedule on batch.
For now, we will only do this for 100 CPUs.
1. Add a cron job to run the scheduler every 15 minutes.
2. Improve region handling in batch (still far from complete).
3. Add function for bulk adding of tasks to queue for use by scheduler.
4. Make fuzz tasks less of a priority on batch than others.

Author: jonathanmetzman

configs/test/batch/batch.yaml

@@ -20,7 +20,6 @@ mapping:
     disk_size_gb: 110
     disk_type: pd-standard
     service_account_email: test-clusterfuzz-service-account-email
-    subnetwork: null
     gce_region: 'gce-region'
     gce_zone: 'gce-zone'
     network: 'projects/google.com:clusterfuzz/global/networks/networkname'

infra/k8s/schedule-fuzz.yaml

@@ -0,0 +1,41 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+apiVersion: batch/v1
+kind: CronJob
+metadata:
+  name: schedule-fuzz
+spec:
+  schedule: "*/10 * * * *"
+  concurrencyPolicy: Forbid
+  jobTemplate:
+    spec:
+      activeDeadlineSeconds: 900  # 15 minutes.
+      template:
+        spec:
+          containers:
+          - name: backup
+            image: gcr.io/clusterfuzz-images/base:091c6c2-202409251610
+            imagePullPolicy: Always
+            env:
+            - name: CLUSTERFUZZ_RELEASE
+              value: "prod"
+            - name: RUN_CMD
+              value: "python3.11 $ROOT_DIR/src/python/bot/startup/run_cron.py schedule_fuzz"
+            - name: IS_K8S_ENV
+              value: "true"
+            - name: DISABLE_MOUNTS
+              value: "true"
+          restartPolicy: OnFailure
+      backoffLimit: 3

src/clusterfuzz/_internal/base/concurrency.py

@@ -0,0 +1,31 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tools for concurrency/parallelism."""
+from concurrent import futures
+import contextlib
+import multiprocessing
+
+from clusterfuzz._internal.system import environment
+
+POOL_SIZE = multiprocessing.cpu_count()
+
+
+@contextlib.contextmanager
+def make_pool(pool_size=POOL_SIZE):
+  # Don't use processes on Windows and unittests to avoid hangs.
+  if (environment.get_value('PY_UNITTESTS') or
+      environment.platform() == 'WINDOWS'):
+    yield futures.ThreadPoolExecutor(pool_size)
+  else:
+    yield futures.ProcessPoolExecutor(pool_size)

src/clusterfuzz/_internal/base/tasks/__init__.py

@@ -368,7 +368,8 @@ def __init__(self,
                eta=None,
                is_command_override=False,
                high_end=False,
-               extra_info=None):
+               extra_info=None,
+               is_from_queue=False):
     self.command = command
     self.argument = argument
     self.job = job
@@ -377,6 +378,19 @@ def __init__(self,
     self.high_end = high_end
     self.extra_info = extra_info
 
+    # is_from_queue is a temporary hack to keep track of which fuzz tasks came
+    # from the queue. Previously all fuzz tasks were picked by the bot when
+    # there was nothing on the queue. With the rearchitecture, we want fuzz
+    # tasks that were put on the queue by the schedule_fuzz cron job to be
+    # executed on batch. is_from_queue is used to do this.
+    # TODO(b/378684001): This code is very ugly, get rid of it when no more
+    # fuzz tasks are executed on the bots themselves (i.e. when the rearch
+    # is complete).
+    self.is_from_queue = is_from_queue
+
+  def __repr__(self):
+    return f'Task: {self.command} {self.argument} {self.job}'
+
   def attribute(self, _):
     return None
 
@@ -414,11 +428,13 @@ def lease(self):
 class PubSubTask(Task):
   """A Pub/Sub task."""
 
-  def __init__(self, pubsub_message):
+  def __init__(self, pubsub_message, is_from_queue=False):
     self._pubsub_message = pubsub_message
     super().__init__(
-        self.attribute('command'), self.attribute('argument'),
-        self.attribute('job'))
+        self.attribute('command'),
+        self.attribute('argument'),
+        self.attribute('job'),
+        is_from_queue=is_from_queue)
 
     self.extra_info = {
         key: value
@@ -524,7 +540,7 @@ def initialize_task(message) -> PubSubTask:
   """Creates a task from |messages|."""
 
   if message.attributes.get('eventType') != 'OBJECT_FINALIZE':
-    return PubSubTask(message)
+    return PubSubTask(message, is_from_queue=True)
 
   # Handle postprocess task.
   # The GCS API for pub/sub notifications uses the data field unlike
@@ -533,7 +549,7 @@ def initialize_task(message) -> PubSubTask:
   name = data['name']
   bucket = data['bucket']
   output_url_argument = storage.get_cloud_storage_file_path(bucket, name)
-  return PostprocessPubSubTask(output_url_argument, message)
+  return PostprocessPubSubTask(output_url_argument, message, is_from_queue=True)
 
 
 class PostprocessPubSubTask(PubSubTask):
@@ -542,14 +558,21 @@ class PostprocessPubSubTask(PubSubTask):
   def __init__(self,
                output_url_argument,
                pubsub_message,
-               is_command_override=False):
+               is_command_override=False,
+               is_from_queue=False):
     command = 'postprocess'
     job_type = 'none'
     eta = None
     high_end = False
     grandparent_class = super(PubSubTask, self)
-    grandparent_class.__init__(command, output_url_argument, job_type, eta,
-                               is_command_override, high_end)
+    grandparent_class.__init__(
+        command,
+        output_url_argument,
+        job_type,
+        eta,
+        is_command_override,
+        high_end,
+        is_from_queue=is_from_queue)
     self._pubsub_message = pubsub_message
 
 
@@ -609,18 +632,36 @@ def add_utask_main(command, input_url, job_type, wait_time=None):
       extra_info={'initial_command': initial_command})
 
 
+def bulk_add_tasks(tasks, queue=None, eta_now=False):
+  """Adds |tasks| in bulk to |queue|."""
+
+  # Old testcases may pass in queue=None explicitly, so we must check this here.
+  if queue is None:
+    queue = default_queue()
+
+  # If callers want delays, they must do it themselves, because this function is
+  # meant to be used for batch tasks which don't need this.
+  # Use an ETA of right now for batch because we don't need extra delay, there
+  # is natural delay added by batch, waiting for utask_main_scheduler,
+  # postprocess etc.
+  if eta_now:
+    now = utils.utcnow()
+    for task in tasks:
+      task.eta = now
+
+  pubsub_client = pubsub.PubSubClient()
+  pubsub_messages = [task.to_pubsub_message() for task in tasks]
+  pubsub_client.publish(
+      pubsub.topic_name(utils.get_application_id(), queue), pubsub_messages)
+
+
 def add_task(command,
              argument,
              job_type,
              queue=None,
              wait_time=None,
              extra_info=None):
   """Add a new task to the job queue."""
-  # Old testcases may pass in queue=None explicitly,
-  # so we must check this here.
-  if not queue:
-    queue = default_queue()
-
   if wait_time is None:
     wait_time = random.randint(1, TASK_CREATION_WAIT_INTERVAL)
 
@@ -636,10 +677,8 @@ def add_task(command,
   # Add the task.
   eta = utils.utcnow() + datetime.timedelta(seconds=wait_time)
   task = Task(command, argument, job_type, eta=eta, extra_info=extra_info)
-  pubsub_client = pubsub.PubSubClient()
-  pubsub_client.publish(
-      pubsub.topic_name(utils.get_application_id(), queue),
-      [task.to_pubsub_message()])
+
+  bulk_add_tasks([task], queue=queue)
 
 
 def get_task_lease_timeout():

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

@@ -95,8 +95,6 @@ def cleanup_task_state():
 
   # Call python's garbage collector.
   utils.python_gc()
-  if 'CF_TASK_ID' in os.environ:
-    del os.environ['CF_TASK_ID']
 
 
 def is_supported_cpu_arch_for_job():
@@ -255,8 +253,10 @@ def process_command(task):
     logs.error('Empty task received.')
     return None
 
+  # TODO(b/378684001): Remove is_from_queue kludge.
   return process_command_impl(task.command, task.argument, task.job,
-                              task.high_end, task.is_command_override)
+                              task.high_end, task.is_command_override,
+                              task.is_from_queue)
 
 
 def _get_task_id(task_name, task_argument, job_name):
@@ -267,12 +267,13 @@ def _get_task_id(task_name, task_argument, job_name):
 # TODO(mbarbella): Rewrite this function to avoid nesting issues.
 @set_task_payload
 def process_command_impl(task_name, task_argument, job_name, high_end,
-                         is_command_override):
+                         is_command_override, is_from_queue):
   """Implementation of process_command."""
   uworker_env = None
   environment.set_value('TASK_NAME', task_name)
   environment.set_value('TASK_ARGUMENT', task_argument)
   environment.set_value('JOB_NAME', job_name)
+  environment.set_value('IS_FROM_QUEUE', is_from_queue)
   if task_name in {'uworker_main', 'postprocess'}:
     # We want the id of the task we are processing, not "uworker_main", or
     # "postprocess".
@@ -452,3 +453,7 @@ def process_command_impl(task_name, task_argument, job_name, high_end,
   finally:
     # Final clean up.
     cleanup_task_state()
+    if 'CF_TASK_ID' in os.environ:
+      del os.environ['CF_TASK_ID']
+    if 'IS_FROM_QUEUE' in os.environ:
+      del os.environ['IS_FROM_QUEUE']

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

@@ -480,7 +480,6 @@ def update_data_bundle(
   data_bundle = uworker_io.entity_from_protobuf(data_bundle_corpus.data_bundle,
                                                 data_types.DataBundle)
   logs.info('Setting up data bundle %s.' % data_bundle)
-
   data_bundle_directory = _prepare_update_data_bundle(fuzzer, data_bundle)
 
   if not _should_update_data_bundle(data_bundle, data_bundle_directory):

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

@@ -156,6 +156,19 @@ def preprocess(self, task_argument, job_type, uworker_env):
     return download_url
 
 
+# TODO(b/378684001): Remove this, it's needed for testing but is otherwise a bad
+# design.
+class UTaskMostlyLocalExecutor(UTask):
+
+  @staticmethod
+  def is_execution_remote(command=None):
+    del command
+    if environment.get_value('IS_FROM_QUEUE'):
+      logs.info('IS FROM QUEUE')
+      return True
+    return False
+
+
 class PostprocessTask(BaseTask):
   """Represents postprocessing of an untrusted task."""
 
@@ -198,7 +211,7 @@ def execute(self, task_argument, job_type, uworker_env):
     'analyze': UTask,
     'blame': TrustedTask,
     'corpus_pruning': UTask,
-    'fuzz': UTaskLocalExecutor,
+    'fuzz': UTaskMostlyLocalExecutor,
     'impact': TrustedTask,
     'minimize': UTask,
     'progression': UTask,