chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[gRPC] Add GetLoads RPC for comprehensive load metrics (#17087)

Browse Source
This commit is contained in:
Simo Lin
2026-01-14 12:05:10 -08:00
committed by GitHub
parent 5b1215d9da
commit f091858304
4 changed files with 439 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
python/sglang/srt/entrypoints/grpc_server.py
View File

@@ -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(

98
python/sglang/srt/grpc/grpc_request_manager.py
View File

@@ -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:

100
python/sglang/srt/grpc/sglang_scheduler.proto
View File

@@ -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;
}