Remove stale load collection from output streaming hot path (#28408)
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
@@ -1179,8 +1179,6 @@ class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
|
|||||||
# The trainer step id. Used to know which step's weights are used for sampling.
|
# The trainer step id. Used to know which step's weights are used for sampling.
|
||||||
token_steps: List[List[int]] = None
|
token_steps: List[List[int]] = None
|
||||||
|
|
||||||
# Load for DP balance
|
|
||||||
load: GetLoadsReqOutput = None
|
|
||||||
# Customized info
|
# Customized info
|
||||||
customized_info: Optional[Dict[str, List[Any]]] = None
|
customized_info: Optional[Dict[str, List[Any]]] = None
|
||||||
# Detailed breakdown of cached tokens by source (device/host/storage)
|
# Detailed breakdown of cached tokens by source (device/host/storage)
|
||||||
@@ -1249,9 +1247,6 @@ class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
|
|||||||
# The trainer step id. Used to know which step's weights are used for sampling.
|
# The trainer step id. Used to know which step's weights are used for sampling.
|
||||||
token_steps: List[List[int]] = None
|
token_steps: List[List[int]] = None
|
||||||
|
|
||||||
# Load for DP balance
|
|
||||||
load: GetLoadsReqOutput = None
|
|
||||||
|
|
||||||
# Customized info
|
# Customized info
|
||||||
customized_info: Optional[Dict[str, List[Any]]] = None
|
customized_info: Optional[Dict[str, List[Any]]] = None
|
||||||
# Detailed breakdown of cached tokens by source (device/host/storage)
|
# Detailed breakdown of cached tokens by source (device/host/storage)
|
||||||
@@ -2156,11 +2151,6 @@ class GetLoadsReqOutput(BaseReq):
|
|||||||
queues: Optional[QueueMetrics] = None
|
queues: Optional[QueueMetrics] = None
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class WatchLoadUpdateReq(BaseReq):
|
|
||||||
loads: List[GetLoadsReqOutput]
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class SetInjectDumpMetadataReqInput(BaseReq):
|
class SetInjectDumpMetadataReqInput(BaseReq):
|
||||||
dump_metadata: Dict[str, Any]
|
dump_metadata: Dict[str, Any]
|
||||||
|
|||||||
@@ -1786,7 +1786,6 @@ class Scheduler(
|
|||||||
spec_algorithm=self.spec_algorithm,
|
spec_algorithm=self.spec_algorithm,
|
||||||
disaggregation_mode=self.disaggregation_mode,
|
disaggregation_mode=self.disaggregation_mode,
|
||||||
enable_hicache_storage=lambda: self.enable_hicache_storage,
|
enable_hicache_storage=lambda: self.enable_hicache_storage,
|
||||||
load_inquirer_get_loads=lambda req: self.load_inquirer.get_loads(req),
|
|
||||||
)
|
)
|
||||||
|
|
||||||
def init_batch_result_processor(self) -> None:
|
def init_batch_result_processor(self) -> None:
|
||||||
|
|||||||
@@ -18,7 +18,6 @@ from sglang.srt.environ import envs
|
|||||||
from sglang.srt.managers.io_struct import (
|
from sglang.srt.managers.io_struct import (
|
||||||
BatchEmbeddingOutput,
|
BatchEmbeddingOutput,
|
||||||
BatchTokenIDOutput,
|
BatchTokenIDOutput,
|
||||||
GetLoadsReqInput,
|
|
||||||
)
|
)
|
||||||
from sglang.srt.managers.schedule_batch import (
|
from sglang.srt.managers.schedule_batch import (
|
||||||
BaseFinishReason,
|
BaseFinishReason,
|
||||||
@@ -44,7 +43,6 @@ class SchedulerOutputStreamer:
|
|||||||
spec_algorithm: SpeculativeAlgorithm
|
spec_algorithm: SpeculativeAlgorithm
|
||||||
disaggregation_mode: DisaggregationMode
|
disaggregation_mode: DisaggregationMode
|
||||||
enable_hicache_storage: Callable[[], bool]
|
enable_hicache_storage: Callable[[], bool]
|
||||||
load_inquirer_get_loads: Callable[..., Any]
|
|
||||||
_test_stream_output_count: int = 0
|
_test_stream_output_count: int = 0
|
||||||
|
|
||||||
def _get_storage_backend_type(self) -> str:
|
def _get_storage_backend_type(self) -> str:
|
||||||
@@ -144,8 +142,6 @@ class SchedulerOutputStreamer:
|
|||||||
default_force_stream_interval=DEFAULT_FORCE_STREAM_INTERVAL,
|
default_force_stream_interval=DEFAULT_FORCE_STREAM_INTERVAL,
|
||||||
get_cached_tokens_details=self.get_cached_tokens_details,
|
get_cached_tokens_details=self.get_cached_tokens_details,
|
||||||
)
|
)
|
||||||
load = self.load_inquirer_get_loads(GetLoadsReqInput(include=["core"]))
|
|
||||||
|
|
||||||
for req in reqs:
|
for req in reqs:
|
||||||
if req is skip_req:
|
if req is skip_req:
|
||||||
continue
|
continue
|
||||||
@@ -159,7 +155,6 @@ class SchedulerOutputStreamer:
|
|||||||
|
|
||||||
# Send to detokenizer
|
# Send to detokenizer
|
||||||
payload = acc.to_payload(
|
payload = acc.to_payload(
|
||||||
load=load,
|
|
||||||
dp_rank=self.ps.dp_rank,
|
dp_rank=self.ps.dp_rank,
|
||||||
is_idle_batch=is_idle_batch,
|
is_idle_batch=is_idle_batch,
|
||||||
has_reqs=bool(reqs),
|
has_reqs=bool(reqs),
|
||||||
@@ -498,7 +493,7 @@ class _GenerationStreamAccumulator:
|
|||||||
self.customized_info[k].append(v[send_token_offset : len(output_ids_)])
|
self.customized_info[k].append(v[send_token_offset : len(output_ids_)])
|
||||||
|
|
||||||
def to_payload(
|
def to_payload(
|
||||||
self, *, load, dp_rank: int, is_idle_batch: bool, has_reqs: bool
|
self, *, dp_rank: int, is_idle_batch: bool, has_reqs: bool
|
||||||
) -> Optional[BatchTokenIDOutput]:
|
) -> Optional[BatchTokenIDOutput]:
|
||||||
if not (has_reqs or is_idle_batch):
|
if not (has_reqs or is_idle_batch):
|
||||||
return None
|
return None
|
||||||
@@ -546,6 +541,5 @@ class _GenerationStreamAccumulator:
|
|||||||
placeholder_tokens_idx=None,
|
placeholder_tokens_idx=None,
|
||||||
placeholder_tokens_val=None,
|
placeholder_tokens_val=None,
|
||||||
retraction_counts=self.retraction_counts,
|
retraction_counts=self.retraction_counts,
|
||||||
load=load,
|
|
||||||
dp_ranks=dp_ranks,
|
dp_ranks=dp_ranks,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -1854,19 +1854,6 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
|
|||||||
"num_retractions": recv_obj.retraction_counts[i],
|
"num_retractions": recv_obj.retraction_counts[i],
|
||||||
}
|
}
|
||||||
|
|
||||||
# Surface scheduler load info on each response so clients can do
|
|
||||||
# response-based flow control without polling /v1/loads. The
|
|
||||||
# scheduler already piggy-backs the per-DP-rank load on
|
|
||||||
# BatchStrOutput / BatchTokenIDOutput via the ``load`` field.
|
|
||||||
load = getattr(recv_obj, "load", None)
|
|
||||||
if load is not None:
|
|
||||||
num_running_reqs = getattr(load, "num_running_reqs", None)
|
|
||||||
num_waiting_reqs = getattr(load, "num_waiting_reqs", None)
|
|
||||||
if num_running_reqs is not None:
|
|
||||||
meta_info["num_running_reqs"] = num_running_reqs
|
|
||||||
if num_waiting_reqs is not None:
|
|
||||||
meta_info["num_waiting_reqs"] = num_waiting_reqs
|
|
||||||
|
|
||||||
if self.enable_metrics:
|
if self.enable_metrics:
|
||||||
if recv_obj.time_stats is not None:
|
if recv_obj.time_stats is not None:
|
||||||
scheduler_time_stats = recv_obj.time_stats[i]
|
scheduler_time_stats = recv_obj.time_stats[i]
|
||||||
|
|||||||
Reference in New Issue
Block a user