diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index 87a5f92e8..60bdcee7e 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -111,7 +111,7 @@ class SchedulerMetricsMixin: self.spec_num_forward_ct += bs self.num_generated_tokens += num_accepted_tokens - def reset_metrics(self): + def reset_metrics(self: Scheduler): self.forward_ct_decode = 0 self.num_generated_tokens = 0 self.spec_num_accepted_tokens = 0 @@ -512,7 +512,7 @@ class SchedulerMetricsMixin: except Exception as e: logger.warning(f"Failed to update LoRA metrics: {e}") - def calculate_utilization(self): + def calculate_utilization(self: Scheduler): if self.disaggregation_mode == DisaggregationMode.PREFILL: self.stats.utilization = -1 else: @@ -556,7 +556,7 @@ class SchedulerMetricsMixin: ) @contextmanager - def record_forward_metrics(self: Scheduler, batch): + def record_forward_metrics(self: Scheduler, batch: ScheduleBatch): if not (self.enable_metrics and ENABLE_METRICS_DEVICE_TIMER): yield return diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index d47a67c31..fffc49911 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -666,7 +666,7 @@ class SchedulerPPMixin: f"Target latency: {self.length_predictor.target_latency:.2f}ms" ) - def predict_next_chunk_size(self: "Scheduler", history_len: int) -> Optional[int]: + def predict_next_chunk_size(self: Scheduler, history_len: int) -> Optional[int]: """ Predict next chunk size dynamically based on current history length. diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py index 0aadc9fcb..d208c9f8c 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py @@ -2,7 +2,7 @@ import logging import os import time from pathlib import Path -from typing import List, Optional +from typing import TYPE_CHECKING, List, Optional import torch @@ -14,6 +14,10 @@ from sglang.srt.utils import is_npu from sglang.srt.utils.profile_merger import ProfileMerger from sglang.srt.utils.profile_utils import ProfileManager +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.managers.scheduler import Scheduler + _is_npu = is_npu() if _is_npu: import torch_npu @@ -29,7 +33,7 @@ logger = logging.getLogger(__name__) class SchedulerProfilerMixin: - def init_profiler(self): + def init_profiler(self: Scheduler): if envs.SGLANG_PROFILE_V2.get(): self._profile_manager = ProfileManager( tp_rank=self.tp_rank, @@ -59,7 +63,7 @@ class SchedulerProfilerMixin: self.rpd_profiler = None def init_profile( - self, + self: Scheduler, output_dir: Optional[str], start_step: Optional[int], num_steps: Optional[int], @@ -130,7 +134,7 @@ class SchedulerProfilerMixin: return ProfileReqOutput(success=True, message="Succeeded") def start_profile( - self, stage: Optional[ForwardMode] = None + self: Scheduler, stage: Optional[ForwardMode] = None ) -> ProfileReqOutput | None: if envs.SGLANG_PROFILE_V2.get(): return self._profile_manager.manual_start() @@ -208,7 +212,7 @@ class SchedulerProfilerMixin: return ProfileReqOutput(success=True, message="Succeeded") - def _merge_profile_traces(self) -> str: + def _merge_profile_traces(self: Scheduler) -> str: if not self.merge_profiles: return "" @@ -241,7 +245,7 @@ class SchedulerProfilerMixin: return merge_message def stop_profile( - self, stage: Optional[ForwardMode] = None + self: Scheduler, stage: Optional[ForwardMode] = None ) -> ProfileReqOutput | None: if envs.SGLANG_PROFILE_V2.get(): return self._profile_manager.manual_stop() @@ -328,7 +332,7 @@ class SchedulerProfilerMixin: return ProfileReqOutput(success=True, message=f"Succeeded.{merge_message}") - def _profile_batch_predicate(self, batch): + def _profile_batch_predicate(self: Scheduler, batch: ScheduleBatch): if envs.SGLANG_PROFILE_V2.get(): self._profile_manager.step(forward_mode=batch.forward_mode) return @@ -368,7 +372,7 @@ class SchedulerProfilerMixin: ): self.start_profile() - def profile(self, recv_req: ProfileReq): + def profile(self: Scheduler, recv_req: ProfileReq): if recv_req.type == ProfileReqType.START_PROFILE: if recv_req.profile_by_stage or recv_req.start_step: return self.init_profile( diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py index f8ebfc1f4..293a84350 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py @@ -43,7 +43,9 @@ logger = logging.getLogger(__name__) class SchedulerUpdateWeightsMixin: - def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): + def update_weights_from_disk( + self: Scheduler, recv_req: UpdateWeightFromDiskReqInput + ): """In-place update of the weights from disk.""" success, message = self.tp_worker.update_weights_from_disk(recv_req) if success: @@ -54,12 +56,16 @@ class SchedulerUpdateWeightsMixin: logger.error(message) return UpdateWeightFromDiskReqOutput(success, message, 0) - def init_weights_update_group(self, recv_req: InitWeightsUpdateGroupReqInput): + def init_weights_update_group( + self: Scheduler, recv_req: InitWeightsUpdateGroupReqInput + ): """Initialize the online model parameter update group.""" success, message = self.tp_worker.init_weights_update_group(recv_req) return InitWeightsUpdateGroupReqOutput(success, message) - def destroy_weights_update_group(self, recv_req: DestroyWeightsUpdateGroupReqInput): + def destroy_weights_update_group( + self: Scheduler, recv_req: DestroyWeightsUpdateGroupReqInput + ): """Destroy the online model parameter update group.""" success, message = self.tp_worker.destroy_weights_update_group(recv_req) return DestroyWeightsUpdateGroupReqOutput(success, message) @@ -78,7 +84,9 @@ class SchedulerUpdateWeightsMixin: logger.error(message) return UpdateWeightsFromDistributedReqOutput(success, message) - def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput): + def update_weights_from_tensor( + self: Scheduler, recv_req: UpdateWeightsFromTensorReqInput + ): """Update the online model parameter from tensors.""" worker = self.draft_worker or self.tp_worker success, message = worker.update_weights_from_tensor(recv_req) @@ -92,7 +100,9 @@ class SchedulerUpdateWeightsMixin: torch.distributed.barrier(group=self.tp_cpu_group) return UpdateWeightsFromTensorReqOutput(success, message) - def update_weights_from_ipc(self, recv_req: UpdateWeightsFromIPCReqInput): + def update_weights_from_ipc( + self: Scheduler, recv_req: UpdateWeightsFromIPCReqInput + ): """Update the online model parameter from IPC for checkpoint-engine integration.""" success, message = self.tp_worker.update_weights_from_ipc(recv_req) if success: @@ -104,7 +114,7 @@ class SchedulerUpdateWeightsMixin: torch.distributed.barrier(group=self.tp_cpu_group) return UpdateWeightsFromIPCReqOutput(success, message) - def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput): + def get_weights_by_name(self: Scheduler, recv_req: GetWeightsByNameReqInput): parameter = self.tp_worker.get_weights_by_name(recv_req) return GetWeightsByNameReqOutput(parameter) diff --git a/python/sglang/srt/multiplex/multiplexing_mixin.py b/python/sglang/srt/multiplex/multiplexing_mixin.py index e328b8186..5ee6add1e 100644 --- a/python/sglang/srt/multiplex/multiplexing_mixin.py +++ b/python/sglang/srt/multiplex/multiplexing_mixin.py @@ -3,6 +3,7 @@ Mixin class providing multiplexing scheduling logic """ import logging +from typing import TYPE_CHECKING import torch import torch.distributed as dist @@ -19,12 +20,15 @@ from sglang.srt.multiplex.pdmux_context import ( set_current_stream_idx, ) +if TYPE_CHECKING: + from sglang.srt.managers.scheduler import Scheduler + logger = logging.getLogger(__name__) class SchedulerMultiplexMixin: - def init_pdmux(self): + def init_pdmux(self: Scheduler): # for pd_multiplexing, Init stream_groups, exclude normal stream for prefill only and decode only self.pdmux_config = load_pdmux_config(self.server_args.pdmux_config_path) initialize_stream_groups(self.gpu_id, self.pdmux_config) @@ -36,7 +40,9 @@ class SchedulerMultiplexMixin: ) # TODO(jason-fxz): This is a temporary demo - def adjust_stream_groups(self) -> tuple[int, tuple[ExternalStream, ExternalStream]]: + def adjust_stream_groups( + self: Scheduler, + ) -> tuple[int, tuple[ExternalStream, ExternalStream]]: if not self.running_batch.is_empty() and self.split_prefill_batch: decode_bs = self.running_batch.batch_size() manual_divisions = self.pdmux_config.manual_divisions @@ -66,7 +72,7 @@ class SchedulerMultiplexMixin: self.tp_worker.model_runner.update_decode_attn_backend(stream_idx) return stream_idx, self.stream_groups[stream_idx] - def update_split_prefill_batch(self, sm_count: int) -> bool: + def update_split_prefill_batch(self: Scheduler, sm_count: int) -> bool: if self.split_prefill_batch: return False @@ -81,7 +87,7 @@ class SchedulerMultiplexMixin: return False @torch.inference_mode() - def event_loop_pdmux(self): + def event_loop_pdmux(self: Scheduler): """A scheduler loop for pd multiplexing.""" decode_done = False prefill_done = False