chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Cheng Wan
2025-12-26 16:50:07 -08:00
committed by GitHub
parent 93495dcac9
commit 988b14ca0e
5 changed files with 42 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/managers/scheduler_metrics_mixin.py
View File

@@ -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

2
python/sglang/srt/managers/scheduler_pp_mixin.py
View File

@@ -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.

20
python/sglang/srt/managers/scheduler_profiler_mixin.py
View File

@@ -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(

22
python/sglang/srt/managers/scheduler_update_weights_mixin.py
View File

@@ -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)

14
python/sglang/srt/multiplex/multiplexing_mixin.py
View File

@@ -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