feat(batch): Add dynamic OS image selection for jobs (#4993)

## What this PR does
This PR introduces the core logic in ClusterFuzz to support running jobs
on different base OS images, as outlined in the Ubuntu 24.04 migration
plan. It enables the platform to dynamically schedule jobs on different
OS versions based on project-level configuration.

## Why this PR is important
This change is a prerequisite for migrating ClusterFuzz jobs from the
EOL Ubuntu 20.04 to Ubuntu 24.04. By making the platform OS-aware, we
can perform a gradual, safe migration without disrupting existing
fuzzing jobs. This implementation specifically addresses the
requirements of **b/441792600**.

## How the changes were implemented
This PR includes three main code changes:

1. **`data_types.py`**: The `OssFuzzProject` model is extended with a
`base_os_version` string property. This allows the Datastore to persist
the desired OS version for each OSS-Fuzz project.

2. **`cron/project_setup.py`**: The project synchronization cron job is
updated. It now reads the `base_os_version` field from a project's
`project.yaml` file (in the `oss-fuzz` repository) and saves it to the
corresponding `OssFuzzProject` entity in the Datastore.

3. **`google_cloud_utils/batch.py`**: The batch job creation logic is
refactored to be backward-compatible and OS-aware:
* It now fetches the `base_os_version` for the project associated with a
task.
* It attempts to find a corresponding image in a new
`versioned_docker_images` map in the batch configuration (`batch.yaml`).
* If a versioned image is not found (for legacy projects or during the
transition), it safely falls back to using the existing `docker_image`
key, ensuring no disruption.

## Related Task
*   **Task:** [b/441792600](b/441792600)

Author: hunsche

.gemini/settings.json

@@ -1,3 +1,5 @@
 {
-  "contextFileName": "AGENTS.md"
-}
+  "context": {
+    "fileName": "AGENTS.md"
+  }
+}

src/clusterfuzz/_internal/cron/project_setup.py

@@ -543,6 +543,7 @@ def create_project_settings(project, info, service_account):
 
   ccs = ccs_from_info(info)
   language = info.get('language')
+  base_os_version = info.get('base_os_version')
 
   if oss_fuzz_project:
     if oss_fuzz_project.service_account != service_account['email']:
@@ -556,6 +557,10 @@ def create_project_settings(project, info, service_account):
     if oss_fuzz_project.ccs != ccs:
       oss_fuzz_project.ccs = ccs
       oss_fuzz_project.put()
+
+    if oss_fuzz_project.base_os_version != base_os_version:
+      oss_fuzz_project.base_os_version = base_os_version
+      oss_fuzz_project.put()
   else:
     if language in MEMORY_SAFE_LANGUAGES:
       cpu_weight = OSS_FUZZ_MEMORY_SAFE_LANGUAGE_PROJECT_WEIGHT
@@ -568,7 +573,8 @@ def create_project_settings(project, info, service_account):
         high_end=is_high_end,
         cpu_weight=cpu_weight,
         service_account=service_account['email'],
-        ccs=ccs).put()
+        ccs=ccs,
+        base_os_version=base_os_version).put()
 
 
 def _create_pubsub_topic(name, client):

src/clusterfuzz/_internal/datastore/data_types.py

@@ -948,6 +948,9 @@ class Job(Model):
   # The platform that this job can run on.
   platform = ndb.StringProperty()
 
+  # Base OS version for the job.
+  base_os_version = ndb.StringProperty()
+
   # Blobstore key of the custom binary for this job.
   custom_binary_key = ndb.StringProperty()
 
@@ -1467,6 +1470,9 @@ class OssFuzzProject(Model):
   # CCs for the project.
   ccs = ndb.StringProperty(repeated=True)
 
+  # Base OS version for the project.
+  base_os_version = ndb.StringProperty()
+
 
 class OssFuzzProjectInfo(Model):
   """Set up information for a project (cpu allocation, instance groups, service

src/clusterfuzz/_internal/google_cloud_utils/batch.py

@@ -268,8 +268,19 @@ def _get_config_names(
     platform = job.platform if not utils.is_oss_fuzz() else 'LINUX'
     disk_size_gb = environment.get_value(
         'DISK_SIZE_GB', env=job.get_environment())
+    # Get the OS version from the job, this is the least specific version.
+    base_os_version = job.base_os_version
+
+    # If we are running in the oss-fuzz context, the project-specific config
+    # is more specific and overrides the job-level one.
+    if environment.get_value('PROJECT_NAME') == 'oss-fuzz':
+      project_name = job.project
+      oss_fuzz_project = data_types.OssFuzzProject.get_by_id(project_name)
+      if oss_fuzz_project and oss_fuzz_project.base_os_version:
+        base_os_version = oss_fuzz_project.base_os_version
+
     config_map[(task.command, task.job_type)] = (f'{platform}{suffix}',
-                                                 disk_size_gb)
+                                                 disk_size_gb, base_os_version)
   # TODO(metzman): Come up with a more systematic way for configs to
   # be overridden by jobs.
   return config_map
@@ -308,11 +319,23 @@ def _get_specs_from_config(batch_tasks) -> Dict:
     if (task.command, task.job_type) in specs:
       # Don't repeat work for no reason.
       continue
-    config_name, disk_size_gb = config_map[(task.command, task.job_type)]
+    config_name, disk_size_gb, base_os_version = config_map[(task.command,
+                                                             task.job_type)]
 
     instance_spec = batch_config.get('mapping').get(config_name)
     if instance_spec is None:
       raise ValueError(f'No mapping for {config_name}')
+
+    # Decide which docker image to use.
+    versioned_images_map = instance_spec.get('versioned_docker_images')
+    if (base_os_version and versioned_images_map and
+        base_os_version in versioned_images_map):
+      # New path: Use the versioned image if specified and available.
+      docker_image_uri = versioned_images_map[base_os_version]
+    else:
+      # Fallback/legacy path: Use the original docker_image key.
+      docker_image_uri = instance_spec['docker_image']
+
     project_name = batch_config.get('project')
     clusterfuzz_release = instance_spec.get('clusterfuzz_release', 'prod')
     # Lower numbers are a lower priority, meaning less likely to run From:
@@ -332,7 +355,7 @@ 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]
     spec = BatchWorkloadSpec(
-        docker_image=instance_spec['docker_image'],
+        docker_image=docker_image_uri,
         disk_size_gb=disk_size_gb,
         disk_type=instance_spec['disk_type'],
         user_data=instance_spec['user_data'],

src/clusterfuzz/_internal/tests/appengine/handlers/cron/project_setup_test.py

@@ -630,6 +630,8 @@ def test_execute(self):
         'ccs': [
             'primary@example.com', 'user@example.com', 'user2@googlemail.com'
         ],
+        'base_os_version':
+            None,
     }, lib1_settings.to_dict())
 
     lib2_settings = ndb.Key(data_types.OssFuzzProject, 'lib2').get()
@@ -644,6 +646,7 @@ def test_execute(self):
         'service_account': 'lib3@serviceaccount.com',
         'high_end': False,
         'ccs': ['user@example.com'],
+        'base_os_version': None,
     }, lib3_settings.to_dict())
 
     lib4_settings = ndb.Key(data_types.OssFuzzProject, 'lib4').get()
@@ -655,6 +658,7 @@ def test_execute(self):
         'service_account': 'lib4@serviceaccount.com',
         'high_end': True,
         'ccs': ['user@example.com'],
+        'base_os_version': None,
     }, lib4_settings.to_dict())
 
     old_lib_settings = ndb.Key(data_types.OssFuzzProject, 'old_lib').get()

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
@@ -148,6 +149,63 @@ def test_get_specs_from_config_with_disk_size_override(self):
         [batch.BatchTask('fuzz', job_name, None)])
     self.assertEqual(spec['fuzz', job_name].disk_size_gb, overridden_size)
 
+  @mock.patch('clusterfuzz._internal.google_cloud_utils.batch.environment')
+  @mock.patch(
+      'clusterfuzz._internal.datastore.data_types.OssFuzzProject.get_by_id')
+  @mock.patch('clusterfuzz._internal.datastore.ndb_utils.get_all_from_query')
+  def test_get_config_names_os_version(self, mock_get_jobs, mock_get_project,
+                                       mock_env):
+    """Test the hierarchical logic for determining base_os_version."""
+    # Test Case 1: Internal project, job-level OS version is used.
+    mock_env.get_value.return_value = 'internal-project'
+    job1 = data_types.Job(
+        name='job1', platform='LINUX', base_os_version='job-os-ubuntu-20')
+    mock_get_jobs.return_value = [job1]
+    config_map = batch._get_config_names(
+        [batch.BatchTask('fuzz', 'job1', None)])
+    self.assertEqual(config_map[('fuzz', 'job1')][2], 'job-os-ubuntu-20')
+
+    # Test Case 2: OSS-Fuzz project, project-level version overrides job-level.
+    mock_env.get_value.return_value = 'oss-fuzz'
+    job2 = data_types.Job(
+        name='job2',
+        project='my-project',
+        platform='LINUX',
+        base_os_version='job-os-ubuntu-20')
+    project = data_types.OssFuzzProject(
+        name='my-project', base_os_version='project-os-ubuntu-24')
+    mock_get_jobs.return_value = [job2]
+    mock_get_project.return_value = project
+    config_map = batch._get_config_names(
+        [batch.BatchTask('fuzz', 'job2', None)])
+    self.assertEqual(config_map[('fuzz', 'job2')][2], 'project-os-ubuntu-24')
+
+    # Test Case 3: OSS-Fuzz project, only project-level version exists.
+    job3 = data_types.Job(name='job3', project='my-project', platform='LINUX')
+    mock_get_jobs.return_value = [job3]
+    config_map = batch._get_config_names(
+        [batch.BatchTask('fuzz', 'job3', None)])
+    self.assertEqual(config_map[('fuzz', 'job3')][2], 'project-os-ubuntu-24')
+
+    # Test Case 4: Internal project, no version is set, should be None.
+    mock_env.get_value.return_value = 'internal-project'
+    job4 = data_types.Job(name='job4', platform='LINUX')
+    mock_get_jobs.return_value = [job4]
+    config_map = batch._get_config_names(
+        [batch.BatchTask('fuzz', 'job4', None)])
+    self.assertIsNone(config_map[('fuzz', 'job4')][2])
+
+    # Test Case 5: OSS-Fuzz project, but no versions are set anywhere.
+    mock_env.get_value.return_value = 'oss-fuzz'
+    job5 = data_types.Job(
+        name='job5', project='my-project-no-version', platform='LINUX')
+    project_no_version = data_types.OssFuzzProject(name='my-project-no-version')
+    mock_get_jobs.return_value = [job5]
+    mock_get_project.return_value = project_no_version
+    config_map = batch._get_config_names(
+        [batch.BatchTask('fuzz', 'job5', None)])
+    self.assertIsNone(config_map[('fuzz', 'job5')][2])
+
 
 def _get_spec_from_config(command, job_name):
   return list(