feat: add endpoint to clear all kv blocks in vllm v1 - #1384
Conversation
|
👋 Hi jain-ria! Thank you for contributing to ai-dynamo/dynamo. Just a reminder: The 🚀 |
github-actions
Bot
added
the
external-contribution
jain-ria
changed the title
jain-ria
changed the title
8f932e1 to
4cfa8ab
Compare
4cfa8ab to
5f76390
Compare
5f76390 to
098835e
Compare
098835e to
ae45360
Compare
e9a7528 to
44fafcc
Compare
44fafcc to
148fa5a
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
126-132:⚠️ Potential issue
reset_prefix_cache()must be off-loaded to avoid blocking the event loop
self.engine_client.reset_prefix_cache()is a synchronous call that may involve GPU / NCCL work.
Invoking it directly in anasync defcoroutine risks freezing the entire event-loop under heavy load.- try: - self.engine_client.reset_prefix_cache() + try: + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self.engine_client.reset_prefix_cache)The identical concern was raised in a previous review but has not yet been addressed.
🧹 Nitpick comments (1)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
268-275: Graceful handling of endpoint-serving failures
asyncio.gather()cancels the remaining tasks as soon as the first one raises.
A transient error in eitherserve_endpointwould therefore bring down the whole worker.Consider running the coroutines shielded and restarting the faulty one, or use
return_exceptions=Trueand decide per-endpoint how to proceed.This strengthens availability without changing happy-path behaviour.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
launch/dynamo-run/src/input/http.rs(1 hunks)launch/dynamo-run/src/subprocess/vllm_v1_inc.py(3 hunks)lib/llm/src/http/service/clear_kv_blocks.rs(1 hunks)lib/llm/src/http/service/service_v2.rs(5 hunks)
✅ Files skipped from review due to trivial changes (1)
- lib/llm/src/http/service/clear_kv_blocks.rs
🚧 Files skipped from review as they are similar to previous changes (2)
- launch/dynamo-run/src/input/http.rs
- lib/llm/src/http/service/service_v2.rs
⏰ Context from checks skipped due to timeout of 90000ms (5)
- GitHub Check: Mirror Repository to GitLab
- GitHub Check: pre-merge-rust (lib/runtime/examples)
- GitHub Check: pre-merge-rust (.)
- GitHub Check: pre-merge-rust (lib/bindings/python)
- GitHub Check: Build and Test - vllm
🔇 Additional comments (1)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
188-190: Consider explicit registration for theclear_kv_blocksendpointOnly
generate_endpointis passed toregister_llm.
If any LLM-registry, routing or discovery logic expects all callable endpoints to be registered, the reset route will be invisible.Please double-check whether
register_llmmust also receiveclear_endpoint; if not, add a clarifying comment explaining why.
148fa5a to
e14366f
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (3)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
126-132: 🛠️ Refactor suggestionAvoid blocking the event-loop when clearing KV blocks
self.engine_client.reset_prefix_cache()is still invoked synchronously.
If that call performs GPU/NCCL synchronisation it will stall the entire asyncio loop and every request handler.- self.engine_client.reset_prefix_cache() + # Off-load to a worker thread so we don’t block the event-loop + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self.engine_client.reset_prefix_cache)lib/llm/src/http/service/clear_kv_blocks.rs (2)
190-195:⚠️ Potential issueEndpoint comparison will never match – use
ends_with
instance.endpointcontains the fully-qualified endpoint (e.g.dyn://ns.comp.clear_kv_blocks).
Comparing with bare"clear_kv_blocks"always fails, so no instance will ever be cleared.- .filter(|instance| instance.endpoint == "clear_kv_blocks") + .filter(|instance| instance.endpoint.ends_with("clear_kv_blocks"))
217-239: 🛠️ Refactor suggestionEmpty payload & lack of result validation
router.round_robin(().into())sendsnull– schema can’t evolve.
Send at least{}.You ignore the JSON body and mark success unconditionally when a frame arrives.
Parse the first item and ensurestatus == "success".- match router.round_robin(().into()).await { + match router.round_robin(json!({}).into()).await { ... - Some(response) => { + Some(response) if response.get("status") == Some(&json!("success")) => { ... + Some(err_resp) => { + add_worker_result( + false, + instance_name, + "Instance reported error", + namespace, + component, + Some(err_resp.to_string()), + ); + }
🧹 Nitpick comments (2)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
268-275: Unhandled cancellation / error propagation inasyncio.gatherIf one of the served endpoints raises,
asyncio.gathercancels the remaining task and re-raises.
Consider:await asyncio.gather( generate_endpoint.serve_endpoint(handler.generate), clear_endpoint.serve_endpoint(handler.clear_kv_blocks), return_exceptions=False, # explicit & self-documenting )and add graceful shutdown / cancellation handling so the worker exits cleanly instead of leaving orphaned tasks.
lib/llm/src/http/service/clear_kv_blocks.rs (1)
44-48: Return appropriate HTTP status code for missing workersThe handler currently returns HTTP 200 with a JSON error message.
Consider usingStatusCode::NOT_FOUND(404) to signal that no worker groups are active.return (StatusCode::NOT_FOUND, Json(json!({ "message": ... })))
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
launch/dynamo-run/src/input/http.rs(1 hunks)launch/dynamo-run/src/subprocess/vllm_v1_inc.py(3 hunks)lib/llm/src/http/service/clear_kv_blocks.rs(1 hunks)lib/llm/src/http/service/service_v2.rs(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- launch/dynamo-run/src/input/http.rs
- lib/llm/src/http/service/service_v2.rs
⏰ Context from checks skipped due to timeout of 90000ms (5)
- GitHub Check: Mirror Repository to GitLab
- GitHub Check: pre-merge-rust (.)
- GitHub Check: pre-merge-rust (lib/runtime/examples)
- GitHub Check: pre-merge-rust (lib/bindings/python)
- GitHub Check: Build and Test - vllm
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
126-131:⚠️ Potential issueStill synchronous ‑ this blocks the event loop
reset_prefix_cache()is executed directly inside the async handler.
If the underlying vLLM call performs GPU / NCCL sync the whole event-loop stalls.The exact same issue was raised previously – please off-load the call via
asyncio.to_thread()orloop.run_in_executor().- self.engine_client.reset_prefix_cache() + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self.engine_client.reset_prefix_cache)lib/llm/src/http/service/clear_kv_blocks.rs (1)
190-194:⚠️ Potential issueEndpoint comparison will never match – use
ends_with
instance.endpointholds the fully-qualified name
(e.g.dyn://ns.comp.clear_kv_blocks).
Comparing it to the bare string"clear_kv_blocks"always fails, so every
worker will be reported as unsupported.- .filter(|instance| instance.endpoint == "clear_kv_blocks") + .filter(|instance| instance + .endpoint + .ends_with("clear_kv_blocks"))
🧹 Nitpick comments (1)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
268-275:asyncio.gathererror-propagation can swallow exceptions
When eitherserve_endpointcoroutine exits unexpectedly the other keeps
running and the first exception wins, potentially masking later failures.
Considerasyncio.TaskGroup(3.11+) orgather(..., return_exceptions=True)
and explicit cancellation handling to surface all failures.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
launch/dynamo-run/src/input/http.rs(1 hunks)launch/dynamo-run/src/subprocess/vllm_v1_inc.py(3 hunks)lib/llm/src/http/service/clear_kv_blocks.rs(1 hunks)lib/llm/src/http/service/service_v2.rs(5 hunks)
🚧 Files skipped from review as they are similar to previous changes (2)
- launch/dynamo-run/src/input/http.rs
- lib/llm/src/http/service/service_v2.rs
🧰 Additional context used
🧬 Code Graph Analysis (1)
launch/dynamo-run/src/subprocess/vllm_v1_inc.py (1)
launch/dynamo-run/src/subprocess/vllm_inc.py (2)
RequestHandler(53-130)generate(90-130)
⏰ Context from checks skipped due to timeout of 90000ms (5)
- GitHub Check: Mirror Repository to GitLab
- GitHub Check: Build and Test - vllm
- GitHub Check: pre-merge-rust (.)
- GitHub Check: pre-merge-rust (lib/runtime/examples)
- GitHub Check: pre-merge-rust (lib/bindings/python)
|
The AIQ team asked for a way to directly clear all blocks for all workers for testing, which is why I added this through http for now |
|
qq, what happens if requests are in flight and this is triggered? Is everything ok? Is the guidance when using this to stop all requests, call this function, and then can send requests again? |
|
Based on some quick tests of a bunch of interleaved generate & clear requests, the generate requests seem to finish normally after the kv blocks are cleared |
|
|
||
| for instance in &instances_filtered { | ||
| let instance_name = format!("{}-instance-{}", entry.name, instance.id()); | ||
| match router.round_robin(().into()).await { |
There was a problem hiding this comment.
Why round_robin if we already identified and stored all the instance IDs? Shouldn't we use direct on our explicit list?
| message: Option<String>| { | ||
| let mut result = json!({ | ||
| "name": name, | ||
| "endpoint": format!("{}/{}/clear_kv_blocks", ns, comp), |
There was a problem hiding this comment.
"clear_kv_blocks" should be a constant somewhere rather than hard-coded strings in a few places to avoid errors and improve compiler's ability to detect errors later, similar to load_metrics:
dynamo/lib/llm/src/kv_router.rs
Line 45 in 0e7d4d8
| .enable_cmpl_endpoints(true) | ||
| .enable_embeddings_endpoints(true) | ||
| .with_request_template(template) | ||
| .runtime(Some(Arc::new(distributed_runtime))) |
There was a problem hiding this comment.
This seems like a pretty intrusive change if it adds runtime object that wasn't previously necessary for other functionalities.
I think we'll probably want some kind of broadcast generalizations to support doing some kind of call or operation on all instances, if that is what we're generally going for with this feature.
There was a problem hiding this comment.
CC @ryanolson @paulhendricks @grahamking (Graham on PTO, but for future viz)
There was a problem hiding this comment.
Yea Graham had me remove the DRT from this. I think there should be a way to access things via the ModelState.
Also see #1037 (comment)
ryanolson
left a comment
ryanolson
left a comment
There was a problem hiding this comment.
I failed to get my review done in a timely manner.
I don't think we should hold an http end point to trigger cache resets.
The proper way to do this would be to take a backend out of service, let it drain, reset the cache then, bring it back in service.
I think we need to think harder about how we issue this directive to a fleet of workers.
Overall this both complicates the public http api and likely introduces more complications that it goodness it provides.
I feel we should refer this change at the soonest opportunity.
|
|
||
| mod openai; | ||
|
|
||
| pub mod clear_kv_blocks; |
There was a problem hiding this comment.
Do we want to put this on the public API? I'm not so sure.
|
@jain-ria Could you revert this please? We now realize this isn't what we want to do, and need to put some extra thought into it. Sorry about that, I was on vacation. |
|
Yes! Sorry for the delay, changes are here in this PR |
|
While the HTTP server may not be the right place for it, we do need an API for certain external changes that is separate from what a User would see. I guess I would break it up into User API, Admin API, and Planner API. With there being some overlap between Planner and Admin. For example an Admin could say I want to drain and reset the KV Cache of these workers or I want to scale up to 6 workers, and this would override any directives from Planner. |
7 participants
Overview:
Adds an endpoint to allow users to directly access vLLM v1's clear all kv blocks feature.
Details:
Added a new axum router, handler function, and dynamo endpoint for clearing all kv blocks.
Example: curl -X POST http://localhost:8000/clear_kv_blocks
Where should the reviewer start?
Router and handler function:
Dynamo endpoint:
Corresponding updates in:
Summary by CodeRabbit
Summary by CodeRabbit