[BUG]: Replicas downscale does not drain the requests

[RealNicolasBourbaki]

Describe the Bug

Currently when a replica is shed, all in-flight and queued requests are aborted a fixed 5 s after SIGTERM (5 s = default DYN_GRACEFUL_SHUTDOWN_GRACE_PERIOD_SECS). Clients get HTTP 500 "Failed to fold completions stream". There is no stage that waits for inflight requests to reach zero.

Steps to Reproduce

1. deploy the attached DGD
   dgd-drain-repro.yaml

2. you can use your own load or use mine in the attachment to generate a trace: python3 make-trace.py --out trace.jsonl. This trace has 2 phases, phase 1 with TPS = 8.0 req/s, phase 2 with 0.2 req/s --> perfect to watch dynamo sheds replicas

3. replay the trace against the deployment:
   FE=$(kubectl -n dynamo-system get svc drain-repro-frontend -o jsonpath='{.spec.clusterIP}') aiperf profile --model openai/gpt-oss-120b --tokenizer openai/gpt-oss-120b \ --endpoint-type completions --streaming --url http://$FE:8000 \ --input-file trace.jsonl --custom-dataset-type mooncake_trace \ --fixed-schedule --fixed-schedule-auto-offset \ --extra-inputs "ignore_eos:true" --artifact-dir ./aiperf-out

4. observe the responses of the clients during replica shedding.make-trace.py

Expected Behavior

a de-registered worker finishes its in-flight requests before exiting, bounded by a timeout

Actual Behavior

The worker de-registers from discovery correctly, so no new requests are routed to it. It then sleeps a fixed 5 s and aborts everything that are still generating. Clients receive HTTP 500.

Worker log of the shed replica (Dynamo 1.4.0, gpt-oss-120b):

16:41:50.979  Unregistering endpoint: namespace=dynamo-system-drain-repro ...
16:41:51.004  Successfully unregistered endpoint instance - worker removed from routing
16:41:51.005  Grace period 5.00s before stopping endpoints
16:41:56.009  Initiating runtime shutdown                     # exactly 5.004 s later

On the client side, at that same instant, 18 of 1500 requests fail simultaneously:

{"message":"Failed to fold completions stream for cc14fa97-...","type":"Internal Server Error"}

graceful_shutdown_with_discovery aborts without draining. Here is what happens chronologically during downscale:

step	source (v1.4.0)
unregister from discovery — correct	graceful_shutdown.py#L104
await asyncio.sleep(grace_period_s), default 5.0 s	#L108 / #L15
drain_callback — skipped, vLLM passes None	#L110
shutdown_event.set()	#L129-L130
every in-flight request calls engine_client.abort()	handlers.py#L1827, #L1861
runtime.shutdown()	#L149

Environment

• Dynamo: 1.4.0 same as the planner and vllm-runtime versions
• Inference engine: vllm-runtime 1.4.0 shipped by dynamo, based on vLLM 0.26.0
• model: gpt-oss-120b
• K8s: v1.34.8+k3s1
• OS: Ubuntu 24.04.3 (kernel 6.8.0-71)
• GPU: Nvidia h200

Additional Context

Just 2 FYIs:

1. there is a request draining feature but for dynamo rust runtime shutdown (here). This never kicks in because SIGTERM first calls shutdown_event, aborts in-flight requests and stops background loops. Only after that the rust runtime starts to shutdown, and start draining. But at that time there is no requests to be drained anymore.

2. (UPDATE: This is confirmed to be irrelevant after the discussion below) Dynamo already has a callbakc receiver: DynamoStatLoggerPublisher, whose record() receives a SchedulerStats every engine iteration (here), and that carries num_running_reqs and num_waiting_reqwhich might be useful for draining_callback

Screenshots

No response

[ayaangazali]

I looked at fixing this and concluded it needs a design decision rather than a patch. Three findings, in case they help.

The hook named here is for something else. drain_callback is documented narrowly: "Use this on prefill workers to wait for in-flight NIXL KV transfers to complete, preventing decode workers from segfaulting due to use-after-free on freed GPU memory (#7319)." Wiring general request draining onto it is exactly what #13286 calls out, that TRT-LLM's callback "polls active + queued, conflating request drain with KV-transfer drain". So copying the TRT-LLM shape to vLLM would spread a pattern a maintainer has already flagged.

The shutdown paths on main differ from the v1.4.0 links in the report. SGLang has its own components/src/dynamo/sglang/shutdown.py, which does call graceful_shutdown_with_discovery (and, as you say, without a drain callback). vLLM does not use that helper at all on main: it installs signal handlers via install_signal_handlers and registers endpoints with graceful_shutdown=True, so it drains through the Rust request-plane path in lib/backend-common. Worth re-confirming the symptom against main before the fix is scoped, because the two backends are not in the same position any more.

Where the fix belongs is what #13286 is settling. That issue is explicitly about unifying shutdown ordering and timeout budgets, and it names this gap, that vLLM installs no prefill KV-transfer drain and that the stages do not share one deadline. Adding a vLLM drain callback now would likely be rewritten by it.

I have not opened a PR. Beyond the overlap, I cannot verify draining without a cluster and sustained load, and a shutdown-path change that is wrong fails in the direction of dropping requests or hanging termination.

If maintainers would rather have the narrow fix ahead of #13286, I am happy to write it once someone says which contract it should follow: whether the grace period should cover the inflight drain, or the drain should be a separate stage with its own budget.

[RealNicolasBourbaki]

@ayaangazali thanks, this is useful — point 1 changes what I was framing. You're right that drain_callback is scoped to prefill NIXL KV transfers (#7319) and that TRT-LLM's is gated on DisaggregationMode.PREFILL, so it would never fire on an aggregated worker anyway. I've rewritten the description to drop that framing.

One correction that I think helps your case instead of hurting it: vLLM does still go through graceful_shutdown_with_discovery on main. install_signal_handlers isn't an alternative to it, it's the entry point: here is the installation and in the handler body there is asyncio.create_task(graceful_shutdown_with_discovery(...)).

So I would reframe this issue as the reproducible symptom of #13286 's stage 2 (stop admission, wait inflight == 0, before shutdown_event.set()) rather than a separate fix.