perf(serverless): take a full batch once drained

runpod/serverless/modules/rp_scale.py

@@ -27,6 +27,11 @@
 PRESTART_CLAIM_TIMEOUT_SECONDS = 10
 # Bound best-effort failure reporting well below the shared 600s HTTP timeout.
 PRESTART_FAILURE_REPORT_TIMEOUT_SECONDS = 15
+# A job-take costs ~0.5 s with jobs queued, however many it asks for, so a full
+# worker waits for every slot to free and takes a full batch. This bounds the
+# wait so one long job can't hold the other slots idle, and so set_scale still
+# runs as often as it did under the old fixed sleep.
+FULL_WAIT_SECONDS = 1
 
 
 def _handle_uncaught_exception(exc_type, exc_value, exc_traceback):
@@ -58,6 +63,8 @@ def __init__(self, config: dict[str, Any]):
         self._shutdown_event = asyncio.Event()
         # Set by get_jobs after queuing so run_jobs dispatches without polling.
         self._jobs_available = asyncio.Event()
+        # Set when a job leaves the worker so a full get_jobs re-checks for room.
+        self._slot_freed = asyncio.Event()
         self._prestart_ready = asyncio.Event()
         self._prestart_succeeded = False
         self._prestart_error: dict[str, Any] | None = None
@@ -315,8 +322,8 @@ async def get_jobs(self, session: ClientSession):
 
             jobs_needed = self.current_concurrency - self.current_occupancy()
             if jobs_needed <= 0:
-                log.debug("JobScaler.get_jobs | Queue is full. Retrying soon.")
-                await asyncio.sleep(1)  # don't go rapidly
+                log.debug("JobScaler.get_jobs | Queue is full. Waiting to drain.")
+                await self._wait_until_drained()
                 continue
 
             try:
@@ -371,6 +378,31 @@ async def get_jobs(self, session: ClientSession):
                 # Yield control back to the event loop
                 await asyncio.sleep(0)
 
+    async def _wait_until_drained(self) -> None:
+        """Block until no job is in progress, shutdown begins, or FULL_WAIT_SECONDS pass."""
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + FULL_WAIT_SECONDS
+        while self.is_alive():
+            # Clear, then re-check: a slot freed after this point re-arms the event.
+            self._slot_freed.clear()
+            remaining = deadline - loop.time()
+            # Zero concurrency has no slots to free, so wait out the bound
+            # rather than return without yielding to the event loop.
+            drained = self.current_occupancy() == 0 and self.current_concurrency > 0
+            if drained or remaining <= 0:
+                return
+            freed = asyncio.create_task(self._slot_freed.wait())
+            stopping = asyncio.create_task(self._shutdown_event.wait())
+            try:
+                await asyncio.wait(
+                    {freed, stopping},
+                    timeout=remaining,
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+            finally:
+                freed.cancel()
+                stopping.cancel()
+
     async def run_jobs(self, session: ClientSession):
         """
         Retrieve jobs from the jobs queue and process them concurrently.
@@ -550,6 +582,7 @@ async def handle_job(self, session: ClientSession, job: dict[str, Any]):
             # Job is no longer in progress
             self.job_progress.remove(job)
             self.jobs_tasks.pop(job["id"], None)
+            self._slot_freed.set()
 
             log.debug("Finished Job", job["id"])
             _reset_batch_id(batch_id_token)

tests/test_serverless/test_rp_scale.py

@@ -597,3 +597,121 @@ async def fetcher(_session, jobs_needed):
     assert job_scaler.progress.count == 0
 
     scaler.kill_worker()
+
+
+def _full_worker(job_scaler: PatchScaler, concurrency: int, release: dict[str, asyncio.Event]):
+    """Start get_jobs/run_jobs with `concurrency` jobs that finish when released."""
+    scaler = job_scaler.scaler
+    scaler.current_concurrency = concurrency
+    fetches = []
+    refetched = asyncio.Event()
+
+    async def fetcher(_session, jobs_needed):
+        fetches.append(jobs_needed)
+        if len(fetches) == 1:
+            return [generate_job(job_id) for job_id in release]
+        refetched.set()
+        await asyncio.sleep(10)  # long-poll with nothing queued
+        return []
+
+    async def handler(_session, _config, job):
+        await release[job["id"]].wait()
+
+    scaler.jobs_fetcher = fetcher
+    scaler.jobs_handler = handler
+    tasks = (asyncio.create_task(scaler.get_jobs(AsyncMock())),
+             asyncio.create_task(scaler.run_jobs(None)))
+    return fetches, refetched, tasks
+
+
+async def _stop(job_scaler: PatchScaler, release, tasks):
+    job_scaler.scaler.kill_worker()
+    for event in release.values():
+        event.set()
+    tasks[0].cancel()
+    await asyncio.gather(*tasks, return_exceptions=True)
+
+
+@pytest.mark.asyncio
+async def test_full_worker_takes_a_full_batch_as_soon_as_it_drains(job_scaler: PatchScaler):
+    """A job-take costs ~0.5 s however many jobs it asks for, so a full worker
+    takes again once every slot is free rather than sleeping a fixed second."""
+    release = {f"job-{i}": asyncio.Event() for i in range(4)}
+    fetches, refetched, tasks = _full_worker(job_scaler, 4, release)
+    await asyncio.sleep(0.05)  # all four running; the worker is full
+
+    for event in release.values():
+        event.set()
+
+    await asyncio.wait_for(refetched.wait(), timeout=DISPATCH_DEADLINE)
+    assert fetches == [4, 4]
+    await _stop(job_scaler, release, tasks)
+
+
+@pytest.mark.asyncio
+async def test_full_worker_does_not_take_while_partly_drained(
+    job_scaler: PatchScaler, monkeypatch
+):
+    monkeypatch.setattr(rp_scale, "FULL_WAIT_SECONDS", 10)
+    release = {f"job-{i}": asyncio.Event() for i in range(4)}
+    fetches, refetched, tasks = _full_worker(job_scaler, 4, release)
+    await asyncio.sleep(0.05)
+
+    release["job-0"].set()
+    await asyncio.sleep(0.1)
+
+    assert fetches == [4]
+    await _stop(job_scaler, release, tasks)
+
+
+@pytest.mark.asyncio
+async def test_full_worker_takes_freed_slots_after_the_wait_bound(
+    job_scaler: PatchScaler, monkeypatch
+):
+    """A straggler can't hold the other slots idle longer than the old sleep."""
+    monkeypatch.setattr(rp_scale, "FULL_WAIT_SECONDS", 0.1)
+    release = {f"job-{i}": asyncio.Event() for i in range(4)}
+    fetches, refetched, tasks = _full_worker(job_scaler, 4, release)
+    await asyncio.sleep(0.05)
+
+    for job_id in ("job-0", "job-1", "job-2"):
+        release[job_id].set()
+
+    await asyncio.wait_for(refetched.wait(), timeout=1)
+    assert fetches == [4, 3]
+    await _stop(job_scaler, release, tasks)
+
+
+@pytest.mark.asyncio
+async def test_full_worker_stops_taking_promptly_on_shutdown(job_scaler: PatchScaler):
+    release = {"job-0": asyncio.Event()}
+    fetches, _refetched, tasks = _full_worker(job_scaler, 1, release)
+    await asyncio.sleep(0.05)  # full: get_jobs is waiting to drain
+
+    job_scaler.scaler.kill_worker()
+
+    await asyncio.wait_for(tasks[0], timeout=DISPATCH_DEADLINE)
+    assert fetches == [1]
+    release["job-0"].set()
+    await asyncio.wait_for(tasks[1], timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_zero_concurrency_waits_instead_of_spinning(job_scaler: PatchScaler):
+    """A modifier returning 0 on an empty worker must still yield to the loop."""
+    scaler = job_scaler.scaler
+    calls = []
+
+    def zero_concurrency(_current):
+        calls.append(1)
+        if len(calls) >= 50:
+            scaler.kill_worker()  # bound the spin so the failure is an assert, not a hang
+        return 0
+
+    scaler.concurrency_modifier = zero_concurrency
+    task = asyncio.create_task(scaler.get_jobs(AsyncMock()))
+    await asyncio.sleep(0.05)
+
+    assert len(calls) == 1
+    scaler.kill_worker()
+    await asyncio.wait_for(task, timeout=DISPATCH_DEADLINE)

tests/test_serverless/test_worker.py

@@ -535,14 +535,19 @@ def test_run_worker_multi_processing_scaling_up(
         }
 
         # Let's mock job_scaler.is_alive so that it returns False
-        # when scale_behavior's counter is now 5.
-        def mock_is_alive():
+        # when scale_behavior's counter reaches 10. Stopping must also set the
+        # shutdown event, as a real stop does, or event waiters never wake.
+        def mock_is_alive(job_scaler):
             res = scale_behavior["counter"] < 10
             scale_behavior["counter"] += 1
+            if not res:
+                job_scaler.kill_worker()
             return res
 
         with patch(
-            "runpod.serverless.modules.rp_scale.JobScaler.is_alive", wraps=mock_is_alive
+            "runpod.serverless.modules.rp_scale.JobScaler.is_alive",
+            autospec=True,
+            side_effect=mock_is_alive,
         ):
             runpod.serverless.start(config)