Try scheduling as much as available

src/clusterfuzz/_internal/cron/schedule_fuzz.py

@@ -17,6 +17,7 @@
 import random
 import time
 from typing import Dict
+from typing import List
 
 from googleapiclient import discovery
 
@@ -25,9 +26,13 @@
 from clusterfuzz._internal.config import local_config
 from clusterfuzz._internal.datastore import data_types
 from clusterfuzz._internal.datastore import ndb_utils
+from clusterfuzz._internal.google_cloud_utils import batch
 from clusterfuzz._internal.google_cloud_utils import credentials
 from clusterfuzz._internal.metrics import logs
 
+# TODO(metzman): Actually implement this.
+CPUS_PER_FUZZ_JOB = 2
+
 
 def _get_quotas(project, region):
   gcp_credentials = credentials.get_default()[0]
@@ -36,6 +41,16 @@ def _get_quotas(project, region):
       region=region, project=project).execute()['quotas']
 
 
+def count_unacked(project: str, topic: str):
+  creds, _ = credentials.get_default()
+  subscription_path = f'projects/{project}/subscriptions/{topic}'
+  client = discovery.build('pubsub', 'v1', credentials=creds)
+  sub = client.projects().subscriptions().get(  # pylint: disable=no-member
+      subscription=subscription_path).execute()
+  # TODO(metzman): Not all of these are fuzz_tasks. Deal with that.
+  return float(sub.get('numUnackedMessages', 0))
+
+
 def get_available_cpus_for_region(project: str, region: str) -> int:
   """Returns the number of available CPUs in the current GCE region."""
   quotas = _get_quotas(project, region)
@@ -59,13 +74,15 @@ def get_available_cpus_for_region(project: str, region: str) -> int:
   assert preemptible_quota or cpu_quota
 
   if not preemptible_quota['limit']:
-    # Preemptible quota is not set. Obey the CPU quota since that limitss us.
+    # Preemptible quota is not set. Obey the CPU quota since that limits us.
     quota = cpu_quota
   else:
     quota = preemptible_quota
   assert quota['limit'], quota
 
-  return quota['limit'] - quota['usage']
+  # TODO(metzman): Do this in a more configurable way.
+  limit = min(quota['limit'], 100_000)
+  return limit - quota['usage']
 
 
 class BaseFuzzTaskScheduler:
@@ -79,8 +96,7 @@ def get_fuzz_tasks(self):
 
   def _get_cpus_per_fuzz_job(self, job_name):
     del job_name
-    # TODO(metzman): Actually implement this.
-    return 2
+    return CPUS_PER_FUZZ_JOB
 
 
 class FuzzTaskCandidate:
@@ -186,15 +202,38 @@ def get_batch_regions(batch_config):
   return list(set(config['gce_region'] for config in mapping.values()))
 
 
+def get_available_cpus(project: str, regions: List[str]) -> int:
+  """Returns the available CPUs for fuzz tasks."""
+  # TODO(metzman): This doesn't distinguish between fuzz and non-fuzz
+  # tasks (nor preemptible and non-preemptible CPUs). Fix this.
+  waiting_tasks = sum(
+      batch.count_queued_or_scheduled_tasks(project, region)
+      for region in regions)
+  waiting_tasks += count_unacked(project, 'preprocess')
+  waiting_tasks += count_unacked(project, 'utasks_main')
+  soon_commited_cpus = waiting_tasks * CPUS_PER_FUZZ_JOB
+  logs.info(f'Soon committed CPUs: {soon_commited_cpus}')
+  available_cpus = sum(
+      get_available_cpus_for_region(project, region) for region in regions)
+  available_cpus = max(available_cpus - soon_commited_cpus, 0)
+
+  # Don't schedule more than 7.5K tasks at once. So we don't overload
+  # batch.
+  available_cpus = min(available_cpus, 15_000 * len(regions))
+  return available_cpus
+
+
 def schedule_fuzz_tasks() -> bool:
   """Schedules fuzz tasks."""
   start = time.time()
   batch_config = local_config.BatchConfig()
-  regions = set(get_batch_regions(batch_config))
   project = batch_config.get('project')
-  available_cpus = sum(
-      get_available_cpus_for_region(project, region) for region in regions)
-  available_cpus = min(available_cpus, 3500 * len(regions))
+  regions = set(get_batch_regions(batch_config))
+  available_cpus = get_available_cpus(project, regions)
+  logs.error(f'{available_cpus} available CPUs.')
+  if not available_cpus:
+    return False
+
   fuzz_tasks = get_fuzz_tasks(available_cpus)
   if not fuzz_tasks:
     logs.error('No fuzz tasks found to schedule.')

src/clusterfuzz/_internal/google_cloud_utils/batch.py

@@ -66,10 +66,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)
 
 
@@ -301,6 +298,14 @@ def _get_subconfig(batch_config, instance_spec):
   return all_subconfigs[weighted_subconfig.name]
 
 
+def count_queued_or_scheduled_tasks(project, region):
+  region = f'projects/{project}/locations/{region}'
+  jobs_filter = 'status.state=SCHEDULED OR status.state=QUEUED'
+  req = batch.types.ListJobsRequest(parent=region, filter=jobs_filter)
+  return sum(job.task_groups[0].task_count
+             for job in _batch_client().list_jobs(request=req))
+
+
 def _get_specs_from_config(batch_tasks) -> Dict:
   """Gets the configured specifications for a batch workload."""
   if not batch_tasks: