fixes

Signed-off-by: Javan Lacerda <javanlacerda@google.com>

Author: javanlacerda

src/clusterfuzz/_internal/base/feature_flags.py

@@ -29,6 +29,8 @@ class FeatureFlags(Enum):
   K8S_JOBS_FREQUENCY = 'k8s_jobs_frequency'
   K8S_JOBS_PENDING_LIMIT = 'k8s_jobs_pending_limit'
 
+  UTASK_MAIN_QUEUE_LIMIT = 'utask_main_queue_limit'
+
   @property
   def flag(self):
     """Get the feature flag."""

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

@@ -17,6 +17,7 @@
 
 from clusterfuzz._internal import swarming
 from clusterfuzz._internal.base import errors
+from clusterfuzz._internal.base import feature_flags
 from clusterfuzz._internal.base import tasks
 from clusterfuzz._internal.base.tasks import task_utils
 from clusterfuzz._internal.batch import service as batch_service
@@ -25,6 +26,8 @@
 from clusterfuzz._internal.metrics import logs
 from clusterfuzz._internal.system import environment
 
+UTASK_MAIN_QUEUE_LIMIT_DEFAULT = 10000
+
 
 class BaseTask:
   """Base module for tasks."""
@@ -157,7 +160,11 @@ def execute(self, task_argument, job_type, uworker_env):
 
     utask_main_queue_size = tasks.get_utask_main_queue_size()
 
-    if utask_main_queue_size > 10000:
+    utask_main_queue_limit = UTASK_MAIN_QUEUE_LIMIT_DEFAULT
+    utask_flag = feature_flags.FeatureFlags.UTASK_MAIN_QUEUE_LIMIT.flag
+    if utask_flag and utask_flag.enabled:
+      utask_main_queue_limit = utask_flag.content
+    if utask_main_queue_size > utask_main_queue_limit:
       raise errors.QueueLimitReachedError(utask_main_queue_size)
 
     logs.info('Preprocessing utask.')

src/clusterfuzz/_internal/k8s/service.py

@@ -16,7 +16,6 @@
 import collections
 import ipaddress
 import os
-import typing
 import uuid
 
 import google.auth
@@ -53,7 +52,7 @@
 ])
 
 
-def _get_config_names(remote_tasks: typing.List[remote_task_types.RemoteTask]):
+def _get_config_names(remote_tasks: list[remote_task_types.RemoteTask]):
   """Gets the name of the configs for each batch_task. Returns a dict
 
   that is indexed by command and job_type for efficient lookup."""
@@ -90,9 +89,8 @@ def _get_config_names(remote_tasks: typing.List[remote_task_types.RemoteTask]):
   return config_map
 
 
-def _get_k8s_job_configs(
-    remote_tasks: typing.List[remote_task_types.RemoteTask]
-) -> typing.Dict[typing.Tuple[str, str], KubernetesJobConfig]:
+def _get_k8s_job_configs(remote_tasks: list[remote_task_types.RemoteTask]
+                        ) -> dict[tuple[str, str], KubernetesJobConfig]:
   """Gets the configured specifications for a batch workload."""
 
   if not remote_tasks:
@@ -286,7 +284,7 @@ def _get_pending_jobs_count(self) -> int:
 
   def create_utask_main_job(self, module: str, job_type: str,
                             input_download_url: str):
-    """Creates a single batch job for a uworker main task."""
+    """Creates a single Kubernetes job for a uworker main task."""
 
     command = task_utils.get_command_from_module(module)
     batch_tasks = [
@@ -295,13 +293,13 @@ def create_utask_main_job(self, module: str, job_type: str,
     uncreated_tasks = self.create_utask_main_jobs(batch_tasks)
     return uncreated_tasks
 
-  def create_utask_main_jobs(
-      self, remote_tasks: typing.List[remote_task_types.RemoteTask]):
-    """Creates a batch job for a list of uworker main tasks.
+  def create_utask_main_jobs(self,
+                             remote_tasks: list[remote_task_types.RemoteTask]):
+    """Creates Kubernetes jobs 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.
+    This method groups the tasks by their workload specification to efficiently
+    schedule them. It then iterates through the groups and creates a separate
+    Kubernetes job for each task.
     """
 
     flag = feature_flags.FeatureFlags.K8S_JOBS_PENDING_LIMIT.flag
@@ -311,7 +309,8 @@ def create_utask_main_jobs(
     pending_jobs_count = self._get_pending_jobs_count()
     if pending_jobs_count >= limit:
       logs.warning(
-          f'Pending jobs count {pending_jobs_count} reached limit {limit}.')
+          f'Pending jobs count {pending_jobs_count} reached limit {limit} '
+          f'for k8s.')
       return remote_tasks
 
     job_specs = collections.defaultdict(list)
@@ -321,7 +320,7 @@ def create_utask_main_jobs(
           f'Scheduling {remote_task.command}, {remote_task.job_type} in K8s.')
       config = configs[(remote_task.command, remote_task.job_type)]
       job_specs[config].append(remote_task.input_download_url)
-    logs.info('Creating GCP batch jobs.')
+    logs.info('Creating Kubernetes jobs.')
     for config, input_urls in job_specs.items():
       for input_url in input_urls:
         self.create_job(config, input_url)

src/clusterfuzz/_internal/remote_task/remote_task_gate.py

@@ -21,10 +21,7 @@
 
 import collections
 import random
-import typing
 
-from clusterfuzz._internal.base.tasks.task_utils import is_testcase_based_task
-from clusterfuzz._internal.bot.tasks.utasks import uworker_io
 from clusterfuzz._internal.metrics import logs
 from clusterfuzz._internal.remote_task import remote_task_adapters
 from clusterfuzz._internal.remote_task import remote_task_types
@@ -107,44 +104,14 @@ def get_job_frequency(self):
     logs.info('Job frequencies', frequencies=frequencies)
     return frequencies
 
-  def prepare_unscheduled_tasks(
-      self, unscheduled_remote_tasks: typing.List[remote_task_types.RemoteTask]
-  ) -> typing.List[remote_task_types.RemoteTask]:
-    """Prepares the unscheduled remote tasks to be sent back to the 
-      preprocess queue.
-
-      The messages in the preprocess queue expects tasks with 
-      the properties followng task_name, fuzzer, job, eta. 
-      The utasks has the signed url as argument instead of the fuzzer,
-      then for sending it back to the preprocess, it recovers its fuzzer
-      and override the argument.
-    """
-
-    prepared_tasks = []
-    if not unscheduled_remote_tasks:
-      return prepared_tasks
-
-    for task in unscheduled_remote_tasks:
-      try:
-        uworker_input = uworker_io.download_and_deserialize_uworker_input(
-            task.input_download_url)
-        if is_testcase_based_task(task.command):
-          task.argument = uworker_input.testcase_id
-        else:
-          task.argument = uworker_input.fuzzer_name
-        prepared_tasks.append(task)
-      except Exception:
-        logs.error('Could not prepare the task due to an error.')
-
-    return prepared_tasks
-
   def create_utask_main_job(self, module, job_type, input_download_url):
+    """Creates a single remote task, selecting a backend dynamically."""
     adapter_id = self._get_adapter()
     service = self._service_map[adapter_id]
     return service.create_utask_main_job(module, job_type, input_download_url)
 
-  def create_utask_main_jobs(
-      self, remote_tasks: typing.List[remote_task_types.RemoteTask]):
+  def create_utask_main_jobs(self,
+                             remote_tasks: list[remote_task_types.RemoteTask]):
     """Creates a batch of remote tasks, distributing them across backends.
 
     This method handles two cases:
@@ -161,6 +128,7 @@ def create_utask_main_jobs(
       # For a single task, use a random distribution.
       adapter_id = self._get_adapter()
       tasks_by_adapter[adapter_id].extend(remote_tasks)
+      unscheduled_tasks = []
     else:
       # For multiple tasks, use deterministic slicing to ensure the
       # distribution precisely matches the frequency configuration.
@@ -172,14 +140,17 @@ def create_utask_main_jobs(
             remote_tasks[start_index:start_index + count])
         start_index += count
 
-      # Distribute any remainder tasks (due to rounding) one by one. This
-      # ensures that all tasks are assigned to a backend.
       remaining_tasks = remote_tasks[start_index:]
-      for i, task in enumerate(remaining_tasks):
-        adapter_id = list(frequencies.keys())[i % len(frequencies)]
-        tasks_by_adapter[adapter_id].append(task)
+      if sum(frequencies.values()) >= 0.999:
+        # Distribute any remainder tasks (due to rounding) one by one. This
+        # ensures that all tasks are assigned to a backend.
+        for i, task in enumerate(remaining_tasks):
+          adapter_id = list(frequencies.keys())[i % len(frequencies)]
+          tasks_by_adapter[adapter_id].append(task)
+        unscheduled_tasks = []
+      else:
+        unscheduled_tasks = list(remaining_tasks)
 
-    unscheduled_tasks = []
     for adapter_id, tasks in tasks_by_adapter.items():
       if tasks:
         try:

src/clusterfuzz/_internal/remote_task/remote_task_types.py

@@ -14,7 +14,6 @@
 """Remote task types."""
 
 import abc
-import typing
 
 from clusterfuzz._internal.base import tasks
 
@@ -53,7 +52,7 @@ def create_utask_main_job(self, module: str, job_type: str,
 
   @abc.abstractmethod
   def create_utask_main_jobs(
-      self, remote_tasks: typing.List[RemoteTask]) -> typing.List[RemoteTask]:
+      self, remote_tasks: list[RemoteTask]) -> list[RemoteTask]:
     """Creates many remote tasks for uworker main tasks.
        Returns the tasks that couldn't be created.
     """

src/clusterfuzz/_internal/tests/core/base/tasks/queue_size_test.py

@@ -1,4 +1,4 @@
-# Copyright 2023 Google LLC
+# Copyright 2026 Google LLC
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

src/clusterfuzz/_internal/tests/core/batch/batch_service_test.py

@@ -157,7 +157,7 @@ def setUp(self):
     self.batch_service = batch_service.GcpBatchService()
     self.mock.uuid4.side_effect = [uuid.UUID(u) for u in UUIDS]
 
-  def test_create_uworker_main_batch_jobs(self):
+  def test_create_utask_main_jobs(self):
     """Tests that create_utask_main_jobs works as expected."""
     # Create mock data.
     spec1 = batch_service.BatchWorkloadSpec(
@@ -217,7 +217,7 @@ def test_create_uworker_main_batch_jobs(self):
           mock.call(expected_create_request_2),
       ])
 
-  def test_create_uworker_main_batch_job(self):
+  def test_create_utask_main_job(self):
     """Tests that create_utask_main_job works as expected."""
     # Create mock data.
     spec1 = batch_service.BatchWorkloadSpec(

src/clusterfuzz/_internal/tests/core/bot/startup/run_bot_test.py

@@ -135,7 +135,6 @@ def setUp(self):
     helpers.patch(self, [
         'clusterfuzz._internal.base.tasks.get_utask_mains',
         'clusterfuzz._internal.remote_task.remote_task_gate.RemoteTaskGate',
-        'clusterfuzz._internal.base.tasks.bulk_add_tasks',
     ])
 
   def test_schedule_tasks_requeue_uncreated(self):
@@ -157,9 +156,6 @@ def test_schedule_tasks_requeue_uncreated(self):
     self.mock.RemoteTaskGate.return_value.create_utask_main_jobs.assert_called_once(
     )
 
-    # Check that bulk_add_tasks was NOT called.
-    self.mock.bulk_add_tasks.assert_not_called()
-
     # Verify that cancel_lease_ack was called on the pubsub task.
     mock_task.cancel_lease_ack.assert_called_once()
 
@@ -179,4 +175,3 @@ def test_schedule_tasks_success(self):
 
     self.mock.RemoteTaskGate.return_value.create_utask_main_jobs.assert_called_once(
     )
-    self.mock.bulk_add_tasks.assert_not_called()

src/clusterfuzz/_internal/tests/core/bot/tasks/task_types_test.py

@@ -137,6 +137,7 @@ def test_event_emit_during_exception(self):
     self.mock.emit.assert_any_call(event_finished)
 
 
+@test_utils.with_cloud_emulators('datastore')
 class UTaskExecuteTest(unittest.TestCase):
   """Tests for UTask execution."""
 

src/clusterfuzz/_internal/tests/core/k8s/k8s_service_e2e_test.py

@@ -150,8 +150,8 @@ def test_create_job(self):
   @mock.patch('clusterfuzz._internal.k8s.service._get_k8s_job_configs')
   @mock.patch(
       'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module')
-  def test_create_uworker_main_batch_job(self, mock_get_command_from_module,
-                                         mock_get_k8s_job_configs):
+  def test_create_utask_main_job(self, mock_get_command_from_module,
+                                 mock_get_k8s_job_configs):
     """Tests creating a single uworker main batch job."""
     mock_get_command_from_module.return_value = 'fuzz'
     config = KubernetesJobConfig(
@@ -175,8 +175,8 @@ def test_create_uworker_main_batch_job(self, mock_get_command_from_module,
   @mock.patch('clusterfuzz._internal.k8s.service._get_k8s_job_configs')
   @mock.patch(
       'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module')
-  def test_create_uworker_main_batch_jobs(self, mock_get_command_from_module,
-                                          mock_get_k8s_job_configs):
+  def test_create_utask_main_jobs(self, mock_get_command_from_module,
+                                  mock_get_k8s_job_configs):
     """Tests creating multiple uworker main batch jobs."""
     mock_get_command_from_module.return_value = 'fuzz'
     config1 = KubernetesJobConfig(