[batch] resolve a broken user share counter in the scheduler

.github/workflows/static-and-unit-tests.yml

@@ -267,7 +267,8 @@ jobs:
             --instafail \
             --durations=50 \
             batch/test/test_utils.py \
-            batch/test/test_time_limited_max_size_cache.py
+            batch/test/test_time_limited_max_size_cache.py \
+            batch/test/test_pool_scheduler.py
 
   check-ui:
     name: Check UI

batch/batch/driver/instance_collection/pool.py

@@ -587,40 +587,9 @@ def allocate_cores(user, mark):
         result = dict(sorted(result.items(), key=lambda item: item[1]['allocated_cores_mcpu'], reverse=True))
         return result
 
-    async def schedule_loop_body(self):
-        if self.app['frozen']:
-            log.info(f'not scheduling any jobs for {self.pool}; batch is frozen')
-            return True
-
-        start = time_msecs()
-        SCHEDULING_LOOP_RUNS.labels(pool_name=self.pool.name).inc()
-
-        n_scheduled = 0
-
-        scheduled_jobs_per_region: Dict[str, int] = defaultdict(int)
-        scheduled_cores_mcpu_per_region: Dict[str, int] = defaultdict(int)
-        unscheduled_jobs_per_region: Dict[str, float] = defaultdict(float)
-        unscheduled_cores_mcpu_per_region: Dict[str, int] = defaultdict(int)
-
-        user_resources = await self.compute_fair_share()
-
-        total = sum(resources['allocated_cores_mcpu'] for resources in user_resources.values())
-        if not total:
-            should_wait = True
-            for region in self.pool.all_supported_regions:
-                SCHEDULING_LOOP_CORES.labels(pool_name=self.pool.name, region=region, scheduled=True).set(0)
-                SCHEDULING_LOOP_CORES.labels(pool_name=self.pool.name, region=region, scheduled=False).set(0)
-                SCHEDULING_LOOP_JOBS.labels(pool_name=self.pool.name, region=region, scheduled=True).set(0)
-                SCHEDULING_LOOP_JOBS.labels(pool_name=self.pool.name, region=region, scheduled=False).set(0)
-            return should_wait
-        user_share = {
-            user: max(int(300 * resources['allocated_cores_mcpu'] / total + 0.5), 20)
-            for user, resources in user_resources.items()
-        }
-
-        async def user_runnable_jobs(user):
-            async for job_group in self.db.select_and_fetchall(
-                """
+    async def _user_runnable_jobs(self, user):
+        async for job_group in self.db.select_and_fetchall(
+            """
 SELECT job_groups.batch_id, job_groups.job_group_id, t.cancelled IS NOT NULL AS cancelled, userdata, job_groups.user, format_version
 FROM job_groups
 LEFT JOIN batches ON job_groups.batch_id = batches.id
@@ -636,11 +605,11 @@ async def user_runnable_jobs(user):
 WHERE job_groups.user = %s AND job_groups.`state` = 'running'
 ORDER BY job_groups.batch_id, job_groups.job_group_id;
 """,
-                (user,),
-                "user_runnable_jobs__select_running_batches",
-            ):
-                async for record in self.db.select_and_fetchall(
-                    """
+            (user,),
+            "user_runnable_jobs__select_running_batches",
+        ):
+            async for record in self.db.select_and_fetchall(
+                """
 SELECT jobs.job_id, spec, cores_mcpu, regions_bits_rep, time_ready, job_group_id, n_max_attempts, count(attempts.job_id) AS n_prior_attempts
 FROM jobs FORCE INDEX(jobs_batch_id_ic_state_ar_n_regions_bits_rep_job_group_id)
 LEFT JOIN jobs_telemetry ON jobs.batch_id = jobs_telemetry.batch_id AND jobs.job_id = jobs_telemetry.job_id
@@ -650,18 +619,18 @@ async def user_runnable_jobs(user):
 ORDER BY jobs.batch_id, jobs.job_group_id, inst_coll, state, always_run, -n_regions DESC, regions_bits_rep, jobs.job_id
 LIMIT 300;
 """,
-                    (job_group['batch_id'], job_group['job_group_id'], self.pool.name),
-                    "user_runnable_jobs__select_ready_always_run_jobs",
-                ):
-                    record['batch_id'] = job_group['batch_id']
-                    record['job_group_id'] = job_group['job_group_id']
-                    record['userdata'] = job_group['userdata']
-                    record['user'] = job_group['user']
-                    record['format_version'] = job_group['format_version']
-                    yield record
-                if not job_group['cancelled']:
-                    async for record in self.db.select_and_fetchall(
-                        """
+                (job_group['batch_id'], job_group['job_group_id'], self.pool.name),
+                "user_runnable_jobs__select_ready_always_run_jobs",
+            ):
+                record['batch_id'] = job_group['batch_id']
+                record['job_group_id'] = job_group['job_group_id']
+                record['userdata'] = job_group['userdata']
+                record['user'] = job_group['user']
+                record['format_version'] = job_group['format_version']
+                yield record
+            if not job_group['cancelled']:
+                async for record in self.db.select_and_fetchall(
+                    """
 SELECT jobs.job_id, spec, cores_mcpu, regions_bits_rep, time_ready, job_group_id, n_max_attempts, count(attempts.attempt_id) AS n_prior_attempts
 FROM jobs FORCE INDEX(jobs_batch_id_ic_state_ar_n_regions_bits_rep_job_id)
 LEFT JOIN jobs_telemetry ON jobs.batch_id = jobs_telemetry.batch_id AND jobs.job_id = jobs_telemetry.job_id
@@ -671,15 +640,46 @@ async def user_runnable_jobs(user):
 ORDER BY jobs.batch_id, inst_coll, state, always_run, n_regions, regions_bits_rep, jobs.job_id
 LIMIT 300;
 """,
-                        (job_group['batch_id'], job_group['job_group_id'], self.pool.name),
-                        "user_runnable_jobs__select_ready_jobs_batch_not_cancelled",
-                    ):
-                        record['batch_id'] = job_group['batch_id']
-                        record['job_group_id'] = job_group['job_group_id']
-                        record['userdata'] = job_group['userdata']
-                        record['user'] = job_group['user']
-                        record['format_version'] = job_group['format_version']
-                        yield record
+                    (job_group['batch_id'], job_group['job_group_id'], self.pool.name),
+                    "user_runnable_jobs__select_ready_jobs_batch_not_cancelled",
+                ):
+                    record['batch_id'] = job_group['batch_id']
+                    record['job_group_id'] = job_group['job_group_id']
+                    record['userdata'] = job_group['userdata']
+                    record['user'] = job_group['user']
+                    record['format_version'] = job_group['format_version']
+                    yield record
+
+    async def schedule_loop_body(self):
+        if self.app['frozen']:
+            log.info(f'not scheduling any jobs for {self.pool}; batch is frozen')
+            return True
+
+        start = time_msecs()
+        SCHEDULING_LOOP_RUNS.labels(pool_name=self.pool.name).inc()
+
+        n_scheduled = 0
+
+        scheduled_jobs_per_region: Dict[str, int] = defaultdict(int)
+        scheduled_cores_mcpu_per_region: Dict[str, int] = defaultdict(int)
+        unscheduled_jobs_per_region: Dict[str, float] = defaultdict(float)
+        unscheduled_cores_mcpu_per_region: Dict[str, int] = defaultdict(int)
+
+        user_resources = await self.compute_fair_share()
+
+        total = sum(resources['allocated_cores_mcpu'] for resources in user_resources.values())
+        if not total:
+            should_wait = True
+            for region in self.pool.all_supported_regions:
+                SCHEDULING_LOOP_CORES.labels(pool_name=self.pool.name, region=region, scheduled=True).set(0)
+                SCHEDULING_LOOP_CORES.labels(pool_name=self.pool.name, region=region, scheduled=False).set(0)
+                SCHEDULING_LOOP_JOBS.labels(pool_name=self.pool.name, region=region, scheduled=True).set(0)
+                SCHEDULING_LOOP_JOBS.labels(pool_name=self.pool.name, region=region, scheduled=False).set(0)
+            return should_wait
+        user_share = {
+            user: max(int(300 * resources['allocated_cores_mcpu'] / total + 0.5), 20)
+            for user, resources in user_resources.items()
+        }
 
         waitable_pool = WaitableSharedPool(self.async_worker_pool)
 
@@ -691,8 +691,9 @@ async def user_runnable_jobs(user):
 
             scheduled_cores_mcpu = 0
             share = user_share[user]
+            n_attempted_for_user = 0
 
-            async for record in user_runnable_jobs(user):
+            async for record in self._user_runnable_jobs(user):
                 attempt_id = secret_alnum_string(6)
                 record['attempt_id'] = attempt_id
 
@@ -767,8 +768,11 @@ async def schedule_with_error_handling(app, record, instance):
                         unscheduled_cores_mcpu_per_region[region] += record['cores_mcpu'] / n_regions
                         unscheduled_jobs_per_region[region] += 1 / n_regions
 
-                if n_scheduled >= share:
-                    should_wait = False
+                n_attempted_for_user += 1
+                if n_attempted_for_user >= share:
+                    # Rerun straight away only if this user's jobs are being placed; otherwise wait for a state change.
+                    if scheduled_cores_mcpu > 0:
+                        should_wait = False
                     break
 
         await waitable_pool.wait()

batch/test/test_pool_scheduler.py

@@ -0,0 +1,126 @@
+import os
+from types import SimpleNamespace
+from typing import Any, Dict, List
+
+import pytest
+
+# The driver modules read these at import time.
+for _name, _value in {
+    'CLOUD': 'gcp',
+    'HAIL_DEFAULT_NAMESPACE': 'default',
+    'HAIL_SCOPE': 'test',
+    'HAIL_DOCKER_ROOT_IMAGE': 'ubuntu:24.04',
+    'HAIL_DOCKER_PREFIX': 'docker.io',
+    'KUBERNETES_SERVER_URL': 'https://kubernetes.default',
+    'INTERNAL_GATEWAY_IP': '10.0.0.1',
+    'HAIL_BATCH_STORAGE_URI': 'gs://test-bucket',
+    'HAIL_SHA': 'test-sha',
+    'HAIL_QUERY_STORAGE_URI': 'gs://test-bucket',
+    'HAIL_QUERY_ACCEPTABLE_JAR_SUBFOLDER': '/jars',
+    'HAIL_QUERY_SPARK_ARCHIVE_SUBFOLDER': '/spark-archives',
+}.items():
+    os.environ.setdefault(_name, _value)
+
+from batch.driver.instance_collection import pool as pool_module  # noqa: E402
+from batch.driver.instance_collection.pool import PoolScheduler  # noqa: E402
+from batch.utils import ExceededSharesCounter  # noqa: E402
+
+REGION = 'us-central1'
+INSTANCE_CORES_MCPU = 16_000
+
+
+class FakeInstance:
+    region = REGION
+    state = 'active'
+
+    def adjust_free_cores_in_memory(self, delta_mcpu):
+        pass
+
+
+class FakePool:
+    name = 'standard'
+
+    def __init__(self):
+        self.all_supported_regions = [REGION]
+        self.inst_coll_manager = SimpleNamespace(regions=[REGION])
+
+    def get_instance(self, cores_mcpu, _regions):
+        return FakeInstance() if cores_mcpu <= INSTANCE_CORES_MCPU else None
+
+
+class InlineWorkerPool:
+    async def call(self, f, *args, **kwargs):
+        await f(*args, **kwargs)
+
+
+def make_scheduler(job_cores_mcpu_by_user: Dict[str, int], n_jobs_per_user: int):
+    """Returns the scheduler and a per-user count of jobs it has read from the queue."""
+    n_read = {user: 0 for user in job_cores_mcpu_by_user}
+
+    async def compute_fair_share():
+        return {user: {'allocated_cores_mcpu': 1_000_000} for user in job_cores_mcpu_by_user}
+
+    async def user_runnable_jobs(user):
+        for job_id in range(1, n_jobs_per_user + 1):
+            n_read[user] += 1
+            yield {
+                'batch_id': 1,
+                'job_group_id': 0,
+                'job_id': job_id,
+                'user': user,
+                'cores_mcpu': job_cores_mcpu_by_user[user],
+                'regions_bits_rep': None,
+                'n_prior_attempts': 0,
+                'n_max_attempts': 5,
+            }
+
+    scheduler: Any = PoolScheduler.__new__(PoolScheduler)
+    scheduler.app = {'frozen': False}
+    scheduler.pool = FakePool()
+    scheduler.async_worker_pool = InlineWorkerPool()
+    scheduler.exceeded_shares_counter = ExceededSharesCounter()
+    scheduler.compute_fair_share = compute_fair_share
+    scheduler._user_runnable_jobs = user_runnable_jobs
+    return scheduler, n_read
+
+
+@pytest.fixture
+def scheduled_users(monkeypatch):
+    users: List[str] = []
+
+    async def schedule_job(_app, record, _instance):
+        users.append(record['user'])
+
+    monkeypatch.setattr(pool_module, 'schedule_job', schedule_job)
+    return users
+
+
+async def test_each_user_gets_their_share_in_one_pass(scheduled_users):
+    # Equal allocations give each user a share of 150 jobs per pass.
+    scheduler, _ = make_scheduler({'alice': 1000, 'bob': 1000}, n_jobs_per_user=200)
+
+    should_wait = await scheduler.schedule_loop_body()
+
+    assert scheduled_users.count('alice') == 150
+    assert scheduled_users.count('bob') == 150
+    assert should_wait is False
+
+
+async def test_unplaceable_jobs_are_bounded_by_share(scheduled_users):
+    # alice's jobs don't fit on any instance.
+    scheduler, n_read = make_scheduler({'alice': 2 * INSTANCE_CORES_MCPU, 'bob': 1000}, n_jobs_per_user=200)
+
+    await scheduler.schedule_loop_body()
+
+    assert n_read['alice'] == 150
+    assert scheduled_users.count('bob') == 150
+
+
+async def test_waits_when_nothing_could_be_placed(scheduled_users):
+    # A lone user's share is 300, so alice reaches it without placing anything.
+    scheduler, _ = make_scheduler({'alice': 2 * INSTANCE_CORES_MCPU}, n_jobs_per_user=400)
+
+    should_wait = await scheduler.schedule_loop_body()
+
+    assert scheduled_users == []
+    assert should_wait is True