[API] Add /v1/loads endpoint for load metrics (#16976)
This commit is contained in:
@@ -359,6 +359,11 @@ app.add_middleware(
|
|||||||
allow_headers=["*"],
|
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)
|
@app.exception_handler(HTTPException)
|
||||||
async def validation_exception_handler(request: Request, exc: HTTPException):
|
async def validation_exception_handler(request: Request, exc: HTTPException):
|
||||||
@@ -601,6 +606,11 @@ async def server_info():
|
|||||||
|
|
||||||
@app.get("/get_load")
|
@app.get("/get_load")
|
||||||
async def 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()
|
return await _global_state.tokenizer_manager.get_load()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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),
|
||||||
|
}
|
||||||
@@ -1697,6 +1697,147 @@ class GetLoadReqOutput(BaseReq):
|
|||||||
ts_tic: float
|
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
|
@dataclass
|
||||||
class WatchLoadUpdateReq(BaseReq):
|
class WatchLoadUpdateReq(BaseReq):
|
||||||
loads: List[GetLoadReqOutput]
|
loads: List[GetLoadReqOutput]
|
||||||
|
|||||||
@@ -87,6 +87,7 @@ from sglang.srt.managers.io_struct import (
|
|||||||
GetInternalStateReq,
|
GetInternalStateReq,
|
||||||
GetInternalStateReqOutput,
|
GetInternalStateReqOutput,
|
||||||
GetLoadReqInput,
|
GetLoadReqInput,
|
||||||
|
GetLoadsReqInput,
|
||||||
GetWeightsByNameReqInput,
|
GetWeightsByNameReqInput,
|
||||||
HealthCheckOutput,
|
HealthCheckOutput,
|
||||||
InitWeightsSendGroupForRemoteInstanceReqInput,
|
InitWeightsSendGroupForRemoteInstanceReqInput,
|
||||||
@@ -1046,6 +1047,7 @@ class Scheduler(
|
|||||||
),
|
),
|
||||||
(UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
|
(UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
|
||||||
(GetLoadReqInput, self.get_load),
|
(GetLoadReqInput, self.get_load),
|
||||||
|
(GetLoadsReqInput, self.get_loads),
|
||||||
(PauseGenerationReqInput, self.pause_generation),
|
(PauseGenerationReqInput, self.pause_generation),
|
||||||
(ContinueGenerationReqInput, self.continue_generation),
|
(ContinueGenerationReqInput, self.continue_generation),
|
||||||
]
|
]
|
||||||
|
|||||||
@@ -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.kv_events import EventPublisherFactory, KVEventBatch
|
||||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||||
from sglang.srt.environ import envs
|
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.schedule_policy import PrefillAdder
|
||||||
from sglang.srt.managers.scheduler import Req, ScheduleBatch
|
from sglang.srt.managers.scheduler import Req, ScheduleBatch
|
||||||
from sglang.srt.managers.utils import GenerationBatchResult
|
from sglang.srt.managers.utils import GenerationBatchResult
|
||||||
@@ -584,6 +594,144 @@ class SchedulerMetricsMixin:
|
|||||||
ts_tic=time.perf_counter(),
|
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
|
@contextmanager
|
||||||
def record_forward_metrics(self: Scheduler, batch: ScheduleBatch):
|
def record_forward_metrics(self: Scheduler, batch: ScheduleBatch):
|
||||||
if not (self.enable_metrics and ENABLE_METRICS_DEVICE_TIMER):
|
if not (self.enable_metrics and ENABLE_METRICS_DEVICE_TIMER):
|
||||||
|
|||||||
@@ -39,6 +39,8 @@ from sglang.srt.managers.io_struct import (
|
|||||||
GetInternalStateReqOutput,
|
GetInternalStateReqOutput,
|
||||||
GetLoadReqInput,
|
GetLoadReqInput,
|
||||||
GetLoadReqOutput,
|
GetLoadReqOutput,
|
||||||
|
GetLoadsReqInput,
|
||||||
|
GetLoadsReqOutput,
|
||||||
GetWeightsByNameReqInput,
|
GetWeightsByNameReqInput,
|
||||||
GetWeightsByNameReqOutput,
|
GetWeightsByNameReqOutput,
|
||||||
InitWeightsSendGroupForRemoteInstanceReqInput,
|
InitWeightsSendGroupForRemoteInstanceReqInput,
|
||||||
@@ -218,6 +220,9 @@ class TokenizerCommunicatorMixin:
|
|||||||
self.get_load_communicator = _Communicator(
|
self.get_load_communicator = _Communicator(
|
||||||
self.send_to_scheduler, server_args.dp_size, mode="watching"
|
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()
|
self._result_dispatcher += self._get_communicator_dispatcher()
|
||||||
|
|
||||||
@@ -304,6 +309,10 @@ class TokenizerCommunicatorMixin:
|
|||||||
GetLoadReqOutput,
|
GetLoadReqOutput,
|
||||||
self.get_load_communicator.handle_recv,
|
self.get_load_communicator.handle_recv,
|
||||||
),
|
),
|
||||||
|
(
|
||||||
|
GetLoadsReqOutput,
|
||||||
|
self.get_loads_communicator.handle_recv,
|
||||||
|
),
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -787,6 +796,33 @@ class TokenizerCommunicatorMixin:
|
|||||||
req = GetLoadReqInput()
|
req = GetLoadReqInput()
|
||||||
return await self.get_load_communicator(req)
|
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(
|
async def open_session(
|
||||||
self, obj: OpenSessionReqInput, request: Optional[fastapi.Request] = None
|
self, obj: OpenSessionReqInput, request: Optional[fastapi.Request] = None
|
||||||
):
|
):
|
||||||
|
|||||||
Reference in New Issue
Block a user