From 2b882d4482ce087cd006ed1c81ae25e1b1e0c65b Mon Sep 17 00:00:00 2001
From: Vitor Guidi <vitorguidi@gmail.com>
Date: Mon, 26 Aug 2024 19:46:30 -0300
Subject: [PATCH] Adding CLUSTERFUZZ_RELEASE env var to batch docker container
 (#4210)

### Motivation

The last move to upgrade our fleet to 3.11 is switch batch. This PR adds
the CLUSTERFUZZ_RELEASE env var to batch config, so we can roll
forward/back between 3.7 and 3.11 on batch by maneuvering between the
appropriate docker image and code release

This PR belongs to this
[initiative](https://github.com/google/clusterfuzz/issues/4059)

### Deployment strategy

This code should be deployed to uworker instances, which are the actual
ones that create batch jobs. Deploy this using --release=candidate, with
the appropriate clusterfuzz-config
---
 configs/test/batch/batch.yaml                                | 2 ++
 src/clusterfuzz/_internal/google_cloud_utils/batch.py        | 5 +++++
 .../_internal/tests/core/google_cloud_utils/batch_test.py    | 2 ++
 3 files changed, 9 insertions(+)

diff --git a/configs/test/batch/batch.yaml b/configs/test/batch/batch.yaml
index 639eae1b45..ed03d2d7f3 100644
--- a/configs/test/batch/batch.yaml
+++ b/configs/test/batch/batch.yaml
@@ -24,6 +24,7 @@ mapping:
     preemptible: false
     machine_type: n1-standard-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
@@ -34,6 +35,7 @@ mapping:
     preemptible: false
     machine_type: n1-standard-1
   LINUX-PREEMPTIBLE:
+    clusterfuzz_release: 'candidate'
     docker_image: 'gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654'
     user_data: 'file://linux-init.yaml'
     disk_size_gb: 75
diff --git a/src/clusterfuzz/_internal/google_cloud_utils/batch.py b/src/clusterfuzz/_internal/google_cloud_utils/batch.py
index 18b297b7b3..cd2a9160c3 100644
--- a/src/clusterfuzz/_internal/google_cloud_utils/batch.py
+++ b/src/clusterfuzz/_internal/google_cloud_utils/batch.py
@@ -45,6 +45,7 @@
 MAX_CONCURRENT_VMS_PER_JOB = 1000
 
 BatchWorkloadSpec = collections.namedtuple('BatchWorkloadSpec', [
+    'clusterfuzz_release',
     'disk_size_gb',
     'disk_type',
     'docker_image',
@@ -158,9 +159,11 @@ def _get_task_spec(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']
@@ -289,9 +292,11 @@ def _get_spec_from_config(command, job_name):
   project_name = batch_config.get('project')
   docker_image = instance_spec['docker_image']
   user_data = instance_spec['user_data']
+  clusterfuzz_release = instance_spec.get('clusterfuzz_release', 'prod')
   # TODO(https://github.com/google/clusterfuzz/issues/3008): Make this use a
   # low-privilege account.
   spec = BatchWorkloadSpec(
+      clusterfuzz_release=clusterfuzz_release,
       docker_image=docker_image,
       user_data=user_data,
       disk_size_gb=instance_spec['disk_size_gb'],
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 e48ef9ead1..9381d18070 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
@@ -34,6 +34,7 @@ def test_nonpreemptible_get_spec_from_config(self):
     job.put()
     spec = batch._get_spec_from_config('corpus_pruning', job.name)  # pylint: disable=protected-access
     expected_spec = batch.BatchWorkloadSpec(
+        clusterfuzz_release='prod',
         docker_image='gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654',
         user_data='file://linux-init.yaml',
         disk_size_gb=110,
@@ -53,6 +54,7 @@ def test_preemptible_get_spec_from_config(self):
     job.put()
     spec = batch._get_spec_from_config('fuzz', job.name)  # pylint: disable=protected-access
     expected_spec = batch.BatchWorkloadSpec(
+        clusterfuzz_release='candidate',
         docker_image='gcr.io/clusterfuzz-images/base:a2f4dd6-202202070654',
         user_data='file://linux-init.yaml',
         disk_size_gb=75,