Merge /get_load into /v1/loads (#23010)

This commit is contained in:
Liangsheng Yin
2026-04-17 13:36:51 -07:00
committed by GitHub
parent 44e67c6835
commit 573e12a7fc
13 changed files with 235 additions and 84 deletions
+19 -2
View File
@@ -647,12 +647,29 @@ async def server_info():
@app.get("/get_load")
async def get_load():
"""Get load metrics (deprecated - use /v1/loads instead)."""
"""Get load metrics (deprecated - use /v1/loads instead).
Legacy shim backed by /v1/loads. Projects GetLoadsReqOutput down to the
historical field shape (dp_rank, num_reqs, num_waiting_reqs, num_tokens,
num_pending_tokens, ts_tic) so existing clients keep working.
"""
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()
load_results = await _global_state.tokenizer_manager.get_loads(include=["core"])
ts = time.perf_counter()
return [
{
"dp_rank": r.dp_rank,
"num_reqs": r.num_running_reqs + r.num_waiting_reqs,
"num_waiting_reqs": r.num_waiting_reqs,
"num_tokens": r.num_total_tokens,
"num_pending_tokens": r.num_total_tokens - r.num_used_tokens,
"ts_tic": ts,
}
for r in load_results
]
# example usage:
@@ -65,6 +65,8 @@ def _compute_aggregate(load_dicts: list) -> dict:
"total_running_reqs": 0,
"total_waiting_reqs": 0,
"total_reqs": 0,
"total_used_tokens": 0,
"total_tokens": 0,
"avg_token_usage": 0.0,
"avg_throughput": 0.0,
"avg_utilization": 0.0,
@@ -77,6 +79,8 @@ def _compute_aggregate(load_dicts: list) -> dict:
"total_reqs": sum(
d["num_running_reqs"] + d["num_waiting_reqs"] for d in load_dicts
),
"total_used_tokens": sum(d["num_used_tokens"] for d in load_dicts),
"total_tokens": sum(d["num_total_tokens"] 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),
@@ -93,8 +93,10 @@ class DPBudget:
def update_budget(self, load_update: WatchLoadUpdateReq):
"""Update the budget."""
for load in load_update.loads:
self.total_requests[load.dp_rank] = load.num_reqs
self.total_tokens[load.dp_rank] = load.num_tokens
self.total_requests[load.dp_rank] = (
load.num_running_reqs + load.num_waiting_reqs
)
self.total_tokens[load.dp_rank] = load.num_total_tokens
def dispatch(self, method: LoadBalanceMethod):
if method == LoadBalanceMethod.TOTAL_REQUESTS:
+8 -18
View File
@@ -1087,7 +1087,7 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
token_steps: List[List[int]] = None
# Load for DP balance
load: GetLoadReqOutput = None
load: GetLoadsReqOutput = None
# Customized info
customized_info: Optional[Dict[str, List[Any]]] = None
# Detailed breakdown of cached tokens by source (device/host/storage)
@@ -1149,7 +1149,7 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
token_steps: List[List[int]] = None
# Load for DP balance
load: GetLoadReqOutput = None
load: GetLoadsReqOutput = None
# Customized info
customized_info: Optional[Dict[str, List[Any]]] = None
@@ -1860,21 +1860,6 @@ class BlockReqInput(BaseReq):
type: BlockReqType
@dataclass
class GetLoadReqInput(BaseReq):
pass
@dataclass
class GetLoadReqOutput(BaseReq):
dp_rank: int
num_reqs: int
num_waiting_reqs: int
num_tokens: int
num_pending_tokens: int
ts_tic: float
@dataclass
class MemoryMetrics:
"""Memory breakdown metrics."""
@@ -1992,6 +1977,11 @@ class GetLoadsReqOutput(BaseReq):
num_used_tokens: int = field(
metadata={"metric": ("gauge", "Number of tokens in use")}
)
# num_used_tokens + pending prefill tokens (waiting-queue seqlen, incl.
# disagg bootstrap/prealloc/transfer queues). Used for DP balance.
num_total_tokens: int = field(
metadata={"metric": ("gauge", "Used tokens plus pending prefill tokens")}
)
max_total_num_tokens: int = field(
metadata={"metric": ("gauge", "Maximum token capacity")}
)
@@ -2020,7 +2010,7 @@ class GetLoadsReqOutput(BaseReq):
@dataclass
class WatchLoadUpdateReq(BaseReq):
loads: List[GetLoadReqOutput]
loads: List[GetLoadsReqOutput]
@dataclass
+1 -3
View File
@@ -109,7 +109,6 @@ from sglang.srt.managers.io_struct import (
FreezeGCReq,
GetInternalStateReq,
GetInternalStateReqOutput,
GetLoadReqInput,
GetLoadsReqInput,
GetWeightsByNameReqInput,
HealthCheckOutput,
@@ -1321,7 +1320,6 @@ class Scheduler(
self.load_lora_adapter_from_tensors,
),
(UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
(GetLoadReqInput, self.get_load),
(GetLoadsReqInput, self.get_loads),
(PauseGenerationReqInput, self.pause_generation),
(ContinueGenerationReqInput, self.continue_generation),
@@ -2360,7 +2358,7 @@ class Scheduler(
# For prefill-only batch, filter out finished requests since they
# won't go through the decode step. This keeps running_batch accurate
# for load reporting (num_running_reqs via /get_load).
# for load reporting (num_running_reqs via /v1/loads).
# Runs outside the last_batch block so stale requests are cleaned
# even when no new batches arrive (e.g. traffic stops).
if self.running_batch.is_prefill_only:
@@ -13,6 +13,7 @@ from sglang.srt.managers.io_struct import (
AbortReq,
BatchEmbeddingOutput,
BatchTokenIDOutput,
GetLoadsReqInput,
)
from sglang.srt.managers.schedule_batch import (
BaseFinishReason,
@@ -964,7 +965,7 @@ class SchedulerOutputProcessorMixin:
spec_acceptance_histogram = []
retraction_counts = []
output_hidden_states = None
load = self.get_load()
load = self.get_loads(GetLoadsReqInput(include=["core"]))
routed_experts = None
customized_info = {}
@@ -39,8 +39,6 @@ from sglang.srt.managers.io_struct import (
FlushCacheReqOutput,
GetInternalStateReq,
GetInternalStateReqOutput,
GetLoadReqInput,
GetLoadReqOutput,
GetLoadsReqInput,
GetLoadsReqOutput,
GetWeightsByNameReqInput,
@@ -121,7 +119,6 @@ _COMMUNICATOR_SPECS = [
("set_internal_state", SetInternalStateReqOutput),
("expert_distribution", ExpertDistributionReqOutput),
("update_lora_adapter", LoRAUpdateOutput),
("get_load", GetLoadReqOutput, "watching"),
("get_loads", GetLoadsReqOutput, "watching"),
("dumper_control", DumperControlReqOutput),
]
@@ -804,11 +801,6 @@ class TokenizerControlMixin:
self.auto_create_handle_loop()
return await self.dumper_control_communicator(obj)
async def get_load(self: TokenizerManager) -> List[GetLoadReqOutput]:
self.auto_create_handle_loop()
req = GetLoadReqInput()
return await self.get_load_communicator(req)
async def get_loads(
self: TokenizerManager,
include: Optional[List[str]] = None,
@@ -12,8 +12,6 @@ from sglang.srt.disaggregation.utils import DisaggregationMode
from sglang.srt.environ import envs
from sglang.srt.managers.io_struct import (
DisaggregationMetrics,
GetLoadReqInput,
GetLoadReqOutput,
GetLoadsReqInput,
GetLoadsReqOutput,
LoRAMetrics,
@@ -776,11 +774,11 @@ class SchedulerMetricsMixin:
Args:
chunk_deduct: extra tokens to subtract from the chunked request's
remaining count. At batch-scheduling time the current chunk
remaining count. At batch-scheduling time the current chunk
has been planned but ``prefix_indices`` does not yet include it,
so callers pass ``extend_input_len`` here. At query time
(``get_load``) ``prefix_indices`` is already up-to-date, so
the default 0 is correct.
so callers pass ``extend_input_len`` here. At load-reporting
time ``prefix_indices`` is already up-to-date, so the default
0 is correct.
"""
num_pending_tokens = sum(req.seqlen for req in self.waiting_queue)
if self.chunked_req is not None:
@@ -788,31 +786,6 @@ class SchedulerMetricsMixin:
num_pending_tokens += req.seqlen - len(req.prefix_indices) - chunk_deduct
return num_pending_tokens
def get_load(self: Scheduler, _: GetLoadReqInput = None) -> GetLoadReqOutput:
num_tokens, _ = self.get_pool_stats().get_kv_token_stats()
num_pending_tokens = self._get_num_pending_tokens()
# Tokens and request count in waiting queue, bootstrap queue, prealloc queue
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_tokens += sum(req.seqlen for queue in waiting_queues for req in queue)
num_waiting_reqs = sum(len(queue) for queue in waiting_queues)
return GetLoadReqOutput(
dp_rank=self.dp_rank,
num_reqs=len(self.running_batch.reqs) + num_waiting_reqs,
num_waiting_reqs=num_waiting_reqs,
num_tokens=num_tokens,
num_pending_tokens=num_pending_tokens,
ts_tic=time.perf_counter(),
)
def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput:
"""
Get comprehensive load metrics for /v1/loads endpoint.
@@ -841,6 +814,9 @@ class SchedulerMetricsMixin:
num_waiting_reqs = sum(len(queue) for queue in waiting_queues)
num_used_tokens, kv_token_usage = self.get_pool_stats().get_kv_token_stats()
num_total_tokens = num_used_tokens + sum(
req.seqlen for queue in waiting_queues for req in queue
)
memory = None
if include_all or "memory" in include:
@@ -925,6 +901,7 @@ class SchedulerMetricsMixin:
num_running_reqs=num_running_reqs,
num_waiting_reqs=num_waiting_reqs,
num_used_tokens=num_used_tokens,
num_total_tokens=num_total_tokens,
max_total_num_tokens=self.max_total_num_tokens,
token_usage=round(kv_token_usage, 4),
gen_throughput=round(self.stats.gen_throughput, 2),