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 ff82f8575a6c..a50184615fb6 100644 --- a/launch/dynamo-run/src/subprocess/vllm_v1_inc.py +++ b/launch/dynamo-run/src/subprocess/vllm_v1_inc.py @@ -116,7 +116,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): @@ -124,6 +124,13 @@ 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): + try: + self.engine_client.reset_prefix_cache() + yield {"status": "success", "message": "KV cache cleared"} + except Exception as e: + yield {"status": "error", "message": str(e)} + async def generate(self, request): request_id = str(uuid.uuid4().hex) @@ -175,13 +182,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, @@ -249,16 +259,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 9c4081f6efec..b36c867d73b5 100644 --- a/lib/llm/src/http/service.rs +++ b/lib/llm/src/http/service.rs @@ -20,6 +20,7 @@ mod openai; +pub mod clear_kv_blocks; pub mod error; pub mod health; pub mod metrics; 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..48bcee70312e --- /dev/null +++ b/lib/llm/src/http/service/clear_kv_blocks.rs @@ -0,0 +1,258 @@ +// 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 dynamo_runtime::{pipeline::PushRouter, stream::StreamExt}; + +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 { + let model_entries = state.manager().get_model_entries(); + + // if there are no active workers + if model_entries.is_empty() { + return Json(serde_json::json!({ + "message": "No active worker groups found" + })); + } + + let distributed = match state.runtime() { + Some(runtime) => runtime, + None => { + return Json(serde_json::json!({ + "message": "Failed to create distributed runtime", + })); + } + }; + + 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) => { + add_worker_result( + false, + entry_name, + "Failed to get namespace", + namespace, + component, + Some(e.to_string()), + ); + continue; + } + }; + + let component_obj = match namespace_obj.component(component) { + Ok(comp) => comp, + Err(e) => { + add_worker_result( + false, + entry_name, + "Failed to get component", + namespace, + component, + Some(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) => { + add_worker_result( + false, + entry_name, + "Failed to get client", + namespace, + component, + Some(e.to_string()), + ); + continue; + } + }; + + let router = match PushRouter::<(), serde_json::Value>::from_client( + client.clone(), + Default::default(), + ) + .await + { + Ok(r) => r, + Err(e) => { + add_worker_result( + false, + entry_name, + "Failed to create router", + namespace, + component, + Some(e.to_string()), + ); + continue; + } + }; + + let instances = match component_obj.list_instances().await { + Ok(instances) => instances, + Err(e) => { + add_worker_result( + false, + entry_name, + "Failed to get instances for worker group", + namespace, + component, + Some(e.to_string()), + ); + continue; + } + }; + + if instances.is_empty() { + add_worker_result( + false, + entry_name, + "No instances found for worker group", + namespace, + component, + None, + ); + 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(); + 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) => 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) => { + add_worker_result( + false, + instance_name, + "Failed to send request for instance", + namespace, + component, + Some(e.to_string()), + ); + } + } + } + } + + Json(serde_json::json!({ + "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 e381fd04b7c2..ac465ead7240 100644 --- a/lib/llm/src/http/service/service_v2.rs +++ b/lib/llm/src/http/service/service_v2.rs @@ -11,6 +11,7 @@ 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; @@ -18,6 +19,7 @@ use tokio_util::sync::CancellationToken; 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,6 +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_kv_blocks::clear_kv_blocks_router(state.clone(), None), ]; if config.enable_chat_endpoints {