Backport batch changes to master from oss-fuzz. (#4689)

The current version of batch on master is non-functional. We need to get
it working to start using master in Chrome.

Author: jonathanmetzman

configs/test/batch/batch.yaml

@@ -20,10 +20,13 @@ mapping:
     disk_size_gb: 110
     disk_type: pd-standard
     service_account_email: test-clusterfuzz-service-account-email
-    gce_region: 'gce-region'
-    gce_zone: 'gce-zone'
-    network: 'projects/google.com:clusterfuzz/global/networks/networkname'
-    subnetwork: 'projects/google.com:clusterfuzz/regions/gce-region/subnetworks/subnetworkname'
+    subconfigs:
+      -
+        name: central1-network1
+        weight: .5
+      -
+        name: central1-network2
+        weight: .5
     preemptible: false
     machine_type: n1-standard-1
   LINUX-NONPREEMPTIBLE-UNPRIVILEGED:
@@ -33,37 +36,56 @@ mapping:
     disk_size_gb: 110
     disk_type: pd-standard
     service_account_email: test-unpriv-clusterfuzz-service-account-email
-    gce_region: 'gce-region'
-    gce_zone: 'gce-zone'
-    network: 'projects/google.com:clusterfuzz/global/networks/networkname'
-    subnetwork: 'projects/google.com:clusterfuzz/regions/gce-region/subnetworks/subnetworkname'
     preemptible: false
     machine_type: n1-standard-1
     retry: true
+    subconfigs:
+      -
+        name: central1-network1
+        weight: .2
+      -
+        name: central1-network2
+        weight: .3
+      -
+        name: east4-network2
+        weight: .5
   LINUX-PREEMPTIBLE:
     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-clusterfuzz-service-account-email
-    gce_region: 'gce-region'
-    gce_zone: 'gce-zone'
-    network: 'projects/google.com:clusterfuzz/global/networks/networkname'
-    subnetwork: 'projects/google.com:clusterfuzz/regions/gce-region/subnetworks/subnetworkname'
     preemptible: true
     machine_type: n1-standard-1
+    subconfigs:
+      -
+        name: east4-network2
+        weight: 1
   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
-    gce_region: 'gce-region'
-    gce_zone: 'gce-zone'
-    network: 'projects/google.com:clusterfuzz/global/networks/networkname'
-    subnetwork: 'projects/google.com:clusterfuzz/regions/gce-region/subnetworks/subnetworkname'
     preemptible: true
     machine_type: n1-standard-1
+    subconfigs:
+      -
+        name: east4-network2
+        weight: 1
 project: 'test-clusterfuzz'
+subconfigs:
+  central1-network1:
+    region: 'us-central1'
+    network: 'projects/project_name/global/networks/networkname'
+    subnetwork: 'projects/project_name/regions/us-central1/subnetworks/subnetworkname'
+  central1-network2:
+    region: 'us-central1'
+    network: 'projects/project_name/global/networks/networkname2'
+    subnetwork: 'projects/project_name/regions/us-central1/subnetworks/subnetworkname2'
+  east4-network2:
+    region: 'us-east4'
+    network: 'projects/project_name/global/networks/networkname2'
+    subnetwork: 'projects/project_name/regions/us-east4/subnetworks/subnetworkname2'

src/clusterfuzz/_internal/google_cloud_utils/batch.py

@@ -14,7 +14,9 @@
 """Cloud Batch helpers."""
 import collections
 import threading
+from typing import Dict
 from typing import List
+from typing import Tuple
 import uuid
 
 from google.cloud import batch_v1 as batch
@@ -25,18 +27,16 @@
 from clusterfuzz._internal.base.tasks import task_utils
 from clusterfuzz._internal.config import local_config
 from clusterfuzz._internal.datastore import data_types
+from clusterfuzz._internal.datastore import ndb_utils
 from clusterfuzz._internal.metrics import logs
+from clusterfuzz._internal.system import environment
 
-# TODO(metzman): Change to from . import credentials when we are done
-# developing.
 from . import credentials
 
 _local = threading.local()
 
 DEFAULT_RETRY_COUNT = 0
 
-TASK_BUNCH_SIZE = 20
-
 # Controls how many containers (ClusterFuzz tasks) can run on a single VM.
 # THIS SHOULD BE 1 OR THERE WILL BE SECURITY PROBLEMS.
 TASK_COUNT_PER_NODE = 1
@@ -54,7 +54,6 @@
     'subnetwork',
     'preemptible',
     'project',
-    'gce_zone',
     'machine_type',
     'network',
     'gce_region',
@@ -66,10 +65,7 @@
 
 def _create_batch_client_new():
   """Creates a batch client."""
-  creds, project = credentials.get_default()
-  if not project:
-    project = utils.get_application_id()
-
+  creds, _ = credentials.get_default()
   return batch.BatchServiceClient(credentials=creds)
 
 
@@ -108,9 +104,10 @@ def create_uworker_main_batch_job(module, job_type, input_download_url):
 def create_uworker_main_batch_jobs(batch_tasks: List[BatchTask]):
   """Creates batch jobs."""
   job_specs = collections.defaultdict(list)
+  specs = _get_specs_from_config(batch_tasks)
   for batch_task in batch_tasks:
     logs.info(f'Scheduling {batch_task.command}, {batch_task.job_type}.')
-    spec = _get_spec_from_config(batch_task.command, batch_task.job_type)
+    spec = specs[(batch_task.command, batch_task.job_type)]
     job_specs[spec].append(batch_task.input_download_url)
 
   logs.info('Creating batch jobs.')
@@ -119,7 +116,7 @@ def create_uworker_main_batch_jobs(batch_tasks: List[BatchTask]):
   logs.info('Batching utask_mains.')
   for spec, input_urls in job_specs.items():
     for input_urls_portion in utils.batched(input_urls,
-                                            MAX_CONCURRENT_VMS_PER_JOB):
+                                            MAX_CONCURRENT_VMS_PER_JOB - 1):
       jobs.append(_create_job(spec, input_urls_portion))
 
   return jobs
@@ -209,7 +206,6 @@ def _create_job(spec, input_urls):
   job = batch.Job()
   job.task_groups = [task_group]
   job.allocation_policy = _get_allocation_policy(spec)
-  job.labels = {'env': 'testing', 'type': 'container'}
   job.logs_policy = batch.LogsPolicy()
   job.logs_policy.destination = batch.LogsPolicy.Destination.CLOUD_LOGGING
   job.priority = spec.priority
@@ -239,86 +235,134 @@ def _get_batch_config():
   return local_config.BatchConfig()
 
 
-def _get_job(job_name):
-  """Returns the Job entity named by |job_name|. This function was made to make
-  mocking easier."""
-  return data_types.Job.query(data_types.Job.name == job_name).get()
-
-
 def is_no_privilege_workload(command, job_name):
   return is_remote_task(command, job_name)
 
 
 def is_remote_task(command, job_name):
   try:
-    _get_spec_from_config(command, job_name)
+    _get_specs_from_config([BatchTask(command, job_name, None)])
     return True
   except ValueError:
     return False
 
 
-def _get_config_name(command, job_name):
-  """Returns the name of the config for |command| and |job_name|."""
-  job = _get_job(job_name)
-  # As of this writing, batch only supports LINUX.
-  if utils.is_oss_fuzz():
-    # TODO(b/377885331): In OSS-Fuzz, the platform can't be used because, as of
-    # it includes the project name.
-    config_name = 'LINUX'
-  else:
-    config_name = job.platform
-  if command == 'fuzz':
-    config_name += '-PREEMPTIBLE-UNPRIVILEGED'
-  else:
-    config_name += '-NONPREEMPTIBLE-UNPRIVILEGED'
-  return config_name
+def _get_config_names(
+    batch_tasks: List[BatchTask]) -> Dict[Tuple[str, str], str]:
+  """"Gets the name of the configs for each batch_task. Returns a dict
+  that is indexed by command and job_type for efficient lookup."""
+  job_names = {task.job_type for task in batch_tasks}
+  query = data_types.Job.query(data_types.Job.name.IN(list(job_names)))
+  jobs = ndb_utils.get_all_from_query(query)
+  job_map = {job.name: job for job in jobs}
+  config_map = {}
+  for task in batch_tasks:
+    if task.job_type not in job_map:
+      logs.error(f'{task.job_type} doesn\'t exist.')
+      continue
+    if task.command == 'fuzz':
+      suffix = '-PREEMPTIBLE-UNPRIVILEGED'
+    else:
+      suffix = '-NONPREEMPTIBLE-UNPRIVILEGED'
+    job = job_map[task.job_type]
+    platform = job.platform if not utils.is_oss_fuzz() else 'LINUX'
+    disk_size_gb = environment.get_value(
+        'DISK_SIZE_GB', env=job.get_environment())
+    config_map[(task.command, task.job_type)] = (f'{platform}{suffix}',
+                                                 disk_size_gb)
+  # TODO(metzman): Come up with a more systematic way for configs to
+  # be overridden by jobs.
+  return config_map
 
 
 def _get_task_duration(command):
   return tasks.TASK_LEASE_SECONDS_BY_COMMAND.get(command,
                                                  tasks.TASK_LEASE_SECONDS)
 
 
-def _get_spec_from_config(command, job_name):
+WeightedSubconfig = collections.namedtuple('WeightedSubconfig',
+                                           ['name', 'weight'])
+
+
+def _get_subconfig(batch_config, instance_spec):
+  # TODO(metzman): Make this pick one at random or based on conditions.
+  all_subconfigs = batch_config.get('subconfigs', {})
+  instance_subconfigs = instance_spec['subconfigs']
+  weighted_subconfigs = [
+      WeightedSubconfig(subconfig['name'], subconfig['weight'])
+      for subconfig in instance_subconfigs
+  ]
+  weighted_subconfig = utils.random_weighted_choice(weighted_subconfigs)
+  return all_subconfigs[weighted_subconfig.name]
+
+
+def _get_specs_from_config(batch_tasks) -> Dict:
   """Gets the configured specifications for a batch workload."""
-  config_name = _get_config_name(command, job_name)
+  if not batch_tasks:
+    return {}
   batch_config = _get_batch_config()
-  instance_spec = batch_config.get('mapping').get(config_name, None)
-  if instance_spec is None:
-    raise ValueError(f'No mapping for {config_name}')
-  project_name = batch_config.get('project')
-  docker_image = instance_spec['docker_image']
-  user_data = instance_spec['user_data']
-  should_retry = instance_spec.get('retry', False)
-  clusterfuzz_release = instance_spec.get('clusterfuzz_release', 'prod')
-
-  # Lower numbers are lower priority. From:
-  # https://cloud.google.com/batch/docs/reference/rest/v1/projects.locations.jobs
-  low_priority = command == 'fuzz'
-  priority = 0 if low_priority else 1
-
-  max_run_duration = f'{_get_task_duration(command)}s'
-  if command == 'corpus_pruning':
-    should_retry = False  # It is naturally retried the next day.
-
-  spec = BatchWorkloadSpec(
-      clusterfuzz_release=clusterfuzz_release,
-      docker_image=docker_image,
-      user_data=user_data,
-      disk_size_gb=instance_spec['disk_size_gb'],
-      disk_type=instance_spec['disk_type'],
-      service_account_email=instance_spec['service_account_email'],
-      # TODO(metzman): Get rid of zone so that we can more easily run in
-      # multiple regions.
-      gce_zone=instance_spec['gce_zone'],
-      gce_region=instance_spec['gce_region'],
-      project=project_name,
-      network=instance_spec['network'],
-      subnetwork=instance_spec['subnetwork'],
-      preemptible=instance_spec['preemptible'],
-      machine_type=instance_spec['machine_type'],
-      priority=priority,
-      max_run_duration=max_run_duration,
-      retry=should_retry,
-  )
-  return spec
+  config_map = _get_config_names(batch_tasks)
+  specs = {}
+  subconfig_map = {}
+  for task in batch_tasks:
+    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)]
+
+    instance_spec = batch_config.get('mapping').get(config_name)
+    if instance_spec is None:
+      raise ValueError(f'No mapping for {config_name}')
+    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:
+    # https://cloud.google.com/batch/docs/reference/rest/v1/projects.locations.jobs
+    priority = 0 if task.command == 'fuzz' else 1
+    max_run_duration = f'{_get_task_duration(task.command)}s'
+    # This saves us time and reduces fragementation, e.g. every linux fuzz task
+    # run in this call will run in the same zone.
+    if config_name not in subconfig_map:
+      subconfig = _get_subconfig(batch_config, instance_spec)
+      subconfig_map[config_name] = subconfig
+
+    should_retry = instance_spec.get('retry', False)
+    if should_retry and task.command == 'corpus_pruning':
+      should_retry = False  # It is naturally retried the next day.
+
+    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'],
+        disk_size_gb=disk_size_gb,
+        disk_type=instance_spec['disk_type'],
+        user_data=instance_spec['user_data'],
+        service_account_email=instance_spec['service_account_email'],
+        preemptible=instance_spec['preemptible'],
+        machine_type=instance_spec['machine_type'],
+        gce_region=subconfig['region'],
+        network=subconfig['network'],
+        subnetwork=subconfig['subnetwork'],
+        project=project_name,
+        clusterfuzz_release=clusterfuzz_release,
+        priority=priority,
+        max_run_duration=max_run_duration,
+        retry=should_retry,
+    )
+    specs[(task.command, task.job_type)] = spec
+  return specs
+
+
+def count_queued_or_scheduled_tasks(project: str,
+                                    region: str) -> Tuple[int, int]:
+  """Counts the number of queued and scheduled tasks."""
+  region = f'projects/{project}/locations/{region}'
+  jobs_filter = 'Status.State="SCHEDULED" OR Status.State="QUEUED"'
+  req = batch.types.ListJobsRequest(parent=region, filter=jobs_filter)
+  queued = 0
+  scheduled = 0
+  for job in _batch_client().list_jobs(request=req):
+    if job.status.state == batch.JobStatus.State.SCHEDULED:
+      scheduled += job.task_groups[0].task_count
+    elif job.status.state == batch.JobStatus.State.QUEUED:
+      queued += job.task_groups[0].task_count
+  return (queued, scheduled)