[API] Add /v1/loads endpoint for load metrics (#16976)

2026-01-14 09:13:20 -08:00
parent de94d793ad
commit b3a3f51320
6 changed files with 513 additions and 1 deletions
--- 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()


--- a/python/sglang/srt/entrypoints/v1_loads.py
+++ 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),
+    }
--- 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]
--- 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),
            ]
--- 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):
--- 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
    ):