From 0f467fd61f165b07afc297fa4d2abdf3851d5d52 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:33:57 -0700 Subject: [PATCH] feat(grpc): add FlushCache and profile RPCs with worker-abstracted admin ops Add KV-cache flush and profiler start/stop support for gRPC-mode engines, dispatched per worker behind the Worker trait: - proto: FlushCache/StartProfile/StopProfile messages in common.proto (shared across engines, same pattern as GetTokenizer); RPCs on SglangScheduler - grpc_client: impl_admin_ops! macro shared by engine clients with local deadlines and trace injection; wired for sglang - Worker trait: flush_cache/start_profile/stop_profile default impls dispatch on connection mode (HTTP endpoint vs gRPC RPC), mirroring check_health_async - WorkerManager: flush_cache_all becomes one uniform fan-out, so gRPC workers are no longer silently skipped; start_profile_all and stop_profile_all support optional worker-URL targeting for PD mode - server: /start_profile and /stop_profile admin routes - grpc_servicer (sglang): FlushCache/StartProfile/StopProfile handlers, communicator plumbing, and the send_communicator_req + serve_grpc(on_request_manager_ready) surface that sglang's gRPC entrypoint requires since sgl-project/sglang#22500. The communicator response socket only binds tokenizer_ipc_name when recv_from_scheduler has not already bound it (--skip-tokenizer-init mode, #1591) Supersedes #1088 Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- crates/grpc_client/proto/common.proto | 36 +++ .../grpc_client/proto/sglang_scheduler.proto | 9 + crates/grpc_client/src/lib.rs | 94 ++++++ crates/grpc_client/src/sglang_scheduler.rs | 1 + crates/protocols/src/worker.rs | 92 ++++++ e2e_test/router/test_admin_ops.py | 143 +++++++++ .../sglang/request_manager.py | 125 +++++++- .../smg_grpc_servicer/sglang/server.py | 31 +- .../smg_grpc_servicer/sglang/servicer.py | 124 ++++++++ model_gateway/src/routers/grpc/client.rs | 45 ++- model_gateway/src/server.rs | 30 +- model_gateway/src/worker/error.rs | 7 + model_gateway/src/worker/manager.rs | 298 ++++++++++++++++-- model_gateway/src/worker/worker.rs | 191 ++++++++++- 14 files changed, 1187 insertions(+), 39 deletions(-) create mode 100644 e2e_test/router/test_admin_ops.py diff --git a/crates/grpc_client/proto/common.proto b/crates/grpc_client/proto/common.proto index cffed74a26..cadca35edb 100644 --- a/crates/grpc_client/proto/common.proto +++ b/crates/grpc_client/proto/common.proto @@ -56,3 +56,39 @@ message KvBlocksRemoved { } message KvCacheCleared {} + +// ===================== +// Admin Operations +// ===================== + +message FlushCacheRequest { + // Seconds to wait for the scheduler to go idle before flushing. + // 0 = flush immediately (fails if requests are in flight). + float timeout_s = 1; +} + +message FlushCacheResponse { + bool success = 1; + string message = 2; +} + +message StartProfileRequest { + // Directory where profiler traces are written (backend default if unset). + optional string output_dir = 1; + // Forward step at which profiling starts (immediately if unset). + optional int32 start_step = 2; + // Number of forward steps to profile (until StopProfile if unset). + optional int32 num_steps = 3; + // Profiler activities, e.g. "CPU", "GPU", "MEM". Backend default if empty. + repeated string activities = 4; + optional bool with_stack = 5; + optional bool record_shapes = 6; + bool profile_by_stage = 7; +} + +message StopProfileRequest {} + +message ProfileResponse { + bool success = 1; + string message = 2; +} diff --git a/crates/grpc_client/proto/sglang_scheduler.proto b/crates/grpc_client/proto/sglang_scheduler.proto index b93a56a5c9..04ed0262e9 100644 --- a/crates/grpc_client/proto/sglang_scheduler.proto +++ b/crates/grpc_client/proto/sglang_scheduler.proto @@ -30,6 +30,15 @@ service SglangScheduler { // Get comprehensive load metrics rpc GetLoads(GetLoadsRequest) returns (GetLoadsResponse); + // Flush the KV cache on all scheduler processes + rpc FlushCache(smg.grpc.common.FlushCacheRequest) returns (smg.grpc.common.FlushCacheResponse); + + // Start the profiler on all scheduler processes + rpc StartProfile(smg.grpc.common.StartProfileRequest) returns (smg.grpc.common.ProfileResponse); + + // Stop the profiler and export traces + rpc StopProfile(smg.grpc.common.StopProfileRequest) returns (smg.grpc.common.ProfileResponse); + // Get tokenizer artifacts for remote construction rpc GetTokenizer(smg.grpc.common.GetTokenizerRequest) returns (stream smg.grpc.common.GetTokenizerChunk); diff --git a/crates/grpc_client/src/lib.rs b/crates/grpc_client/src/lib.rs index a9f5474b32..8729aac9bf 100644 --- a/crates/grpc_client/src/lib.rs +++ b/crates/grpc_client/src/lib.rs @@ -57,6 +57,100 @@ macro_rules! impl_get_tokenizer { } pub(crate) use impl_get_tokenizer; +/// Extra local-deadline margin for `flush_cache` on top of the timeout +/// forwarded to the backend. The servicer bounds its own scheduler +/// round-trip at `max(30, timeout_s + 10)` seconds, so the margin must +/// cover that budget plus transport overhead. +pub const FLUSH_RPC_DEADLINE_MARGIN: std::time::Duration = std::time::Duration::from_secs(45); + +/// Local deadline for profile start/stop RPCs. Stopping a profile can take +/// a long time while the backend serializes large traces. +pub const PROFILE_RPC_DEADLINE: std::time::Duration = std::time::Duration::from_secs(630); + +/// Shared admin-op implementations (`flush_cache`, `start_profile`, +/// `stop_profile`) for engine clients whose protos expose the common +/// admin RPCs (request/response messages live in `common.proto`). +/// +/// Every call enforces a local deadline so an unresponsive backend cannot +/// hang the gateway, and injects trace context for distributed tracing. +macro_rules! impl_admin_ops { + () => { + /// Flush the KV cache on the backend scheduler. + /// + /// `timeout_s` is forwarded to the backend: 0 = flush immediately + /// (fails if requests are in flight), >0 = wait up to that many + /// seconds for the scheduler to go idle first. + pub async fn flush_cache( + &self, + timeout_s: f32, + ) -> Result<$crate::common_proto::FlushCacheResponse, tonic::Status> { + tracing::debug!("Requesting cache flush (timeout_s={timeout_s})"); + let mut request = + tonic::Request::new($crate::common_proto::FlushCacheRequest { timeout_s }); + if let Err(e) = self.trace_injector.inject(request.metadata_mut()) { + tracing::warn!("Failed to inject trace context: {}", e); + } + let deadline = std::time::Duration::from_secs_f32(timeout_s.max(0.0)) + + $crate::FLUSH_RPC_DEADLINE_MARGIN; + let mut client = self.client.clone(); + let response = tokio::time::timeout(deadline, client.flush_cache(request)) + .await + .map_err(|_| { + tonic::Status::deadline_exceeded(format!( + "FlushCache did not complete within {deadline:?}" + )) + })??; + Ok(response.into_inner()) + } + + /// Start the profiler on the backend scheduler. + pub async fn start_profile( + &self, + req: $crate::common_proto::StartProfileRequest, + ) -> Result<$crate::common_proto::ProfileResponse, tonic::Status> { + tracing::debug!("Requesting profile start"); + let mut request = tonic::Request::new(req); + if let Err(e) = self.trace_injector.inject(request.metadata_mut()) { + tracing::warn!("Failed to inject trace context: {}", e); + } + let mut client = self.client.clone(); + let response = + tokio::time::timeout($crate::PROFILE_RPC_DEADLINE, client.start_profile(request)) + .await + .map_err(|_| { + tonic::Status::deadline_exceeded(format!( + "StartProfile did not complete within {:?}", + $crate::PROFILE_RPC_DEADLINE + )) + })??; + Ok(response.into_inner()) + } + + /// Stop the profiler on the backend scheduler and export traces. + pub async fn stop_profile( + &self, + ) -> Result<$crate::common_proto::ProfileResponse, tonic::Status> { + tracing::debug!("Requesting profile stop"); + let mut request = tonic::Request::new($crate::common_proto::StopProfileRequest {}); + if let Err(e) = self.trace_injector.inject(request.metadata_mut()) { + tracing::warn!("Failed to inject trace context: {}", e); + } + let mut client = self.client.clone(); + let response = + tokio::time::timeout($crate::PROFILE_RPC_DEADLINE, client.stop_profile(request)) + .await + .map_err(|_| { + tonic::Status::deadline_exceeded(format!( + "StopProfile did not complete within {:?}", + $crate::PROFILE_RPC_DEADLINE + )) + })??; + Ok(response.into_inner()) + } + }; +} +pub(crate) use impl_admin_ops; + /// Shared `subscribe_kv_events()` implementation for all engine clients. /// /// Each engine's generated proto client has a `subscribe_kv_events` RPC method diff --git a/crates/grpc_client/src/sglang_scheduler.rs b/crates/grpc_client/src/sglang_scheduler.rs index 531208c5cb..0a49cebe04 100644 --- a/crates/grpc_client/src/sglang_scheduler.rs +++ b/crates/grpc_client/src/sglang_scheduler.rs @@ -141,6 +141,7 @@ impl SglangSchedulerClient { crate::impl_get_tokenizer!(); crate::impl_subscribe_kv_events!(); + crate::impl_admin_ops!(); /// Build a single SGLang EmbedRequest #[expect( diff --git a/crates/protocols/src/worker.rs b/crates/protocols/src/worker.rs index 2c69e278f3..3729f9f85c 100644 --- a/crates/protocols/src/worker.rs +++ b/crates/protocols/src/worker.rs @@ -995,9 +995,64 @@ pub struct FlushCacheResult { pub failed: Vec<(String, String)>, pub total_workers: usize, pub http_workers: usize, + #[serde(default)] + pub grpc_workers: usize, + pub message: String, +} + +/// Options for starting a profiling run on workers. +/// +/// Mirrors the engines' native profile parameters: serialized verbatim as +/// the JSON body for HTTP workers and mapped to the `StartProfile` RPC for +/// gRPC workers. Unset fields fall back to backend defaults. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct ProfileOptions { + #[serde(skip_serializing_if = "Option::is_none")] + pub output_dir: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub start_step: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub num_steps: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub activities: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub with_stack: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub record_shapes: Option, + pub profile_by_stage: bool, +} + +/// Result from profile start/stop operations across workers +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ProfileResult { + pub successful: Vec, + pub failed: Vec<(String, String)>, + pub total_workers: usize, pub message: String, } +/// Request body for the gateway `/start_profile` route: profile options +/// plus an optional worker URL to target a single worker (e.g. one +/// PD-disaggregation role). All workers are profiled when `url` is unset. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct StartProfileRequest { + #[serde(skip_serializing_if = "Option::is_none")] + pub url: Option, + #[serde(flatten)] + pub options: ProfileOptions, +} + +/// Request body for the gateway `/stop_profile` route: optional worker URL +/// to target a single worker. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] +pub struct StopProfileRequest { + #[serde(skip_serializing_if = "Option::is_none")] + pub url: Option, +} + /// Result from getting worker loads #[derive(Debug, Clone, Deserialize, Serialize)] pub struct WorkerLoadsResult { @@ -1101,6 +1156,43 @@ impl IntoResponse for FlushCacheResult { "message": self.message, "workers_flushed": self.successful.len(), "total_http_workers": self.http_workers, + "total_grpc_workers": self.grpc_workers, + "total_workers": self.total_workers + }); + + if !self.failed.is_empty() { + body["successful"] = json!(self.successful); + body["failed"] = json!(self + .failed + .into_iter() + .map(|(url, err)| json!({"worker": url, "error": err})) + .collect::>()); + } + + (status, Json(body)).into_response() + } +} + +#[cfg(feature = "axum")] +impl IntoResponse for ProfileResult { + fn into_response(self) -> Response { + let status = if self.total_workers == 0 { + StatusCode::NOT_FOUND + } else if self.failed.is_empty() { + StatusCode::OK + } else { + StatusCode::PARTIAL_CONTENT + }; + + let status_str = match status { + StatusCode::OK => "success", + StatusCode::PARTIAL_CONTENT => "partial_success", + _ => "error", + }; + let mut body = json!({ + "status": status_str, + "message": self.message, + "workers_profiled": self.successful.len(), "total_workers": self.total_workers }); diff --git a/e2e_test/router/test_admin_ops.py b/e2e_test/router/test_admin_ops.py new file mode 100644 index 0000000000..93f271bd80 --- /dev/null +++ b/e2e_test/router/test_admin_ops.py @@ -0,0 +1,143 @@ +"""Tests for gateway admin operations (cache flush, profiling). + +Covers the worker-abstracted admin fan-out behind /flush_cache, +/start_profile, and /stop_profile: the gateway dispatches per worker on +connection mode (HTTP endpoint vs gRPC RPC), so both modes are exercised +with the same assertions, including that profiling actually exports trace +artifacts to the requested output directory. + +Usage: + pytest e2e_test/router/test_admin_ops.py -v +""" + +from __future__ import annotations + +import logging +import time +from pathlib import Path + +import httpx +import pytest + +logger = logging.getLogger(__name__) + +FLUSH_TIMEOUT = 60.0 +PROFILE_START_TIMEOUT = 120.0 +PROFILE_STOP_TIMEOUT = 300.0 +TRACE_EXPORT_WAIT = 60.0 + + +def _wait_for_trace_artifacts(output_dir: Path) -> list[Path]: + """Poll for profiler artifacts — trace export may lag the RPC reply.""" + deadline = time.time() + TRACE_EXPORT_WAIT + while time.time() < deadline: + if output_dir.is_dir(): + files = [f for f in output_dir.rglob("*") if f.is_file() and f.stat().st_size > 0] + if files: + return files + time.sleep(1.0) + return [] + + +class AdminOpsBehavior: + """Shared admin-endpoint assertions, parametrized per engine below.""" + + def test_flush_cache_full_cycle(self, setup_backend): + backend, model, client, gateway = setup_backend + + # Populate the prefix cache so there is something to flush. + client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "Say hi"}], + max_tokens=8, + ) + + resp = httpx.post(f"{gateway.base_url}/flush_cache", timeout=FLUSH_TIMEOUT) + assert resp.status_code == 200, resp.text + body = resp.json() + logger.info("flush_cache response: %s", body) + assert body["status"] == "success" + assert body["workers_flushed"] == 1 + assert body["total_workers"] == 1 + if backend == "grpc": + assert body["total_grpc_workers"] == 1 + else: + assert body["total_http_workers"] == 1 + + # Generation still works against the flushed cache. + completion = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "Say hi again"}], + max_tokens=8, + ) + assert completion.choices, "generation after cache flush should succeed" + + def test_start_and_stop_profile_produces_traces(self, setup_backend, tmp_path): + _backend, model, client, gateway = setup_backend + + output_dir = tmp_path / "traces" + output_dir.mkdir() + resp = httpx.post( + f"{gateway.base_url}/start_profile", + json={"output_dir": str(output_dir)}, + timeout=PROFILE_START_TIMEOUT, + ) + assert resp.status_code == 200, resp.text + body = resp.json() + logger.info("start_profile response: %s", body) + assert body["status"] == "success" + assert body["workers_profiled"] == 1 + + # Give the profiler some forward steps to record. + client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "Say hi"}], + max_tokens=8, + ) + + resp = httpx.post(f"{gateway.base_url}/stop_profile", timeout=PROFILE_STOP_TIMEOUT) + assert resp.status_code == 200, resp.text + body = resp.json() + logger.info("stop_profile response: %s", body) + assert body["status"] == "success" + assert body["workers_profiled"] == 1 + + artifacts = _wait_for_trace_artifacts(output_dir) + assert artifacts, ( + f"profiler reported success but produced no artifacts in {output_dir} " + f"within {TRACE_EXPORT_WAIT}s" + ) + logger.info( + "profiler produced %d artifact(s): %s", + len(artifacts), + [f.name for f in artifacts[:5]], + ) + + def test_profile_url_filter_without_match_returns_404(self, setup_backend): + _backend, _model, _client, gateway = setup_backend + + resp = httpx.post( + f"{gateway.base_url}/start_profile", + json={"url": "http://nonexistent:9999"}, + timeout=30.0, + ) + assert resp.status_code == 404, resp.text + + def test_start_profile_rejects_malformed_json(self, setup_backend): + _backend, _model, _client, gateway = setup_backend + + resp = httpx.post( + f"{gateway.base_url}/start_profile", + content=b"{not json", + headers={"content-type": "application/json"}, + timeout=30.0, + ) + assert resp.status_code == 400, resp.text + + +@pytest.mark.engine("sglang") +@pytest.mark.gpu(1) +@pytest.mark.e2e +@pytest.mark.parametrize("setup_backend", ["grpc", "http"], indirect=True) +class TestAdminOps(AdminOpsBehavior): + """SGLang serves both connection modes behind the gateway.""" diff --git a/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py b/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py index 65dc26641c..235430b313 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/request_manager.py @@ -23,9 +23,12 @@ AbortReq, BatchEmbeddingOutput, BatchTokenIDOutput, + FlushCacheReqOutput, + GetInternalStateReqOutput, GetLoadsReqInput, GetLoadsReqOutput, HealthCheckOutput, + ProfileReqOutput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, WatchLoadUpdateReq, @@ -194,6 +197,19 @@ def __init__( self.context, zmq.PULL, self.recv_ipc_name, bind=True ) + # Communicator responses (ProfileReqOutput, FlushCacheReqOutput, ...) + # always arrive on tokenizer_ipc_name via the scheduler's + # send_to_tokenizer path. In skip_tokenizer_init mode + # recv_from_scheduler already binds that endpoint, so generation + # outputs and communicator responses share one socket; otherwise a + # second socket is needed for communicator responses. + if server_args.skip_tokenizer_init: + self.recv_from_tokenizer = None + else: + self.recv_from_tokenizer = get_zmq_socket( + self.context, zmq.PULL, port_args.tokenizer_ipc_name, bind=True + ) + # Socket for sending requests to scheduler self.send_to_scheduler = get_zmq_socket( self.context, zmq.PUSH, port_args.scheduler_input_ipc_name, bind=True @@ -228,6 +244,15 @@ def __init__( self.get_loads_communicator = _GrpcCommunicator( self.send_to_scheduler, fan_out=server_args.dp_size ) + self.profile_communicator = _GrpcCommunicator( + self.send_to_scheduler, fan_out=server_args.dp_size + ) + self.flush_cache_communicator = _GrpcCommunicator( + self.send_to_scheduler, fan_out=server_args.dp_size + ) + self.get_internal_state_communicator = _GrpcCommunicator( + self.send_to_scheduler, fan_out=server_args.dp_size + ) logger.info( f"GrpcRequestManager initialized with ZMQ IPC: " @@ -508,9 +533,10 @@ async def handle_loop(self): await self._handle_health_check_output(recv_obj) elif isinstance(recv_obj, AbortReq): await self._handle_abort_req(recv_obj) - elif isinstance(recv_obj, GetLoadsReqOutput): - # Route to communicator for request/response pattern - self.get_loads_communicator.handle_recv(recv_obj) + elif self._dispatch_communicator_output(recv_obj): + # Communicator responses arrive here in skip_tokenizer_init + # mode, where this socket is bound to tokenizer_ipc_name. + pass else: logger.warning(f"Unknown output type: {type(recv_obj)}") @@ -531,6 +557,52 @@ async def handle_loop(self): if self.gracefully_exit: break + def _dispatch_communicator_output(self, recv_obj) -> bool: + """Route a communicator response from the scheduler to its communicator. + + Returns True when the object was a recognized communicator response. + """ + if isinstance(recv_obj, ProfileReqOutput): + self.profile_communicator.handle_recv(recv_obj) + elif isinstance(recv_obj, FlushCacheReqOutput): + self.flush_cache_communicator.handle_recv(recv_obj) + elif isinstance(recv_obj, GetInternalStateReqOutput): + self.get_internal_state_communicator.handle_recv(recv_obj) + elif isinstance(recv_obj, GetLoadsReqOutput): + self.get_loads_communicator.handle_recv(recv_obj) + else: + return False + return True + + async def _handle_tokenizer_loop(self): + """Process communicator responses from the scheduler's send_to_tokenizer path. + + The scheduler sends communicator responses (ProfileReqOutput, + FlushCacheReqOutput, ...) directly to tokenizer_ipc_name, not through + the detokenizer. This loop only runs when recv_from_tokenizer is a + dedicated socket — in skip_tokenizer_init mode those responses arrive + on recv_from_scheduler and are dispatched by handle_loop instead. + """ + while not self.gracefully_exit: + try: + recv_obj = await self.recv_from_tokenizer.recv_pyobj() + if not self._dispatch_communicator_output(recv_obj): + logger.warning(f"Unknown type on tokenizer socket: {type(recv_obj)}") + except zmq.error.Again: + if self.gracefully_exit: + break + continue + except zmq.error.ZMQError as e: + if self.gracefully_exit: + logger.debug(f"ZMQ recv interrupted during shutdown: {e}") + break + logger.error(f"ZMQ error in tokenizer loop: {e}\n{get_exception_traceback()}") + break + except Exception as e: + logger.error(f"Tokenizer loop error: {e}\n{get_exception_traceback()}") + if self.gracefully_exit: + break + def _convert_logprob_style( self, state: GrpcReqState, @@ -874,6 +946,8 @@ async def shutdown(self): # Close ZMQ sockets self.recv_from_scheduler.close() + if self.recv_from_tokenizer is not None: + self.recv_from_tokenizer.close() self.send_to_scheduler.close() # Terminate the ZMQ context - this is critical for asyncio loop to exit cleanly @@ -914,6 +988,47 @@ async def get_loads( return results + # Communicators that callers may address through send_communicator_req. + _PUBLIC_COMMUNICATORS = frozenset( + { + "profile_communicator", + "flush_cache_communicator", + "get_internal_state_communicator", + "get_loads_communicator", + } + ) + + async def send_communicator_req(self, req, communicator_name: str, timeout: float = 30.0): + """Send a request to the scheduler via a named communicator and return responses. + + This is the generic transport method for request/response patterns that + go through the scheduler's send_to_tokenizer path (profile, flush_cache, + get_internal_state, ...). Business logic (request construction, response + interpretation) belongs in the caller — sglang's HTTP sidecar relies on + this exact signature, so treat it as a public contract. + + Args: + req: The request object to send to the scheduler. + communicator_name: Attribute name of the communicator + (e.g. "profile_communicator"). + timeout: Timeout in seconds for the scheduler round-trip. + + Returns: + List of response objects from the scheduler(s), one per DP rank. + + Raises: + ValueError: If communicator_name is not a known communicator. + TimeoutError: If the scheduler does not respond within timeout. + """ + if communicator_name not in self._PUBLIC_COMMUNICATORS: + raise ValueError( + f"Unknown communicator '{communicator_name}'. " + f"Allowed: {sorted(self._PUBLIC_COMMUNICATORS)}" + ) + self.auto_create_handle_loop() + communicator = getattr(self, communicator_name) + return await communicator(req, timeout=timeout) + def auto_create_handle_loop(self): """Automatically create and start the handle_loop task, matching TokenizerManager pattern.""" if self.no_create_loop: @@ -922,6 +1037,10 @@ def auto_create_handle_loop(self): self.no_create_loop = True loop = get_or_create_event_loop() self.asyncio_tasks.add(loop.create_task(print_exception_wrapper(self.handle_loop))) + if self.recv_from_tokenizer is not None: + self.asyncio_tasks.add( + loop.create_task(print_exception_wrapper(self._handle_tokenizer_loop)) + ) self.event_loop = loop diff --git a/grpc_servicer/smg_grpc_servicer/sglang/server.py b/grpc_servicer/smg_grpc_servicer/sglang/server.py index b9309397c1..943ed06a61 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/server.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/server.py @@ -12,6 +12,7 @@ import signal import threading import time +from collections.abc import Callable from concurrent import futures import grpc @@ -36,8 +37,20 @@ async def serve_grpc( server_args: ServerArgs, model_info: dict | None = None, + on_request_manager_ready: Callable | None = None, ): - """Start the standalone gRPC server with integrated scheduler.""" + """Start the standalone gRPC server with integrated scheduler. + + Args: + server_args: Server configuration. + model_info: Optional model metadata override. + on_request_manager_ready: Optional async callback invoked with + ``(request_manager, server_args, scheduler_info)`` once the + ``GrpcRequestManager`` has been created but before the gRPC + server starts accepting requests. sglang's HTTP sidecar uses + this to wire its admin endpoints to the scheduler, so the + callback signature is a public contract. + """ # Start bootstrap server BEFORE launching scheduler processes (only in PREFILL mode) # This ensures the bootstrap server is ready when prefill schedulers try to register @@ -104,6 +117,22 @@ async def serve_grpc( bootstrap_server=bootstrap_server, ) + if on_request_manager_ready is not None: + try: + await on_request_manager_ready(request_manager, server_args, scheduler_info) + except Exception: + # The shutdown guard below only covers the post-start path, so + # clean up the scheduler processes and sockets before re-raising. + logger.exception("on_request_manager_ready callback failed, shutting down") + await request_manager.shutdown() + for proc in scheduler_procs: + if proc.is_alive(): + proc.terminate() + proc.join(timeout=2.0) + if proc.is_alive(): + proc.kill() + raise + # Create gRPC server server = grpc.aio.server( futures.ThreadPoolExecutor(max_workers=10), diff --git a/grpc_servicer/smg_grpc_servicer/sglang/servicer.py b/grpc_servicer/smg_grpc_servicer/sglang/servicer.py index 51b19829d5..20de8cb791 100644 --- a/grpc_servicer/smg_grpc_servicer/sglang/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/sglang/servicer.py @@ -37,13 +37,17 @@ ) from sglang.srt.disaggregation.utils import FAKE_BOOTSTRAP_HOST, DisaggregationMode from sglang.srt.managers.io_struct import ( + FlushCacheReqInput, GetLoadsReqOutput, + ProfileReq, + ProfileReqType, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, ) from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem, MultimodalInputs from sglang.srt.sampling.sampling_params import SamplingParams as SGLSamplingParams from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import get_bool_env_var from sglang.utils import get_exception_traceback from smg_grpc_proto import sglang_scheduler_pb2, sglang_scheduler_pb2_grpc from smg_grpc_proto.generated import common_pb2 @@ -55,6 +59,20 @@ logger = logging.getLogger(__name__) HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20)) +# Profile round-trips include trace serialization, which can take minutes. +PROFILE_COMM_TIMEOUT = 600.0 + + +def _aggregate_communicator_results(results: list, ok_message: str) -> tuple[bool, str]: + """Aggregate per-DP-rank communicator outputs into (success, message).""" + if not results: + return False, "No response from scheduler" + failures = [r for r in results if not r.success] + if failures: + return False, " | ".join(r.message or "failed" for r in failures) + return True, ok_message + + SAMPLING_DEFAULT_KEYS = ( "temperature", "top_p", @@ -602,6 +620,112 @@ async def GetLoads( aggregate=_compute_aggregate_protobuf(loads), ) + async def FlushCache( + self, + request: common_pb2.FlushCacheRequest, + context: grpc.aio.ServicerContext, + ) -> common_pb2.FlushCacheResponse: + """Flush the KV cache on all scheduler processes.""" + logger.debug("Receive flush cache request (timeout_s=%.1f)", request.timeout_s) + timeout_s = request.timeout_s + comm_timeout = max(30.0, timeout_s + 10.0) + try: + results = await self.request_manager.send_communicator_req( + FlushCacheReqInput(timeout_s=timeout_s), + "flush_cache_communicator", + timeout=comm_timeout, + ) + except TimeoutError: + context.set_code(grpc.StatusCode.DEADLINE_EXCEEDED) + context.set_details( + f"Flush cache timed out after {comm_timeout}s. " + "The scheduler may be unresponsive or under heavy load." + ) + return common_pb2.FlushCacheResponse( + success=False, message=f"Flush cache timed out after {comm_timeout}s" + ) + except Exception as e: + logger.error(f"FlushCache failed: {e}\n{get_exception_traceback()}") + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details(f"Flush cache failed: {e}") + return common_pb2.FlushCacheResponse(success=False, message=f"Flush cache failed: {e}") + + success, message = _aggregate_communicator_results(results, "Cache flushed successfully") + if not results: + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details(message) + return common_pb2.FlushCacheResponse(success=success, message=message) + + async def StartProfile( + self, + request: common_pb2.StartProfileRequest, + context: grpc.aio.ServicerContext, + ) -> common_pb2.ProfileResponse: + """Start the profiler on all scheduler processes.""" + logger.debug("Receive start profile request") + + # Env-var overrides mirror sglang's TokenizerCommunicatorMixin defaults. + with_stack = request.with_stack if request.HasField("with_stack") else None + with_stack = (with_stack is not False) and get_bool_env_var( + "SGLANG_PROFILE_WITH_STACK", "true" + ) + record_shapes = request.record_shapes if request.HasField("record_shapes") else None + record_shapes = (record_shapes is not False) and get_bool_env_var( + "SGLANG_PROFILE_RECORD_SHAPES", "true" + ) + + req = ProfileReq( + type=ProfileReqType.START_PROFILE, + output_dir=request.output_dir if request.HasField("output_dir") else None, + start_step=request.start_step if request.HasField("start_step") else None, + num_steps=request.num_steps if request.HasField("num_steps") else None, + activities=list(request.activities) if request.activities else None, + with_stack=with_stack, + record_shapes=record_shapes, + profile_by_stage=request.profile_by_stage, + profile_id=str(time.time()), + ) + return await self._run_profile_req(req, "Start profiling", context) + + async def StopProfile( + self, + request: common_pb2.StopProfileRequest, + context: grpc.aio.ServicerContext, + ) -> common_pb2.ProfileResponse: + """Stop the profiler and export traces.""" + logger.debug("Receive stop profile request") + req = ProfileReq(type=ProfileReqType.STOP_PROFILE) + return await self._run_profile_req(req, "Stop profiling", context) + + async def _run_profile_req( + self, + req: ProfileReq, + op_name: str, + context: grpc.aio.ServicerContext, + ) -> common_pb2.ProfileResponse: + """Send a ProfileReq to the scheduler and aggregate per-rank results.""" + try: + results = await self.request_manager.send_communicator_req( + req, "profile_communicator", timeout=PROFILE_COMM_TIMEOUT + ) + except TimeoutError: + context.set_code(grpc.StatusCode.DEADLINE_EXCEEDED) + context.set_details(f"{op_name} timed out after {PROFILE_COMM_TIMEOUT}s") + return common_pb2.ProfileResponse( + success=False, message=f"{op_name} timed out after {PROFILE_COMM_TIMEOUT}s" + ) + except Exception as e: + logger.error(f"{op_name} failed: {e}\n{get_exception_traceback()}") + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details(f"{op_name} failed: {e}") + return common_pb2.ProfileResponse(success=False, message=f"{op_name} failed: {e}") + + success, message = _aggregate_communicator_results(results, f"{op_name} succeeded") + if not results: + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details(message) + return common_pb2.ProfileResponse(success=success, message=message) + async def GetTokenizer( self, request: common_pb2.GetTokenizerRequest, diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index b254e3f7ed..06d6af277a 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -7,8 +7,8 @@ use openai_protocol::{ messages::CreateMessageRequest, worker::WorkerLoadResponse, }; use smg_grpc_client::{ - tokenizer_bundle, tokenizer_bundle::StreamBundle, MlxEngineClient, SglangSchedulerClient, - TokenSpeedSchedulerClient, TrtllmServiceClient, VllmEngineClient, + common_proto, tokenizer_bundle, tokenizer_bundle::StreamBundle, MlxEngineClient, + SglangSchedulerClient, TokenSpeedSchedulerClient, TrtllmServiceClient, VllmEngineClient, }; use crate::routers::grpc::{ @@ -257,12 +257,51 @@ impl GrpcClient { } } + /// Flush the KV cache on the backend. Returns `Unimplemented` for + /// backends without the FlushCache RPC. + pub async fn flush_cache( + &self, + timeout_s: f32, + ) -> Result { + match self { + Self::Sglang(client) => client.flush_cache(timeout_s).await, + Self::Vllm(_) | Self::Trtllm(_) | Self::Mlx(_) | Self::TokenSpeed(_) => Err( + tonic::Status::unimplemented("FlushCache RPC not supported for this backend"), + ), + } + } + + /// Start the profiler on the backend. Returns `Unimplemented` for + /// backends without the StartProfile RPC. + pub async fn start_profile( + &self, + req: common_proto::StartProfileRequest, + ) -> Result { + match self { + Self::Sglang(client) => client.start_profile(req).await, + Self::Vllm(_) | Self::Trtllm(_) | Self::Mlx(_) | Self::TokenSpeed(_) => Err( + tonic::Status::unimplemented("StartProfile RPC not supported for this backend"), + ), + } + } + + /// Stop the profiler on the backend. Returns `Unimplemented` for + /// backends without the StopProfile RPC. + pub async fn stop_profile(&self) -> Result { + match self { + Self::Sglang(client) => client.stop_profile().await, + Self::Vllm(_) | Self::Trtllm(_) | Self::Mlx(_) | Self::TokenSpeed(_) => Err( + tonic::Status::unimplemented("StopProfile RPC not supported for this backend"), + ), + } + } + /// Subscribe to KV cache events. Returns `Unimplemented` on backends /// without KV-event streaming. pub async fn subscribe_kv_events( &self, start_seq: u64, - ) -> Result, tonic::Status> { + ) -> Result, tonic::Status> { match self { Self::Sglang(client) => client.subscribe_kv_events(start_seq).await, Self::Vllm(client) => client.subscribe_kv_events(start_seq).await, diff --git a/model_gateway/src/server.rs b/model_gateway/src/server.rs index a8466f376a..e61dc6aad0 100644 --- a/model_gateway/src/server.rs +++ b/model_gateway/src/server.rs @@ -32,7 +32,7 @@ use openai_protocol::{ responses::ResponsesRequest, tokenize::{AddTokenizerRequest, DetokenizeRequest, TokenizeRequest}, validated::ValidatedJson, - worker::{WorkerSpec, WorkerUpdateRequest}, + worker::{StartProfileRequest, StopProfileRequest, WorkerSpec, WorkerUpdateRequest}, }; use rustls::crypto::ring; use serde::Deserialize; @@ -575,7 +575,31 @@ async fn v1_realtime_transcription_session( } async fn flush_cache(State(state): State>, _req: Request) -> Response { - WorkerManager::flush_cache_all(&state.context.worker_registry, &state.context.client) + WorkerManager::flush_cache_all(&state.context.worker_registry) + .await + .into_response() +} + +async fn start_profile( + State(state): State>, + body: Option>, +) -> Response { + let body = body.map_or_else(StartProfileRequest::default, |Json(body)| body); + WorkerManager::start_profile_all( + &state.context.worker_registry, + &body.options, + body.url.as_deref(), + ) + .await + .into_response() +} + +async fn stop_profile( + State(state): State>, + body: Option>, +) -> Response { + let body = body.map_or_else(StopProfileRequest::default, |Json(body)| body); + WorkerManager::stop_profile_all(&state.context.worker_registry, body.url.as_deref()) .await .into_response() } @@ -907,6 +931,8 @@ pub fn build_app( // Build admin routes with control plane auth if configured, otherwise use simple API key auth let admin_routes = Router::new() .route("/flush_cache", post(flush_cache)) + .route("/start_profile", post(start_profile)) + .route("/stop_profile", post(stop_profile)) .route("/get_loads", get(get_loads)) .route("/parse/function_call", post(parse_function_call)) .route("/parse/reasoning", post(parse_reasoning)) diff --git a/model_gateway/src/worker/error.rs b/model_gateway/src/worker/error.rs index b17360a18c..81397479ca 100644 --- a/model_gateway/src/worker/error.rs +++ b/model_gateway/src/worker/error.rs @@ -19,6 +19,13 @@ pub enum WorkerError { #[error("Connection failed for worker {url}: {reason}")] ConnectionFailed { url: String, reason: String }, + + #[error("{operation} failed for worker {url}: {reason}")] + OperationFailed { + url: String, + operation: String, + reason: String, + }, } /// Result type for worker operations diff --git a/model_gateway/src/worker/manager.rs b/model_gateway/src/worker/manager.rs index 138e9e266e..dfaac95324 100644 --- a/model_gateway/src/worker/manager.rs +++ b/model_gateway/src/worker/manager.rs @@ -17,7 +17,8 @@ use futures::{ }; use http::StatusCode; use openai_protocol::worker::{ - FlushCacheResult, HealthCheckConfig, WorkerLoadInfo, WorkerLoadsResult, WorkerStatus, + FlushCacheResult, HealthCheckConfig, ProfileOptions, ProfileResult, WorkerLoadInfo, + WorkerLoadsResult, WorkerStatus, }; use tokio::{ sync::{broadcast, Notify}, @@ -33,7 +34,7 @@ use crate::{ monitor::WorkerMonitor, registry::{WorkerDescriptor, WorkerId}, worker::WorkerTypeExt, - ConnectionMode, Worker, WorkerRegistry, WorkerType, + ConnectionMode, Worker, WorkerRegistry, WorkerResult, WorkerType, }, workflow::{Job, JobQueue}, }; @@ -218,7 +219,7 @@ struct ProbeCompletion { expected_revision: u64, launched_status: WorkerStatus, health_config: HealthCheckConfig, - probe_result: crate::worker::WorkerResult<()>, + probe_result: WorkerResult<()>, } struct RemovalCandidate { @@ -691,55 +692,149 @@ impl WorkerManager { .collect() } - pub async fn flush_cache_all( + /// Fan an admin operation out to workers in parallel, collecting + /// successful worker URLs and per-worker failure messages. + /// + /// The connection-mode dispatch (HTTP endpoint vs gRPC RPC) lives in + /// the [`Worker`] admin methods, so the fan-out is uniform. + async fn admin_fan_out( + workers: Vec>, + op: F, + ) -> (Vec, Vec<(String, String)>) + where + F: Fn(Arc) -> Fut, + Fut: Future>, + { + let futures: Vec<_> = workers + .into_iter() + .map(|worker| { + let url = worker.url().to_string(); + let fut = op(worker); + async move { (url, fut.await) } + }) + .collect(); + + let results: Vec<(String, WorkerResult<()>)> = stream::iter(futures) + .buffer_unordered(MAX_CONCURRENT) + .collect() + .await; + + let mut successful = Vec::new(); + let mut failed = Vec::new(); + for (url, result) in results { + match result { + Ok(()) => successful.push(url), + Err(e) => failed.push((url, e.to_string())), + } + } + (successful, failed) + } + + /// Workers targeted by an admin op: all registered workers, or only + /// the worker(s) whose URL matches the filter. + fn admin_target_workers( worker_registry: &WorkerRegistry, - client: &reqwest::Client, - ) -> FlushCacheResult { + url_filter: Option<&str>, + ) -> Vec> { let workers = worker_registry.get_all(); - let total_workers = workers.len(); + match url_filter { + Some(url) => workers.into_iter().filter(|w| w.url() == url).collect(), + None => workers, + } + } - let http_workers: Vec<_> = workers - .into_iter() + pub async fn flush_cache_all(worker_registry: &WorkerRegistry) -> FlushCacheResult { + let workers = worker_registry.get_all(); + let total_workers = workers.len(); + let http_workers = workers + .iter() .filter(|w| matches!(w.connection_mode(), ConnectionMode::Http)) - .collect(); + .count(); + let grpc_workers = total_workers - http_workers; - if http_workers.is_empty() { + if workers.is_empty() { return FlushCacheResult { successful: vec![], failed: vec![], total_workers, - http_workers: 0, - message: "No HTTP workers available for cache flush".to_string(), + http_workers, + grpc_workers, + message: "No workers available for cache flush".to_string(), }; } info!( - "Flushing cache on {} HTTP workers (out of {} total)", - http_workers.len(), - total_workers + "Flushing cache on {} workers ({} HTTP, {} gRPC)", + total_workers, http_workers, grpc_workers ); - let responses = fan_out(&http_workers, client, "flush_cache", reqwest::Method::POST).await; + let (successful, failed) = + Self::admin_fan_out(workers, |w| async move { w.flush_cache().await }).await; - let mut successful = Vec::new(); - let mut failed = Vec::new(); + let message = if failed.is_empty() { + format!( + "Successfully flushed cache on all {} workers", + successful.len() + ) + } else { + format!( + "Cache flush: {} succeeded, {} failed", + successful.len(), + failed.len() + ) + }; - for resp in responses { - match resp.result { - Ok(r) if r.status().is_success() => successful.push(resp.url), - Ok(r) => failed.push((resp.url, format!("HTTP {}", r.status()))), - Err(e) => failed.push((resp.url, e.to_string())), - } + info!("{}", message); + + FlushCacheResult { + successful, + failed, + total_workers, + http_workers, + grpc_workers, + message, + } + } + + /// Start a profiling run on all workers, or on the single worker + /// matching `worker_url`. + pub async fn start_profile_all( + worker_registry: &WorkerRegistry, + options: &ProfileOptions, + worker_url: Option<&str>, + ) -> ProfileResult { + let workers = Self::admin_target_workers(worker_registry, worker_url); + let total_workers = workers.len(); + + if workers.is_empty() { + return ProfileResult { + successful: vec![], + failed: vec![], + total_workers, + message: match worker_url { + Some(url) => format!("No worker matching url '{url}'"), + None => "No workers available for profiling".to_string(), + }, + }; } + info!("Starting profiling on {} workers", total_workers); + + let options = options.clone(); + let (successful, failed) = Self::admin_fan_out(workers, move |w| { + let options = options.clone(); + async move { w.start_profile(&options).await } + }) + .await; + let message = if failed.is_empty() { format!( - "Successfully flushed cache on all {} HTTP workers", + "Successfully started profiling on all {} workers", successful.len() ) } else { format!( - "Cache flush: {} succeeded, {} failed", + "Profile start: {} succeeded, {} failed", successful.len(), failed.len() ) @@ -747,11 +842,59 @@ impl WorkerManager { info!("{}", message); - FlushCacheResult { + ProfileResult { + successful, + failed, + total_workers, + message, + } + } + + /// Stop the in-flight profiling run on all workers, or on the single + /// worker matching `worker_url`. + pub async fn stop_profile_all( + worker_registry: &WorkerRegistry, + worker_url: Option<&str>, + ) -> ProfileResult { + let workers = Self::admin_target_workers(worker_registry, worker_url); + let total_workers = workers.len(); + + if workers.is_empty() { + return ProfileResult { + successful: vec![], + failed: vec![], + total_workers, + message: match worker_url { + Some(url) => format!("No worker matching url '{url}'"), + None => "No workers available for profiling".to_string(), + }, + }; + } + + info!("Stopping profiling on {} workers", total_workers); + + let (successful, failed) = + Self::admin_fan_out(workers, |w| async move { w.stop_profile().await }).await; + + let message = if failed.is_empty() { + format!( + "Successfully stopped profiling on all {} workers", + successful.len() + ) + } else { + format!( + "Profile stop: {} succeeded, {} failed", + successful.len(), + failed.len() + ) + }; + + info!("{}", message); + + ProfileResult { successful, failed, total_workers, - http_workers: http_workers.len(), message, } } @@ -1141,4 +1284,101 @@ mod tests { WorkerStatus::Ready ); } + + /// Spawn a loopback HTTP server stubbing the engine admin endpoints, + /// answering every admin POST with the given status. + async fn spawn_admin_stub(status: StatusCode) -> String { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let app = axum::Router::new() + .route( + "/flush_cache", + axum::routing::post(move || async move { status }), + ) + .route( + "/start_profile", + axum::routing::post(move || async move { status }), + ) + .route( + "/stop_profile", + axum::routing::post(move || async move { status }), + ); + #[expect( + clippy::disallowed_methods, + reason = "test stub server lives for the duration of the test process" + )] + tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + format!("http://{addr}") + } + + #[tokio::test] + async fn test_flush_cache_all_aggregates_mixed_results() { + let ok_url = spawn_admin_stub(StatusCode::OK).await; + let fail_url = spawn_admin_stub(StatusCode::INTERNAL_SERVER_ERROR).await; + + let registry = WorkerRegistry::new(); + registry.register(make_worker(&ok_url, 1, 1)).unwrap(); + registry.register(make_worker(&fail_url, 1, 1)).unwrap(); + + let result = WorkerManager::flush_cache_all(®istry).await; + + assert_eq!(result.total_workers, 2); + assert_eq!(result.http_workers, 2); + assert_eq!(result.grpc_workers, 0); + assert_eq!(result.successful, vec![ok_url]); + assert_eq!(result.failed.len(), 1); + assert_eq!(result.failed[0].0, fail_url); + assert!( + result.failed[0].1.contains("500"), + "failure reason should carry the HTTP status: {}", + result.failed[0].1 + ); + } + + #[tokio::test] + async fn test_flush_cache_all_no_workers() { + let registry = WorkerRegistry::new(); + let result = WorkerManager::flush_cache_all(®istry).await; + assert_eq!(result.total_workers, 0); + assert!(result.successful.is_empty()); + assert!(result.failed.is_empty()); + assert_eq!(result.message, "No workers available for cache flush"); + } + + #[tokio::test] + async fn test_profile_all_targets_single_worker_via_url_filter() { + let url_a = spawn_admin_stub(StatusCode::OK).await; + let url_b = spawn_admin_stub(StatusCode::OK).await; + + let registry = WorkerRegistry::new(); + registry.register(make_worker(&url_a, 1, 1)).unwrap(); + registry.register(make_worker(&url_b, 1, 1)).unwrap(); + + let options = ProfileOptions::default(); + let result = + WorkerManager::start_profile_all(®istry, &options, Some(url_a.as_str())).await; + assert_eq!(result.total_workers, 1); + assert_eq!(result.successful, vec![url_a.clone()]); + assert!(result.failed.is_empty()); + + let result = WorkerManager::stop_profile_all(®istry, Some(url_a.as_str())).await; + assert_eq!(result.total_workers, 1); + assert_eq!(result.successful, vec![url_a]); + } + + #[tokio::test] + async fn test_profile_all_url_filter_without_match() { + let registry = WorkerRegistry::new(); + registry.register(make_worker("http://w:1", 1, 1)).unwrap(); + + let options = ProfileOptions::default(); + let result = + WorkerManager::start_profile_all(®istry, &options, Some("http://absent:1")).await; + assert_eq!(result.total_workers, 0); + assert!(result.successful.is_empty()); + assert!(result.failed.is_empty()); + assert!(result.message.contains("No worker matching")); + } } diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index 05ae8c583b..7d520903d5 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -12,12 +12,13 @@ use arc_swap::ArcSwap; use async_trait::async_trait; use axum::body::Body; // Re-export protocol types as the canonical types for the gateway -pub use openai_protocol::worker::{ConnectionMode, RuntimeType, WorkerType}; +pub use openai_protocol::worker::{ConnectionMode, ProfileOptions, RuntimeType, WorkerType}; use openai_protocol::{ model_card::ModelCard, model_type::{Endpoint, ModelType}, worker::{HealthCheckConfig, ProviderType, WorkerInfo, WorkerModels, WorkerSpec, WorkerStatus}, }; +use smg_grpc_client::common_proto; use tokio::{sync::OnceCell, time}; use super::{CircuitBreaker, ResolvedResilience, WorkerError, WorkerResult, UNKNOWN_MODEL_ID}; @@ -29,6 +30,15 @@ use crate::{ /// Default HTTP client timeout for worker requests (in seconds) pub const DEFAULT_WORKER_HTTP_TIMEOUT_SECS: u64 = 30; +/// Timeout for worker HTTP `flush_cache` requests. Matches the gRPC +/// client's local flush deadline. +const FLUSH_HTTP_TIMEOUT: Duration = Duration::from_secs(45); + +/// Timeout for worker HTTP profile requests. Stopping a profile can take +/// a long time while the backend serializes large traces. Matches the +/// gRPC client's profile deadline. +const PROFILE_HTTP_TIMEOUT: Duration = Duration::from_secs(630); + /// Default bootstrap port for PD disaggregation (used by SGLang and vLLM Mooncake) pub const DEFAULT_BOOTSTRAP_PORT: u16 = 8998; @@ -38,6 +48,77 @@ pub const MOONCAKE_CONNECTOR: &str = "MooncakeConnector"; /// vLLM NIXL KV connector name pub const NIXL_CONNECTOR: &str = "NixlConnector"; +/// POST an admin endpoint on an HTTP worker and map the outcome to a +/// [`WorkerResult`]. +async fn admin_http_post( + client: &reqwest::Client, + url: String, + api_key: Option<&String>, + body: Option, + operation: &str, + timeout: Duration, +) -> WorkerResult<()> { + let mut req = client.post(&url).timeout(timeout); + if let Some(key) = api_key { + req = req.bearer_auth(key); + } + if let Some(body) = body { + req = req.json(&body); + } + match req.send().await { + Ok(resp) if resp.status().is_success() => Ok(()), + Ok(resp) => Err(WorkerError::OperationFailed { + url, + operation: operation.to_string(), + reason: format!("HTTP {}", resp.status()), + }), + Err(e) => Err(WorkerError::OperationFailed { + url, + operation: operation.to_string(), + reason: e.to_string(), + }), + } +} + +/// Unwrap the optional gRPC client for an admin op, erroring when absent. +fn require_grpc_client( + url: &str, + operation: &str, + client: Option>, +) -> WorkerResult> { + client.ok_or_else(|| WorkerError::OperationFailed { + url: url.to_string(), + operation: operation.to_string(), + reason: "no gRPC client available".to_string(), + }) +} + +/// Map a gRPC admin-op outcome (`success` flag plus message) to a +/// [`WorkerResult`]. +fn admin_grpc_result( + url: &str, + operation: &str, + result: Result<(bool, String), tonic::Status>, +) -> WorkerResult<()> { + match result { + Ok((true, _)) => Ok(()), + Ok((false, message)) => Err(WorkerError::OperationFailed { + url: url.to_string(), + operation: operation.to_string(), + reason: if message.is_empty() { + "backend reported failure".to_string() + } else { + message + }, + }), + Err(status) => Err(WorkerError::OperationFailed { + url: url.to_string(), + operation: operation.to_string(), + reason: status.to_string(), + }), + } +} + pub struct WorkerRoutingKeyLoad { url: String, active_routing_keys: dashmap::DashMap, @@ -392,6 +473,114 @@ pub trait Worker: Send + Sync + fmt::Debug + 'static { } async fn grpc_health_check(&self) -> WorkerResult; async fn http_health_check(&self) -> WorkerResult; + + // ── Admin operations ──────────────────────────────────────────── + // + // Dual-path dispatch on connection mode, mirroring + // `check_health_async`: HTTP workers receive a POST to the engine's + // native admin endpoint, gRPC workers receive the corresponding RPC. + // Backends without the RPC surface the dispatcher's `Unimplemented` + // status as an `OperationFailed` error. + + /// Flush the KV cache on this worker's backend. + async fn flush_cache(&self) -> WorkerResult<()> { + match self.connection_mode() { + ConnectionMode::Http => { + admin_http_post( + self.http_client(), + self.endpoint_url("/flush_cache"), + self.api_key(), + None, + "flush_cache", + FLUSH_HTTP_TIMEOUT, + ) + .await + } + ConnectionMode::Grpc => { + let client = + require_grpc_client(self.url(), "flush_cache", self.get_grpc_client().await?)?; + let result = client.flush_cache(0.0).await; + admin_grpc_result( + self.url(), + "flush_cache", + result.map(|r| (r.success, r.message)), + ) + } + } + } + + /// Start a profiling run on this worker's backend. + async fn start_profile(&self, options: &ProfileOptions) -> WorkerResult<()> { + match self.connection_mode() { + ConnectionMode::Http => { + let body = + serde_json::to_value(options).map_err(|e| WorkerError::OperationFailed { + url: self.url().to_string(), + operation: "start_profile".to_string(), + reason: format!("failed to serialize profile options: {e}"), + })?; + admin_http_post( + self.http_client(), + self.endpoint_url("/start_profile"), + self.api_key(), + Some(body), + "start_profile", + PROFILE_HTTP_TIMEOUT, + ) + .await + } + ConnectionMode::Grpc => { + let client = require_grpc_client( + self.url(), + "start_profile", + self.get_grpc_client().await?, + )?; + let req = common_proto::StartProfileRequest { + output_dir: options.output_dir.clone(), + start_step: options.start_step, + num_steps: options.num_steps, + activities: options.activities.clone().unwrap_or_default(), + with_stack: options.with_stack, + record_shapes: options.record_shapes, + profile_by_stage: options.profile_by_stage, + }; + let result = client.start_profile(req).await; + admin_grpc_result( + self.url(), + "start_profile", + result.map(|r| (r.success, r.message)), + ) + } + } + } + + /// Stop the in-flight profiling run on this worker's backend and + /// export traces. + async fn stop_profile(&self) -> WorkerResult<()> { + match self.connection_mode() { + ConnectionMode::Http => { + admin_http_post( + self.http_client(), + self.endpoint_url("/stop_profile"), + self.api_key(), + None, + "stop_profile", + PROFILE_HTTP_TIMEOUT, + ) + .await + } + ConnectionMode::Grpc => { + let client = + require_grpc_client(self.url(), "stop_profile", self.get_grpc_client().await?)?; + let result = client.stop_profile().await; + admin_grpc_result( + self.url(), + "stop_profile", + result.map(|r| (r.success, r.message)), + ) + } + } + } } /// Extension trait for model_gateway-specific ConnectionMode methods.