[router] add grpc client get and set (#10955)

This commit is contained in:
Simo Lin
2025-09-26 06:07:05 -04:00
committed by GitHub
parent a5095d6262
commit 1e57b9472d
4 changed files with 111 additions and 37 deletions

View File

@@ -202,12 +202,23 @@ impl GrpcRouter {
debug!("Selected worker: {}", worker.url());
// Step 2: Get gRPC client for worker (fail fast if can't connect)
// TODO(CahterineSue): manage grpc connection in worker. (it should be simpler here)
let client = match self.get_or_create_grpc_client(worker.url()).await {
Ok(c) => c,
// Step 2: Get gRPC client from worker
let client = match worker.get_grpc_client().await {
Ok(Some(client_arc)) => {
// Clone the client from inside the Arc<Mutex<>>
let client = client_arc.lock().await.clone();
client
}
Ok(None) => {
error!("Selected worker is not a gRPC worker");
return (
StatusCode::INTERNAL_SERVER_ERROR,
"Selected worker is not configured for gRPC",
)
.into_response();
}
Err(e) => {
error!("Failed to get gRPC client: {}", e);
error!("Failed to get gRPC client from worker: {}", e);
return (
StatusCode::INTERNAL_SERVER_ERROR,
format!("Failed to get gRPC client: {}", e),
@@ -552,18 +563,6 @@ impl GrpcRouter {
None
}
/// Get or create a gRPC client for the worker
async fn get_or_create_grpc_client(
&self,
worker_url: &str,
) -> Result<SglangSchedulerClient, String> {
// TODO: move to worker
debug!("Creating new gRPC client for worker: {}", worker_url);
SglangSchedulerClient::connect(worker_url)
.await
.map_err(|e| format!("Failed to connect to gRPC server: {}", e))
}
/// Placeholder for streaming handler (to be implemented in Phase 2)
async fn handle_streaming_chat(
&self,