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:
@@ -295,8 +293,8 @@ 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]):
+  def create_utask_main_jobs(self,
+                             remote_tasks: list[remote_task_types.RemoteTask]):
     """Creates a batch job for a list of uworker main tasks.
 
     This method groups the tasks by their workload specification and creates a
@@ -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,13 @@ 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):
     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 +127,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 +139,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/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/remote_task/remote_task_gate_test.py

@@ -55,12 +55,6 @@ def setUp(self):
     self.patcher.start()
     self.addCleanup(self.patcher.stop)
 
-    self.mock_prepare_unscheduled_tasks = mock.patch.object(
-        remote_task_gate.RemoteTaskGate,
-        'prepare_unscheduled_tasks',
-        side_effect=lambda x: x).start()
-    self.addCleanup(mock.patch.stopall)
-
   def test_init(self):
     """Tests that the RemoteTaskGate initializes correctly and creates
     service map."""
@@ -161,11 +155,8 @@ def test_create_utask_main_jobs_remainder_distribution(
         remote_task_types.RemoteTask('c', 'j', 'u3'),
     ]
 
-    # 33/33/33 split - one task will be a remainder
-    mock_get_job_frequency.return_value = {
-        'kubernetes': 0.33,
-        'gcp_batch': 0.33
-    }
+    # 50/50 split - one task will be a remainder
+    mock_get_job_frequency.return_value = {'kubernetes': 0.5, 'gcp_batch': 0.5}
 
     gate = remote_task_gate.RemoteTaskGate()
     gate.create_utask_main_jobs(tasks)
@@ -177,6 +168,38 @@ def test_create_utask_main_jobs_remainder_distribution(
     self.mock_gcp_batch_service.create_utask_main_jobs.assert_called_once_with(
         [tasks[1]])
 
+  @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency')
+  def test_create_utask_main_jobs_unscheduled(self, mock_get_job_frequency):
+    """Tests that create_utask_main_jobs returns remainder as unscheduled
+    when sum < 1.0."""
+    tasks = [
+        remote_task_types.RemoteTask('c', 'j', 'u1'),
+        remote_task_types.RemoteTask('c', 'j', 'u2'),
+        remote_task_types.RemoteTask('c', 'j', 'u3'),
+        remote_task_types.RemoteTask('c', 'j', 'u4'),
+    ]
+
+    # 0.25 each. Sum 0.5.
+    mock_get_job_frequency.return_value = {
+        'kubernetes': 0.25,
+        'gcp_batch': 0.25
+    }
+
+    gate = remote_task_gate.RemoteTaskGate()
+    result = gate.create_utask_main_jobs(tasks)
+
+    # 4 * 0.25 = 1 task each.
+    # Total assigned = 2.
+    # Total unscheduled = 2.
+
+    self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with(
+        [tasks[0]])
+    self.mock_gcp_batch_service.create_utask_main_jobs.assert_called_once_with(
+        [tasks[1]])
+
+    # Result should contain unscheduled (tasks[2], tasks[3]).
+    self.assertEqual(result, [tasks[2], tasks[3]])
+
   @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency')
   def test_create_utask_main_jobs_full_kubernetes(self, mock_get_job_frequency):
     """Tests that all tasks are routed to Kubernetes when frequency is 1.0."""
@@ -222,7 +245,6 @@ def test_create_utask_main_jobs_returns_unscheduled_tasks(
     result = gate.create_utask_main_jobs(tasks)
 
     self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with(tasks)
-    self.mock_prepare_unscheduled_tasks.assert_not_called()
     self.assertEqual(result, unscheduled_tasks)
 
 
@@ -250,47 +272,3 @@ def setUp(self):
     self.patcher.start()
     self.addCleanup(self.patcher.stop)
     self.gate = remote_task_gate.RemoteTaskGate()
-
-  @mock.patch(
-      'clusterfuzz._internal.bot.tasks.utasks.uworker_io.download_and_deserialize_uworker_input'
-  )
-  def test_prepare_unscheduled_tasks(self, mock_download):
-    """Tests that prepare_unscheduled_tasks correctly updates task arguments."""
-
-    def side_effect(url):
-      mock_input = mock.Mock()
-      if url == 'signed_url_1':
-        mock_input.testcase_id = None
-        mock_input.fuzzer_name = 'fuzzer1'
-      elif url == 'signed_url_2':
-        mock_input.testcase_id = None
-        mock_input.fuzzer_name = 'fuzzer2'
-      elif url == 'signed_url_3':
-        mock_input.testcase_id = '12345'
-        mock_input.fuzzer_name = 'fuzzer3'
-      return mock_input
-
-    mock_download.side_effect = side_effect
-
-    task1 = remote_task_types.RemoteTask(
-        command='fuzz', job_type='job1', input_download_url='signed_url_1')
-    task2 = remote_task_types.RemoteTask(
-        command='fuzz', job_type='job2', input_download_url='signed_url_2')
-    task3 = remote_task_types.RemoteTask(
-        command='progression',
-        job_type='job3',
-        input_download_url='signed_url_3')
-
-    tasks = [task1, task2, task3]
-
-    updated_tasks = self.gate.prepare_unscheduled_tasks(tasks)
-
-    self.assertEqual(len(updated_tasks), 3)
-    self.assertEqual(updated_tasks[0].argument, 'fuzzer1')
-    self.assertEqual(updated_tasks[1].argument, 'fuzzer2')
-    self.assertEqual(updated_tasks[2].argument, '12345')
-
-  def test_prepare_unscheduled_tasks_empty(self):
-    """Tests that prepare_unscheduled_tasks handles empty list of tasks."""
-    updated_tasks = self.gate.prepare_unscheduled_tasks([])
-    self.assertEqual(updated_tasks, [])

src/clusterfuzz/_internal/tests/core/remote_task/remote_task_test.py

@@ -47,14 +47,6 @@ def setUp(self):
     self.addCleanup(patcher.stop)
     patcher.start()
 
-    # Mock prepare_unscheduled_tasks to avoid NDB context issues
-    patcher = mock.patch.object(
-        remote_task_gate.RemoteTaskGate,
-        'prepare_unscheduled_tasks',
-        side_effect=lambda x: x)
-    self.addCleanup(patcher.stop)
-    patcher.start()
-
     self.gate = remote_task_gate.RemoteTaskGate()
 
   @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency')