From 11128be4fa3d183348e8baa174f97dda83f9f5fb Mon Sep 17 00:00:00 2001 From: Ivan Barba Date: Fri, 25 Sep 2026 00:41:12 +0000 Subject: [PATCH] [Swarming][Scheduling] Splits swarming tasks in scheduler --- .../_internal/remote_task/remote_task_gate.py | 25 ++- .../_internal/swarming/__init__.py | 19 +- .../core/remote_task/remote_task_gate_test.py | 162 ++++++++++++++---- 3 files changed, 163 insertions(+), 43 deletions(-) diff --git a/src/clusterfuzz/_internal/remote_task/remote_task_gate.py b/src/clusterfuzz/_internal/remote_task/remote_task_gate.py index 0425f29c348..c7c85796430 100644 --- a/src/clusterfuzz/_internal/remote_task/remote_task_gate.py +++ b/src/clusterfuzz/_internal/remote_task/remote_task_gate.py @@ -62,7 +62,7 @@ def _is_swarming_applicable(self): return feature_flags.FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled def _is_swarming_task(self, job_type): - return swarming.is_swarming_task(job_type) + return swarming.is_swarming_task(job_type, ignore_feature_flag=True) def _handle_swarming_job(self, module, job_type, input_download_url): return self._service_map['swarming'].create_utask_main_job( @@ -72,6 +72,19 @@ def _handle_swarming_jobs(self, remote_tasks: list[remote_task_types.RemoteTask]): return self._service_map['swarming'].create_utask_main_jobs(remote_tasks) + def _split_remote_tasks( + self, remote_tasks: list[remote_task_types.RemoteTask]) -> tuple[list[ + remote_task_types.RemoteTask], list[remote_task_types.RemoteTask]]: + """Splits remote tasks into swarming and non-swarming tasks.""" + swarming_tasks = [] + other_tasks = [] + for task in remote_tasks: + if self._is_swarming_task(task.job_type): + swarming_tasks.append(task) + else: + other_tasks.append(task) + return swarming_tasks, other_tasks + def get_job_frequency(self): """Returns the frequency distribution for all remote task adapters. @@ -153,9 +166,15 @@ def create_utask_main_jobs(self, tasks_by_adapter = collections.defaultdict(list) unscheduled_tasks = [] + swarming_tasks, remote_tasks = self._split_remote_tasks(remote_tasks) if self._is_swarming_applicable(): - logs.info(f'[Swarming] enabled, pushing {len(remote_tasks)} tasks.') - remote_tasks = self._handle_swarming_jobs(remote_tasks) + logs.info(f'[Swarming] scheduling {len(swarming_tasks)} tasks.') + unscheduled_tasks = self._handle_swarming_jobs(swarming_tasks) + else: + logs.info( + f'[Swarming] Swarming flag not enabled, {len(swarming_tasks)} tasks' + ' unscheduled.') + unscheduled_tasks.extend(swarming_tasks) if not remote_tasks: pass diff --git a/src/clusterfuzz/_internal/swarming/__init__.py b/src/clusterfuzz/_internal/swarming/__init__.py index 2c0755e9c91..15a8720800a 100644 --- a/src/clusterfuzz/_internal/swarming/__init__.py +++ b/src/clusterfuzz/_internal/swarming/__init__.py @@ -36,9 +36,22 @@ def has_swarming_env_vars(job_environment: dict) -> bool: job_environment.get('SWARMING_DIMENSIONS')) -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: +def is_swarming_task(job_name: str, + job: data_types.Job | None = None, + ignore_feature_flag: bool = False) -> bool: + """Validates that the current job and environment can send a task to swarming. + + Args: + job_name: The name of the job. + job: The job object, use if available to avoid querying datastore. + ignore_feature_flag: So that we check the job even if the feature flag is + not enabled. + + Returns: + True if the task is supposed to run on swarming. + """ + if not (ignore_feature_flag or + FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled): logs.info('[DEBUG] Flag is disabled', job_name=job_name) return False if job is None: diff --git a/src/clusterfuzz/_internal/tests/core/remote_task/remote_task_gate_test.py b/src/clusterfuzz/_internal/tests/core/remote_task/remote_task_gate_test.py index e8545bfd68e..564201d37b2 100644 --- a/src/clusterfuzz/_internal/tests/core/remote_task/remote_task_gate_test.py +++ b/src/clusterfuzz/_internal/tests/core/remote_task/remote_task_gate_test.py @@ -63,6 +63,12 @@ def setUp(self): }) self.patcher.start() self.addCleanup(self.patcher.stop) + is_swarming_task_patcher = mock.patch.object( + remote_task_gate.RemoteTaskGate, + '_is_swarming_task', + return_value=False) + is_swarming_task_patcher.start() + self.addCleanup(is_swarming_task_patcher.stop) def test_init(self): """Tests that the RemoteTaskGate initializes correctly and creates @@ -351,21 +357,22 @@ def test_create_utask_main_jobs_returns_unscheduled_tasks( self.assertEqual(result, unscheduled_tasks) @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_applicable') + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_task') @mock.patch.object(remote_task_gate.RemoteTaskGate, '_handle_swarming_jobs') @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency') def test_create_utask_main_jobs_swarming_remote_execution_enabled( self, mock_get_job_frequency, mock_handle_swarming_jobs, - mock_is_swarming_applicable): - """Tests that create_utask_main_jobs passes tasks to swarming service when the - feature flag is enabled.""" - tasks = [ - remote_task_types.RemoteTask('swarming_cmd', 'job1', 'url1'), - remote_task_types.RemoteTask('regular_cmd', 'job2', 'url2'), - ] + mock_is_swarming_task, mock_is_swarming_applicable): + """Tests that create_utask_main_jobs passes swarming tasks to swarming + service when the feature flag is enabled.""" + swarming_task = remote_task_types.RemoteTask('swarming_cmd', 'swarming_job', + 'url1') + other_task = remote_task_types.RemoteTask('regular_cmd', 'regular_job', + 'url2') mock_is_swarming_applicable.return_value = True - # _handle_swarming_jobs should process the swarming task and return the regular task. - mock_handle_swarming_jobs.return_value = [tasks[1]] + mock_is_swarming_task.side_effect = lambda job_type: job_type == 'swarming_job' + mock_handle_swarming_jobs.return_value = [] mock_get_job_frequency.return_value = { 'kubernetes': 1.0, @@ -374,22 +381,19 @@ def test_create_utask_main_jobs_swarming_remote_execution_enabled( } gate = remote_task_gate.RemoteTaskGate() - result = gate.create_utask_main_jobs(tasks) - - # ALL tasks should be sent to SwarmingService initially via _handle_swarming_jobs. - mock_handle_swarming_jobs.assert_called_once_with(tasks) + unscheduled_tasks = gate.create_utask_main_jobs([swarming_task, other_task]) - # The regular task returned by _handle_swarming_jobs should be routed to Kubernetes. + mock_handle_swarming_jobs.assert_called_once_with([swarming_task]) self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with( - [tasks[1]]) - - # No tasks should be unscheduled. - self.assertEqual(result, []) + [other_task]) + self.assertEqual(unscheduled_tasks, []) @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_applicable') + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_task') @mock.patch.object(remote_task_gate.RemoteTaskGate, '_handle_swarming_jobs') def test_create_utask_main_jobs_swarming_remote_execution_all_swarming( - self, mock_handle_swarming_jobs, mock_is_swarming_applicable): + self, mock_handle_swarming_jobs, mock_is_swarming_task, + mock_is_swarming_applicable): """Tests that create_utask_main_jobs handles the case where all tasks are swarming tasks.""" tasks = [ @@ -398,6 +402,7 @@ def test_create_utask_main_jobs_swarming_remote_execution_all_swarming( ] mock_is_swarming_applicable.return_value = True + mock_is_swarming_task.return_value = True # All tasks successfully scheduled as swarming. mock_handle_swarming_jobs.return_value = [] @@ -413,22 +418,24 @@ def test_create_utask_main_jobs_swarming_remote_execution_all_swarming( self.assertEqual(result, []) @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_applicable') + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_task') @mock.patch.object(remote_task_gate.RemoteTaskGate, '_handle_swarming_jobs') @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency') def test_create_utask_main_jobs_swarming_failure_preservation( self, mock_get_job_frequency, mock_handle_swarming_jobs, - mock_is_swarming_applicable): + mock_is_swarming_task, mock_is_swarming_applicable): """Tests that failed swarming tasks are correctly included in unscheduled_tasks.""" - tasks = [ - remote_task_types.RemoteTask('swarming_cmd1', 'job1', 'url1'), - remote_task_types.RemoteTask('swarming_cmd2', 'job2', 'url2'), - remote_task_types.RemoteTask('regular_cmd', 'job3', 'url3'), - ] + swarming_task = remote_task_types.RemoteTask('cmd', 'swarming_job', 'url1') + failed_swarming_task = remote_task_types.RemoteTask('cmd2', 'swarming_job', + 'url2') + other_task = remote_task_types.RemoteTask('cmd', 'job3', 'url3') + tasks = [swarming_task, failed_swarming_task, other_task] mock_is_swarming_applicable.return_value = True - # Mock one success, one failure (by returning it as unscheduled) and one regular task. - mock_handle_swarming_jobs.return_value = [tasks[1], tasks[2]] + mock_is_swarming_task.side_effect = ( + lambda job_type: job_type == 'swarming_job') + mock_handle_swarming_jobs.return_value = [failed_swarming_task] mock_get_job_frequency.return_value = { 'kubernetes': 1.0, @@ -437,17 +444,13 @@ def test_create_utask_main_jobs_swarming_failure_preservation( } gate = remote_task_gate.RemoteTaskGate() - result = gate.create_utask_main_jobs(tasks) + unscheduled_tasks = gate.create_utask_main_jobs(tasks) - # All tasks sent to _handle_swarming_jobs. - mock_handle_swarming_jobs.assert_called_once_with(tasks) - - # The tasks returned by _handle_swarming_jobs (failed swarming + regular) sent to k8s. + mock_handle_swarming_jobs.assert_called_once_with( + [swarming_task, failed_swarming_task]) self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with( - [tasks[1], tasks[2]]) - - # Both should be successfully processed by K8s (mocked to return [] by default). - self.assertEqual(result, []) + [other_task]) + self.assertEqual(unscheduled_tasks, [failed_swarming_task]) @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_applicable') @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency') @@ -474,6 +477,90 @@ def test_create_utask_main_jobs_swarming_remote_execution_disabled( # Should be routed normally to Kubernetes. self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with(tasks) + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_applicable') + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_task') + def test_unscheduled_swarming_tasks_dont_get_tried_on_other_services( + self, mock_is_swarming_task, mock_is_swarming_applicable): + """Tests that if swarming is unable to schedule a task, it doesn't get + tried on other services.""" + swarming_tasks = [ + remote_task_types.RemoteTask('swarming_cmd', 'swarming_job', 'url'), + remote_task_types.RemoteTask('swarming_cmd', 'swarming_job', 'url') + ] + self.mock_swarming_service.create_utask_main_jobs.return_value = swarming_tasks + mock_is_swarming_applicable.return_value = True + mock_is_swarming_task.side_effect = ( + lambda job_type: job_type == 'swarming_job') + + gate = remote_task_gate.RemoteTaskGate() + unscheduled_tasks = gate.create_utask_main_jobs(swarming_tasks) + + self.assertCountEqual(unscheduled_tasks, swarming_tasks) + self.mock_k8s_service.create_utask_main_jobs.assert_not_called() + + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_applicable') + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_task') + @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency') + def test_swarming_tasks_dont_try_on_other_services_when_feature_flag_disabled( + self, mock_get_job_frequency, mock_is_swarming_task, + mock_is_swarming_applicable): + """Tests that when a swarming task is pulled, and the feature flag is + disabled, we don't try to schedule it on other services.""" + swarming_tasks = [ + remote_task_types.RemoteTask('swarming_cmd', 'swarming_job', 'url'), + remote_task_types.RemoteTask('swarming_cmd', 'swarming_job', 'url') + ] + k8s_task = remote_task_types.RemoteTask('k8s_cmd', 'k8s_job', 'url') + tasks = swarming_tasks + [k8s_task] + mock_is_swarming_applicable.return_value = False + mock_is_swarming_task.side_effect = ( + lambda job_type: job_type == 'swarming_job') + mock_get_job_frequency.return_value = { + 'kubernetes': 1.0, + 'gcp_batch': 0.0, + 'swarming': 0.0 + } + + gate = remote_task_gate.RemoteTaskGate() + unscheduled_tasks = gate.create_utask_main_jobs(tasks) + + self.assertCountEqual(unscheduled_tasks, swarming_tasks) + self.mock_swarming_service.create_utask_main_jobs.assert_not_called() + self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with( + [k8s_task]) + + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_applicable') + @mock.patch.object(remote_task_gate.RemoteTaskGate, '_is_swarming_task') + @mock.patch.object(remote_task_gate.RemoteTaskGate, 'get_job_frequency') + def test_swarming_service_only_recieves_swarming_tasks( + self, mock_get_job_frequency, mock_is_swarming_task, + mock_is_swarming_applicable): + """Tests that the gate filters non swarming task so that the swarming + service only receives swarming tasks.""" + swarming_tasks = [ + remote_task_types.RemoteTask('swarming_cmd', 'swarming_job', 'url'), + remote_task_types.RemoteTask('swarming_cmd', 'swarming_job', 'url') + ] + k8s_task = remote_task_types.RemoteTask('k8s_cmd', 'k8s_job', 'url') + tasks = swarming_tasks + [k8s_task] + mock_is_swarming_applicable.return_value = True + mock_is_swarming_task.side_effect = ( + lambda job_type: job_type == 'swarming_job') + mock_get_job_frequency.return_value = { + 'kubernetes': 1.0, + 'gcp_batch': 0.0, + 'swarming': 0.0 + } + + gate = remote_task_gate.RemoteTaskGate() + unscheduled_tasks = gate.create_utask_main_jobs(tasks) + + self.assertCountEqual(unscheduled_tasks, []) + self.mock_k8s_service.create_utask_main_jobs.assert_called_once_with( + [k8s_task]) + self.mock_swarming_service.create_utask_main_jobs.assert_called_once_with( + swarming_tasks) + class RemoteTaskGateProcessingTest(unittest.TestCase): """Tests for logic in RemoteTaskGate that doesn't require full service mocking.""" @@ -524,7 +611,8 @@ def test_is_swarming_task(self, mock_swarming): mock_swarming.is_swarming_task.return_value = True self.assertTrue(self.gate._is_swarming_task('job')) - mock_swarming.is_swarming_task.assert_called_once_with('job') + mock_swarming.is_swarming_task.assert_called_once_with( + 'job', ignore_feature_flag=True) def test_handle_swarming_job(self): """Tests _handle_swarming_job."""