Make the scheduler track the published weight version (#35925)

python/sglang/srt/entrypoints/http_server.py

@@ -1442,9 +1442,7 @@ async def update_weight_version(
     # Use a simple approach without the complex lock mechanism for now
     # since weight_version update is a simple operation that doesn't affect model weights
     try:
-        _global_state.tokenizer_manager.record_config_updates(
-            "http.update_weight_version", weight_version=obj.new_version
-        )
+        await _global_state.tokenizer_manager.update_weight_version(obj)
 
         return ORJSONResponse(
             {

python/sglang/srt/managers/io_struct.py

@@ -1923,6 +1923,10 @@ class UpdateWeightVersionReqInput(BaseReq, kw_only=True):
     abort_all_requests: bool = True
 
 
+class UpdateWeightVersionReqOutput(BaseReq, kw_only=True):
+    pass
+
+
 class GetWeightsByNameReqInput(BaseReq, kw_only=True):
     name: str
     truncate_size: int = 100

python/sglang/srt/managers/scheduler.py

@@ -179,6 +179,8 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromIPCReqInput,
     UpdateWeightsFromTensorReqInput,
+    UpdateWeightVersionReqInput,
+    UpdateWeightVersionReqOutput,
     sock_send,
 )
 from sglang.srt.managers.load_snapshot import create_load_snapshot_writer
@@ -1608,6 +1610,10 @@ class Scheduler(
                     UpdateWeightsFromIPCReqInput,
                     self.weight_updater.update_weights_from_ipc,
                 ),
+                (
+                    UpdateWeightVersionReqInput,
+                    self.handle_update_weight_version,
+                ),
                 (
                     GetWeightsByNameReqInput,
                     self.weight_updater.get_weights_by_name,
@@ -4568,6 +4574,20 @@ class Scheduler(
         barrier(group=self.tp_group.cpu_group)
         return RpcReqOutput(success=success, message="" if not exec else str(exec))
 
+    def handle_update_weight_version(
+        self, recv_req: UpdateWeightVersionReqInput
+    ) -> UpdateWeightVersionReqOutput:
+        self.record_weight_version_change(new_version=recv_req.new_version)
+        return UpdateWeightVersionReqOutput()
+
+    def record_weight_version_change(self, new_version: Optional[str]) -> None:
+        if new_version is None or new_version == get_serving().weight_version:
+            return
+
+        old_version = get_serving().weight_version
+        get_context().override("scheduler.weight_version", weight_version=new_version)
+        logger.info(f"Weight version changed. {old_version=} {new_version=}")
+
     def collect_inflight_reqs(self) -> Set[Req]:
         if self.ps.pp_size == 1:
             inflight_batches = [self.running_batch, self.last_batch]

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

@@ -107,6 +107,9 @@ class SchedulerWeightUpdaterManager:
             )
             assert flush_cache_success, "Cache flush failed after updating weights"
 
+    def record_weight_version_after_update(self, weight_version: Optional[str]) -> None:
+        self.scheduler.record_weight_version_change(new_version=weight_version)
+
     def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
         """In-place update of the weights from disk."""
         with self._observe_weight_load("disk"):
@@ -116,7 +119,9 @@ class SchedulerWeightUpdaterManager:
                 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:
+            if success:
+                self.record_weight_version_after_update(recv_req.weight_version)
+            else:
                 logger.error(message)
             return UpdateWeightFromDiskReqOutput(
                 success=success, message=message, num_paused_requests=0
@@ -144,6 +149,7 @@ class SchedulerWeightUpdaterManager:
             success, message = self.tp_worker.update_weights_from_distributed(recv_req)
             if success:
                 self.flush_cache_after_weight_update(recv_req)
+                self.record_weight_version_after_update(recv_req.weight_version)
             else:
                 logger.error(message)
             return UpdateWeightsFromDistributedReqOutput(
@@ -160,6 +166,7 @@ class SchedulerWeightUpdaterManager:
             success, message = worker.update_weights_from_tensor(recv_req)
             if success:
                 self.flush_cache_after_weight_update(recv_req)
+                self.record_weight_version_after_update(recv_req.weight_version)
             else:
                 logger.error(message)
             torch.distributed.barrier(group=self.tp_cpu_group)
@@ -174,7 +181,9 @@ class SchedulerWeightUpdaterManager:
                 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:
+            if success:
+                self.record_weight_version_after_update(recv_req.weight_version)
+            else:
                 logger.error(message)
             torch.distributed.barrier(group=self.tp_cpu_group)
             return UpdateWeightsFromIPCReqOutput(success=success, message=message)

python/sglang/srt/managers/tokenizer_control_mixin.py

@@ -72,6 +72,8 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightsFromIPCReqOutput,
     UpdateWeightsFromTensorReqInput,
     UpdateWeightsFromTensorReqOutput,
+    UpdateWeightVersionReqInput,
+    UpdateWeightVersionReqOutput,
 )
 from sglang.srt.managers.load_snapshot import LoadSnapshot
 from sglang.srt.runtime_context import (
@@ -106,6 +108,7 @@ _COMMUNICATOR_SPECS = [
     ("send_weights_to_remote_instance", SendWeightsToRemoteInstanceReqOutput),
     ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput),
     ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput),
+    ("update_weight_version", UpdateWeightVersionReqOutput),
     ("get_weights_by_name", GetWeightsByNameReqOutput),
     ("release_memory_occupation", ReleaseMemoryOccupationReqOutput),
     ("resume_memory_occupation", ResumeMemoryOccupationReqOutput),
@@ -929,6 +932,13 @@ class TokenizerControlMixin:
     ):
         await self._async_dispatch_to_scheduler(obj)
 
+    async def update_weight_version(
+        self: TokenizerManager, obj: UpdateWeightVersionReqInput
+    ) -> None:
+        self.auto_create_handle_loop()
+        await self.update_weight_version_communicator(obj)
+        self._update_weight_version_if_provided(obj.new_version)
+
     def _update_weight_version_if_provided(
         self: TokenizerManager, weight_version: Optional[str]
     ) -> None: