From f0918583042a59a99ea3a548f3d47093529296c9 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Wed, 14 Jan 2026 12:05:10 -0800 Subject: [PATCH] [gRPC] Add GetLoads RPC for comprehensive load metrics (#17087) --- python/sglang/srt/entrypoints/grpc_server.py | 141 ++++++++++++++++++ .../sglang/srt/grpc/grpc_request_manager.py | 98 ++++++++++++ python/sglang/srt/grpc/sglang_scheduler.proto | 100 +++++++++++++ .../src/proto/sglang_scheduler.proto | 100 +++++++++++++ 4 files changed, 439 insertions(+) diff --git a/python/sglang/srt/entrypoints/grpc_server.py b/python/sglang/srt/entrypoints/grpc_server.py index 10513e000..61b7157f3 100644 --- a/python/sglang/srt/entrypoints/grpc_server.py +++ b/python/sglang/srt/entrypoints/grpc_server.py @@ -12,6 +12,7 @@ import signal import threading import time from concurrent import futures +from datetime import datetime, timezone from typing import AsyncIterator, Dict, Optional import grpc @@ -30,6 +31,7 @@ from sglang.srt.grpc.health_servicer import SGLangHealthServicer from sglang.srt.grpc.scheduler_launcher import launch_scheduler_process_only from sglang.srt.managers.disagg_service import start_disagg_service from sglang.srt.managers.io_struct import ( + GetLoadsReqOutput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, ) @@ -42,6 +44,100 @@ logger = logging.getLogger(__name__) HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20)) +def _convert_loads_to_protobuf( + result: GetLoadsReqOutput, +) -> sglang_scheduler_pb2.SchedulerLoad: + """Convert GetLoadsReqOutput dataclass to protobuf SchedulerLoad message.""" + scheduler_load = sglang_scheduler_pb2.SchedulerLoad( + dp_rank=result.dp_rank, + num_running_reqs=result.num_running_reqs, + num_waiting_reqs=result.num_waiting_reqs, + num_total_reqs=result.num_running_reqs + result.num_waiting_reqs, + num_used_tokens=result.num_used_tokens, + max_total_num_tokens=result.max_total_num_tokens, + token_usage=result.token_usage, + gen_throughput=result.gen_throughput, + cache_hit_rate=result.cache_hit_rate, + utilization=result.utilization, + max_running_requests=result.max_running_requests, + ) + + # Add optional sections using CopyFrom for proper protobuf assignment + if result.memory: + scheduler_load.memory.CopyFrom( + sglang_scheduler_pb2.MemoryMetrics( + weight_gb=result.memory.weight_gb, + kv_cache_gb=result.memory.kv_cache_gb, + graph_gb=result.memory.graph_gb, + token_capacity=result.memory.token_capacity, + ) + ) + + if result.speculative: + scheduler_load.speculative.CopyFrom( + sglang_scheduler_pb2.SpeculativeMetrics( + accept_length=result.speculative.accept_length, + accept_rate=result.speculative.accept_rate, + ) + ) + + if result.lora: + scheduler_load.lora.CopyFrom( + sglang_scheduler_pb2.LoRAMetrics( + slots_used=result.lora.slots_used, + slots_total=result.lora.slots_total, + utilization=result.lora.utilization, + ) + ) + + if result.disaggregation: + scheduler_load.disaggregation.CopyFrom( + sglang_scheduler_pb2.DisaggregationMetrics( + mode=result.disaggregation.mode, + prefill_prealloc_queue_reqs=result.disaggregation.prefill_prealloc_queue_reqs, + prefill_inflight_queue_reqs=result.disaggregation.prefill_inflight_queue_reqs, + decode_prealloc_queue_reqs=result.disaggregation.decode_prealloc_queue_reqs, + decode_transfer_queue_reqs=result.disaggregation.decode_transfer_queue_reqs, + decode_retracted_queue_reqs=result.disaggregation.decode_retracted_queue_reqs, + kv_transfer_speed_gb_s=result.disaggregation.kv_transfer_speed_gb_s, + kv_transfer_latency_ms=result.disaggregation.kv_transfer_latency_ms, + ) + ) + + if result.queues: + scheduler_load.queues.CopyFrom( + sglang_scheduler_pb2.QueueMetrics( + waiting=result.queues.waiting, + grammar=result.queues.grammar, + paused=result.queues.paused, + retracted=result.queues.retracted, + ) + ) + + return scheduler_load + + +def _compute_aggregate_protobuf( + loads: list, +) -> sglang_scheduler_pb2.AggregateMetrics: + """Compute aggregate metrics from list of SchedulerLoad protobuf messages.""" + if not loads: + return sglang_scheduler_pb2.AggregateMetrics() + + n = len(loads) + total_running = sum(load.num_running_reqs for load in loads) + total_waiting = sum(load.num_waiting_reqs for load in loads) + + return sglang_scheduler_pb2.AggregateMetrics( + total_running_reqs=total_running, + total_waiting_reqs=total_waiting, + total_reqs=total_running + total_waiting, + avg_token_usage=round(sum(load.token_usage for load in loads) / n, 4), + avg_throughput=round(sum(load.gen_throughput for load in loads) / n, 2), + avg_utilization=round(sum(load.utilization for load in loads) / n, 4), + ) + + class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer): """ Standalone gRPC service implementation using GrpcRequestManager. @@ -392,6 +488,51 @@ class SGLangSchedulerServicer(sglang_scheduler_pb2_grpc.SglangSchedulerServicer) start_time=start_timestamp, ) + async def GetLoads( + self, + request: sglang_scheduler_pb2.GetLoadsRequest, + context: grpc.aio.ServicerContext, + ) -> sglang_scheduler_pb2.GetLoadsResponse: + """ + Get comprehensive load metrics for all DP ranks. + + Uses the communicator pattern to fetch real-time metrics, + providing full parity with the HTTP /v1/loads endpoint. + """ + logger.debug("Receive get loads request") + + include = list(request.include) if request.include else ["all"] + dp_rank = request.dp_rank if request.HasField("dp_rank") else None + + try: + results = await self.request_manager.get_loads( + include=include, dp_rank=dp_rank + ) + except ValueError as e: + # Validation error (e.g., invalid include sections) + context.set_code(grpc.StatusCode.INVALID_ARGUMENT) + context.set_details(str(e)) + return sglang_scheduler_pb2.GetLoadsResponse() + except asyncio.TimeoutError: + context.set_code(grpc.StatusCode.DEADLINE_EXCEEDED) + context.set_details("Timeout waiting for scheduler response") + return sglang_scheduler_pb2.GetLoadsResponse() + except Exception as e: + logger.error(f"GetLoads failed: {e}\n{get_exception_traceback()}") + context.set_code(grpc.StatusCode.INTERNAL) + context.set_details(f"Failed to get load metrics: {e}") + return sglang_scheduler_pb2.GetLoadsResponse() + + loads = [_convert_loads_to_protobuf(r) for r in results] + + return sglang_scheduler_pb2.GetLoadsResponse( + timestamp=datetime.now(timezone.utc).isoformat(), + version=sglang.__version__, + dp_rank_count=len(loads), + loads=loads, + aggregate=_compute_aggregate_protobuf(loads), + ) + # Helper methods for request/response conversion def _convert_generate_request( diff --git a/python/sglang/srt/grpc/grpc_request_manager.py b/python/sglang/srt/grpc/grpc_request_manager.py index cc1e9f64a..dc0c27539 100644 --- a/python/sglang/srt/grpc/grpc_request_manager.py +++ b/python/sglang/srt/grpc/grpc_request_manager.py @@ -23,6 +23,8 @@ from sglang.srt.managers.io_struct import ( AbortReq, BatchEmbeddingOutput, BatchTokenIDOutput, + GetLoadsReqInput, + GetLoadsReqOutput, HealthCheckOutput, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, @@ -34,6 +36,68 @@ from sglang.utils import get_exception_traceback logger = logging.getLogger(__name__) +class _GrpcCommunicator: + """ + Communicator for request/response patterns with scheduler. + + Thread-safe and handles the async request/response cycle with proper + timeout handling to prevent hangs if the scheduler becomes unresponsive. + """ + + DEFAULT_TIMEOUT = 30.0 # seconds + + def __init__(self, sender: zmq.Socket, fan_out: int = 1): + self._sender = sender + self._fan_out = fan_out + self._result_event: Optional[asyncio.Event] = None + self._result_values: Optional[List[Any]] = None + self._lock = asyncio.Lock() + + async def __call__(self, obj, timeout: float = DEFAULT_TIMEOUT) -> List[Any]: + """ + Send request and wait for response(s). + + Args: + obj: Request object to send to scheduler + timeout: Maximum time to wait for response (seconds) + + Returns: + List of response objects from scheduler(s) + + Raises: + asyncio.TimeoutError: If no response within timeout + """ + async with self._lock: + # Initialize state BEFORE sending to avoid race condition + self._result_event = asyncio.Event() + self._result_values = [] + + # Send request to scheduler + if obj: + self._sender.send_pyobj(obj) + + try: + # Wait for response(s) with timeout + await asyncio.wait_for(self._result_event.wait(), timeout=timeout) + return self._result_values + finally: + # Always clean up state + self._result_event = None + self._result_values = None + + def handle_recv(self, recv_obj: Any): + """ + Handle received response from scheduler. + + Called by handle_loop when a matching response type is received. + Safe to call even if no request is pending (will be ignored). + """ + if self._result_values is not None and self._result_event is not None: + self._result_values.append(recv_obj) + if len(self._result_values) >= self._fan_out: + self._result_event.set() + + class GrpcSignalHandler: """Minimal signal handler for gRPC server - delegates real crash handling to scheduler.""" @@ -154,6 +218,12 @@ class GrpcRequestManager: # Bootstrap server (passed from serve_grpc, not started here) self.bootstrap_server = bootstrap_server + # Communicators for request/response patterns with scheduler + # Note: These must be initialized after send_to_scheduler socket is created + self.get_loads_communicator = _GrpcCommunicator( + self.send_to_scheduler, fan_out=server_args.dp_size + ) + logger.info( f"GrpcRequestManager initialized with ZMQ IPC: " f"recv={port_args.detokenizer_ipc_name}, " @@ -462,6 +532,9 @@ class GrpcRequestManager: 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) else: logger.warning(f"Unknown output type: {type(recv_obj)}") @@ -872,6 +945,31 @@ class GrpcRequestManager: "last_receive_time": self.last_receive_tstamp, } + async def get_loads( + self, include: List[str], dp_rank: Optional[int] = None + ) -> List[GetLoadsReqOutput]: + """ + Get comprehensive load metrics from the scheduler. + + This method uses the communicator pattern to send GetLoadsReqInput to the + scheduler and wait for GetLoadsReqOutput responses. + + Args: + include: List of metric sections to include (core, memory, spec, lora, disagg, queues, all) + dp_rank: Optional DP rank filter (None for all ranks) + + Returns: + List of GetLoadsReqOutput objects, one per scheduler/DP rank + """ + req = GetLoadsReqInput(include=include, dp_rank=dp_rank) + results = await self.get_loads_communicator(req) + + # Filter by dp_rank if specified + if dp_rank is not None: + results = [r for r in results if r.dp_rank == dp_rank] + + return results + def auto_create_handle_loop(self): """Automatically create and start the handle_loop task, matching TokenizerManager pattern.""" if self.no_create_loop: diff --git a/python/sglang/srt/grpc/sglang_scheduler.proto b/python/sglang/srt/grpc/sglang_scheduler.proto index 85ee64a08..578fc3317 100644 --- a/python/sglang/srt/grpc/sglang_scheduler.proto +++ b/python/sglang/srt/grpc/sglang_scheduler.proto @@ -26,6 +26,9 @@ service SglangScheduler { // Get server information rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse); + // Get comprehensive load metrics + rpc GetLoads(GetLoadsRequest) returns (GetLoadsResponse); + } // ===================== @@ -464,3 +467,100 @@ message GetServerInfoResponse { // bidirectional communicator infrastructure not available in gRPC. // Use HTTP /get_server_info if scheduler internal state is needed. } + +// ===================== +// Load Metrics (v1/loads) +// ===================== + +message GetLoadsRequest { + // Optional: filter to specific DP rank + optional int32 dp_rank = 1; + + // Sections to include: core, memory, spec, lora, disagg, queues, all + repeated string include = 2; +} + +message GetLoadsResponse { + // ISO 8601 timestamp + string timestamp = 1; + + // SGLang version + string version = 2; + + // Number of DP ranks + int32 dp_rank_count = 3; + + // Per-DP-rank load metrics + repeated SchedulerLoad loads = 4; + + // Aggregate metrics across all DP ranks + AggregateMetrics aggregate = 5; +} + +message SchedulerLoad { + int32 dp_rank = 1; + + // Core metrics (always included) + int32 num_running_reqs = 2; + int32 num_waiting_reqs = 3; + int32 num_total_reqs = 4; + int32 num_used_tokens = 5; + int32 max_total_num_tokens = 6; + double token_usage = 7; + double gen_throughput = 8; + double cache_hit_rate = 9; + double utilization = 10; + int32 max_running_requests = 11; + + // Optional sections + optional MemoryMetrics memory = 12; + optional SpeculativeMetrics speculative = 13; + optional LoRAMetrics lora = 14; + optional DisaggregationMetrics disaggregation = 15; + optional QueueMetrics queues = 16; +} + +message MemoryMetrics { + double weight_gb = 1; + double kv_cache_gb = 2; + double graph_gb = 3; + int32 token_capacity = 4; +} + +message SpeculativeMetrics { + double accept_length = 1; + double accept_rate = 2; +} + +message LoRAMetrics { + int32 slots_used = 1; + int32 slots_total = 2; + double utilization = 3; +} + +message DisaggregationMetrics { + string mode = 1; // "prefill", "decode", or "null" + int32 prefill_prealloc_queue_reqs = 2; + int32 prefill_inflight_queue_reqs = 3; + int32 decode_prealloc_queue_reqs = 4; + int32 decode_transfer_queue_reqs = 5; + int32 decode_retracted_queue_reqs = 6; + double kv_transfer_speed_gb_s = 7; + double kv_transfer_latency_ms = 8; +} + +message QueueMetrics { + int32 waiting = 1; + int32 grammar = 2; + int32 paused = 3; + int32 retracted = 4; +} + +message AggregateMetrics { + int32 total_running_reqs = 1; + int32 total_waiting_reqs = 2; + int32 total_reqs = 3; + double avg_token_usage = 4; + double avg_throughput = 5; + double avg_utilization = 6; +} diff --git a/sgl-model-gateway/src/proto/sglang_scheduler.proto b/sgl-model-gateway/src/proto/sglang_scheduler.proto index 85ee64a08..578fc3317 100644 --- a/sgl-model-gateway/src/proto/sglang_scheduler.proto +++ b/sgl-model-gateway/src/proto/sglang_scheduler.proto @@ -26,6 +26,9 @@ service SglangScheduler { // Get server information rpc GetServerInfo(GetServerInfoRequest) returns (GetServerInfoResponse); + // Get comprehensive load metrics + rpc GetLoads(GetLoadsRequest) returns (GetLoadsResponse); + } // ===================== @@ -464,3 +467,100 @@ message GetServerInfoResponse { // bidirectional communicator infrastructure not available in gRPC. // Use HTTP /get_server_info if scheduler internal state is needed. } + +// ===================== +// Load Metrics (v1/loads) +// ===================== + +message GetLoadsRequest { + // Optional: filter to specific DP rank + optional int32 dp_rank = 1; + + // Sections to include: core, memory, spec, lora, disagg, queues, all + repeated string include = 2; +} + +message GetLoadsResponse { + // ISO 8601 timestamp + string timestamp = 1; + + // SGLang version + string version = 2; + + // Number of DP ranks + int32 dp_rank_count = 3; + + // Per-DP-rank load metrics + repeated SchedulerLoad loads = 4; + + // Aggregate metrics across all DP ranks + AggregateMetrics aggregate = 5; +} + +message SchedulerLoad { + int32 dp_rank = 1; + + // Core metrics (always included) + int32 num_running_reqs = 2; + int32 num_waiting_reqs = 3; + int32 num_total_reqs = 4; + int32 num_used_tokens = 5; + int32 max_total_num_tokens = 6; + double token_usage = 7; + double gen_throughput = 8; + double cache_hit_rate = 9; + double utilization = 10; + int32 max_running_requests = 11; + + // Optional sections + optional MemoryMetrics memory = 12; + optional SpeculativeMetrics speculative = 13; + optional LoRAMetrics lora = 14; + optional DisaggregationMetrics disaggregation = 15; + optional QueueMetrics queues = 16; +} + +message MemoryMetrics { + double weight_gb = 1; + double kv_cache_gb = 2; + double graph_gb = 3; + int32 token_capacity = 4; +} + +message SpeculativeMetrics { + double accept_length = 1; + double accept_rate = 2; +} + +message LoRAMetrics { + int32 slots_used = 1; + int32 slots_total = 2; + double utilization = 3; +} + +message DisaggregationMetrics { + string mode = 1; // "prefill", "decode", or "null" + int32 prefill_prealloc_queue_reqs = 2; + int32 prefill_inflight_queue_reqs = 3; + int32 decode_prealloc_queue_reqs = 4; + int32 decode_transfer_queue_reqs = 5; + int32 decode_retracted_queue_reqs = 6; + double kv_transfer_speed_gb_s = 7; + double kv_transfer_latency_ms = 8; +} + +message QueueMetrics { + int32 waiting = 1; + int32 grammar = 2; + int32 paused = 3; + int32 retracted = 4; +} + +message AggregateMetrics { + int32 total_running_reqs = 1; + int32 total_waiting_reqs = 2; + int32 total_reqs = 3; + double avg_token_usage = 4; + double avg_throughput = 5; + double avg_utilization = 6; +}