refactor: add type hints to scheduler mixins (#15913)

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

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.
 

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(

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)
 

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