From f1a0fd23c586ce8e150b3bb23cf635437ccd661e Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 4 Jun 2025 13:54:58 -0700 Subject: [PATCH 1/5] Add clear_all_blocks endpoint --- lib/llm/src/http/service.rs | 1 + lib/llm/src/http/service/clear_all_blocks.rs | 57 ++++++++++++++++++++ lib/llm/src/http/service/service_v2.rs | 1 + 3 files changed, 59 insertions(+) create mode 100644 lib/llm/src/http/service/clear_all_blocks.rs diff --git a/lib/llm/src/http/service.rs b/lib/llm/src/http/service.rs index 9c4081f6efec..02863afebb9a 100644 --- a/lib/llm/src/http/service.rs +++ b/lib/llm/src/http/service.rs @@ -20,6 +20,7 @@ mod openai; +pub mod clear_all_blocks; pub mod error; pub mod health; pub mod metrics; diff --git a/lib/llm/src/http/service/clear_all_blocks.rs b/lib/llm/src/http/service/clear_all_blocks.rs new file mode 100644 index 000000000000..4183bef47bb2 --- /dev/null +++ b/lib/llm/src/http/service/clear_all_blocks.rs @@ -0,0 +1,57 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{service_v2, RouteDoc}; +use axum::{http::Method, http::StatusCode, response::IntoResponse, routing::get, Json, Router}; +use serde_json::json; +use std::sync::Arc; + +pub fn clear_all_blocks_router( + state: Arc, + path: Option, +) -> (Vec, Router) { + let path = path.unwrap_or_else(|| "/clear_all_blocks".to_string()); + + let docs: Vec = vec![RouteDoc::new(Method::GET, &path)]; + + let router = Router::new() + .route(&path, get(clear_all_blocks_handler)) + .with_state(state); + + (docs, router) +} + +async fn clear_all_blocks_handler( + axum::extract::State(_state): axum::extract::State>, +) -> impl IntoResponse { + Json(serde_json::json!({ + "message": "hello" + })) +} diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index e381fd04b7c2..508fd0a99436 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -148,6 +148,7 @@ impl HttpServiceConfigBuilder { metrics::router(registry, None), super::openai::list_models_router(state.clone(), None), super::health::health_check_router(state.clone(), None), + super::clear_all_blocks::clear_all_blocks_router(state.clone(), None), ]; if config.enable_chat_endpoints { From 19d59e43c3ce9db39783193f6a3600800f03b2ca Mon Sep 17 00:00:00 2001 From: jain-ria Date: Sun, 8 Jun 2025 20:55:44 -0700 Subject: [PATCH 2/5] Updated clear_kv_blocks integration with vllm v1 exmaple --- .../dynamo-run/src/subprocess/vllm_v1_inc.py | 50 ++-- lib/llm/src/http/service.rs | 2 +- lib/llm/src/http/service/clear_all_blocks.rs | 57 ---- lib/llm/src/http/service/clear_kv_blocks.rs | 243 ++++++++++++++++++ lib/llm/src/http/service/service_v2.rs | 32 ++- 5 files changed, 308 insertions(+), 76 deletions(-) delete mode 100644 lib/llm/src/http/service/clear_all_blocks.rs create mode 100644 lib/llm/src/http/service/clear_kv_blocks.rs diff --git a/launch/dynamo-run/src/subprocess/vllm_v1_inc.py b/launch/dynamo-run/src/subprocess/vllm_v1_inc.py index 04732e11f5e8..f6a3b35f3fb9 100644 --- a/launch/dynamo-run/src/subprocess/vllm_v1_inc.py +++ b/launch/dynamo-run/src/subprocess/vllm_v1_inc.py @@ -22,15 +22,6 @@ from typing import Optional import uvloop -from vllm.config import VllmConfig -from vllm.distributed.kv_events import KVEventsConfig -from vllm.engine.arg_utils import AsyncEngineArgs -from vllm.inputs import TokensPrompt -from vllm.sampling_params import SamplingParams -from vllm.usage.usage_lib import UsageContext -from vllm.v1.engine.async_llm import AsyncLLM -from vllm.v1.metrics.loggers import StatLoggerBase -from vllm.v1.metrics.stats import IterationStats, SchedulerStats from dynamo.llm import ( ModelType, @@ -40,6 +31,15 @@ register_llm, ) from dynamo.runtime import Component, DistributedRuntime, dynamo_worker +from vllm.config import VllmConfig +from vllm.distributed.kv_events import KVEventsConfig +from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.inputs import TokensPrompt +from vllm.sampling_params import SamplingParams +from vllm.usage.usage_lib import UsageContext +from vllm.v1.engine.async_llm import AsyncLLM +from vllm.v1.metrics.loggers import StatLoggerBase +from vllm.v1.metrics.stats import IterationStats, SchedulerStats # Only used if you run it manually from the command line DEFAULT_ENDPOINT = "dyn://dynamo.backend.generate" @@ -115,7 +115,7 @@ def __call__(self, vllm_config: VllmConfig, dp_rank: int) -> StatLoggerBase: class RequestHandler: """ - Request handler for the generate endpoint + Request handler for the generate and clear_kv_blocks endpoints. """ def __init__(self, component, engine, default_sampling_params): @@ -123,6 +123,16 @@ def __init__(self, component, engine, default_sampling_params): self.engine_client = engine self.default_sampling_params = default_sampling_params + async def clear_kv_blocks(self, request=None): + logger.info("clear_kv_blocks endpoint called") + try: + self.engine_client.reset_prefix_cache() + logger.info("Successfully reset prefix cache") + yield {"status": "success", "message": "KV cache cleared"} + except Exception as e: + logger.error(f"Error clearing KV cache: {e}") + yield {"status": "error", "message": str(e)} + async def generate(self, request): request_id = str(uuid.uuid4().hex) @@ -174,13 +184,16 @@ async def init(runtime: DistributedRuntime, config: Config): """ Instantiate and serve """ + component = runtime.namespace(config.namespace).component(config.component) await component.create_service() - endpoint = component.endpoint(config.endpoint) + generate_endpoint = component.endpoint(config.endpoint) + clear_endpoint = component.endpoint("clear_kv_blocks") + await register_llm( ModelType.Backend, - endpoint, + generate_endpoint, config.model_path, config.model_name, kv_cache_block_size=config.kv_block_size, @@ -248,16 +261,21 @@ async def init(runtime: DistributedRuntime, config: Config): logger.info("VllmWorker has been initialized") zmq_config = ZmqKvEventPublisherConfig( - worker_id=endpoint.lease_id(), kv_block_size=engine_args.block_size + worker_id=generate_endpoint.lease_id(), kv_block_size=engine_args.block_size ) _ = ZmqKvEventPublisher(component=component, config=zmq_config) handler = RequestHandler(component, engine_client, default_sampling_params) - # the server will gracefully shutdown (i.e., keep opened TCP streams finishes) - # after the lease is revoked - await endpoint.serve_endpoint(handler.generate) + try: + await asyncio.gather( + generate_endpoint.serve_endpoint(handler.generate), + clear_endpoint.serve_endpoint(handler.clear_kv_blocks), + ) + except Exception as e: + logger.error(f"Failed to serve endpoints: {e}") + raise def cmd_line_args(): diff --git a/lib/llm/src/http/service.rs b/lib/llm/src/http/service.rs index 02863afebb9a..b36c867d73b5 100644 --- a/lib/llm/src/http/service.rs +++ b/lib/llm/src/http/service.rs @@ -20,7 +20,7 @@ mod openai; -pub mod clear_all_blocks; +pub mod clear_kv_blocks; pub mod error; pub mod health; pub mod metrics; diff --git a/lib/llm/src/http/service/clear_all_blocks.rs b/lib/llm/src/http/service/clear_all_blocks.rs deleted file mode 100644 index 4183bef47bb2..000000000000 --- a/lib/llm/src/http/service/clear_all_blocks.rs +++ /dev/null @@ -1,57 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -use super::{service_v2, RouteDoc}; -use axum::{http::Method, http::StatusCode, response::IntoResponse, routing::get, Json, Router}; -use serde_json::json; -use std::sync::Arc; - -pub fn clear_all_blocks_router( - state: Arc, - path: Option, -) -> (Vec, Router) { - let path = path.unwrap_or_else(|| "/clear_all_blocks".to_string()); - - let docs: Vec = vec![RouteDoc::new(Method::GET, &path)]; - - let router = Router::new() - .route(&path, get(clear_all_blocks_handler)) - .with_state(state); - - (docs, router) -} - -async fn clear_all_blocks_handler( - axum::extract::State(_state): axum::extract::State>, -) -> impl IntoResponse { - Json(serde_json::json!({ - "message": "hello" - })) -} diff --git a/lib/llm/src/http/service/clear_kv_blocks.rs b/lib/llm/src/http/service/clear_kv_blocks.rs new file mode 100644 index 000000000000..a0975acbec6d --- /dev/null +++ b/lib/llm/src/http/service/clear_kv_blocks.rs @@ -0,0 +1,243 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use super::{service_v2, RouteDoc}; +use axum::{http::Method, response::IntoResponse, routing::post, Json, Router}; +use serde_json::json; +use std::sync::Arc; +use serde::{Deserialize, Serialize}; + +use dynamo_runtime::{ + logging, pipeline::PushRouter, protocols::annotated::Annotated, stream::StreamExt, + DistributedRuntime, Result, Runtime, Worker, +}; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ClearKvBlocksRequest { + pub message: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ClearKvBlocksResponse { + pub success: bool, + pub message: String, +} + +pub fn clear_kv_blocks_router( + state: Arc, + path: Option, +) -> (Vec, Router) { + let path = path.unwrap_or_else(|| "/clear_kv_blocks".to_string()); + + let docs: Vec = vec![RouteDoc::new(Method::POST, &path)]; + + let router = Router::new() + .route(&path, post(clear_kv_blocks_handler)) + .with_state(state); + + (docs, router) +} + +async fn clear_kv_blocks_handler( + axum::extract::State(state): axum::extract::State>, +) -> impl IntoResponse { + + tracing::trace!("Received clear all KV blocks request"); + + let model_entries = state.manager().get_model_entries(); + + tracing::debug!("Found {} model entries:", model_entries.len()); + for entry in &model_entries { + tracing::debug!("Entry: name={}, namespace={}, component={}", + entry.name, entry.endpoint.namespace, entry.endpoint.component); + } + + // if there are no active workers + if model_entries.is_empty() { + return Json(serde_json::json!({ + "message": "No active worker groups found" + })); + } + + let mut cleared_workers = Vec::new(); + let mut failed_workers = Vec::new(); + + let distributed = match state.runtime() { + Some(runtime) => runtime, + None => { + return Json(serde_json::json!({ + "message": "Failed to create distributed runtime", + })); + } + }; + + // create client for each model entry + for entry in &model_entries { + let namespace = &entry.endpoint.namespace; + let component = &entry.endpoint.component; + + tracing::debug!("Processing worker group: {}/{}", namespace, component); + + let namespace_obj = match distributed.namespace(namespace) { + Ok(ns) => ns, + Err(e) => { + failed_workers.push(json!({ + "name": entry.name, + "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), + "status": "failed to get namespace", + "error": e.to_string() + })); + continue; + } + }; + + let component_obj = match namespace_obj.component(component) { + Ok(comp) => comp, + Err(e) => { + failed_workers.push(json!({ + "name": entry.name, + "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), + "status": "failed to get component", + "error": e.to_string() + })); + continue; + } + }; + + + let endpoint: dynamo_runtime::component::Endpoint = component_obj.endpoint("clear_kv_blocks"); + + let client = match endpoint.client().await { + Ok(c) => { + c + }, + Err(e) => { + failed_workers.push(json!({ + "name": entry.name, + "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), + "status": "failed to create client", + "error": e.to_string() + })); + continue; } + }; + + let router = match PushRouter::<(), serde_json::Value>::from_client(client.clone(), Default::default()).await { + Ok(r) => r, + Err(e) => { + failed_workers.push(json!({ + "name": entry.name, + "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), + "status": "failed to create router", + "error": e.to_string() + })); + continue; + } + }; + + let instances = match component_obj.list_instances().await { + Ok(instances) => { + instances + } + Err(e) => { + failed_workers.push(json!({ + "name": entry.name, + "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), + "status": "Failed to get instances for worker group", + "error": e.to_string() + })); + continue; + } + }; + + if instances.is_empty() { + failed_workers.push(json!({ + "name": entry.name, + "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), + "status": "No instances found for worker group", + })); + continue; + } + + let instances_filtered = instances.clone().into_iter().filter(|instance| instance.endpoint == "clear_kv_blocks").collect::>(); + + if instances_filtered.is_empty() { + let found_endpoints: Vec = instances.iter().map(|instance| instance.endpoint.clone()).collect(); + failed_workers.push(json!({ + "name": entry.name, + "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), + "status": format!("Worker group doesn't support clear_kv_blocks. Supported endpoints: {}", found_endpoints.join(", ")), + })); + continue; + } + + for instance in &instances_filtered { + match router.round_robin(().into()).await { + Ok(mut stream) => { + // Successfully sent request, now process the response + match stream.next().await { + Some(response) => { + // Instance successfully cleared its KV blocks + cleared_workers.push(json!({ + "name": format!("{}-instance-{}", entry.name, instance.id()), + "endpoint": format!("{}/{}/clear_kv_blocks", entry.endpoint.namespace, entry.endpoint.component), + "status": "successfully cleared kv blocks for instance", + "response": response.to_string() + })); + } + None => { + // No response from instance + failed_workers.push(json!({ + "name": format!("{}-instance-{}", entry.name, instance.id()), + "endpoint": format!("{}/{}/clear_kv_blocks", entry.endpoint.namespace, entry.endpoint.component), + "status": "no response from instance", + })); + } + } + } + Err(e) => { + // Failed to send request to this instance + failed_workers.push(json!({ + "name": format!("{}-instance-{}", entry.name, instance.id()), + "endpoint": format!("{}/{}/clear_kv_blocks", entry.endpoint.namespace, entry.endpoint.component), + "status": "failed to send request for instance", + "error": e.to_string() + })); + } + } + } + } + + Json(serde_json::json!({ + "message": format!("Cleared prefix cache on {} out of {} worker groups", + cleared_workers.len(), model_entries.len()), + "cleared_workers": cleared_workers, + "failed_workers": failed_workers + })) +} diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index 508fd0a99436..3f49731f3ddc 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -13,11 +13,13 @@ use anyhow::Result; use derive_builder::Builder; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; +use dynamo_runtime::DistributedRuntime; /// HTTP service shared state pub struct State { metrics: Arc, manager: Arc, + runtime: Option>, } impl State { @@ -25,6 +27,15 @@ impl State { Self { manager, metrics: Arc::new(Metrics::default()), + runtime: None, + } + } + + pub fn with_runtime(manager: Arc, runtime: Arc) -> Self { + Self { + manager, + metrics: Arc::new(Metrics::default()), + runtime: Some(runtime), } } @@ -41,6 +52,11 @@ impl State { self.manager.clone() } + /// Get the DistributedRuntime if available + pub fn runtime(&self) -> Option<&DistributedRuntime> { + self.runtime.as_ref().map(|r| r.as_ref()) + } + // TODO pub fn sse_keep_alive(&self) -> Option { None @@ -80,6 +96,9 @@ pub struct HttpServiceConfig { #[builder(default = "None")] request_template: Option, + + #[builder(default = "None")] + runtime: Option>, } impl HttpService { @@ -134,7 +153,11 @@ impl HttpServiceConfigBuilder { let config: HttpServiceConfig = self.build_internal()?; let model_manager = Arc::new(ModelManager::new()); - let state = Arc::new(State::new(model_manager)); + let state = if let Some(runtime) = config.runtime { + Arc::new(State::with_runtime(model_manager, runtime)) + } else { + Arc::new(State::new(model_manager)) + }; // enable prometheus metrics let registry = metrics::Registry::new(); @@ -148,7 +171,7 @@ impl HttpServiceConfigBuilder { metrics::router(registry, None), super::openai::list_models_router(state.clone(), None), super::health::health_check_router(state.clone(), None), - super::clear_all_blocks::clear_all_blocks_router(state.clone(), None), + super::clear_kv_blocks::clear_kv_blocks_router(state.clone(), None), ]; if config.enable_chat_endpoints { @@ -190,4 +213,9 @@ impl HttpServiceConfigBuilder { self.request_template = Some(request_template); self } + + pub fn with_runtime(mut self, runtime: Arc) -> Self { + self.runtime = Some(Some(runtime)); + self + } } From 44aa1eeb1a125ba8ad5613994efd18e20025c8ba Mon Sep 17 00:00:00 2001 From: jain-ria Date: Mon, 9 Jun 2025 09:19:50 -0700 Subject: [PATCH 3/5] fixed formatting --- lib/llm/src/http/service/clear_kv_blocks.rs | 47 +++++++++++++-------- lib/llm/src/http/service/service_v2.rs | 2 +- 2 files changed, 30 insertions(+), 19 deletions(-) diff --git a/lib/llm/src/http/service/clear_kv_blocks.rs b/lib/llm/src/http/service/clear_kv_blocks.rs index a0975acbec6d..7a3c00ada2f4 100644 --- a/lib/llm/src/http/service/clear_kv_blocks.rs +++ b/lib/llm/src/http/service/clear_kv_blocks.rs @@ -30,13 +30,12 @@ use super::{service_v2, RouteDoc}; use axum::{http::Method, response::IntoResponse, routing::post, Json, Router}; +use serde::{Deserialize, Serialize}; use serde_json::json; use std::sync::Arc; -use serde::{Deserialize, Serialize}; use dynamo_runtime::{ - logging, pipeline::PushRouter, protocols::annotated::Annotated, stream::StreamExt, - DistributedRuntime, Result, Runtime, Worker, + pipeline::PushRouter, stream::StreamExt, }; #[derive(Debug, Clone, Serialize, Deserialize)] @@ -68,15 +67,18 @@ pub fn clear_kv_blocks_router( async fn clear_kv_blocks_handler( axum::extract::State(state): axum::extract::State>, ) -> impl IntoResponse { - tracing::trace!("Received clear all KV blocks request"); let model_entries = state.manager().get_model_entries(); tracing::debug!("Found {} model entries:", model_entries.len()); for entry in &model_entries { - tracing::debug!("Entry: name={}, namespace={}, component={}", - entry.name, entry.endpoint.namespace, entry.endpoint.component); + tracing::debug!( + "Entry: name={}, namespace={}, component={}", + entry.name, + entry.endpoint.namespace, + entry.endpoint.component + ); } // if there are no active workers @@ -131,13 +133,11 @@ async fn clear_kv_blocks_handler( } }; - - let endpoint: dynamo_runtime::component::Endpoint = component_obj.endpoint("clear_kv_blocks"); + let endpoint: dynamo_runtime::component::Endpoint = + component_obj.endpoint("clear_kv_blocks"); let client = match endpoint.client().await { - Ok(c) => { - c - }, + Ok(c) => c, Err(e) => { failed_workers.push(json!({ "name": entry.name, @@ -145,10 +145,16 @@ async fn clear_kv_blocks_handler( "status": "failed to create client", "error": e.to_string() })); - continue; } + continue; + } }; - let router = match PushRouter::<(), serde_json::Value>::from_client(client.clone(), Default::default()).await { + let router = match PushRouter::<(), serde_json::Value>::from_client( + client.clone(), + Default::default(), + ) + .await + { Ok(r) => r, Err(e) => { failed_workers.push(json!({ @@ -162,9 +168,7 @@ async fn clear_kv_blocks_handler( }; let instances = match component_obj.list_instances().await { - Ok(instances) => { - instances - } + Ok(instances) => instances, Err(e) => { failed_workers.push(json!({ "name": entry.name, @@ -185,10 +189,17 @@ async fn clear_kv_blocks_handler( continue; } - let instances_filtered = instances.clone().into_iter().filter(|instance| instance.endpoint == "clear_kv_blocks").collect::>(); + let instances_filtered = instances + .clone() + .into_iter() + .filter(|instance| instance.endpoint == "clear_kv_blocks") + .collect::>(); if instances_filtered.is_empty() { - let found_endpoints: Vec = instances.iter().map(|instance| instance.endpoint.clone()).collect(); + let found_endpoints: Vec = instances + .iter() + .map(|instance| instance.endpoint.clone()) + .collect(); failed_workers.push(json!({ "name": entry.name, "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index 3f49731f3ddc..9014b15258fa 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -11,9 +11,9 @@ use crate::discovery::ModelManager; use crate::request_template::RequestTemplate; use anyhow::Result; use derive_builder::Builder; +use dynamo_runtime::DistributedRuntime; use tokio::task::JoinHandle; use tokio_util::sync::CancellationToken; -use dynamo_runtime::DistributedRuntime; /// HTTP service shared state pub struct State { From c5f0f54430bc02757865fc02d09710f0a6f37c4e Mon Sep 17 00:00:00 2001 From: jain-ria Date: Mon, 9 Jun 2025 09:39:42 -0700 Subject: [PATCH 4/5] More formatting --- .../dynamo-run/src/subprocess/vllm_v1_inc.py | 18 +++++++++--------- lib/llm/src/http/service/clear_kv_blocks.rs | 4 +--- 2 files changed, 10 insertions(+), 12 deletions(-) diff --git a/launch/dynamo-run/src/subprocess/vllm_v1_inc.py b/launch/dynamo-run/src/subprocess/vllm_v1_inc.py index f6a3b35f3fb9..6f2cc0b41d9e 100644 --- a/launch/dynamo-run/src/subprocess/vllm_v1_inc.py +++ b/launch/dynamo-run/src/subprocess/vllm_v1_inc.py @@ -22,15 +22,6 @@ from typing import Optional import uvloop - -from dynamo.llm import ( - ModelType, - WorkerMetricsPublisher, - ZmqKvEventPublisher, - ZmqKvEventPublisherConfig, - register_llm, -) -from dynamo.runtime import Component, DistributedRuntime, dynamo_worker from vllm.config import VllmConfig from vllm.distributed.kv_events import KVEventsConfig from vllm.engine.arg_utils import AsyncEngineArgs @@ -41,6 +32,15 @@ from vllm.v1.metrics.loggers import StatLoggerBase from vllm.v1.metrics.stats import IterationStats, SchedulerStats +from dynamo.llm import ( + ModelType, + WorkerMetricsPublisher, + ZmqKvEventPublisher, + ZmqKvEventPublisherConfig, + register_llm, +) +from dynamo.runtime import Component, DistributedRuntime, dynamo_worker + # Only used if you run it manually from the command line DEFAULT_ENDPOINT = "dyn://dynamo.backend.generate" DEFAULT_MODEL = "Qwen/Qwen3-0.6B" diff --git a/lib/llm/src/http/service/clear_kv_blocks.rs b/lib/llm/src/http/service/clear_kv_blocks.rs index 7a3c00ada2f4..6ae510e33001 100644 --- a/lib/llm/src/http/service/clear_kv_blocks.rs +++ b/lib/llm/src/http/service/clear_kv_blocks.rs @@ -34,9 +34,7 @@ use serde::{Deserialize, Serialize}; use serde_json::json; use std::sync::Arc; -use dynamo_runtime::{ - pipeline::PushRouter, stream::StreamExt, -}; +use dynamo_runtime::{pipeline::PushRouter, stream::StreamExt}; #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ClearKvBlocksRequest { From 15135dc2b605bf3482d3cb18ed62c00ab64f38c0 Mon Sep 17 00:00:00 2001 From: jain-ria Date: Wed, 11 Jun 2025 09:40:01 -0700 Subject: [PATCH 5/5] Cleaned up clear_kv_blocks.rs --- launch/dynamo-run/src/input/http.rs | 2 + .../dynamo-run/src/subprocess/vllm_v1_inc.py | 3 - lib/llm/src/http/service/clear_kv_blocks.rs | 230 +++++++++--------- lib/llm/src/http/service/service_v2.rs | 5 - 4 files changed, 120 insertions(+), 120 deletions(-) diff --git a/launch/dynamo-run/src/input/http.rs b/launch/dynamo-run/src/input/http.rs index 8172d245ce1b..4227a04244c8 100644 --- a/launch/dynamo-run/src/input/http.rs +++ b/launch/dynamo-run/src/input/http.rs @@ -29,12 +29,14 @@ pub async fn run( engine_config: EngineConfig, template: Option, ) -> anyhow::Result<()> { + let distributed_runtime = DistributedRuntime::from_settings(runtime.clone()).await?; let http_service = service_v2::HttpService::builder() .port(flags.http_port) .enable_chat_endpoints(true) .enable_cmpl_endpoints(true) .enable_embeddings_endpoints(true) .with_request_template(template) + .runtime(Some(Arc::new(distributed_runtime))) .build()?; match engine_config { EngineConfig::Dynamic => { diff --git a/launch/dynamo-run/src/subprocess/vllm_v1_inc.py b/launch/dynamo-run/src/subprocess/vllm_v1_inc.py index 6f2cc0b41d9e..1ed9a822b92c 100644 --- a/launch/dynamo-run/src/subprocess/vllm_v1_inc.py +++ b/launch/dynamo-run/src/subprocess/vllm_v1_inc.py @@ -124,13 +124,10 @@ def __init__(self, component, engine, default_sampling_params): self.default_sampling_params = default_sampling_params async def clear_kv_blocks(self, request=None): - logger.info("clear_kv_blocks endpoint called") try: self.engine_client.reset_prefix_cache() - logger.info("Successfully reset prefix cache") yield {"status": "success", "message": "KV cache cleared"} except Exception as e: - logger.error(f"Error clearing KV cache: {e}") yield {"status": "error", "message": str(e)} async def generate(self, request): diff --git a/lib/llm/src/http/service/clear_kv_blocks.rs b/lib/llm/src/http/service/clear_kv_blocks.rs index 6ae510e33001..48bcee70312e 100644 --- a/lib/llm/src/http/service/clear_kv_blocks.rs +++ b/lib/llm/src/http/service/clear_kv_blocks.rs @@ -13,40 +13,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - use super::{service_v2, RouteDoc}; use axum::{http::Method, response::IntoResponse, routing::post, Json, Router}; -use serde::{Deserialize, Serialize}; use serde_json::json; use std::sync::Arc; use dynamo_runtime::{pipeline::PushRouter, stream::StreamExt}; -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ClearKvBlocksRequest { - pub message: String, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct ClearKvBlocksResponse { - pub success: bool, - pub message: String, -} - pub fn clear_kv_blocks_router( state: Arc, path: Option, @@ -65,20 +38,8 @@ pub fn clear_kv_blocks_router( async fn clear_kv_blocks_handler( axum::extract::State(state): axum::extract::State>, ) -> impl IntoResponse { - tracing::trace!("Received clear all KV blocks request"); - let model_entries = state.manager().get_model_entries(); - tracing::debug!("Found {} model entries:", model_entries.len()); - for entry in &model_entries { - tracing::debug!( - "Entry: name={}, namespace={}, component={}", - entry.name, - entry.endpoint.namespace, - entry.endpoint.component - ); - } - // if there are no active workers if model_entries.is_empty() { return Json(serde_json::json!({ @@ -86,9 +47,6 @@ async fn clear_kv_blocks_handler( })); } - let mut cleared_workers = Vec::new(); - let mut failed_workers = Vec::new(); - let distributed = match state.runtime() { Some(runtime) => runtime, None => { @@ -98,22 +56,53 @@ async fn clear_kv_blocks_handler( } }; + let mut cleared_workers = Vec::new(); + let mut failed_workers = Vec::new(); + + // update cleared and failed workers + let mut add_worker_result = |success: bool, + name: String, + status: &str, + ns: &str, + comp: &str, + message: Option| { + let mut result = json!({ + "name": name, + "endpoint": format!("{}/{}/clear_kv_blocks", ns, comp), + "status": status, + }); + if success { + if let Some(m) = message { + result["response"] = json!(m); + } + cleared_workers.push(result); + } else { + if let Some(m) = message { + result["error"] = json!(m); + } + failed_workers.push(result); + } + }; + // create client for each model entry for entry in &model_entries { let namespace = &entry.endpoint.namespace; let component = &entry.endpoint.component; + let entry_name = entry.name.to_string(); tracing::debug!("Processing worker group: {}/{}", namespace, component); let namespace_obj = match distributed.namespace(namespace) { Ok(ns) => ns, Err(e) => { - failed_workers.push(json!({ - "name": entry.name, - "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), - "status": "failed to get namespace", - "error": e.to_string() - })); + add_worker_result( + false, + entry_name, + "Failed to get namespace", + namespace, + component, + Some(e.to_string()), + ); continue; } }; @@ -121,12 +110,14 @@ async fn clear_kv_blocks_handler( let component_obj = match namespace_obj.component(component) { Ok(comp) => comp, Err(e) => { - failed_workers.push(json!({ - "name": entry.name, - "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), - "status": "failed to get component", - "error": e.to_string() - })); + add_worker_result( + false, + entry_name, + "Failed to get component", + namespace, + component, + Some(e.to_string()), + ); continue; } }; @@ -137,12 +128,14 @@ async fn clear_kv_blocks_handler( let client = match endpoint.client().await { Ok(c) => c, Err(e) => { - failed_workers.push(json!({ - "name": entry.name, - "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), - "status": "failed to create client", - "error": e.to_string() - })); + add_worker_result( + false, + entry_name, + "Failed to get client", + namespace, + component, + Some(e.to_string()), + ); continue; } }; @@ -155,12 +148,14 @@ async fn clear_kv_blocks_handler( { Ok(r) => r, Err(e) => { - failed_workers.push(json!({ - "name": entry.name, - "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), - "status": "failed to create router", - "error": e.to_string() - })); + add_worker_result( + false, + entry_name, + "Failed to create router", + namespace, + component, + Some(e.to_string()), + ); continue; } }; @@ -168,22 +163,27 @@ async fn clear_kv_blocks_handler( let instances = match component_obj.list_instances().await { Ok(instances) => instances, Err(e) => { - failed_workers.push(json!({ - "name": entry.name, - "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), - "status": "Failed to get instances for worker group", - "error": e.to_string() - })); + add_worker_result( + false, + entry_name, + "Failed to get instances for worker group", + namespace, + component, + Some(e.to_string()), + ); continue; } }; if instances.is_empty() { - failed_workers.push(json!({ - "name": entry.name, - "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), - "status": "No instances found for worker group", - })); + add_worker_result( + false, + entry_name, + "No instances found for worker group", + namespace, + component, + None, + ); continue; } @@ -198,54 +198,60 @@ async fn clear_kv_blocks_handler( .iter() .map(|instance| instance.endpoint.clone()) .collect(); - failed_workers.push(json!({ - "name": entry.name, - "endpoint": format!("{}/{}/clear_kv_blocks", namespace, component), - "status": format!("Worker group doesn't support clear_kv_blocks. Supported endpoints: {}", found_endpoints.join(", ")), - })); + add_worker_result( + false, + entry_name, + &format!( + "Worker group doesn't support clear_kv_blocks. Supported endpoints: {}", + found_endpoints.join(", ") + ), + namespace, + component, + None, + ); continue; } for instance in &instances_filtered { + let instance_name = format!("{}-instance-{}", entry.name, instance.id()); match router.round_robin(().into()).await { - Ok(mut stream) => { - // Successfully sent request, now process the response - match stream.next().await { - Some(response) => { - // Instance successfully cleared its KV blocks - cleared_workers.push(json!({ - "name": format!("{}-instance-{}", entry.name, instance.id()), - "endpoint": format!("{}/{}/clear_kv_blocks", entry.endpoint.namespace, entry.endpoint.component), - "status": "successfully cleared kv blocks for instance", - "response": response.to_string() - })); - } - None => { - // No response from instance - failed_workers.push(json!({ - "name": format!("{}-instance-{}", entry.name, instance.id()), - "endpoint": format!("{}/{}/clear_kv_blocks", entry.endpoint.namespace, entry.endpoint.component), - "status": "no response from instance", - })); - } + Ok(mut stream) => match stream.next().await { + Some(response) => { + add_worker_result( + true, + instance_name, + "Successfully cleared kv blocks for instance", + namespace, + component, + Some(response.to_string()), + ); } - } + None => { + add_worker_result( + false, + instance_name, + "No response from instance", + namespace, + component, + None, + ); + } + }, Err(e) => { - // Failed to send request to this instance - failed_workers.push(json!({ - "name": format!("{}-instance-{}", entry.name, instance.id()), - "endpoint": format!("{}/{}/clear_kv_blocks", entry.endpoint.namespace, entry.endpoint.component), - "status": "failed to send request for instance", - "error": e.to_string() - })); + add_worker_result( + false, + instance_name, + "Failed to send request for instance", + namespace, + component, + Some(e.to_string()), + ); } } } } Json(serde_json::json!({ - "message": format!("Cleared prefix cache on {} out of {} worker groups", - cleared_workers.len(), model_entries.len()), "cleared_workers": cleared_workers, "failed_workers": failed_workers })) diff --git a/lib/llm/src/http/service/service_v2.rs b/lib/llm/src/http/service/service_v2.rs index 9014b15258fa..ac465ead7240 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -213,9 +213,4 @@ impl HttpServiceConfigBuilder { self.request_template = Some(request_template); self } - - pub fn with_runtime(mut self, runtime: Arc) -> Self { - self.runtime = Some(Some(runtime)); - self - } }