Swarming: Removes Unnecessary DB calls (#5223)

We are adding back the last batch of the [changes that were
rollbacked](c964710)
which make sure that we weren't making a lot of unnecessary calls into
the DB. By consequence the signature of multiple methods were changed to
simplified.


I also added a new tests to avoid the previous issue (the one that
caused the rollback) from ever appearing again. This tests are in its
own test class and its purpose is to check that the BadConfigError is
correctly handled when no config is found(which is the expected
bahaviour in the external and google envs)

## Tests performed in `dev` environment
This changes have been present since March 25th. Its been over one week
since then and metrics seem consistent
<img width="957" height="337" alt="image"
src="https://github.com/user-attachments/assets/1ffcdba5-9adc-434b-831e-905c4653357e"
/>
Note: Down times in fuzzing hours matches when we deployed a change into
the env.

Also no new error groups seem to be appearing in swarming, the only
notable error is one relating missing permissions for uploading a test
case, which seems unrelated

This changes have been in `dev` for quite some time and we have
successfully tested their functionality without compromising any
functionality. For example [this
Task](https://chrome-swarming.appspot.com/task?id=7743c89f08030410) was
scheduled with this changes.

---------

Co-authored-by: Diego Jardon <37823380+jardondiego@users.noreply.github.com>
Co-authored-by: Javan Lacerda <javanlacerda@google.com>

Author: 3 people

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

@@ -115,8 +115,8 @@ def is_remote_utask(command, job):
     # Return True even if we can't query the db.
     return True
 
-  return batch_service.is_remote_task(
-      command, job) or swarming.is_swarming_task(command, job)
+  return batch_service.is_remote_task(command,
+                                      job) or swarming.is_swarming_task(job)
 
 
 def task_main_runs_on_uworker():
@@ -178,11 +178,8 @@ def execute(self, task_argument, job_type, uworker_env):
       return
 
     logs.info('Queueing utask for remote execution.', download_url=download_url)
-    if batch_service.is_remote_task(command, job_type):
-      tasks.add_utask_main(command, download_url, job_type)
-    else:
-      assert swarming.is_swarming_task(command, job_type)
-      swarming.push_swarming_task(command, download_url, job_type)
+    assert batch_service.is_remote_task(command, job_type)
+    tasks.add_utask_main(command, download_url, job_type)
 
   @logs.task_stage_context(logs.Stage.PREPROCESS)
   def preprocess(self, task_argument, job_type, uworker_env):

src/clusterfuzz/_internal/bot/tasks/utasks/__init__.py

@@ -63,10 +63,9 @@ def _timestamp_now() -> Timestamp:
   return ts
 
 
-def _get_execution_mode(utask_module, job_type):
+def _get_execution_mode(job_type):
   """Determines whether this task in executed on swarming or batch."""
-  command = task_utils.get_command_from_module(utask_module.__name__)
-  if swarming.is_swarming_task(command, job_type):
+  if swarming.is_swarming_task(job_type):
     return Mode.SWARMING
   return Mode.BATCH
 
@@ -410,7 +409,7 @@ def tworker_preprocess(utask_module, task_argument, job_type, uworker_env):
   signed download URL for the uworker's input and the (unsigned) download URL
   for its output."""
   with _MetricRecorder(_Subtask.PREPROCESS) as recorder:
-    execution_mode = _get_execution_mode(utask_module, job_type)
+    execution_mode = _get_execution_mode(job_type)
     uworker_input = _preprocess(utask_module, task_argument, job_type,
                                 uworker_env, recorder, execution_mode)
     if not uworker_input:
@@ -501,8 +500,7 @@ def tworker_postprocess(output_download_url) -> None:
     task_utils.reset_task_stage_env()
 
     utask_module = get_utask_module(uworker_output.uworker_input.module_name)
-    execution_mode = _get_execution_mode(utask_module,
-                                         uworker_output.uworker_input.job_type)
+    execution_mode = _get_execution_mode(uworker_output.uworker_input.job_type)
     recorder.set_task_details(
         utask_module,
         uworker_output.uworker_input.job_type,

src/clusterfuzz/_internal/metrics/logs.py

@@ -1013,6 +1013,7 @@ class Stage(enum.Enum):
   MAIN = 'main'
   POSTPROCESS = 'postprocess'
   UNKNOWN = 'unknown'
+  SCHEDULER = 'scheduler'
   NA = 'n/a'
 
 

src/clusterfuzz/_internal/remote_task/remote_task_gate.py

@@ -24,7 +24,6 @@
 
 from clusterfuzz._internal import swarming
 from clusterfuzz._internal.base import feature_flags
-from clusterfuzz._internal.base.tasks import task_utils
 from clusterfuzz._internal.metrics import logs
 from clusterfuzz._internal.remote_task import remote_task_adapters
 from clusterfuzz._internal.remote_task import remote_task_types
@@ -61,9 +60,8 @@ def _get_adapter(self) -> str:
   def _is_swarming_applicable(self):
     return feature_flags.FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled
 
-  def _is_swarming_task(self, module, job_type):
-    return swarming.is_swarming_task(
-        task_utils.get_command_from_module(module), job_type)
+  def _is_swarming_task(self, job_type):
+    return swarming.is_swarming_task(job_type)
 
   def _handle_swarming_job(self, module, job_type, input_download_url):
     return self._service_map['swarming'].create_utask_main_job(
@@ -124,8 +122,7 @@ def get_job_frequency(self):
 
   def create_utask_main_job(self, module, job_type, input_download_url):
     """Creates a single remote task, selecting a backend dynamically."""
-    if self._is_swarming_applicable() and self._is_swarming_task(
-        module, job_type):
+    if self._is_swarming_applicable() and self._is_swarming_task(job_type):
       return self._handle_swarming_job(module, job_type, input_download_url)
 
     adapter_id = self._get_adapter()
@@ -156,6 +153,7 @@ def create_utask_main_jobs(self,
     unscheduled_tasks = []
 
     if self._is_swarming_applicable():
+      logs.info(f'[Swarming] enabled, pushing {len(remote_tasks)} tasks.')
       remote_tasks = self._handle_swarming_jobs(remote_tasks)
 
     if not remote_tasks:

src/clusterfuzz/_internal/swarming/__init__.py

@@ -21,6 +21,7 @@
 from google.protobuf import json_format
 
 from clusterfuzz._internal.base import utils
+from clusterfuzz._internal.base.errors import BadConfigError
 from clusterfuzz._internal.base.feature_flags import FeatureFlags
 from clusterfuzz._internal.config import local_config
 from clusterfuzz._internal.datastore import data_types
@@ -35,38 +36,62 @@
 ]
 
 
-def is_swarming_task(command: str, job_name: str):
+def is_swarming_task(job_name: str, job: data_types.Job | None = None) -> bool:
   """Returns True if the task is supposed to run on swarming."""
   if not FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled:
+    logs.info('[DEBUG] Flag is disabled', job_name=job_name)
     return False
-  job = data_types.Job.query(data_types.Job.name == job_name).get()
-  if not job:
-    return False
+  if job is None:
+    job = data_types.Job.query(data_types.Job.name == job_name).get()
+    if not job:
+      logs.info('[Swarming DEBUG] Job not found', job_name=job_name)
+      return False
 
   job_environment = job.get_environment()
-  if not utils.string_is_true(job_environment.get('IS_SWARMING_JOB')):
+  if not utils.string_is_true(job_environment.get(
+      'IS_SWARMING_JOB')) and not job_environment.get('SWARMING_DIMENSIONS'):
+    logs.info('[Swarming DEBUG] No swarming env var', job_name=job_name)
     return False
 
-  try:
-    _get_new_task_spec(command, job_name, '')
-    return True
-  except ValueError:
+  swarming_config = _get_swarming_config()
+  if swarming_config is None:
+    logs.warning(
+        """[Swarming DEBUG] current task is not suitable for swarming. 
+    'Reason: failed to retrieve config.""",
+        job_name=job_name)
     return False
 
+  return _get_instance_spec(swarming_config, job) is not None
+
 
-def _get_task_name():
-  return 't-' + str(uuid.uuid4()).lower()
+def _get_instance_spec(swarming_config: local_config.SwarmingConfig,
+                       job: data_types.Job) -> dict | None:
+  return swarming_config.get('mapping').get(job.platform, None)
 
 
-def _get_swarming_config():
+def _get_task_name(job_name: str):
+  return f't-{str(uuid.uuid4()).lower()}-{job_name}'
+
+
+def _get_swarming_config() -> local_config.SwarmingConfig | None:
   """Returns the swarming config."""
-  return local_config.SwarmingConfig()
+  try:
+    return local_config.SwarmingConfig()
+  except (BadConfigError, ValueError) as e:
+    logs.error(f'[Swarming] Failed to retrieve config: {e}')
+    return None
 
 
 def _get_task_dimensions(job: data_types.Job, platform_specific_dimensions: list
                         ) -> list[swarming_pb2.StringPair]:  # pylint: disable=no-member
   """ Gets all swarming dimensions for a task.
   Job dimensions have more precedence than static dimensions"""
+  swarming_config = _get_swarming_config()
+  if not swarming_config:
+    logs.error(
+        '[Swarming] No dimensions set. Reason: failed to retrieve config')
+    return []
+
   unique_dimensions = {}
   unique_dimensions['os'] = str(job.platform).capitalize()
   unique_dimensions['pool'] = _get_swarming_config().get('swarming_pool')
@@ -100,16 +125,27 @@ def _env_vars_to_json(
       value=json.dumps(env_vars_dict))
 
 
-def _get_new_task_spec(command: str, job_name: str,
-                       download_url: str) -> swarming_pb2.NewTaskRequest:  # pylint: disable=no-member
-  """Gets the configured specifications for a swarming task."""
+def create_new_task_request(command: str, job_name: str, download_url: str
+                           ) -> swarming_pb2.NewTaskRequest | None:  # pylint: disable=no-member
+  """Gets the configured specifications for a swarming task. 
+  Returns None if the task should'nt be executed on swarming 
+  or if the SWARMING_REMOTE_EXECUTION flag is disabled."""
+  if not FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled:
+    return None
+
   job = data_types.Job.query(data_types.Job.name == job_name).get()
-  config_name = job.platform
+  if job is None:
+    return None
+
   swarming_config = _get_swarming_config()
-  instance_spec = swarming_config.get('mapping').get(config_name, None)
+  if not swarming_config:
+    return None
+
+  instance_spec = _get_instance_spec(swarming_config, job)
   if instance_spec is None:
-    raise ValueError(f'No mapping for {config_name}')
-  swarming_realm = swarming_config.get('swarming_realm')
+    return None
+
+  swarming_realm = swarming_config.get('swarming_realm',)
   logs_project_id = swarming_config.get('logs_project_id')
   priority = instance_spec['priority']
   startup_command = instance_spec['command']
@@ -155,7 +191,7 @@ def _get_new_task_spec(command: str, job_name: str,
   cas_input_root = instance_spec.get('cas_input_root', {})
 
   new_task_request = swarming_pb2.NewTaskRequest(  # pylint: disable=no-member
-      name=_get_task_name(),
+      name=_get_task_name(job_name),
       priority=priority,
       realm=swarming_realm,
       service_account=service_account,
@@ -176,13 +212,13 @@ def _get_new_task_spec(command: str, job_name: str,
   return new_task_request
 
 
-def push_swarming_task(command, download_url, job_type):
+def push_swarming_task(task_request: swarming_pb2.NewTaskRequest):  # pylint: disable=no-member
   """Schedules a task on swarming."""
-  job = data_types.Job.query(data_types.Job.name == job_type).get()
-  if not job:
-    raise ValueError('invalid job_name')
-
-  task_spec = _get_new_task_spec(command, job_type, download_url)
+  swarming_config = _get_swarming_config()
+  if not swarming_config:
+    logs.error(
+        '[Swarming] Failed to push task into swarming. Reason: No config.')
+    return
   creds = credentials.get_scoped_service_account_credentials(_SWARMING_SCOPES)
   if not creds:
     logs.error(
@@ -199,11 +235,11 @@ def push_swarming_task(command, download_url, job_type):
   }
   swarming_server = _get_swarming_config().get('swarming_server')
   url = f'https://{swarming_server}/prpc/swarming.v2.Tasks/NewTask'
-  message_body = json_format.MessageToJson(task_spec)
+  message_body = json_format.MessageToJson(task_request)
   logs.info(
-      f"""[Swarming] Pushing task for {job_type}
+      f"""[Swarming] Pushing task {task_request.name}
             as {creds.service_account_email}""",
       url=url,
       body=message_body)
   response = utils.post_url(url=url, data=message_body, headers=headers)
-  logs.info(f'[Swarming] Response from {job_type}', response=response)
+  logs.info(f'[Swarming] Response from {task_request.name}', response=response)

src/clusterfuzz/_internal/swarming/service.py

@@ -35,21 +35,23 @@ def create_utask_main_job(self, module: str, job_type: str,
 
     return result[0]
 
+  @logs.task_stage_context(logs.Stage.SCHEDULER)
   def create_utask_main_jobs(self,
                              remote_tasks: list[remote_task_types.RemoteTask]
                             ) -> list[remote_task_types.RemoteTask]:
     """Creates many remote tasks for uworker main tasks.
        Returns the tasks that couldn't be created.
     """
     unscheduled_tasks = []
+    logs.info(f'[Swarming] Pushing {len(remote_tasks)} tasks trough service.')
     for task in remote_tasks:
       try:
-        if not swarming.is_swarming_task(task.command, task.job_type):
+        if not swarming.is_swarming_task(task.job_type):
           unscheduled_tasks.append(task)
           continue
-
-        swarming.push_swarming_task(task.command, task.input_download_url,
-                                    task.job_type)
+        if request := swarming.create_new_task_request(
+            task.command, task.job_type, task.argument):
+          swarming.push_swarming_task(request)
       except Exception:  # pylint: disable=broad-except
         logs.error(
             f'Failed to push task to Swarming: {task.command}, {task.job_type}.'

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

@@ -518,15 +518,13 @@ def test_is_swarming_applicable(self, mock_swarming_flag):
     mock_swarming_flag.return_value = False
     self.assertFalse(self.gate._is_swarming_applicable())
 
-  @mock.patch('clusterfuzz._internal.remote_task.remote_task_gate.task_utils')
   @mock.patch('clusterfuzz._internal.remote_task.remote_task_gate.swarming')
-  def test_is_swarming_task(self, mock_swarming, mock_task_utils):
+  def test_is_swarming_task(self, mock_swarming):
     """Tests _is_swarming_task."""
-    mock_task_utils.get_command_from_module.return_value = 'fuzz'
     mock_swarming.is_swarming_task.return_value = True
 
-    self.assertTrue(self.gate._is_swarming_task('module', 'job'))
-    mock_swarming.is_swarming_task.assert_called_once_with('fuzz', 'job')
+    self.assertTrue(self.gate._is_swarming_task('job'))
+    mock_swarming.is_swarming_task.assert_called_once_with('job')
 
   def test_handle_swarming_job(self):
     """Tests _handle_swarming_job."""

src/clusterfuzz/_internal/tests/core/swarming/service_test.py

@@ -28,10 +28,12 @@ def setUp(self):
     helpers.patch(self, [
         'clusterfuzz._internal.swarming.is_swarming_task',
         'clusterfuzz._internal.swarming.push_swarming_task',
+        'clusterfuzz._internal.swarming.create_new_task_request',
         'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module',
         'clusterfuzz._internal.metrics.logs.error',
     ])
     self.service = service.SwarmingService()
+    self.mock.create_new_task_request.return_value = 'fake_request'
 
   def test_create_utask_main_job_success(self):
     """Test creating a single task successfully."""
@@ -44,8 +46,7 @@ def test_create_utask_main_job_success(self):
     # Success returns None in this interface (consistent with GcpBatchService)
     self.assertIsNone(result)
 
-    self.mock.push_swarming_task.assert_called_once_with(
-        'fuzz', 'http://url', 'job_type')
+    self.mock.push_swarming_task.assert_called_once_with('fake_request')
 
   def test_create_utask_main_job_failure(self):
     """Test creating a single task that is not a swarming task."""
@@ -78,8 +79,8 @@ def test_create_utask_main_jobs_mixed_results(self):
 
     self.assertEqual(self.mock.push_swarming_task.call_count, 2)
     self.mock.push_swarming_task.assert_has_calls([
-        mock.call('fuzz', 'url1', 'job1'),
-        mock.call('fuzz', 'url3', 'job3'),
+        mock.call('fake_request'),
+        mock.call('fake_request'),
     ])
 
   def test_create_utask_main_jobs_all_success(self):