From 19593359971ebc3582a74f000bf285488d993362 Mon Sep 17 00:00:00 2001 From: cctry Date: Thu, 9 Jul 2026 05:58:10 -0700 Subject: [PATCH] refactor(load-snapshot): build LoadSnapshot directly, drop legacy get_loads IPC (#30525) Co-authored-by: cctry --- python/sglang/srt/entrypoints/http_server.py | 2 +- python/sglang/srt/managers/io_struct.py | 96 -------- python/sglang/srt/managers/load_snapshot.py | 230 ++++++------------ python/sglang/srt/managers/scheduler.py | 10 +- .../scheduler_components/load_inquirer.py | 167 ++++++------- .../srt/managers/tokenizer_control_mixin.py | 2 - .../entrypoints/test_v1_loads_aggregate.py | 14 +- .../utils/test_type_based_dispatcher.py | 3 - 8 files changed, 153 insertions(+), 371 deletions(-) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index ce061978a..fe6f26db6 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -751,7 +751,7 @@ async def server_info(): async def get_load(): """Get load metrics (deprecated - use /v1/loads instead). - Legacy shim backed by /v1/loads. Projects GetLoadsReqOutput down to the + Legacy shim backed by /v1/loads. Projects the load snapshot 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. """ diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 8435da43c..60985e55e 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1980,102 +1980,6 @@ class BlockReqInput(BaseReq, kw_only=True): req_type: BlockReqType -class MemoryMetrics(msgspec.Struct, array_like=True): - """Memory breakdown metrics.""" - - weight_gb: float - kv_cache_gb: float - graph_gb: float - token_capacity: int - - -class SpeculativeMetrics(msgspec.Struct, array_like=True): - """Speculative decoding metrics.""" - - accept_length: float - accept_rate: float - - -class LoRAMetrics(msgspec.Struct, array_like=True): - """LoRA adapter pool metrics.""" - - slots_used: int - slots_total: int - utilization: float - - -class DisaggregationMetrics(msgspec.Struct, array_like=True): - """PD disaggregation metrics.""" - - mode: str # "prefill", "decode", or "null" - prefill_bootstrap_queue_reqs: int = 0 - prefill_inflight_queue_reqs: int = 0 - decode_prealloc_queue_reqs: int = 0 - decode_transfer_queue_reqs: int = 0 - decode_retracted_queue_reqs: int = 0 - kv_transfer_speed_gb_s: float = 0.0 - kv_transfer_latency_ms: float = 0.0 - - -class QueueMetrics(msgspec.Struct, array_like=True): - """Detailed queue info breakdown.""" - - waiting: int - grammar: int - paused: int - retracted: int - - -class GetLoadsReqInput(BaseReq, kw_only=True): - """Request for /v1/loads endpoint.""" - - VALID_SECTIONS = frozenset( - {"core", "memory", "spec", "lora", "disagg", "queues", "all"} - ) - - include: List[str] = msgspec.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)}" - ) - - -class GetLoadsReqOutput(BaseReq, kw_only=True): - """Per-DP-rank load metrics for /v1/loads endpoint.""" - - dp_rank: int - timestamp: float - - num_running_reqs: int - num_waiting_reqs: int - num_waiting_uncached_tokens: int - num_used_tokens: int - # num_used_tokens plus pending tokens not already allocated in the KV pool. - # Used for DP balance. - num_total_tokens: int - max_total_num_tokens: int - # FIXME: token_usage is actually max usage across all pools (KV, SWA, mamba), - # not just KV token usage. Rename requires API deprecation. - token_usage: float - gen_throughput: float - cache_hit_rate: float - utilization: float - max_running_requests: int - - memory: Optional[MemoryMetrics] = None - speculative: Optional[SpeculativeMetrics] = None - lora: Optional[LoRAMetrics] = None - disaggregation: Optional[DisaggregationMetrics] = None - queues: Optional[QueueMetrics] = None - - class SetInjectDumpMetadataReqInput(BaseReq, kw_only=True): dump_metadata: Dict[str, Any] diff --git a/python/sglang/srt/managers/load_snapshot.py b/python/sglang/srt/managers/load_snapshot.py index ff0ce49bb..fd0f53246 100644 --- a/python/sglang/srt/managers/load_snapshot.py +++ b/python/sglang/srt/managers/load_snapshot.py @@ -45,7 +45,7 @@ import mmap import os import struct from contextlib import contextmanager -from typing import TYPE_CHECKING, Optional +from typing import Optional import msgspec import msgspec.msgpack @@ -54,25 +54,12 @@ import msgspec.structs from sglang.srt.environ import envs from sglang.srt.utils.network import is_zmq_endpoint_ipv6 -if TYPE_CHECKING: - from sglang.srt.managers.io_struct import GetLoadsReqOutput - logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- -DISAGG_MODE_TO_INT = {"null": 0, "prefill": 1, "decode": 2} -INT_TO_DISAGG_MODE = {v: k for k, v in DISAGG_MODE_TO_INT.items()} - - -def _native(v): - """Coerce numpy scalars to Python int/float for msgpack encoding.""" - if hasattr(v, "item"): - return v.item() - return v - def should_use_zmq(server_args) -> bool: """Whether to use zmq PUSH/PULL instead of shared memory for load snapshots. @@ -139,7 +126,54 @@ def zmq_reader_owner(server_args, caller: str) -> bool: # LoadSnapshot data class # --------------------------------------------------------------------------- -CORE_METRIC_FIELDS = ( + +class MemoryMetrics(msgspec.Struct, array_like=True): + """Memory breakdown metrics.""" + + weight_gb: float + kv_cache_gb: float + graph_gb: float + token_capacity: int + + +class SpeculativeMetrics(msgspec.Struct, array_like=True): + """Speculative decoding metrics.""" + + accept_length: float + accept_rate: float + + +class LoRAMetrics(msgspec.Struct, array_like=True): + """LoRA adapter pool metrics.""" + + slots_used: int + slots_total: int + utilization: float + + +class DisaggregationMetrics(msgspec.Struct, array_like=True): + """PD disaggregation metrics.""" + + mode: str # "prefill", "decode", or "null" + prefill_bootstrap_queue_reqs: int = 0 + prefill_inflight_queue_reqs: int = 0 + decode_prealloc_queue_reqs: int = 0 + decode_transfer_queue_reqs: int = 0 + decode_retracted_queue_reqs: int = 0 + kv_transfer_speed_gb_s: float = 0.0 + kv_transfer_latency_ms: float = 0.0 + + +class QueueMetrics(msgspec.Struct, array_like=True): + """Detailed queue info breakdown.""" + + waiting: int + grammar: int + paused: int + retracted: int + + +_CORE_KEYS = ( "timestamp", "dp_rank", "num_running_reqs", @@ -154,67 +188,11 @@ CORE_METRIC_FIELDS = ( "cache_hit_rate", "utilization", ) -SECTION_FIELDS = ( - ( - "memory", - "memory", - "has_memory", - ( - ("weight_gb", "memory_weight_gb"), - ("kv_cache_gb", "memory_kv_cache_gb"), - ("graph_gb", "memory_graph_gb"), - ("token_capacity", "memory_token_capacity"), - ), - ), - ( - "spec", - "speculative", - "has_speculative", - ( - ("accept_length", "speculative_accept_length"), - ("accept_rate", "speculative_accept_rate"), - ), - ), - ( - "lora", - "lora", - "has_lora", - ( - ("slots_used", "lora_slots_used"), - ("slots_total", "lora_slots_total"), - ("utilization", "lora_utilization"), - ), - ), - ( - "disagg", - "disaggregation", - "has_disaggregation", - ( - ("mode", "disagg_mode"), - ("prefill_bootstrap_queue_reqs", "prefill_bootstrap_queue_reqs"), - ("prefill_inflight_queue_reqs", "prefill_inflight_queue_reqs"), - ("decode_prealloc_queue_reqs", "decode_prealloc_queue_reqs"), - ("decode_transfer_queue_reqs", "decode_transfer_queue_reqs"), - ("decode_retracted_queue_reqs", "decode_retracted_queue_reqs"), - ("kv_transfer_speed_gb_s", "kv_transfer_speed_gb_s"), - ("kv_transfer_latency_ms", "kv_transfer_latency_ms"), - ), - ), - ( - "queues", - "queues", - "has_queues", - ( - ("waiting", "queue_waiting"), - ("grammar", "queue_grammar"), - ("paused", "queue_paused"), - ("retracted", "queue_retracted"), - ), - ), -) class LoadSnapshot(msgspec.Struct, omit_defaults=True): + """Per-DP-rank load metrics: the SHM/zmq wire format and the /v1/loads source.""" + timestamp: float = 0.0 dp_rank: int = 0 num_running_reqs: int = 0 @@ -229,81 +207,18 @@ class LoadSnapshot(msgspec.Struct, omit_defaults=True): cache_hit_rate: float = 0.0 utilization: float = 0.0 - has_memory: int = 0 - memory_weight_gb: float = 0.0 - memory_kv_cache_gb: float = 0.0 - memory_graph_gb: float = 0.0 - memory_token_capacity: int = 0 - - has_speculative: int = 0 - speculative_accept_length: float = 0.0 - speculative_accept_rate: float = 0.0 - - has_lora: int = 0 - lora_slots_used: int = 0 - lora_slots_total: int = 0 - lora_utilization: float = 0.0 - - has_disaggregation: int = 0 - disagg_mode: int = 0 - prefill_bootstrap_queue_reqs: int = 0 - prefill_inflight_queue_reqs: int = 0 - decode_prealloc_queue_reqs: int = 0 - decode_transfer_queue_reqs: int = 0 - decode_retracted_queue_reqs: int = 0 - kv_transfer_speed_gb_s: float = 0.0 - kv_transfer_latency_ms: float = 0.0 - - has_queues: int = 0 - queue_waiting: int = 0 - queue_grammar: int = 0 - queue_paused: int = 0 - queue_retracted: int = 0 - - @classmethod - def from_get_loads_output(cls, output: GetLoadsReqOutput) -> LoadSnapshot: - snapshot: dict = {} - for name in CORE_METRIC_FIELDS: - value = getattr(output, name) - if name == "dp_rank": - snapshot[name] = int(value) if value is not None else 0 - else: - snapshot[name] = _native(value) - - for _, section_name, present_attr, attrs in SECTION_FIELDS: - section = getattr(output, section_name, None) - snapshot[present_attr] = int(section is not None) - if section is None: - continue - for section_attr, snapshot_attr in attrs: - value = getattr(section, section_attr) - if snapshot_attr == "disagg_mode": - value = DISAGG_MODE_TO_INT.get(value, 0) - else: - value = _native(value) - snapshot[snapshot_attr] = value - - return cls(**snapshot) + memory: Optional[MemoryMetrics] = None + speculative: Optional[SpeculativeMetrics] = None + lora: Optional[LoRAMetrics] = None + disaggregation: Optional[DisaggregationMetrics] = None + queues: Optional[QueueMetrics] = None VALID_SECTIONS = frozenset( {"core", "memory", "spec", "lora", "disagg", "queues", "all"} ) def to_dict(self, include: Optional[set[str]] = None) -> dict: - load = { - "dp_rank": self.dp_rank, - "num_running_reqs": self.num_running_reqs, - "num_waiting_reqs": self.num_waiting_reqs, - "num_waiting_uncached_tokens": self.num_waiting_uncached_tokens, - "num_used_tokens": self.num_used_tokens, - "num_total_tokens": self.num_total_tokens, - "max_total_num_tokens": self.max_total_num_tokens, - "max_running_requests": self.max_running_requests, - "token_usage": self.token_usage, - "gen_throughput": self.gen_throughput, - "cache_hit_rate": self.cache_hit_rate, - "utilization": self.utilization, - } + load = {key: getattr(self, key) for key in _CORE_KEYS} if include is None or "all" in include: include_all = True @@ -317,24 +232,29 @@ class LoadSnapshot(msgspec.Struct, omit_defaults=True): return load include_all = False - for include_key, section_name, present_attr, attrs in SECTION_FIELDS: - if not getattr(self, present_attr): + for field, include_name, section in ( + ("memory", "memory", self.memory), + ("speculative", "spec", self.speculative), + ("lora", "lora", self.lora), + ("disaggregation", "disagg", self.disaggregation), + ("queues", "queues", self.queues), + ): + if section is None or (not include_all and include_name not in include): continue - if not include_all and include_key not in include: - continue - - section = {} - for section_attr, snapshot_attr in attrs: - value = getattr(self, snapshot_attr) - if snapshot_attr == "disagg_mode": - value = INT_TO_DISAGG_MODE.get(value, "null") - section[section_attr] = value - load[section_name] = section + load[field] = msgspec.structs.asdict(section) return load -snapshot_encoder = msgspec.msgpack.Encoder() +def _enc_hook(obj): + """Coerce numpy scalars to native Python; msgpack has no numpy types.""" + to_item = getattr(obj, "item", None) + if to_item is not None: + return to_item() + raise NotImplementedError(f"cannot encode {type(obj).__name__} in load snapshot") + + +snapshot_encoder = msgspec.msgpack.Encoder(enc_hook=_enc_hook) snapshot_decoder = msgspec.msgpack.Decoder(LoadSnapshot) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 97d3a21df..a872dbbcb 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -109,7 +109,6 @@ from sglang.srt.managers.io_struct import ( FreezeGCReq, GetInternalStateReq, GetInternalStateReqOutput, - GetLoadsReqInput, GetWeightsByNameReqInput, HealthCheckOutput, InitWeightsSendGroupForRemoteInstanceReqInput, @@ -147,7 +146,7 @@ from sglang.srt.managers.io_struct import ( UpdateWeightsFromTensorReqInput, sock_send, ) -from sglang.srt.managers.load_snapshot import LoadSnapshot, create_load_snapshot_writer +from sglang.srt.managers.load_snapshot import create_load_snapshot_writer from sglang.srt.managers.min_free_slots_delayer import ( MinFreeSlotsDelayer, resolve_min_free_slots, @@ -653,14 +652,10 @@ class Scheduler( return writer.publish_counter = 0 try: - result = self.load_inquirer.get_loads(GetLoadsReqInput(include=["all"])) - writer.write(LoadSnapshot.from_get_loads_output(result)) + writer.write(self.load_inquirer.get_loads()) except Exception as e: logger.warning("load snapshot publish failed: %s", e) - def handle_get_loads_req(self, req: GetLoadsReqInput): - return self.load_inquirer.get_loads(req) - def init_tokenizer(self): server_args = self.server_args self.is_generation = self.model_config.is_generation @@ -1444,7 +1439,6 @@ class Scheduler( self.load_lora_adapter_from_tensors, ), (UnloadLoRAAdapterReqInput, self.unload_lora_adapter), - (GetLoadsReqInput, self.handle_get_loads_req), (PauseGenerationReqInput, self.pause_generation), (ContinueGenerationReqInput, self.continue_generation), (ConfigureLoggingReq, self.configure_logging), diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index b49409fc7..e8619a2f7 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/python/sglang/srt/managers/scheduler_components/load_inquirer.py @@ -6,10 +6,9 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Callable from sglang.srt.disaggregation.utils import DisaggregationMode -from sglang.srt.managers.io_struct import ( +from sglang.srt.managers.load_snapshot import ( DisaggregationMetrics, - GetLoadsReqInput, - GetLoadsReqOutput, + LoadSnapshot, LoRAMetrics, MemoryMetrics, QueueMetrics, @@ -85,22 +84,9 @@ class SchedulerLoadInquirer: num_tokens += max(0, cr.seqlen - len(cr.prefix_indices)) return num_tokens - def get_loads(self, 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 - + def get_loads(self) -> LoadSnapshot: + """Build the per-DP-rank load snapshot for DP balancing and /v1/loads.""" + stats = self.get_stats() num_running_reqs = len(self.get_running_batch().reqs) waiting_queues = [self.get_waiting_queue()] @@ -124,7 +110,6 @@ class SchedulerLoadInquirer: pending_token_queues = [decode_prealloc_queue, decode_retracted_queue] num_waiting_reqs = sum(len(queue) for queue in waiting_queues) - num_waiting_uncached_tokens = self.get_num_waiting_uncached_tokens() num_used_tokens, kv_token_usage = ( self.pool_stats_observer.get_pool_stats().get_kv_token_stats() ) @@ -133,99 +118,85 @@ class SchedulerLoadInquirer: ) 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}") + 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, TypeError) 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.get_spec_total_num_forward_ct() > 0 - ): - speculative = SpeculativeMetrics( - accept_length=( - self.get_spec_total_num_accept_tokens() - / self.get_spec_total_num_forward_ct() - ), - accept_rate=self.get_stats().spec_accept_rate, - ) + if ( + not self.spec_algorithm.is_none() + and self.get_spec_total_num_forward_ct() > 0 + ): + speculative = SpeculativeMetrics( + accept_length=( + self.get_spec_total_num_accept_tokens() + / self.get_spec_total_num_forward_ct() + ), + accept_rate=stats.spec_accept_rate, + ) lora = None - if include_all or "lora" in include: - if self.server_args.enable_lora: - lora = LoRAMetrics( - slots_used=self.get_stats().lora_pool_slots_used, - slots_total=self.get_stats().lora_pool_slots_total, - utilization=self.get_stats().lora_pool_utilization, - ) - - disaggregation = None - if include_all or "disagg" in include: - mode_str = "null" - prefill_bootstrap = 0 - prefill_inflight = 0 - decode_prealloc = 0 - decode_transfer = 0 - decode_retracted = 0 - - if self.disaggregation_mode == DisaggregationMode.PREFILL: - mode_str = "prefill" - prefill_bootstrap = len(self.get_disagg_prefill_bootstrap_queue().queue) - prefill_inflight = len(self.get_disagg_prefill_inflight_queue()) - elif self.disaggregation_mode == DisaggregationMode.DECODE: - mode_str = "decode" - decode_prealloc = len(self.get_disagg_decode_prealloc_queue().queue) - decode_transfer = len(self.get_disagg_decode_transfer_queue().queue) - decode_retracted = len( - self.get_disagg_decode_prealloc_queue().retracted_queue - ) - - disaggregation = DisaggregationMetrics( - mode=mode_str, - prefill_bootstrap_queue_reqs=prefill_bootstrap, - 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.get_stats().kv_transfer_speed_gb_s, - kv_transfer_latency_ms=self.get_stats().kv_transfer_latency_ms, + if self.server_args.enable_lora: + lora = LoRAMetrics( + slots_used=stats.lora_pool_slots_used, + slots_total=stats.lora_pool_slots_total, + utilization=stats.lora_pool_utilization, ) - queues = None - if include_all or "queues" in include: - queues = QueueMetrics( - waiting=len(self.get_waiting_queue()), - grammar=self.get_stats().num_grammar_queue_reqs, - paused=self.get_stats().num_paused_reqs, - retracted=self.get_stats().num_retracted_reqs, + mode_str = "null" + prefill_bootstrap = prefill_inflight = 0 + decode_prealloc = decode_transfer = decode_retracted = 0 + if self.disaggregation_mode == DisaggregationMode.PREFILL: + mode_str = "prefill" + prefill_bootstrap = len(self.get_disagg_prefill_bootstrap_queue().queue) + prefill_inflight = len(self.get_disagg_prefill_inflight_queue()) + elif self.disaggregation_mode == DisaggregationMode.DECODE: + mode_str = "decode" + decode_prealloc = len(self.get_disagg_decode_prealloc_queue().queue) + decode_transfer = len(self.get_disagg_decode_transfer_queue().queue) + decode_retracted = len( + self.get_disagg_decode_prealloc_queue().retracted_queue ) + disaggregation = DisaggregationMetrics( + mode=mode_str, + prefill_bootstrap_queue_reqs=prefill_bootstrap, + 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=stats.kv_transfer_speed_gb_s, + kv_transfer_latency_ms=stats.kv_transfer_latency_ms, + ) - return GetLoadsReqOutput( - dp_rank=self.ps.dp_rank, + queues = QueueMetrics( + waiting=len(self.get_waiting_queue()), + grammar=stats.num_grammar_queue_reqs, + paused=stats.num_paused_reqs, + retracted=stats.num_retracted_reqs, + ) + + return LoadSnapshot( + dp_rank=int(self.ps.dp_rank) if self.ps.dp_rank is not None else 0, timestamp=time.time(), num_running_reqs=num_running_reqs, num_waiting_reqs=num_waiting_reqs, - num_waiting_uncached_tokens=num_waiting_uncached_tokens, + num_waiting_uncached_tokens=self.get_num_waiting_uncached_tokens(), 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.get_stats().gen_throughput, 2), - cache_hit_rate=round(self.get_stats().cache_hit_rate, 4), - utilization=round(self.get_stats().utilization, 4), max_running_requests=self.max_running_requests, + token_usage=round(kv_token_usage, 4), + gen_throughput=round(stats.gen_throughput, 2), + cache_hit_rate=round(stats.cache_hit_rate, 4), + utilization=round(stats.utilization, 4), memory=memory, speculative=speculative, lora=lora, diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py index dcc937c7b..24b608c0c 100644 --- a/python/sglang/srt/managers/tokenizer_control_mixin.py +++ b/python/sglang/srt/managers/tokenizer_control_mixin.py @@ -33,7 +33,6 @@ from sglang.srt.managers.io_struct import ( FlushCacheReqOutput, GetInternalStateReq, GetInternalStateReqOutput, - GetLoadsReqOutput, GetWeightsByNameReqInput, GetWeightsByNameReqOutput, InitWeightsSendGroupForRemoteInstanceReqInput, @@ -116,7 +115,6 @@ _COMMUNICATOR_SPECS = [ ("set_internal_state", SetInternalStateReqOutput), ("expert_distribution", ExpertDistributionReqOutput), ("update_lora_adapter", LoRAUpdateOutput), - ("get_loads", GetLoadsReqOutput, "watching"), ("dumper_control", DumperControlReqOutput), ] diff --git a/test/registered/unit/entrypoints/test_v1_loads_aggregate.py b/test/registered/unit/entrypoints/test_v1_loads_aggregate.py index 12967a42c..8b325c890 100644 --- a/test/registered/unit/entrypoints/test_v1_loads_aggregate.py +++ b/test/registered/unit/entrypoints/test_v1_loads_aggregate.py @@ -15,7 +15,9 @@ from sglang.srt.managers.load_snapshot import ( SLOT_LEN_STRUCT, SLOT_SIZE, VERSION, + DisaggregationMetrics, LoadSnapshot, + QueueMetrics, ShmLoadSnapshotReader, ShmLoadSnapshotWriter, slot_offset, @@ -150,14 +152,10 @@ class TestGetLoads(CustomTestCase): cache_hit_rate=0.75, utilization=0.5, max_running_requests=128, - has_disaggregation=1, - disagg_mode=2, - decode_transfer_queue_reqs=4, - has_queues=1, - queue_waiting=2, - queue_grammar=1, - queue_paused=0, - queue_retracted=3, + disaggregation=DisaggregationMetrics( + mode="decode", decode_transfer_queue_reqs=4 + ), + queues=QueueMetrics(waiting=2, grammar=1, paused=0, retracted=3), ) ) diff --git a/test/registered/utils/test_type_based_dispatcher.py b/test/registered/utils/test_type_based_dispatcher.py index a4c29a71a..911cda058 100644 --- a/test/registered/utils/test_type_based_dispatcher.py +++ b/test/registered/utils/test_type_based_dispatcher.py @@ -34,7 +34,6 @@ class TestTypeBasedDispatcher(unittest.TestCase): FlushCacheReqInput, FreezeGCReq, GetInternalStateReq, - GetLoadsReqInput, GetWeightsByNameReqInput, InitWeightsSendGroupForRemoteInstanceReqInput, InitWeightsUpdateGroupReqInput, @@ -114,7 +113,6 @@ class TestTypeBasedDispatcher(unittest.TestCase): (ExpertDistributionReq, lambda req: "expert_distribution_handled"), (LoadLoRAAdapterReqInput, lambda req: "load_lora_adapter_handled"), (UnloadLoRAAdapterReqInput, lambda req: "unload_lora_adapter_handled"), - (GetLoadsReqInput, lambda req: "get_loads_handled"), ] # Create requests that conforms to the real distribution @@ -209,7 +207,6 @@ class TestTypeBasedDispatcher(unittest.TestCase): test_requests.append(GetWeightsByNameReqInput(name="")) test_requests.append(ReleaseMemoryOccupationReqInput()) test_requests.append(RpcReqInput(method="")) - test_requests.append(GetLoadsReqInput()) dispatcher = TypeBasedDispatcher(mapping)