From 583cced69d5f962ae433a307f5f5532aa91e9a87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Fri, 13 Nov 2026 09:03:27 -0800 Subject: [PATCH 1/2] perf(apps): vendor fastapi without the [standard] extra --- runpod/apps/build.py | 8 ++++++-- tests/test_apps/test_build.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/runpod/apps/build.py b/runpod/apps/build.py index c9cb31e0..f392cdec 100644 --- a/runpod/apps/build.py +++ b/runpod/apps/build.py @@ -63,8 +63,12 @@ "tqdm-loggable", ) -# api workers additionally serve over fastapi/uvicorn -API_RUNTIME_REQUIREMENTS = ("fastapi[standard]", "uvicorn>=0.30") +# api workers additionally serve over fastapi/uvicorn. no [standard] extra: +# the runtime imports only fastapi and uvicorn, and the extra's closure +# (fastapi-cloud-cli, sentry, opentelemetry, rich, ...) is ~4,000 files and +# ~18 MB compressed that every cold worker downloads and unpacks. uvicorn +# still picks up uvloop/httptools from the api images, which ship them. +API_RUNTIME_REQUIREMENTS = ("fastapi", "uvicorn>=0.30") # runpod's serverless limit for build artifacts MAX_ARTIFACT_MB = 1500 diff --git a/tests/test_apps/test_build.py b/tests/test_apps/test_build.py index b100387b..68e6460d 100644 --- a/tests/test_apps/test_build.py +++ b/tests/test_apps/test_build.py @@ -476,6 +476,22 @@ def test_gpu_builtin_image_prunes_transitive_torch(self, tmp_path): assert result.excluded == ["torch"] assert _installed(result.env_dir) == ["peft"] + def test_api_runtime_vendors_fastapi_without_the_standard_extra(self, tmp_path): + app = App("a") + + @app.api(name="web", cpu="cpu3c-2-4") + def web(): + pass + + seen = [] + with patch( + "runpod.apps.build.vendor", + side_effect=lambda env_dir, requirements, *a, **k: seen.extend(requirements), + ): + build_environment(app, tmp_path, scratch_dir=tmp_path / "s") + assert "fastapi" in seen and "uvicorn>=0.30" in seen + assert not any(r.startswith("fastapi[") for r in seen) + def test_runpod_in_dependencies_fails_before_vendoring(self, tmp_path): app = App("a") From fad9e64273359dc07103c6cc391f95ac77f5fb5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dean=20Qui=C3=B1anola?= Date: Fri, 13 Nov 2026 09:03:27 -0800 Subject: [PATCH 2/2] perf(serverless): take the next job while a result uploads --- runpod/serverless/modules/rp_job.py | 16 ++- runpod/serverless/modules/rp_scale.py | 59 ++++++-- tests/test_serverless/test_result_overlap.py | 138 +++++++++++++++++++ 3 files changed, 200 insertions(+), 13 deletions(-) create mode 100644 tests/test_serverless/test_result_overlap.py diff --git a/runpod/serverless/modules/rp_job.py b/runpod/serverless/modules/rp_job.py index 99addbc7..e780c648 100644 --- a/runpod/serverless/modules/rp_job.py +++ b/runpod/serverless/modules/rp_job.py @@ -194,7 +194,18 @@ async def get_job( return jobs -async def handle_job(session: ClientSession, config: Dict[str, Any], job) -> dict: +async def handle_job( + session: ClientSession, + config: Dict[str, Any], + job, + on_result_ready: Optional[Callable[[Dict[str, Any]], None]] = None, +) -> dict: + """Run a job and send its result. + + on_result_ready, if given, is called with the final result just before it + is sent, so the caller can free the job's slot and take the next job + while the result upload is still in flight. + """ if is_generator(config["handler"]): is_stream = True generator_output = run_job_generator(config["handler"], job) @@ -237,6 +248,9 @@ async def handle_job(session: ClientSession, config: Dict[str, Any], job) -> dic log.debug("rp_debugger | Flag not set, skipping debugger output.", job["id"]) rp_debugger.clear_debugger_output() + if on_result_ready is not None: + on_result_ready(job_result) + # Send the job result back to JOB_DONE_URL await send_result(session, job_result, job, is_stream=is_stream) return job_result diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 8537ebab..7ed9eec5 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/runpod/serverless/modules/rp_scale.py @@ -80,6 +80,9 @@ def __init__(self, config: dict[str, Any]): # maps in-progress job ids to their running tasks so individual jobs # can be stopped without killing the whole worker self.jobs_tasks: dict[str, asyncio.Task[Any]] = {} + # Jobs whose handler finished and whose result is still being sent. + # They no longer hold a slot, so the next job-take overlaps the upload. + self._sending_results: set[str] = set() self.stop_signals_fetcher = get_stop_signals self.stop_signals_fetcher_timeout = 90 @@ -416,8 +419,12 @@ async def run_jobs(self, session: ClientSession): # Clear before draining: a job queued after this point re-arms it. self._jobs_available.clear() - # Fetch as many jobs as the concurrency allows - while len(tasks) < self.current_concurrency and not self.jobs_queue.empty(): + # Fetch as many jobs as the concurrency allows; a task that is only + # sending its result no longer counts against it + while ( + len(tasks) - len(self._sending_results) < self.current_concurrency + and not self.jobs_queue.empty() + ): job = await self.jobs_queue.get() # Create a new task for each job and track it by job id task = asyncio.create_task(self.handle_job(session, job)) @@ -562,15 +569,16 @@ async def handle_job(self, session: ClientSession, job: dict[str, Any]): ) return - job_result = await self.jobs_handler(session, self.config, job) - - # Stop intake before the finally frees this slot: the platform - # restarts the worker on stopPod, killing any job taken meanwhile. - if self.config.get("refresh_worker", False) or ( - isinstance(job_result, dict) and job_result.get("stopPod") - ): - log.info("Refresh requested, taking no more jobs.", job["id"]) - self.kill_worker() + if self.jobs_handler is handle_job: + job_result = await self.jobs_handler( + session, + self.config, + job, + on_result_ready=lambda result: self._release_slot(job, result), + ) + else: + job_result = await self.jobs_handler(session, self.config, job) + self._stop_intake_if_refreshing(job, job_result) except asyncio.CancelledError: log.info("Job stopped.", job["id"]) @@ -585,13 +593,40 @@ async def handle_job(self, session: ClientSession, job: dict[str, Any]): self.jobs_queue.task_done() # Job is no longer in progress - self.job_progress.remove(job) + if job["id"] in self._sending_results: + self._sending_results.discard(job["id"]) + else: + 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) + def _stop_intake_if_refreshing(self, job: dict[str, Any], job_result: Any) -> None: + """Stop intake before this job's slot is freed: the platform restarts + the worker on stopPod, killing any job taken meanwhile.""" + if not self.is_alive(): + return + if self.config.get("refresh_worker", False) or ( + isinstance(job_result, dict) and job_result.get("stopPod") + ): + log.info("Refresh requested, taking no more jobs.", job["id"]) + self.kill_worker() + + def _release_slot(self, job: dict[str, Any], job_result: Any) -> None: + """Free a finished job's slot before its result is sent. + + The result upload takes a full round trip (~0.3 s measured on + serverless), and a single-slot worker used to wait it out before + asking for the next job. Taking the next job while the upload is in + flight hides it for every job queued behind another. + """ + self._stop_intake_if_refreshing(job, job_result) + self._sending_results.add(job["id"]) + self.job_progress.remove(job) + self._slot_freed.set() + async def _wait_for_prestart(self) -> None: """Wait for prestart to finish or shutdown to begin.""" ready = asyncio.create_task(self._prestart_ready.wait()) diff --git a/tests/test_serverless/test_result_overlap.py b/tests/test_serverless/test_result_overlap.py new file mode 100644 index 00000000..51fb4ed1 --- /dev/null +++ b/tests/test_serverless/test_result_overlap.py @@ -0,0 +1,138 @@ +""" +A finished job frees its slot before its result is sent, so the next job-take +overlaps the result upload instead of waiting for it. + +The upload is a full round trip to ai-api (~0.3 s measured on serverless), and +a single-slot worker used to wait it out before asking for the next job: every +job queued behind another paid it again. The refresh guarantee (SLS-900) still +holds: a job asking for a refresh stops intake before its slot is freed. +""" + +import asyncio +import time + +import pytest + +from runpod.serverless.modules import rp_http, rp_job, rp_scale + +from .test_refresh_worker import DEADLINE, FakeAiApi, _sync_handler, generate_job + +UPLOAD_S = 0.3 + + +class SlowUploadAiApi(FakeAiApi): + """ai-api whose job-done answers only after UPLOAD_S, recording when each + job-take arrived and when each upload finished.""" + + def __init__(self, jobs): + super().__init__(jobs) + self.take_at: dict[str, float] = {} + self.done_at: dict[str, float] = {} + + async def job_take(self, request): + response = await super().job_take(request) + if self.takes and self.takes[-1] is not None: + self.take_at.setdefault(self.takes[-1], time.monotonic()) + return response + + async def job_done(self, request): + await asyncio.sleep(UPLOAD_S) + response = await super().job_done(request) + self.done_at[request.match_info["job_id"]] = time.monotonic() + return response + + +@pytest.fixture +async def slow_upload_worker(monkeypatch): + servers = [] + + async def make(jobs): + server = SlowUploadAiApi(jobs) + await server.start() + servers.append(server) + monkeypatch.setattr(rp_job, "JOB_GET_URL", server.job_take_url.replace("$ID", "pod-1")) + monkeypatch.setattr( + rp_http, "JOB_DONE_URL", server.job_done_url.replace("$RUNPOD_POD_ID", "pod-1") + ) + return server + + yield make + for server in servers: + await server.close() + + +async def _serve_until_drained(server, scaler): + run = asyncio.create_task(scaler.run()) + try: + await asyncio.wait_for(server.polled_empty.wait(), timeout=DEADLINE) + # the last upload may still be in flight when the queue reads empty + deadline = time.monotonic() + DEADLINE + while len(server.done) < len(server.take_at) and time.monotonic() < deadline: + await asyncio.sleep(0.01) + finally: + scaler.kill_worker() + run.cancel() + await asyncio.gather(run, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_next_job_is_taken_while_the_previous_result_uploads(slow_upload_worker): + server = await slow_upload_worker([generate_job(j, n=i) for i, j in enumerate("ABC")]) + await _serve_until_drained(server, rp_scale.JobScaler({"handler": _sync_handler})) + + assert server.takes[:3] == ["A", "B", "C"] + assert server.done == {j: {"output": {"echo": i}} for i, j in enumerate("ABC")} + assert server.take_at["B"] < server.done_at["A"] + assert server.take_at["C"] < server.done_at["B"] + + +@pytest.mark.asyncio +async def test_uploads_no_longer_add_up_across_queued_jobs(slow_upload_worker): + jobs = [generate_job(f"J{i}", n=i) for i in range(5)] + server = await slow_upload_worker(jobs) + await _serve_until_drained(server, rp_scale.JobScaler({"handler": _sync_handler})) + + elapsed = server.done_at["J4"] - server.take_at["J0"] + # serial: 5 uploads back to back (1.5 s); overlapped: about one upload + # plus the takes, which this fake answers at once + assert elapsed < 3 * UPLOAD_S, f"uploads were serialized: {elapsed:.2f}s" + + +@pytest.mark.asyncio +async def test_refresh_still_leaves_the_next_job_queued(slow_upload_worker): + server = await slow_upload_worker( + [generate_job("A", n=1, refresh=True), generate_job("B", n=2)] + ) + scaler = rp_scale.JobScaler({"handler": _sync_handler}) + run = asyncio.create_task(scaler.run()) + try: + _, pending = await asyncio.wait({run}, timeout=DEADLINE) + assert server.takes == ["A"], "job B was taken after stopPod and would be lost" + assert [job["id"] for job in server.queue] == ["B"] + assert server.done == {"A": {"output": {"echo": 1}, "stopPod": True}} + assert not pending + finally: + scaler.kill_worker() + run.cancel() + await asyncio.gather(run, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_custom_jobs_handler_keeps_the_slot_until_it_returns(): + """Local-test jobs handlers have no on_result_ready; they keep the old contract.""" + seen = [] + + async def handler(session, config, job): + seen.append(job["id"]) + return {"output": job["id"]} + + scaler = rp_scale.JobScaler({"handler": _sync_handler}) + scaler.jobs_handler = handler + job = generate_job("A") + scaler.job_progress.add(job) + await scaler.jobs_queue.put(job) + await scaler.jobs_queue.get() + await scaler.handle_job(None, job) + assert seen == ["A"] + assert scaler.current_occupancy() == 0 + assert scaler._sending_results == set()