Swarming: Adds dynamic Job specific bot dimensions

src/clusterfuzz/_internal/swarming/__init__.py

@@ -62,6 +62,30 @@ def _get_swarming_config():
   return local_config.SwarmingConfig()
 
 
+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"""
+  unique_dimensions = {}
+  unique_dimensions['os'] = job.platform
+  unique_dimensions['pool'] = _get_swarming_config().get('swarming_pool')
+
+  for dimension in platform_specific_dimensions:
+    unique_dimensions[dimension['key'].lower()] = dimension['value']
+
+  swarming_dimensions = environment.get_value('SWARMING_DIMENSIONS')
+  if isinstance(swarming_dimensions, dict):
+    for key, value in swarming_dimensions.items():
+      unique_dimensions[key.lower()] = value
+
+  task_dimensions = []
+  for dimension, value in unique_dimensions.items():
+    task_dimensions.append(
+        swarming_pb2.StringPair(  # pylint: disable=no-member
+            key=dimension, value=value))
+  return task_dimensions
+
+
 def _env_vars_to_json(
     env_vars: list[swarming_pb2.StringPair]) -> swarming_pb2.StringPair:  # pylint: disable=no-member
   """
@@ -84,7 +108,6 @@ def _get_new_task_spec(command: str, job_name: str,
   instance_spec = swarming_config.get('mapping').get(config_name, None)
   if instance_spec is None:
     raise ValueError(f'No mapping for {config_name}')
-  swarming_pool = swarming_config.get('swarming_pool')
   swarming_realm = swarming_config.get('swarming_realm')
   logs_project_id = swarming_config.get('logs_project_id')
   priority = instance_spec['priority']
@@ -127,18 +150,7 @@ def _get_new_task_spec(command: str, job_name: str,
   swarming_bot_environment.append(_env_vars_to_json(default_task_environment))
   swarming_bot_environment.extend(default_task_environment)
 
-  task_dimensions = [
-      swarming_pb2.StringPair(key='os', value=job.platform),  # pylint: disable=no-member
-      swarming_pb2.StringPair(key='pool', value=swarming_pool)  # pylint: disable=no-member
-  ]
-
-  dimensions = instance_spec.get('dimensions', None)
-  if dimensions:
-    for dimension in dimensions:
-      task_dimensions.append(
-          swarming_pb2.StringPair(  # pylint: disable=no-member
-              key=dimension['key'],
-              value=dimension['value']))
+  dimensions = instance_spec.get('dimensions', [])
 
   cas_input_root = instance_spec.get('cas_input_root', {})
 
@@ -152,7 +164,7 @@ def _get_new_task_spec(command: str, job_name: str,
               expiration_secs=expiration_secs,
               properties=swarming_pb2.TaskProperties(  # pylint: disable=no-member
                   command=startup_command,
-                  dimensions=task_dimensions,
+                  dimensions=_get_task_dimensions(job, dimensions),
                   cipd_input=cipd_input,
                   cas_input_root=cas_input_root,
                   execution_timeout_secs=execution_timeout_secs,

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

@@ -21,6 +21,7 @@
 from clusterfuzz._internal import swarming
 from clusterfuzz._internal.datastore import data_types
 from clusterfuzz._internal.protos import swarming_pb2
+from clusterfuzz._internal.system import environment
 from clusterfuzz._internal.tests.test_libs import helpers
 from clusterfuzz._internal.tests.test_libs import test_utils
 
@@ -37,6 +38,7 @@ def setUp(self):
         'google.auth.transport.requests.Request',
         'clusterfuzz._internal.swarming.FeatureFlags',
     ])
+    helpers.patch_environ(self)
     self.mock._get_task_name.return_value = 'task_name'  # pylint: disable=protected-access
     self.mock.FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled = True
     self.maxDiff = None
@@ -336,3 +338,41 @@ def test_is_swarming_task(self):
     job.environment_string = ''
     job.put()
     self.assertFalse(swarming.is_swarming_task('fuzz', job.name))
+
+  def test_get_task_dimensions_with_env_var(self):
+    """Tests that _get_task_dimensions handles SWARMING_DIMENSIONS env var."""
+    environment.set_value('SWARMING_DIMENSIONS', {
+        'cpu': 'x86',
+        'os': 'windows'
+    })
+    job = data_types.Job(name='libfuzzer_chrome_asan', platform='LINUX')
+    dimensions = swarming._get_task_dimensions(job, [])  # pylint: disable=protected-access
+
+    expected_dimensions = [
+        swarming_pb2.StringPair(key='os', value='windows'),
+        swarming_pb2.StringPair(key='pool', value='pool-name'),
+        swarming_pb2.StringPair(key='cpu', value='x86'),
+    ]
+    self.assertCountEqual(dimensions, expected_dimensions)
+
+  def test_get_task_dimensions_job_precedence(self):
+    """Tests that job swarming dimensions have more precedence than platform ones."""
+    # Use 'MAC' platform which has static dimensions (key1, key2) in swarming.yaml.
+    job = data_types.Job(name='mac_job', platform='MAC')
+    job.put()
+
+    # Platform dimensions for MAC are: key1: value1, key2: value2.
+    # We set SWARMING_DIMENSIONS in the environment to override key1.
+    environment.set_value('SWARMING_DIMENSIONS', {'key1': 'job_value1'})
+
+    spec = swarming._get_new_task_spec(  # pylint: disable=protected-access
+        'fuzz', job.name, 'https://download_url')
+    dimensions = spec.task_slices[0].properties.dimensions
+
+    expected_dimensions = [
+        swarming_pb2.StringPair(key='os', value='MAC'),
+        swarming_pb2.StringPair(key='pool', value='pool-name'),
+        swarming_pb2.StringPair(key='key1', value='job_value1'),
+        swarming_pb2.StringPair(key='key2', value='value2'),
+    ]
+    self.assertCountEqual(dimensions, expected_dimensions)