From dc47bf579bd7847bd32861f99b4d1582bc835fa1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Thu, 22 Oct 2026 07:36:26 -0700 Subject: [PATCH 1/2] perf(serverless): take a full batch once drained A full worker slept a fixed 1 s before taking again. With jobs queued a job-take costs ~0.5 s server-side whatever its size, so a 32-slot vLLM worker ran take (0.5 s) -> jobs (~0.4 s, inside the sleep) -> idle, one 32-job batch per 1.5 s. It now waits until no job is in progress, at most FULL_WAIT_SECONDS (the old sleep), so it takes the same or a bigger batch no later than before. Taking each freed slot at once was tried and is slower: smaller takes pay the same 0.5 s each. --- runpod/serverless/modules/rp_scale.py | 34 ++++++++- tests/test_serverless/test_rp_scale.py | 97 ++++++++++++++++++++++++++ 2 files changed, 129 insertions(+), 2 deletions(-) diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 11702673..439469d7 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/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,28 @@ 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() + if self.current_occupancy() == 0 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 +579,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) diff --git a/tests/test_serverless/test_rp_scale.py b/tests/test_serverless/test_rp_scale.py index ec6a752a..0db822af 100644 --- a/tests/test_serverless/test_rp_scale.py +++ b/tests/test_serverless/test_rp_scale.py @@ -597,3 +597,100 @@ 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) From bfca5addd634b563d09f3984f63ff2982ce909d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Thu, 22 Oct 2026 09:22:46 -0700 Subject: [PATCH 2/2] fix(serverless): wait when concurrency is zero - _wait_until_drained returned immediately on an empty worker even when concurrency_modifier returned 0. set_scale also returns without awaiting in that state, so get_jobs spun without yielding and starved the event loop, including shutdown. An empty worker now only counts as drained when it has slots; otherwise it waits out FULL_WAIT_SECONDS or shutdown, as the old sleep did. - The scaling_up worker test stopped the worker by mocking is_alive without setting the shutdown event, so run_jobs could wait on it forever (hung on 3.10/3.11). The mock now calls kill_worker() when it starts returning False, as a real stop does. --- runpod/serverless/modules/rp_scale.py | 5 ++++- tests/test_serverless/test_rp_scale.py | 21 +++++++++++++++++++++ tests/test_serverless/test_worker.py | 11 ++++++++--- 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 439469d7..f2079fdc 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -386,7 +386,10 @@ async def _wait_until_drained(self) -> None: # Clear, then re-check: a slot freed after this point re-arms the event. self._slot_freed.clear() remaining = deadline - loop.time() - if self.current_occupancy() == 0 or remaining <= 0: + # 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()) diff --git a/tests/test_serverless/test_rp_scale.py b/tests/test_serverless/test_rp_scale.py index 0db822af..18bebb0e 100644 --- a/tests/test_serverless/test_rp_scale.py +++ b/tests/test_serverless/test_rp_scale.py @@ -694,3 +694,24 @@ async def test_full_worker_stops_taking_promptly_on_shutdown(job_scaler: PatchSc 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) diff --git a/tests/test_serverless/test_worker.py b/tests/test_serverless/test_worker.py index 547bd88f..0349ceb9 100644 --- a/tests/test_serverless/test_worker.py +++ b/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)