[Mohith26]

Summary

This is a concrete option for question 2 in #647: SDK-side take-ahead, opt-in and off by default. I know you're assigned to #647 and may prefer a server-side fix, so please treat this as one possible approach; happy to change it or close it.

With jobs queued, a job-take costs a flat ~0.5s whatever its size, so even after #646 a saturated worker spends about half of every take/run cycle idle, waiting on the take. With take_ahead set, a worker whose last take came back full also asks for up to take_ahead extra jobs and holds them in jobs_queue until slots free. The next take is then in flight while those jobs run, instead of starting after them.

• take_ahead comes from the start() config or RUNPOD_TAKE_AHEAD, default 0. With 0 nothing changes: the queue size, the take sizes and the perf(serverless): take a full batch once drained #646 full-worker wait are the same as on main.
• Deep queue only. Take-ahead asks for extra jobs only after a take returned everything it asked for. A short take means the queue is shallow, so the next take asks only for free slots, as today. That keeps a shallow queue spread across workers, which was the concern in Job-take round trip costs ~500ms with a deep queue, capping saturated worker throughput #647.
• One buffer at a time: while held jobs are waiting, get_jobs waits for run_jobs to empty the buffer before it takes again. Concurrency still caps what runs.
• Stop signals apply to held jobs: stop_job marks a held job and run_jobs drops it before it starts (and returns True, so it is reported as stopped).
• Shutdown drains held jobs like any other queued job.
• Disabled when refresh_worker is set in the config, since the worker exits after its first job. Docs say not to combine it with handlers that return refresh_worker either (same race as fix(serverless): stop taking jobs after a per-job refresh_worker (SLS-900) #663, with more jobs at stake).
• jobs_queue is sized concurrency + take_ahead, so set_scale keeps working.

Docs: new "Take-ahead" section and config row in docs/serverless/worker.md.

Benchmark

I don't have a GPU endpoint, so I couldn't rerun your vLLM setup. Instead scripts/benchmark_job_take.py runs real JobScaler instances against an in-process fake queue that models the latency you measured in #647: a take with jobs queued costs 0.5s regardless of size, a take already long-polling an empty queue returns 50ms after a job arrives. This is a model of the scheduling, not a platform measurement.

32 slots, take_ahead=32, time until 99% of the burst is done, median of 5 runs, Apple Silicon arm64. The main column is the script run against unmodified main.

scenario	main	this PR, take_ahead=32	change
1 worker, 512 burst, 0.5s jobs (the #647 setup)	15.57s	8.07s	-48%
1 worker, 256 burst, 2.0s jobs	19.56s	16.06s	-18%
4 workers, 1024 burst, 0.5s jobs	7.58s	4.07s	-46%
4 workers, 64 burst (shallow)	0.56s	0.55s	same, same split across workers

On the deep-queue runs the job split across workers was identical with and without take-ahead (256 each with 4 workers). With take_ahead=0 the branch matches main (15.58s / 7.57s / 0.55s). The 2.0s-job scenario is noisier: baseline runs ranged from 19.6s to 24.6s across reps, so treat its -18% as approximate.

uv run python scripts/benchmark_job_take.py --reps 5

The part I can't check here is how this behaves next to the autoscaler: a held job could have gone to a worker that was just scaled up. That is why it only kicks in while takes come back full, and why it is opt-in.

Test plan

• New tests/test_serverless/test_take_ahead.py (20 cases): config/env parsing and invalid values, refresh_worker disables it, queue sizing, default worker unchanged, deep queue takes ahead while slots are busy, no new take until the buffer drains, short take asks only for free slots, shallow queue never takes ahead, zero concurrency, stop signal drops a held job, shutdown runs held jobs
• test_rp_scale.py + test_take_ahead.py repeated 12x, no failures
• Full suite: 2340 passed, coverage 93.9%. Two test_apps/test_dispatch.py::TestSyncBridge::test_interrupt_* cases timed out under load on my machine and pass when run on their own; they don't touch the scaler.

AI-assisted; tested locally with the commands above.