From 22dc493dd0e7701c1ef07a1b6e01daecb4aa7693 Mon Sep 17 00:00:00 2001
From: jonathanmetzman <31354670+jonathanmetzman@users.noreply.github.com>
Date: Thu, 19 Dec 2024 10:11:53 -0500
Subject: [PATCH]  Support multiple regions for schedule_fuzz_task.py  (#4525)

---
 src/clusterfuzz/_internal/cron/schedule_fuzz.py | 17 ++++++++++-------
 .../handlers/cron/schedule_fuzz_test.py         | 14 ++++++++------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/clusterfuzz/_internal/cron/schedule_fuzz.py b/src/clusterfuzz/_internal/cron/schedule_fuzz.py
index 80ad36a915..44c4b8278b 100644
--- a/src/clusterfuzz/_internal/cron/schedule_fuzz.py
+++ b/src/clusterfuzz/_internal/cron/schedule_fuzz.py
@@ -36,7 +36,7 @@ def _get_quotas(project, region):
       region=region, project=project).execute()['quotas']
 
 
-def get_available_cpus(project: str, region: str) -> int:
+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)
 
@@ -181,17 +181,20 @@ def get_fuzz_tasks(available_cpus: int) -> [tasks.Task]:
   return fuzz_tasks
 
 
+def get_batch_regions(batch_config):
+  mapping = batch_config.get('mapping')
+  return list(set(config['gce_region'] for config in mapping.values()))
+
+
 def schedule_fuzz_tasks() -> bool:
   """Schedules fuzz tasks."""
-  # TODO(metzman): Remove this when we are ready to run on Chrome.
   start = time.time()
-
   batch_config = local_config.BatchConfig()
+  regions = set(get_batch_regions(batch_config))
   project = batch_config.get('project')
-  # TODO(metzman): Put the CPU-based scheduling in tworkers.
-  available_cpus = get_available_cpus(project, 'us-east4')
-  # TODO(metzman): Remove this as we move from experimental code to production.
-  available_cpus = min(available_cpus, 2750)
+  available_cpus = sum(
+      get_available_cpus_for_region(project, region) for region in regions)
+  available_cpus = min(available_cpus, 3500 * len(regions))
   fuzz_tasks = get_fuzz_tasks(available_cpus)
   if not fuzz_tasks:
     logs.error('No fuzz tasks found to schedule.')
diff --git a/src/clusterfuzz/_internal/tests/appengine/handlers/cron/schedule_fuzz_test.py b/src/clusterfuzz/_internal/tests/appengine/handlers/cron/schedule_fuzz_test.py
index a849875c0b..004768456d 100644
--- a/src/clusterfuzz/_internal/tests/appengine/handlers/cron/schedule_fuzz_test.py
+++ b/src/clusterfuzz/_internal/tests/appengine/handlers/cron/schedule_fuzz_test.py
@@ -78,24 +78,25 @@ def test_get_fuzz_tasks(self):
     self.assertListEqual(comparable_results, expected_results)
 
 
-class TestGetAvailableCpus(unittest.TestCase):
-  """Tests for get_available_cpus."""
+class TestGetAvailableCpusForRegion(unittest.TestCase):
+  """Tests for get_available_cpus_for_region."""
 
   def setUp(self):
     test_helpers.patch(self,
                        ['clusterfuzz._internal.cron.schedule_fuzz._get_quotas'])
 
   def test_usage(self):
-    """Tests that get_available_cpus handles usage properly."""
+    """Tests that get_available_cpus_for_region handles usage properly."""
     self.mock._get_quotas.return_value = [{
         'metric': 'PREEMPTIBLE_CPUS',
         'limit': 5,
         'usage': 2
     }]
-    self.assertEqual(schedule_fuzz.get_available_cpus('project', 'region'), 3)
+    self.assertEqual(
+        schedule_fuzz.get_available_cpus_for_region('project', 'region'), 3)
 
   def test_cpus_and_preemptible_cpus(self):
-    """Tests that get_available_cpus handles usage properly."""
+    """Tests that get_available_cpus_for_region handles usage properly."""
     self.mock._get_quotas.return_value = [{
         'metric': 'PREEMPTIBLE_CPUS',
         'limit': 5,
@@ -105,4 +106,5 @@ def test_cpus_and_preemptible_cpus(self):
         'limit': 5,
         'usage': 5
     }]
-    self.assertEqual(schedule_fuzz.get_available_cpus('region', 'project'), 5)
+    self.assertEqual(
+        schedule_fuzz.get_available_cpus_for_region('region', 'project'), 5)