diff --git a/src/clusterfuzz/_internal/google_cloud_utils/batch.py b/src/clusterfuzz/_internal/google_cloud_utils/batch.py index abd073110cb..df1b80cc4fa 100644 --- a/src/clusterfuzz/_internal/google_cloud_utils/batch.py +++ b/src/clusterfuzz/_internal/google_cloud_utils/batch.py @@ -125,16 +125,13 @@ def create_uworker_main_batch_jobs(batch_tasks: List[BatchTask]): def _get_task_spec(batch_workload_spec): """Gets the task spec based on the batch workload spec.""" runnable = batch.Runnable() - runnable.container = batch.Runnable.Container() - runnable.container.image_uri = batch_workload_spec.docker_image - clusterfuzz_release = batch_workload_spec.clusterfuzz_release - runnable.container.options = ( - '--memory-swappiness=40 --shm-size=1.9g --rm --net=host ' - '-e HOST_UID=1337 -P --privileged --cap-add=all ' - f'-e CLUSTERFUZZ_RELEASE={clusterfuzz_release} ' - '--name=clusterfuzz -e UNTRUSTED_WORKER=False -e UWORKER=True ' - '-e UWORKER_INPUT_DOWNLOAD_URL') - runnable.container.volumes = ['/var/scratch0:/mnt/scratch0'] + runnable.script = batch.Runnable.Script() + runnable.script.text = f""" +#!/bin/bash +{batch_workload_spec.user_data} +docker pull {batch_workload_spec.docker_image} +docker run --name=clusterfuzz -e CLUSTERFUZZ_RELEASE={batch_workload_spec.clusterfuzz_release} -e UNTRUSTED_WORKER=False -e UWORKER=True -e UWORKER_INPUT_DOWNLOAD_URL --memory-swappiness=40 --shm-size=1.9g --rm --net=host -e HOST_UID=1337 -P --privileged --cap-add=all -v /var/scratch0:/mnt/scratch0 {batch_workload_spec.docker_image} +""" task_spec = batch.TaskSpec() task_spec.runnables = [runnable] if batch_workload_spec.retry: @@ -331,11 +328,13 @@ def _get_specs_from_config(batch_tasks) -> Dict: disk_size_gb = (disk_size_gb or instance_spec['disk_size_gb']) subconfig = subconfig_map[config_name] + with open(instance_spec['user_data'].replace('file://', ''), 'r') as f: + user_data = f.read() spec = BatchWorkloadSpec( docker_image=instance_spec['docker_image'], disk_size_gb=disk_size_gb, disk_type=instance_spec['disk_type'], - user_data=instance_spec['user_data'], + user_data=user_data, service_account_email=instance_spec['service_account_email'], preemptible=instance_spec['preemptible'], machine_type=instance_spec['machine_type'], diff --git a/src/clusterfuzz/_internal/tests/core/google_cloud_utils/batch_test.py b/src/clusterfuzz/_internal/tests/core/google_cloud_utils/batch_test.py index 35117daa341..e3123e72c3d 100644 --- a/src/clusterfuzz/_internal/tests/core/google_cloud_utils/batch_test.py +++ b/src/clusterfuzz/_internal/tests/core/google_cloud_utils/batch_test.py @@ -14,6 +14,7 @@ """Batch tests.""" import unittest +from unittest import mock from clusterfuzz._internal.datastore import data_types from clusterfuzz._internal.google_cloud_utils import batch @@ -38,6 +39,53 @@ def setUp(self): name='east4-network2', weight=1, ) + self.mock_open = mock.patch( + 'builtins.open', mock.mock_open(read_data='dummy_user_data')).start() + + self.mock_batch_config = mock.patch( + 'clusterfuzz._internal.google_cloud_utils.batch._get_batch_config').start() + self.mock_batch_config.return_value.get.side_effect = ( + lambda key, default=None: { + 'mapping': { + 'LINUX-PREEMPTIBLE-UNPRIVILEGED': { + 'clusterfuzz_release': 'prod', + 'docker_image': 'gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654', + 'user_data': 'file://linux-init.yaml', + 'disk_size_gb': 75, + 'disk_type': 'pd-standard', + 'service_account_email': 'test-unpriv-clusterfuzz-service-account-email', + 'preemptible': True, + 'machine_type': 'n1-standard-1', + 'retry': False, + 'subconfigs': [{'name': 'east4-network2', 'weight': 1}], + }, + 'LINUX-NONPREEMPTIBLE-UNPRIVILEGED': { + 'clusterfuzz_release': 'prod', + 'docker_image': 'gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654', + 'user_data': 'file://linux-init.yaml', + 'disk_size_gb': 110, + 'disk_type': 'pd-standard', + 'service_account_email': 'test-unpriv-clusterfuzz-service-account-email', + 'preemptible': False, + 'machine_type': 'n1-standard-1', + 'retry': True, + 'subconfigs': [{'name': 'east4-network2', 'weight': 1}], + }, + }, + 'subconfigs': { + 'east4-network2': { + 'region': 'us-east4', + 'network': 'projects/project_name/global/networks/networkname2', + 'subnetwork': 'projects/project_name/regions/us-east4/subnetworks/subnetworkname2', + }, + }, + 'project': 'test-clusterfuzz', + }.get(key, default) + ) + + def tearDown(self): + mock.patch.stopall() + def test_nonpreemptible(self): """Tests that _get_specs_from_config works for non-preemptibles as @@ -46,7 +94,7 @@ def test_nonpreemptible(self): expected_spec = batch.BatchWorkloadSpec( clusterfuzz_release='prod', docker_image='gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654', - user_data='file://linux-init.yaml', + user_data='dummy_user_data', disk_size_gb=110, disk_type='pd-standard', service_account_email='test-unpriv-clusterfuzz-service-account-email', @@ -72,7 +120,7 @@ def test_fuzz_get_specs_from_config(self): expected_spec = batch.BatchWorkloadSpec( clusterfuzz_release='prod', docker_image='gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654', - user_data='file://linux-init.yaml', + user_data='dummy_user_data', disk_size_gb=75, disk_type='pd-standard', service_account_email='test-unpriv-clusterfuzz-service-account-email',