Refactor: Remove batch.py and introduce RemoteTaskGate

This refactor removes the  file and introduces a  that proportionally distributes tasks between  and . This allows for A/B testing and performance comparisons between the two platforms.

Author: javanlacerda

local/bin/kind

@@ -335,7 +335,7 @@ class GcpBatchService(RemoteTaskInterface):
   provides a way to check if a task is configured to run remotely.
   """
 
-  def create_job(self, spec, input_urls: List[str]):
+  def create_job(self, spec: BatchWorkloadSpec, input_urls: List[str]):
     """Creates and starts a batch job from |spec| that executes all tasks.
     
     This method creates a new GCP Batch job with a single task group. The

src/clusterfuzz/_internal/batch/service.py

@@ -12,12 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Kubernetes batch client."""
+import collections
 from typing import List
 import uuid
 
 from kubernetes import client as k8s_client
 from kubernetes import config as k8s_config
 
+from clusterfuzz._internal.base import utils
+from clusterfuzz._internal.base.tasks import task_utils
+from clusterfuzz._internal.batch.service import _get_specs_from_config
+from clusterfuzz._internal.batch.service import MAX_CONCURRENT_VMS_PER_JOB
+from clusterfuzz._internal.metrics import logs
+from clusterfuzz._internal.remote_task import RemoteTask
 from clusterfuzz._internal.remote_task import RemoteTaskInterface
 
 
@@ -88,6 +95,41 @@ def create_job(self, remote_task: RemoteTaskInterface,
     return self._create_job_client_wrapper(remote_task.docker_image, job_spec,
                                            input_urls)
 
+  def create_uworker_main_batch_job(self, module: str, job_type: str,
+                                    input_download_url: str):
+    """Creates a single batch job for a uworker main task."""
+    command = task_utils.get_command_from_module(module)
+    batch_tasks = [RemoteTask(command, job_type, input_download_url)]
+    result = self.create_uworker_main_batch_jobs(batch_tasks)
+    if result is None:
+      return result
+    return result[0]
+
+  def create_uworker_main_batch_jobs(self, batch_tasks: List[RemoteTask]):
+    """Creates a batch job for a list of uworker main tasks.
+    
+    This method groups the tasks by their workload specification and creates a
+    separate batch job for each group. This allows tasks with similar
+    requirements to be processed together, which can improve efficiency.
+    """
+    job_specs = collections.defaultdict(list)
+    specs = _get_specs_from_config(batch_tasks)
+    for batch_task in batch_tasks:
+      logs.info(f'Scheduling {batch_task.command}, {batch_task.job_type}.')
+      spec = specs[(batch_task.command, batch_task.job_type)]
+      job_specs[spec].append(batch_task.input_download_url)
+
+    logs.info('Creating batch jobs.')
+    jobs = []
+
+    logs.info('Batching utask_mains.')
+    for spec, input_urls in job_specs.items():
+      for input_urls_portion in utils.batched(input_urls,
+                                              MAX_CONCURRENT_VMS_PER_JOB - 1):
+        jobs.append(self.create_job(spec, input_urls_portion))
+
+    return jobs
+
   def create_kata_container_job(self, container_image: str,
                                 input_urls: List[str]) -> str:
     """Creates a Kubernetes job that runs in a Kata container."""

src/clusterfuzz/_internal/google_cloud_utils/batch.py

@@ -53,3 +53,69 @@ def create_job(self, remote_task: RemoteTask, input_urls: List[str]):
     and returns a representation of the created job.
     """
     raise NotImplementedError
+
+
+import random
+
+from clusterfuzz._internal.remote_task import job_frequency
+
+
+class RemoteTaskGate:
+  """A gatekeeper for remote task execution.
+  
+  This class is responsible for choosing the remote execution backend (GCP Batch
+  or Kubernetes) for a given task, based on the configured frequencies in the
+  `job_frequency` module.
+  """
+
+  def __init__(self):
+    from clusterfuzz._internal.batch.service import GcpBatchService
+    from clusterfuzz._internal.k8s.service import KubernetesService
+    self._gcp_batch_service = GcpBatchService()
+    self._kubernetes_service = KubernetesService()
+
+  def _should_use_kubernetes(self, job_type: str) -> bool:
+    """Determines whether to use the Kubernetes backend for a given job.
+    
+    The decision is made based on a random roll and the configured frequency
+    for the given job type.
+    """
+    frequencies = job_frequency.get_job_frequency(job_type)
+    return random.random() < frequencies['kubernetes']
+
+  def create_uworker_main_batch_job(self, module: str, job_type: str,
+                                    input_download_url: str):
+    """Creates a batch job on either GCP Batch or Kubernetes.
+    
+    The choice of backend is determined by the `_should_use_kubernetes` method.
+    """
+    if self._should_use_kubernetes(job_type):
+      return self._kubernetes_service.create_uworker_main_batch_job(
+          module, job_type, input_download_url)
+    return self._gcp_batch_service.create_uworker_main_batch_job(
+        module, job_type, input_download_url)
+
+  def create_uworker_main_batch_jobs(self, batch_tasks: List[RemoteTask]):
+    """Creates batch jobs on either GCP Batch or Kubernetes.
+    
+    The tasks are grouped by their target backend (GCP Batch or Kubernetes) and
+    then created in separate batches.
+    """
+    gcp_batch_tasks = []
+    kubernetes_tasks = []
+    for task in batch_tasks:
+      if self._should_use_kubernetes(task.job_type):
+        kubernetes_tasks.append(task)
+      else:
+        gcp_batch_tasks.append(task)
+
+    results = []
+    if gcp_batch_tasks:
+      results.extend(
+          self._gcp_batch_service.create_uworker_main_batch_jobs(
+              gcp_batch_tasks))
+    if kubernetes_tasks:
+      results.extend(
+          self._kubernetes_service.create_uworker_main_batch_jobs(
+              kubernetes_tasks))
+    return results

src/clusterfuzz/_internal/k8s/service.py

@@ -0,0 +1,60 @@
+# Copyright 2025 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.
+"""Configurable job frequencies for remote task execution.
+
+This module provides a way to define how frequently certain jobs are executed
+on different remote backends, such as GCP Batch and Kubernetes. This allows for
+A/B testing and performance comparisons between the two platforms.
+"""
+
+from clusterfuzz._internal.system import environment
+
+# By default, all jobs are sent to the GCP Batch backend. This can be
+# overridden on a per-job basis by setting the `K8S_JOBS_FREQUENCY`
+# environment variable.
+DEFAULT_FREQUENCY = {'gcp_batch': 1.0, 'kubernetes': 0.0}
+
+
+def _get_job_frequencies_from_env():
+  """Parses the `K8S_JOBS_FREQUENCY` environment variable.
+  
+  The variable should be a comma-separated list of key-value pairs, where the
+  key is the job name and the value is the frequency (a float between 0 and 1).
+  For example: `libfuzzer_asan_chrome=0.5,libfuzzer_msan_chrome=0.2`.
+  """
+  job_frequencies = {}
+  frequency_string = environment.get_value('K8S_JOBS_FREQUENCY')
+  if not frequency_string:
+    return {}
+
+  for item in frequency_string.split(','):
+    key, value = item.split('=')
+    job_frequencies[key] = float(value)
+  return job_frequencies
+
+
+def get_job_frequency(job_name):
+  """Returns the frequency for a given job.
+  
+  If the frequency is not explicitly defined in the `K8S_JOBS_FREQUENCY`
+  environment variable, the default frequency is returned.
+  """
+  job_frequencies = _get_job_frequencies_from_env()
+  if job_name in job_frequencies:
+    kubernetes_frequency = job_frequencies[job_name]
+    return {
+        'gcp_batch': 1.0 - kubernetes_frequency,
+        'kubernetes': kubernetes_frequency
+    }
+  return DEFAULT_FREQUENCY

src/clusterfuzz/_internal/remote_task/__init__.py

@@ -146,7 +146,7 @@ class GcpBatchServiceTest(unittest.TestCase):
 
   def setUp(self):
     helpers.patch(self, [
-        'clusterfuzz._internal.batch.batch_service._batch_client',
+        'clusterfuzz._internal.batch.service._batch_client',
         'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module',
         'uuid.uuid4',
     ])
@@ -190,9 +190,8 @@ def test_create_uworker_main_batch_jobs(self):
         priority=0,
         max_run_duration='2s',
         retry=True)
-    with mock.patch(
-        'clusterfuzz._internal.batch.batch_service._get_specs_from_config'
-    ) as mock_get_specs_from_config:
+    with mock.patch('clusterfuzz._internal.batch.service._get_specs_from_config'
+                   ) as mock_get_specs_from_config:
       mock_get_specs_from_config.return_value = {
           ('command1', 'job1'): spec1,
           ('command2', 'job2'): spec2,
@@ -235,9 +234,8 @@ def test_create_uworker_main_batch_job(self):
         priority=1,
         max_run_duration='1s',
         retry=False)
-    with mock.patch(
-        'clusterfuzz._internal.batch.batch_service._get_specs_from_config'
-    ) as mock_get_specs_from_config:
+    with mock.patch('clusterfuzz._internal.batch.service._get_specs_from_config'
+                   ) as mock_get_specs_from_config:
       mock_get_specs_from_config.return_value = {
           ('fuzz', 'job1'): spec1,
       }
@@ -262,7 +260,7 @@ class IsRemoteTaskTest(unittest.TestCase):
 
   def setUp(self):
     helpers.patch(self, [
-        'clusterfuzz._internal.batch.batch_service._get_specs_from_config',
+        'clusterfuzz._internal.batch.service._get_specs_from_config',
     ])
     data_types.Job(name='job', platform='LINUX').put()
 
@@ -408,7 +406,7 @@ def test_get_specs_from_config_with_disk_size_override(self):
         [batch_service.RemoteTask('fuzz', job_name, None)])
     self.assertEqual(spec['fuzz', job_name].disk_size_gb, overridden_size)
 
-  @mock.patch('clusterfuzz._internal.batch.batch_service.utils.is_oss_fuzz')
+  @mock.patch('clusterfuzz._internal.batch.service.utils.is_oss_fuzz')
   @mock.patch('clusterfuzz._internal.datastore.data_types.OssFuzzProject.query')
   @mock.patch('clusterfuzz._internal.datastore.ndb_utils.get_all_from_query')
   def test_get_config_names_os_version(self, mock_get_all_from_query,