Merge /get_load into /v1/loads (#23010)
This commit is contained in:
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
Reference in New Issue
Block a user