diff --git a/src/clusterfuzz/_internal/swarming/__init__.py b/src/clusterfuzz/_internal/swarming/__init__.py index 9196c2a19a..aaff9447b6 100644 --- a/src/clusterfuzz/_internal/swarming/__init__.py +++ b/src/clusterfuzz/_internal/swarming/__init__.py @@ -25,6 +25,7 @@ from clusterfuzz._internal.base.feature_flags import FeatureFlags from clusterfuzz._internal.config import local_config from clusterfuzz._internal.datastore import data_types +from clusterfuzz._internal.google_cloud_utils import compute_metadata from clusterfuzz._internal.google_cloud_utils import credentials from clusterfuzz._internal.metrics import logs from clusterfuzz._internal.protos import swarming_pb2 @@ -112,6 +113,68 @@ def _get_task_dimensions(job: data_types.Job, platform_specific_dimensions: list return task_dimensions +def _append_metadata_env_var( + env_vars: list[swarming_pb2.StringPair], # pylint: disable=no-member + env_var_name: str, + metadata_path: str) -> None: + """Attempts to get a variable from the environment or metadata and appends it. + + Args: + env_vars: The list of string pairs to append the retrieved value to. + env_var_name: The name of the environment variable to check and set. + metadata_path: The path in the compute metadata to check if not set in env. + """ + value = environment.get_value(env_var_name) + if not value: + try: + value = compute_metadata.get(metadata_path) + except Exception: + pass + + if value: + env_vars.append( + swarming_pb2.StringPair( # pylint: disable=no-member + key=env_var_name, value=str(value))) + else: + logs.warning(f'{env_var_name} is not set or cannot be fetched.') + + +def _get_env_vars(logs_project_id: str, + instance_spec: dict) -> list[swarming_pb2.StringPair]: # pylint: disable=no-member + """Retrieve required environment variables from metadata and config.""" + default_task_environment = [ + swarming_pb2.StringPair(key='UWORKER', value='True'), # pylint: disable=no-member + swarming_pb2.StringPair(key='SWARMING_BOT', value='True'), # pylint: disable=no-member + swarming_pb2.StringPair(key='LOG_TO_GCP', value='True'), # pylint: disable=no-member + swarming_pb2.StringPair(key='IS_K8S_ENV', value='True'), # pylint: disable=no-member + swarming_pb2.StringPair( # pylint: disable=no-member + key='LOGGING_CLOUD_PROJECT_ID', + value=logs_project_id or ''), + ] + + _append_metadata_env_var(default_task_environment, 'DEPLOYMENT_BUCKET', + 'project/attributes/deployment-bucket') + _append_metadata_env_var(default_task_environment, 'HOST_JOB_SELECTION', + 'instance/attributes/host-job-selection') + _append_metadata_env_var(default_task_environment, 'DEPLOYMENT_ZIP', + 'project/attributes/deployment-zip') + + env_vars = [] + env_vars.append( + swarming_pb2.StringPair( # pylint: disable=no-member + key='DOCKER_IMAGE', + value=instance_spec.get('docker_image', ''))) + + platform_specific_env = instance_spec.get('env', []) + for var in platform_specific_env: + env_vars.append(swarming_pb2.StringPair(key=var['key'], value=var['value'])) # pylint: disable=no-member + + env_vars.append(_env_vars_to_json(default_task_environment)) + env_vars.extend(default_task_environment) + + return env_vars + + def _env_vars_to_json( env_vars: list[swarming_pb2.StringPair]) -> swarming_pb2.StringPair: # pylint: disable=no-member """ @@ -166,27 +229,7 @@ def create_new_task_request(command: str, job_name: str, download_url: str # env_prefixes allows the modification of existing environment variables by # adding the values as prefixes to the env variable. env_prefixes = instance_spec.get('env_prefixes', {}) - default_task_environment = [ - swarming_pb2.StringPair(key='UWORKER', value='True'), # pylint: disable=no-member - swarming_pb2.StringPair(key='SWARMING_BOT', value='True'), # pylint: disable=no-member - swarming_pb2.StringPair(key='LOG_TO_GCP', value='True'), # pylint: disable=no-member - swarming_pb2.StringPair(key='IS_K8S_ENV', value='True'), # pylint: disable=no-member - swarming_pb2.StringPair( # pylint: disable=no-member - key='LOGGING_CLOUD_PROJECT_ID', - value=logs_project_id), - ] - - platform_specific_env = instance_spec.get('env', []) - swarming_bot_environment = [] - swarming_bot_environment.append( - swarming_pb2.StringPair( # pylint: disable=no-member - key='DOCKER_IMAGE', - value=instance_spec.get('docker_image', ''))) - for var in platform_specific_env: - swarming_bot_environment.append( - swarming_pb2.StringPair(key=var['key'], value=var['value'])) # pylint: disable=no-member - swarming_bot_environment.append(_env_vars_to_json(default_task_environment)) - swarming_bot_environment.extend(default_task_environment) + swarming_bot_environment = _get_env_vars(logs_project_id, instance_spec) dimensions = instance_spec.get('dimensions', []) cas_input_root = instance_spec.get('cas_input_root', {}) diff --git a/src/clusterfuzz/_internal/tests/core/swarming/service_test.py b/src/clusterfuzz/_internal/tests/core/swarming/service_test.py index e80e51b323..234150cf5f 100644 --- a/src/clusterfuzz/_internal/tests/core/swarming/service_test.py +++ b/src/clusterfuzz/_internal/tests/core/swarming/service_test.py @@ -31,9 +31,11 @@ def setUp(self): 'clusterfuzz._internal.swarming.create_new_task_request', 'clusterfuzz._internal.base.tasks.task_utils.get_command_from_module', 'clusterfuzz._internal.metrics.logs.error', + 'clusterfuzz._internal.google_cloud_utils.compute_metadata.get', ]) self.service = service.SwarmingService() self.mock.create_new_task_request.return_value = 'fake_request' + self.mock.get.return_value = None def test_create_utask_main_job_success(self): """Test creating a single task successfully.""" diff --git a/src/clusterfuzz/_internal/tests/core/swarming/swarming_config_error_test.py b/src/clusterfuzz/_internal/tests/core/swarming/swarming_config_error_test.py index 8958d73b08..5d09102709 100644 --- a/src/clusterfuzz/_internal/tests/core/swarming/swarming_config_error_test.py +++ b/src/clusterfuzz/_internal/tests/core/swarming/swarming_config_error_test.py @@ -29,9 +29,11 @@ class SwarmingConfigErrorTest(unittest.TestCase): def setUp(self): helpers.patch(self, [ 'clusterfuzz._internal.swarming.FeatureFlags', + 'clusterfuzz._internal.google_cloud_utils.compute_metadata.get', ]) helpers.patch_environ(self) self.mock.FeatureFlags.SWARMING_REMOTE_EXECUTION.enabled = True + self.mock.get.return_value = None def test_is_swarming_task_bad_config(self): """Tests that is_swarming_task returns False when there's a BadConfigError.""" diff --git a/src/clusterfuzz/_internal/tests/core/swarming/swarming_test.py b/src/clusterfuzz/_internal/tests/core/swarming/swarming_test.py index 020b94c3d8..a8667989a5 100644 --- a/src/clusterfuzz/_internal/tests/core/swarming/swarming_test.py +++ b/src/clusterfuzz/_internal/tests/core/swarming/swarming_test.py @@ -13,6 +13,7 @@ # limitations under the License. """Swarming tests.""" import base64 +import os import unittest from unittest import mock @@ -38,11 +39,17 @@ def setUp(self): 'clusterfuzz._internal.google_cloud_utils.credentials.get_scoped_service_account_credentials', 'google.auth.transport.requests.Request', 'clusterfuzz._internal.swarming.FeatureFlags', + 'clusterfuzz._internal.google_cloud_utils.compute_metadata.get', ]) 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.mock.get.return_value = None self.maxDiff = None + os.environ.pop('DEPLOYMENT_ZIP', None) + os.environ.pop('DEPLOYMENT_BUCKET', None) + os.environ.pop('PROJECT_NAME', None) + os.environ.pop('HOST_JOB_SELECTION', None) def test_get_spec_from_config_with_docker_image(self): """Tests that create_new_task_request works as expected.""" @@ -430,3 +437,37 @@ def test_get_task_dimensions_job_precedence(self): swarming_pb2.StringPair(key='key2', value='value2'), ] self.assertCountEqual(dimensions, expected_dimensions) + + def test_get_env_vars_with_metadata_server(self): + """Tests that _get_env_vars uses values from the metadata server when available.""" + + def metadata_get(path): + if path == 'project/attributes/deployment-bucket': + return 'test-bucket-from-metadata' + return None + + self.mock.get.side_effect = metadata_get + instance_spec = { + "docker_image": "gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654" + } + env = swarming._get_env_vars('project_id', instance_spec) # pylint: disable=protected-access + + expected_env = [ + swarming_pb2.StringPair( + key='DOCKER_IMAGE', + value='gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654'), + swarming_pb2.StringPair( + key='DOCKER_ENV_VARS', + value= + '{"UWORKER": "True", "SWARMING_BOT": "True", "LOG_TO_GCP": "True", "IS_K8S_ENV": "True", "LOGGING_CLOUD_PROJECT_ID": "project_id", "DEPLOYMENT_BUCKET": "test-bucket-from-metadata"}' + ), + swarming_pb2.StringPair(key='UWORKER', value='True'), + swarming_pb2.StringPair(key='SWARMING_BOT', value='True'), + swarming_pb2.StringPair(key='LOG_TO_GCP', value='True'), + swarming_pb2.StringPair(key='IS_K8S_ENV', value='True'), + swarming_pb2.StringPair( + key='LOGGING_CLOUD_PROJECT_ID', value='project_id'), + swarming_pb2.StringPair( + key='DEPLOYMENT_BUCKET', value='test-bucket-from-metadata'), + ] + self.assertEqual(env, expected_env)