Move weight-update RPC handlers to SchedulerWeightUpdaterManager (#25616)

python/sglang/srt/managers/scheduler.py

@@ -186,9 +186,6 @@ from sglang.srt.managers.scheduler_runtime_checker_mixin import (
     SchedulerRuntimeCheckerMixin,
     create_scheduler_watchdog,
 )
-from sglang.srt.managers.scheduler_update_weights_mixin import (
-    SchedulerUpdateWeightsMixin,
-)
 from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
 from sglang.srt.mem_cache import kv_cache_builder
 from sglang.srt.mem_cache.common import maybe_cache_unfinished_req, release_kv_cache
@@ -324,7 +321,6 @@ def validate_dflash_request(req: Req) -> Optional[str]:
 
 class Scheduler(
     SchedulerOutputProcessorMixin,
-    SchedulerUpdateWeightsMixin,
     SchedulerMetricsMixin,
     SchedulerDisaggregationDecodeMixin,
     SchedulerDisaggregationPrefillMixin,
@@ -1312,19 +1308,15 @@ class Scheduler(
                 (CloseSessionReqInput, self.close_session),
                 (
                     UpdateWeightFromDiskReqInput,
-                    lambda req: self.update_weights_from_disk(self.weight_updater, req),
+                    self.weight_updater.update_weights_from_disk,
                 ),
                 (
                     InitWeightsUpdateGroupReqInput,
-                    lambda req: self.init_weights_update_group(
+                    self.weight_updater.init_weights_update_group,
-                        self.weight_updater, req
-                    ),
                 ),
                 (
                     DestroyWeightsUpdateGroupReqInput,
-                    lambda req: self.destroy_weights_update_group(
+                    self.weight_updater.destroy_weights_update_group,
-                        self.weight_updater, req
-                    ),
                 ),
                 (
                     InitWeightsSendGroupForRemoteInstanceReqInput,
@@ -1336,37 +1328,31 @@ class Scheduler(
                 ),
                 (
                     UpdateWeightsFromDistributedReqInput,
-                    lambda req: self.update_weights_from_distributed(
+                    self.weight_updater.update_weights_from_distributed,
-                        self.weight_updater, req
-                    ),
                 ),
                 (
                     UpdateWeightsFromTensorReqInput,
-                    lambda req: self.update_weights_from_tensor(
+                    self.weight_updater.update_weights_from_tensor,
-                        self.weight_updater, req
-                    ),
                 ),
                 (
                     UpdateWeightsFromIPCReqInput,
-                    lambda req: self.update_weights_from_ipc(self.weight_updater, req),
+                    self.weight_updater.update_weights_from_ipc,
                 ),
                 (
                     GetWeightsByNameReqInput,
-                    lambda req: self.get_weights_by_name(self.weight_updater, req),
+                    self.weight_updater.get_weights_by_name,
                 ),
                 (
                     ReleaseMemoryOccupationReqInput,
-                    lambda req: self.release_memory_occupation(
+                    self.weight_updater.release_memory_occupation,
-                        self.weight_updater, req
-                    ),
                 ),
                 (
                     ResumeMemoryOccupationReqInput,
-                    lambda req: self.resume_memory_occupation(self.weight_updater, req),
+                    self.weight_updater.resume_memory_occupation,
                 ),
                 (
                     CheckWeightsReqInput,
-                    lambda req: self.check_weights(self.weight_updater, req),
+                    self.weight_updater.check_weights,
                 ),
                 (SlowDownReqInput, self.slow_down),
                 (
@@ -3289,10 +3275,10 @@ class Scheduler(
         )
 
     def save_remote_model(self, **kwargs):
-        SchedulerUpdateWeightsMixin.save_remote_model(self.weight_updater, kwargs)
+        self.weight_updater.save_remote_model(kwargs)
 
     def save_sharded_model(self, **kwargs):
-        SchedulerUpdateWeightsMixin.save_sharded_model(self.weight_updater, kwargs)
+        self.weight_updater.save_sharded_model(kwargs)
 
     def handle_rpc_request(self, recv_req: RpcReqInput):
         # Handle RPC requests

python/sglang/srt/managers/scheduler_components/weight_updater.py

@@ -1,7 +1,42 @@
 from __future__ import annotations
 
+import logging
+import traceback
 from dataclasses import dataclass, field
-from typing import Any, Callable
+from typing import Any, Callable, Tuple
+
+import torch
+
+from sglang.srt.constants import (
+    GPU_MEMORY_ALL_TYPES,
+    GPU_MEMORY_TYPE_CUDA_GRAPH,
+    GPU_MEMORY_TYPE_KV_CACHE,
+    GPU_MEMORY_TYPE_WEIGHTS,
+)
+from sglang.srt.managers.io_struct import (
+    CheckWeightsReqInput,
+    CheckWeightsReqOutput,
+    DestroyWeightsUpdateGroupReqInput,
+    DestroyWeightsUpdateGroupReqOutput,
+    GetWeightsByNameReqInput,
+    GetWeightsByNameReqOutput,
+    InitWeightsUpdateGroupReqInput,
+    InitWeightsUpdateGroupReqOutput,
+    ReleaseMemoryOccupationReqInput,
+    ReleaseMemoryOccupationReqOutput,
+    ResumeMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqOutput,
+    UpdateWeightFromDiskReqInput,
+    UpdateWeightFromDiskReqOutput,
+    UpdateWeightsFromDistributedReqInput,
+    UpdateWeightsFromDistributedReqOutput,
+    UpdateWeightsFromIPCReqInput,
+    UpdateWeightsFromIPCReqOutput,
+    UpdateWeightsFromTensorReqInput,
+    UpdateWeightsFromTensorReqOutput,
+)
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass(kw_only=True, slots=True)
@@ -14,3 +49,180 @@ class SchedulerWeightUpdaterManager:
     is_fully_idle: Callable[..., bool]
     offload_tags: set = field(default_factory=set)
     stashed_model_static_state: Any = None
+
+    def flush_cache_after_weight_update(self, recv_req) -> None:
+        if recv_req.flush_cache:
+            flush_cache_success = self.flush_cache(
+                empty_cache=recv_req.torch_empty_cache
+            )
+            assert flush_cache_success, "Cache flush failed after updating weights"
+
+    def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
+        """In-place update of the weights from disk."""
+        success, message = self.tp_worker.update_weights_from_disk(recv_req)
+        tp_success = success
+        if success and self.draft_worker is not None:
+            success, message = self.draft_worker.update_weights_from_disk(recv_req)
+        if tp_success:
+            self.flush_cache_after_weight_update(recv_req)
+        if not success:
+            logger.error(message)
+        return UpdateWeightFromDiskReqOutput(success, message, 0)
+
+    def init_weights_update_group(self, 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,
+    ):
+        """Destroy the online model parameter update group."""
+        success, message = self.tp_worker.destroy_weights_update_group(recv_req)
+        return DestroyWeightsUpdateGroupReqOutput(success, message)
+
+    def update_weights_from_distributed(
+        self,
+        recv_req: UpdateWeightsFromDistributedReqInput,
+    ) -> Tuple[bool, str]:
+        """Update the online model parameter."""
+        success, message = self.tp_worker.update_weights_from_distributed(recv_req)
+        if success:
+            self.flush_cache_after_weight_update(recv_req)
+        else:
+            logger.error(message)
+        return UpdateWeightsFromDistributedReqOutput(success, message)
+
+    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
+        """Update the online model parameter from tensors."""
+        if recv_req.disable_draft_model:
+            worker = self.tp_worker
+        else:
+            worker = self.draft_worker or self.tp_worker
+        success, message = worker.update_weights_from_tensor(recv_req)
+        if success:
+            self.flush_cache_after_weight_update(recv_req)
+        else:
+            logger.error(message)
+        torch.distributed.barrier(group=self.tp_cpu_group)
+        return UpdateWeightsFromTensorReqOutput(success, message)
+
+    def update_weights_from_ipc(self, 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)
+        tp_success = success
+        if success and self.draft_worker is not None:
+            success, message = self.draft_worker.update_weights_from_ipc(recv_req)
+        if tp_success:
+            self.flush_cache_after_weight_update(recv_req)
+        if not success:
+            logger.error(message)
+        torch.distributed.barrier(group=self.tp_cpu_group)
+        return UpdateWeightsFromIPCReqOutput(success, message)
+
+    def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
+        parameter = self.tp_worker.get_weights_by_name(recv_req)
+        return GetWeightsByNameReqOutput(parameter)
+
+    def release_memory_occupation(self, recv_req: ReleaseMemoryOccupationReqInput):
+        assert (
+            self.is_fully_idle()
+        ), "release_memory_occupation should be called only when server is idle."
+
+        tags = recv_req.tags
+
+        if tags is None or len(tags) == 0:
+            tags = GPU_MEMORY_ALL_TYPES
+
+        for tag in tags:
+            self.offload_tags.add(tag)
+
+        if GPU_MEMORY_TYPE_KV_CACHE in tags:
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
+            self.flush_cache()
+
+        if GPU_MEMORY_TYPE_WEIGHTS in tags:
+            self.stashed_model_static_state = _export_static_state(
+                self.tp_worker.model_runner.model
+            )
+            torch.distributed.barrier(self.tp_cpu_group)
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
+
+        if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
+            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_CUDA_GRAPH)
+
+        torch.get_device_module().synchronize()
+
+        return ReleaseMemoryOccupationReqOutput()
+
+    def resume_memory_occupation(self, recv_req: ResumeMemoryOccupationReqInput):
+        tags = recv_req.tags
+
+        if tags is None or len(tags) == 0:
+            tags = GPU_MEMORY_ALL_TYPES
+
+        for tag in tags:
+            self.offload_tags.remove(tag)
+
+        if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_CUDA_GRAPH)
+
+        if GPU_MEMORY_TYPE_WEIGHTS in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
+            torch.distributed.barrier(self.tp_cpu_group)
+            _import_static_state(
+                self.tp_worker.model_runner.model,
+                self.stashed_model_static_state,
+            )
+            del self.stashed_model_static_state
+
+        if GPU_MEMORY_TYPE_KV_CACHE in tags:
+            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE)
+
+        return ResumeMemoryOccupationReqOutput()
+
+    def check_weights(self, recv_req: CheckWeightsReqInput):
+        try:
+            payload = self.tp_worker.model_runner.check_weights(action=recv_req.action)
+            return CheckWeightsReqOutput(
+                success=True, message="Success.", payload=payload
+            )
+        except Exception as e:
+            logger.warning(f"check_weights see error: {e}")
+            traceback.print_exc()
+            return CheckWeightsReqOutput(success=False, message=f"{e}")
+
+    def save_remote_model(self, params):
+        url = params["url"]
+
+        self.tp_worker.model_runner.save_remote_model(url)
+
+        if self.draft_worker is not None:
+            draft_url = params.get("draft_url", None)
+            assert (
+                draft_url is not None
+            ), "draft_url must be provided when draft model is enabled"
+            self.draft_worker.model_runner.save_remote_model(draft_url)
+
+    def save_sharded_model(self, params):
+        self.tp_worker.model_runner.save_sharded_model(
+            path=params["path"],
+            pattern=params["pattern"],
+            max_size=params["max_size"],
+        )
+
+
+def _export_static_state(model):
+    return dict(
+        buffers=[
+            (name, buffer.detach().clone()) for name, buffer in model.named_buffers()
+        ]
+    )
+
+
+def _import_static_state(model, static_params):
+    with torch.inference_mode():
+        self_named_buffers = dict(model.named_buffers())
+        for name, tensor in static_params["buffers"]:
+            self_named_buffers[name][...] = tensor

python/sglang/srt/managers/scheduler_update_weights_mixin.py

@@ -1,253 +0,0 @@
-from __future__ import annotations
-
-import logging
-import traceback
-from typing import TYPE_CHECKING, Tuple
-
-import torch
-
-from sglang.srt.constants import (
-    GPU_MEMORY_ALL_TYPES,
-    GPU_MEMORY_TYPE_CUDA_GRAPH,
-    GPU_MEMORY_TYPE_KV_CACHE,
-    GPU_MEMORY_TYPE_WEIGHTS,
-)
-from sglang.srt.managers.io_struct import (
-    CheckWeightsReqInput,
-    CheckWeightsReqOutput,
-    DestroyWeightsUpdateGroupReqInput,
-    DestroyWeightsUpdateGroupReqOutput,
-    GetWeightsByNameReqInput,
-    GetWeightsByNameReqOutput,
-    InitWeightsUpdateGroupReqInput,
-    InitWeightsUpdateGroupReqOutput,
-    ReleaseMemoryOccupationReqInput,
-    ReleaseMemoryOccupationReqOutput,
-    ResumeMemoryOccupationReqInput,
-    ResumeMemoryOccupationReqOutput,
-    UpdateWeightFromDiskReqInput,
-    UpdateWeightFromDiskReqOutput,
-    UpdateWeightsFromDistributedReqInput,
-    UpdateWeightsFromDistributedReqOutput,
-    UpdateWeightsFromIPCReqInput,
-    UpdateWeightsFromIPCReqOutput,
-    UpdateWeightsFromTensorReqInput,
-    UpdateWeightsFromTensorReqOutput,
-)
-
-if TYPE_CHECKING:
-    from sglang.srt.managers.scheduler_components.weight_updater import (
-        SchedulerWeightUpdaterManager,
-    )
-
-logger = logging.getLogger(__name__)
-
-
-class SchedulerUpdateWeightsMixin:
-    @staticmethod
-    def flush_cache_after_weight_update(
-        self: "SchedulerWeightUpdaterManager", recv_req
-    ) -> None:
-        if recv_req.flush_cache:
-            flush_cache_success = self.flush_cache(
-                empty_cache=recv_req.torch_empty_cache
-            )
-            assert flush_cache_success, "Cache flush failed after updating weights"
-
-    @staticmethod
-    def update_weights_from_disk(
-        self: "SchedulerWeightUpdaterManager", recv_req: UpdateWeightFromDiskReqInput
-    ):
-        """In-place update of the weights from disk."""
-        success, message = self.tp_worker.update_weights_from_disk(recv_req)
-        tp_success = success
-        if success and self.draft_worker is not None:
-            success, message = self.draft_worker.update_weights_from_disk(recv_req)
-        if tp_success:
-            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
-        if not success:
-            logger.error(message)
-        return UpdateWeightFromDiskReqOutput(success, message, 0)
-
-    @staticmethod
-    def init_weights_update_group(
-        self: "SchedulerWeightUpdaterManager", 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)
-
-    @staticmethod
-    def destroy_weights_update_group(
-        self: "SchedulerWeightUpdaterManager",
-        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)
-
-    @staticmethod
-    def update_weights_from_distributed(
-        self: "SchedulerWeightUpdaterManager",
-        recv_req: UpdateWeightsFromDistributedReqInput,
-    ) -> Tuple[bool, str]:
-        """Update the online model parameter."""
-        success, message = self.tp_worker.update_weights_from_distributed(recv_req)
-        if success:
-            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
-        else:
-            logger.error(message)
-        return UpdateWeightsFromDistributedReqOutput(success, message)
-
-    @staticmethod
-    def update_weights_from_tensor(
-        self: "SchedulerWeightUpdaterManager", recv_req: UpdateWeightsFromTensorReqInput
-    ):
-        """Update the online model parameter from tensors."""
-        if recv_req.disable_draft_model:
-            worker = self.tp_worker
-        else:
-            worker = self.draft_worker or self.tp_worker
-        success, message = worker.update_weights_from_tensor(recv_req)
-        if success:
-            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
-        else:
-            logger.error(message)
-        torch.distributed.barrier(group=self.tp_cpu_group)
-        return UpdateWeightsFromTensorReqOutput(success, message)
-
-    @staticmethod
-    def update_weights_from_ipc(
-        self: "SchedulerWeightUpdaterManager", 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)
-        tp_success = success
-        if success and self.draft_worker is not None:
-            success, message = self.draft_worker.update_weights_from_ipc(recv_req)
-        if tp_success:
-            SchedulerUpdateWeightsMixin.flush_cache_after_weight_update(self, recv_req)
-        if not success:
-            logger.error(message)
-        torch.distributed.barrier(group=self.tp_cpu_group)
-        return UpdateWeightsFromIPCReqOutput(success, message)
-
-    @staticmethod
-    def get_weights_by_name(
-        self: "SchedulerWeightUpdaterManager", recv_req: GetWeightsByNameReqInput
-    ):
-        parameter = self.tp_worker.get_weights_by_name(recv_req)
-        return GetWeightsByNameReqOutput(parameter)
-
-    @staticmethod
-    def release_memory_occupation(
-        self: "SchedulerWeightUpdaterManager", recv_req: ReleaseMemoryOccupationReqInput
-    ):
-        assert (
-            self.is_fully_idle()
-        ), "release_memory_occupation should be called only when server is idle."
-
-        tags = recv_req.tags
-
-        if tags is None or len(tags) == 0:
-            tags = GPU_MEMORY_ALL_TYPES
-
-        for tag in tags:
-            self.offload_tags.add(tag)
-
-        if GPU_MEMORY_TYPE_KV_CACHE in tags:
-            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
-            self.flush_cache()
-
-        if GPU_MEMORY_TYPE_WEIGHTS in tags:
-            self.stashed_model_static_state = _export_static_state(
-                self.tp_worker.model_runner.model
-            )
-            torch.distributed.barrier(self.tp_cpu_group)
-            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_WEIGHTS)
-
-        if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
-            self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_CUDA_GRAPH)
-
-        torch.get_device_module().synchronize()
-
-        return ReleaseMemoryOccupationReqOutput()
-
-    @staticmethod
-    def resume_memory_occupation(
-        self: "SchedulerWeightUpdaterManager", recv_req: ResumeMemoryOccupationReqInput
-    ):
-        tags = recv_req.tags
-
-        if tags is None or len(tags) == 0:
-            tags = GPU_MEMORY_ALL_TYPES
-
-        for tag in tags:
-            self.offload_tags.remove(tag)
-
-        if GPU_MEMORY_TYPE_CUDA_GRAPH in tags:
-            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_CUDA_GRAPH)
-
-        if GPU_MEMORY_TYPE_WEIGHTS in tags:
-            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_WEIGHTS)
-            torch.distributed.barrier(self.tp_cpu_group)
-            _import_static_state(
-                self.tp_worker.model_runner.model,
-                self.stashed_model_static_state,
-            )
-            del self.stashed_model_static_state
-
-        if GPU_MEMORY_TYPE_KV_CACHE in tags:
-            self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE)
-
-        return ResumeMemoryOccupationReqOutput()
-
-    @staticmethod
-    def check_weights(
-        self: "SchedulerWeightUpdaterManager", recv_req: CheckWeightsReqInput
-    ):
-        try:
-            payload = self.tp_worker.model_runner.check_weights(action=recv_req.action)
-            return CheckWeightsReqOutput(
-                success=True, message="Success.", payload=payload
-            )
-        except Exception as e:
-            logger.warning(f"check_weights see error: {e}")
-            traceback.print_exc()
-            return CheckWeightsReqOutput(success=False, message=f"{e}")
-
-    @staticmethod
-    def save_remote_model(self: "SchedulerWeightUpdaterManager", params):
-        url = params["url"]
-
-        self.tp_worker.model_runner.save_remote_model(url)
-
-        if self.draft_worker is not None:
-            draft_url = params.get("draft_url", None)
-            assert (
-                draft_url is not None
-            ), "draft_url must be provided when draft model is enabled"
-            self.draft_worker.model_runner.save_remote_model(draft_url)
-
-    @staticmethod
-    def save_sharded_model(self: "SchedulerWeightUpdaterManager", params):
-        self.tp_worker.model_runner.save_sharded_model(
-            path=params["path"],
-            pattern=params["pattern"],
-            max_size=params["max_size"],
-        )
-
-
-def _export_static_state(model):
-    return dict(
-        buffers=[
-            (name, buffer.detach().clone()) for name, buffer in model.named_buffers()
-        ]
-    )
-
-
-def _import_static_state(model, static_params):
-    with torch.inference_mode():
-        self_named_buffers = dict(model.named_buffers())
-        for name, tensor in static_params["buffers"]:
-            self_named_buffers[name][...] = tensor