Update logic to identify Swarming task

src/clusterfuzz/_internal/swarming/__init__.py

@@ -26,18 +26,16 @@
 from clusterfuzz._internal.system import environment
 
 
-def _requires_gpu() -> bool:
-  """Checks whether the REQUIRES_GPU env variable is set. This means
-  that the current job needs a gpu enabled device."""
-  requires_gpu = environment.get_value('REQUIRES_GPU')
-  return bool(utils.string_is_true(requires_gpu))
-
-
 def is_swarming_task(command: str, job_name: str):
   """Returns True if the task is supposed to run on swarming."""
   job = data_types.Job.query(data_types.Job.name == job_name).get()
-  if not job or not _requires_gpu():
+  if not job:
     return False
+
+  job_environment = job.get_environment()
+  if not utils.string_is_true(job_environment.get('IS_SWARMING_JOB')):
+    return False
+
   try:
     _get_new_task_spec(command, job_name, '')
     return True

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

@@ -264,3 +264,27 @@ def test_push_swarming_task(self):
         url=expected_url,
         data=json_format.MessageToJson(expected_new_task_request),
         headers=expected_headers)
+
+  def test_is_swarming_task(self):
+    """Tests is_swarming_task."""
+    job = data_types.Job(
+        name='libfuzzer_chrome_asan',
+        platform='LINUX',
+        environment_string='IS_SWARMING_JOB = True\n')
+    job.put()
+
+    self.assertTrue(swarming.is_swarming_task('fuzz', job.name))
+
+    job.environment_string = 'IS_SWARMING_JOB = False\n'
+    job.put()
+    self.assertFalse(swarming.is_swarming_task('fuzz', job.name))
+
+    job.environment_string = ''
+    job.put()
+    self.assertFalse(swarming.is_swarming_task('fuzz', job.name))
+
+    # Test unknown platform (ValueError in _get_new_task_spec)
+    job.platform = 'UNKNOWN'
+    job.environment_string = 'IS_SWARMING_JOB = True\n'
+    job.put()
+    self.assertFalse(swarming.is_swarming_task('fuzz', job.name))