diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 652f5b061..51c090379 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -359,6 +359,11 @@ app.add_middleware( allow_headers=["*"], ) +# Include routers +from sglang.srt.entrypoints.v1_loads import router as v1_loads_router + +app.include_router(v1_loads_router) + @app.exception_handler(HTTPException) async def validation_exception_handler(request: Request, exc: HTTPException): @@ -601,6 +606,11 @@ async def server_info(): @app.get("/get_load") async def get_load(): + """Get load metrics (deprecated - use /v1/loads instead).""" + logger.warning( + "Endpoint '/get_load' is deprecated and will be removed in a future version. " + "Please use '/v1/loads' instead." + ) return await _global_state.tokenizer_manager.get_load() diff --git a/python/sglang/srt/entrypoints/v1_loads.py b/python/sglang/srt/entrypoints/v1_loads.py new file mode 100644 index 000000000..784417010 --- /dev/null +++ b/python/sglang/srt/entrypoints/v1_loads.py @@ -0,0 +1,175 @@ +# Copyright 2023-2024 SGLang Team +# 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. +# ============================================================================== +""" +/v1/loads API endpoint for comprehensive load metrics. + +This module provides the /v1/loads endpoint which returns detailed scheduler +metrics for load balancing, monitoring, and capacity planning. +""" + +import dataclasses +from datetime import datetime, timezone +from typing import Optional + +from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import Response + +from sglang.srt.managers.io_struct import ( + DisaggregationMetrics, + GetLoadsReqOutput, + LoRAMetrics, + MemoryMetrics, + QueueMetrics, + SpeculativeMetrics, +) +from sglang.version import __version__ + +router = APIRouter() + +_OPTIONAL_METRIC_SECTIONS = { + "memory": ("memory", MemoryMetrics), + "speculative": ("spec", SpeculativeMetrics), + "lora": ("lora", LoRAMetrics), + "disaggregation": ("disagg", DisaggregationMetrics), + "queues": ("queues", QueueMetrics), +} + + +def _get_tokenizer_manager(): + """Dependency to get tokenizer_manager from global state.""" + from sglang.srt.entrypoints.http_server import get_global_state + + return get_global_state().tokenizer_manager + + +def _loads_dict_factory(items): + """Factory for dataclasses.asdict() that excludes None values and timestamp.""" + return {k: v for k, v in items if v is not None and k != "timestamp"} + + +def _compute_aggregate(load_dicts: list) -> dict: + """Compute aggregate metrics from load dicts.""" + if not load_dicts: + return { + "total_running_reqs": 0, + "total_waiting_reqs": 0, + "total_reqs": 0, + "avg_token_usage": 0.0, + "avg_throughput": 0.0, + "avg_utilization": 0.0, + } + + n = len(load_dicts) + return { + "total_running_reqs": sum(d["num_running_reqs"] for d in load_dicts), + "total_waiting_reqs": sum(d["num_waiting_reqs"] for d in load_dicts), + "total_reqs": sum( + d["num_running_reqs"] + d["num_waiting_reqs"] for d in load_dicts + ), + "avg_token_usage": round(sum(d["token_usage"] for d in load_dicts) / n, 4), + "avg_throughput": round(sum(d["gen_throughput"] for d in load_dicts) / n, 2), + "avg_utilization": round(sum(d["utilization"] for d in load_dicts) / n, 4), + } + + +def _format_loads_prometheus(load_results) -> Response: + """Format load metrics in Prometheus text exposition format. + + Metrics are derived from dataclass field metadata, providing a single source of truth. + """ + lines = [] + + for f in dataclasses.fields(GetLoadsReqOutput): + if "metric" not in f.metadata: + continue + metric_type, description = f.metadata["metric"] + metric_name = f"sglang_{f.name}" + lines.append(f"# HELP {metric_name} {description}") + lines.append(f"# TYPE {metric_name} {metric_type}") + for load in load_results: + value = getattr(load, f.name, None) + if value is not None: + lines.append(f'{metric_name}{{dp_rank="{load.dp_rank}"}} {value}') + + for attr_name, (prefix, dataclass_type) in _OPTIONAL_METRIC_SECTIONS.items(): + if not any(getattr(load, attr_name, None) for load in load_results): + continue + for f in dataclasses.fields(dataclass_type): + if "metric" not in f.metadata: + continue + metric_type, description = f.metadata["metric"] + metric_name = f"sglang_{prefix}_{f.name}" + lines.append(f"# HELP {metric_name} {description}") + lines.append(f"# TYPE {metric_name} {metric_type}") + for load in load_results: + section = getattr(load, attr_name, None) + if section: + value = getattr(section, f.name, None) + if value is not None: + lines.append( + f'{metric_name}{{dp_rank="{load.dp_rank}"}} {value}' + ) + + return Response( + content="\n".join(lines) + "\n", + media_type="text/plain; version=0.0.4; charset=utf-8", + ) + + +@router.get("/v1/loads") +async def get_loads( + dp_rank: Optional[int] = None, + include: Optional[str] = None, + format: Optional[str] = None, + tokenizer_manager=Depends(_get_tokenizer_manager), +): + """ + Get comprehensive load metrics for all DP ranks. + + Query Parameters: + dp_rank: Filter to specific DP rank (optional) + include: Comma-separated sections to include (optional) + Options: core, memory, spec, lora, disagg, queues, all + Default: all + format: Response format - 'json' (default) or 'prometheus' + + Returns: + JSON response with timestamp, version, dp_rank_count, per-DP-rank loads, and aggregates + """ + include_list = [s.strip() for s in include.split(",")] if include else None + + try: + load_results = await tokenizer_manager.get_loads( + include=include_list, + dp_rank=dp_rank, + ) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + if format == "prometheus": + return _format_loads_prometheus(load_results) + + loads = [] + for load in load_results: + d = dataclasses.asdict(load, dict_factory=_loads_dict_factory) + d["num_total_reqs"] = d["num_running_reqs"] + d["num_waiting_reqs"] + loads.append(d) + + return { + "timestamp": datetime.now(timezone.utc).isoformat(), + "version": __version__, + "dp_rank_count": len(loads), + "loads": loads, + "aggregate": _compute_aggregate(loads), + } diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 18cc3d2aa..3df095559 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1697,6 +1697,147 @@ class GetLoadReqOutput(BaseReq): ts_tic: float +@dataclass +class MemoryMetrics: + """Memory breakdown metrics.""" + + weight_gb: float = field( + metadata={"metric": ("gauge", "Model weight memory in GB")} + ) + kv_cache_gb: float = field(metadata={"metric": ("gauge", "KV cache memory in GB")}) + graph_gb: float = field(metadata={"metric": ("gauge", "CUDA graph memory in GB")}) + token_capacity: int = field( + metadata={"metric": ("gauge", "Max tokens in KV cache")} + ) + + +@dataclass +class SpeculativeMetrics: + """Speculative decoding metrics.""" + + accept_length: float = field( + metadata={"metric": ("gauge", "Avg accepted tokens per step")} + ) + accept_rate: float = field( + metadata={"metric": ("gauge", "Speculative acceptance rate")} + ) + + +@dataclass +class LoRAMetrics: + """LoRA adapter pool metrics.""" + + slots_used: int = field(metadata={"metric": ("gauge", "LoRA adapter slots in use")}) + slots_total: int = field(metadata={"metric": ("gauge", "Total LoRA adapter slots")}) + utilization: float = field( + metadata={"metric": ("gauge", "LoRA pool utilization ratio")} + ) + + +@dataclass +class DisaggregationMetrics: + """PD disaggregation metrics.""" + + mode: str # "prefill", "decode", or "null" - not a metric + prefill_prealloc_queue_reqs: int = field( + default=0, metadata={"metric": ("gauge", "Prefill prealloc queue requests")} + ) + prefill_inflight_queue_reqs: int = field( + default=0, metadata={"metric": ("gauge", "Prefill inflight queue requests")} + ) + decode_prealloc_queue_reqs: int = field( + default=0, metadata={"metric": ("gauge", "Decode prealloc queue requests")} + ) + decode_transfer_queue_reqs: int = field( + default=0, metadata={"metric": ("gauge", "Decode transfer queue requests")} + ) + decode_retracted_queue_reqs: int = field( + default=0, metadata={"metric": ("gauge", "Decode retracted queue requests")} + ) + kv_transfer_speed_gb_s: float = field( + default=0.0, metadata={"metric": ("gauge", "KV transfer speed in GB/s")} + ) + kv_transfer_latency_ms: float = field( + default=0.0, metadata={"metric": ("gauge", "KV transfer latency in ms")} + ) + + +@dataclass +class QueueMetrics: + """Detailed queue breakdown.""" + + waiting: int = field(metadata={"metric": ("gauge", "Main waiting queue size")}) + grammar: int = field( + metadata={"metric": ("gauge", "Grammar compilation queue size")} + ) + paused: int = field( + metadata={"metric": ("gauge", "Requests paused by weight sync")} + ) + retracted: int = field(metadata={"metric": ("gauge", "Retracted requests count")}) + + +@dataclass +class GetLoadsReqInput(BaseReq): + """Request for /v1/loads endpoint.""" + + VALID_SECTIONS = frozenset( + {"core", "memory", "spec", "lora", "disagg", "queues", "all"} + ) + + include: List[str] = field(default_factory=lambda: ["all"]) + dp_rank: Optional[int] = None + + def __post_init__(self): + """Validate include sections.""" + if self.include: + invalid = set(self.include) - self.VALID_SECTIONS + if invalid: + raise ValueError( + f"Invalid include sections: {invalid}. " + f"Valid options: {sorted(self.VALID_SECTIONS)}" + ) + + +@dataclass +class GetLoadsReqOutput(BaseReq): + """Per-DP-rank load metrics for /v1/loads endpoint.""" + + dp_rank: int + timestamp: float + + num_running_reqs: int = field( + metadata={"metric": ("gauge", "Number of running requests")} + ) + num_waiting_reqs: int = field( + metadata={"metric": ("gauge", "Number of waiting requests")} + ) + num_used_tokens: int = field( + metadata={"metric": ("gauge", "Number of tokens in use")} + ) + max_total_num_tokens: int = field( + metadata={"metric": ("gauge", "Maximum token capacity")} + ) + token_usage: float = field(metadata={"metric": ("gauge", "Token pool usage ratio")}) + gen_throughput: float = field( + metadata={"metric": ("gauge", "Generation throughput tokens/sec")} + ) + cache_hit_rate: float = field( + metadata={"metric": ("gauge", "Prefix cache hit rate")} + ) + utilization: float = field( + metadata={"metric": ("gauge", "Overall utilization ratio")} + ) + max_running_requests: int = field( + metadata={"metric": ("gauge", "Maximum running requests capacity")} + ) + + memory: Optional[MemoryMetrics] = None + speculative: Optional[SpeculativeMetrics] = None + lora: Optional[LoRAMetrics] = None + disaggregation: Optional[DisaggregationMetrics] = None + queues: Optional[QueueMetrics] = None + + @dataclass class WatchLoadUpdateReq(BaseReq): loads: List[GetLoadReqOutput] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 1bf294973..80e5e4c0f 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -87,6 +87,7 @@ from sglang.srt.managers.io_struct import ( GetInternalStateReq, GetInternalStateReqOutput, GetLoadReqInput, + GetLoadsReqInput, GetWeightsByNameReqInput, HealthCheckOutput, InitWeightsSendGroupForRemoteInstanceReqInput, @@ -1046,6 +1047,7 @@ class Scheduler( ), (UnloadLoRAAdapterReqInput, self.unload_lora_adapter), (GetLoadReqInput, self.get_load), + (GetLoadsReqInput, self.get_loads), (PauseGenerationReqInput, self.pause_generation), (ContinueGenerationReqInput, self.continue_generation), ] diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index c943d1886..37d47010d 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -9,7 +9,17 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Union from sglang.srt.disaggregation.kv_events import EventPublisherFactory, KVEventBatch from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs -from sglang.srt.managers.io_struct import GetLoadReqInput, GetLoadReqOutput +from sglang.srt.managers.io_struct import ( + DisaggregationMetrics, + GetLoadReqInput, + GetLoadReqOutput, + GetLoadsReqInput, + GetLoadsReqOutput, + LoRAMetrics, + MemoryMetrics, + QueueMetrics, + SpeculativeMetrics, +) from sglang.srt.managers.schedule_policy import PrefillAdder from sglang.srt.managers.scheduler import Req, ScheduleBatch from sglang.srt.managers.utils import GenerationBatchResult @@ -584,6 +594,144 @@ class SchedulerMetricsMixin: ts_tic=time.perf_counter(), ) + def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: + """ + Get comprehensive load metrics for /v1/loads endpoint. + + Args: + req: Request containing include list and optional dp_rank filter + + Returns: + GetLoadsReqOutput with core metrics and optional detailed sections + """ + if req is None: + req = GetLoadsReqInput() + + include = set(req.include) if req.include else {"core"} + include_all = "all" in include + + num_running_reqs = len(self.running_batch.reqs) + + waiting_queues = [self.waiting_queue] + if self.disaggregation_mode == DisaggregationMode.PREFILL: + waiting_queues.append(self.disagg_prefill_bootstrap_queue.queue) + elif self.disaggregation_mode == DisaggregationMode.DECODE: + waiting_queues.append(self.disagg_decode_prealloc_queue.queue) + waiting_queues.append(self.disagg_decode_transfer_queue.queue) + waiting_queues.append(self.disagg_decode_prealloc_queue.retracted_queue) + + num_waiting_reqs = sum(len(queue) for queue in waiting_queues) + + if self.is_hybrid_swa: + full_num_used, swa_num_used, *_ = self._get_swa_token_info() + num_used_tokens = max(full_num_used, swa_num_used) + elif self.is_hybrid_ssm: + num_used_tokens = self._get_mamba_token_info()[0] + else: + num_used_tokens = self._get_token_info()[0] + + token_usage = ( + num_used_tokens / self.max_total_num_tokens + if self.max_total_num_tokens > 0 + else 0.0 + ) + + memory = None + if include_all or "memory" in include: + try: + memory = MemoryMetrics( + weight_gb=round( + self.tp_worker.model_runner.weight_load_mem_usage, 3 + ), + kv_cache_gb=round( + self.token_to_kv_pool_allocator.get_kvcache().mem_usage, 3 + ), + graph_gb=round(self.tp_worker.model_runner.graph_mem_usage, 3), + token_capacity=int(self.max_total_num_tokens), + ) + except AttributeError as e: + logger.debug(f"Memory metrics not available: {e}") + + speculative = None + if include_all or "spec" in include: + if not self.spec_algorithm.is_none() and self.spec_total_num_forward_ct > 0: + speculative = SpeculativeMetrics( + accept_length=( + self.spec_total_num_accepted_tokens + / self.spec_total_num_forward_ct + ), + accept_rate=self.stats.spec_accept_rate, + ) + + lora = None + if include_all or "lora" in include: + if hasattr(self, "lora_scheduler") and self.lora_scheduler is not None: + lora = LoRAMetrics( + slots_used=self.stats.lora_pool_slots_used, + slots_total=self.stats.lora_pool_slots_total, + utilization=self.stats.lora_pool_utilization, + ) + + disaggregation = None + if include_all or "disagg" in include: + mode_str = "null" + prefill_prealloc = 0 + prefill_inflight = 0 + decode_prealloc = 0 + decode_transfer = 0 + decode_retracted = 0 + + if self.disaggregation_mode == DisaggregationMode.PREFILL: + mode_str = "prefill" + prefill_prealloc = len(self.disagg_prefill_bootstrap_queue.queue) + prefill_inflight = len(self.disagg_prefill_inflight_queue) + elif self.disaggregation_mode == DisaggregationMode.DECODE: + mode_str = "decode" + decode_prealloc = len(self.disagg_decode_prealloc_queue.queue) + decode_transfer = len(self.disagg_decode_transfer_queue.queue) + decode_retracted = len( + self.disagg_decode_prealloc_queue.retracted_queue + ) + + disaggregation = DisaggregationMetrics( + mode=mode_str, + prefill_prealloc_queue_reqs=prefill_prealloc, + prefill_inflight_queue_reqs=prefill_inflight, + decode_prealloc_queue_reqs=decode_prealloc, + decode_transfer_queue_reqs=decode_transfer, + decode_retracted_queue_reqs=decode_retracted, + kv_transfer_speed_gb_s=self.stats.kv_transfer_speed_gb_s, + kv_transfer_latency_ms=self.stats.kv_transfer_latency_ms, + ) + + queues = None + if include_all or "queues" in include: + queues = QueueMetrics( + waiting=len(self.waiting_queue), + grammar=self.stats.num_grammar_queue_reqs, + paused=self.stats.num_paused_reqs, + retracted=self.stats.num_retracted_reqs, + ) + + return GetLoadsReqOutput( + dp_rank=self.dp_rank, + timestamp=time.time(), + num_running_reqs=num_running_reqs, + num_waiting_reqs=num_waiting_reqs, + num_used_tokens=num_used_tokens, + max_total_num_tokens=self.max_total_num_tokens, + token_usage=round(token_usage, 4), + gen_throughput=round(self.stats.gen_throughput, 2), + cache_hit_rate=round(self.stats.cache_hit_rate, 4), + utilization=round(self.stats.utilization, 4), + max_running_requests=self.max_running_requests, + memory=memory, + speculative=speculative, + lora=lora, + disaggregation=disaggregation, + queues=queues, + ) + @contextmanager def record_forward_metrics(self: Scheduler, batch: ScheduleBatch): if not (self.enable_metrics and ENABLE_METRICS_DEVICE_TIMER): diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py index 231527143..e25729e71 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -39,6 +39,8 @@ from sglang.srt.managers.io_struct import ( GetInternalStateReqOutput, GetLoadReqInput, GetLoadReqOutput, + GetLoadsReqInput, + GetLoadsReqOutput, GetWeightsByNameReqInput, GetWeightsByNameReqOutput, InitWeightsSendGroupForRemoteInstanceReqInput, @@ -218,6 +220,9 @@ class TokenizerCommunicatorMixin: self.get_load_communicator = _Communicator( self.send_to_scheduler, server_args.dp_size, mode="watching" ) + self.get_loads_communicator = _Communicator( + self.send_to_scheduler, server_args.dp_size + ) self._result_dispatcher += self._get_communicator_dispatcher() @@ -304,6 +309,10 @@ class TokenizerCommunicatorMixin: GetLoadReqOutput, self.get_load_communicator.handle_recv, ), + ( + GetLoadsReqOutput, + self.get_loads_communicator.handle_recv, + ), ] ) @@ -787,6 +796,33 @@ class TokenizerCommunicatorMixin: req = GetLoadReqInput() return await self.get_load_communicator(req) + async def get_loads( + self: TokenizerManager, + include: Optional[List[str]] = None, + dp_rank: Optional[int] = None, + ) -> List[GetLoadsReqOutput]: + """ + Get comprehensive load metrics for /v1/loads endpoint. + + Args: + include: List of sections to include. Options: core, memory, spec, lora, disagg, queues, all + dp_rank: Optional filter for specific DP rank + + Returns: + List of GetLoadsReqOutput, one per scheduler (filtered by dp_rank if specified) + """ + req = GetLoadsReqInput( + include=include if include else ["all"], + 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 + async def open_session( self, obj: OpenSessionReqInput, request: Optional[fastapi.Request] = None ):