[RL] Add /pull_weights: engine-side pull of published weights into a host-local checkpoint

python/pyproject.toml

@@ -87,6 +87,7 @@ dependencies = [
   "uvloop",
   "watchfiles",
   "xgrammar==0.2.1",
+  "xxhash",
   "zstandard",
 ]
 

python/sglang/srt/entrypoints/http_server.py

@@ -130,6 +130,7 @@
     ParseFunctionCallReq,
     PauseGenerationReqInput,
     ProfileReq,
+    PullWeightsReqInput,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
     SendWeightsToRemoteInstanceReqInput,
@@ -1187,6 +1188,18 @@ async def update_weights_from_disk(
         )
 
 
+@app.post("/pull_weights")
+@auth_level(AuthLevel.ADMIN_OPTIONAL)
+async def pull_weights(obj: Annotated[PullWeightsReqInput, Body()], request: Request):
+    """Materialize published weights on every engine host."""
+    success, message = await _global_state.tokenizer_manager.pull_weights(obj, request)
+
+    content = {"success": success, "message": message}
+    return ORJSONResponse(
+        content, status_code=HTTPStatus.OK if success else HTTPStatus.BAD_REQUEST
+    )
+
+
 @app.post("/init_weights_send_group_for_remote_instance")
 @auth_level(AuthLevel.ADMIN_OPTIONAL)
 async def init_weights_send_group_for_remote_instance(

python/sglang/srt/managers/io_struct.py

@@ -1562,6 +1562,19 @@ class UpdateWeightFromDiskReqOutput(BaseReq, kw_only=True):
     num_paused_requests: int = 0
 
 
+class PullWeightsReqInput(BaseReq, kw_only=True):
+    """Request to materialize a published weight version."""
+
+    local_checkpoint_dir: str
+    source_dir: str
+    target_version: int
+
+
+class PullWeightsReqOutput(BaseReq, kw_only=True):
+    success: bool
+    message: str
+
+
 class UpdateWeightsFromDistributedReqInput(BaseReq, kw_only=True):
     names: List[str]
     dtypes: List[str]

python/sglang/srt/managers/scheduler.py

@@ -121,6 +121,7 @@
     OpenSessionReqInput,
     PauseGenerationReqInput,
     ProfileReq,
+    PullWeightsReqInput,
     ReleaseMemoryOccupationReqInput,
     RemoveExternalCorpusReqInput,
     RemoveExternalCorpusReqOutput,
@@ -1392,6 +1393,7 @@ def init_request_dispatcher(self):
                     CheckWeightsReqInput,
                     self.weight_updater.check_weights,
                 ),
+                (PullWeightsReqInput, self.weight_updater.pull_weights),
                 (SlowDownReqInput, self.slow_down),
                 (
                     ProfileReq,

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

@@ -28,6 +28,8 @@
     GetWeightsByNameReqOutput,
     InitWeightsUpdateGroupReqInput,
     InitWeightsUpdateGroupReqOutput,
+    PullWeightsReqInput,
+    PullWeightsReqOutput,
     ReleaseMemoryOccupationReqInput,
     ReleaseMemoryOccupationReqOutput,
     ResumeMemoryOccupationReqInput,
@@ -84,6 +86,7 @@ class SchedulerWeightUpdaterManager:
     metrics_collector: Optional[Any] = None
     offload_tags: set = field(default_factory=set)
     stashed_model_static_state: Any = None
+    _pull_weights_base_dir: Optional[str] = None
 
     @contextmanager
     def _observe_weight_load(self, source: str) -> Iterator[None]:
@@ -122,6 +125,40 @@ def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput):
                 success=success, message=message, num_paused_requests=0
             )
 
+    def pull_weights(self, recv_req: PullWeightsReqInput):
+        """Materialize a published weight version on every host."""
+        from sglang.srt.weight_sync import local_checkpoint
+
+        server_args = self.tp_worker.model_runner.server_args
+        if self._pull_weights_base_dir is None:
+            self._pull_weights_base_dir = server_args.model_path
+        try:
+            local_checkpoint.pull(
+                local_checkpoint_dir=recv_req.local_checkpoint_dir,
+                base_dir=self._pull_weights_base_dir,
+                source_dir=recv_req.source_dir,
+                target_version=recv_req.target_version,
+                pre_read_hook=server_args.custom_pull_weights_pre_read_hook,
+            )
+            success, message = True, "Success."
+        except Exception:
+            success, message = False, traceback.format_exc()
+            logger.error(message)
+
+        world_size = (
+            torch.distributed.get_world_size(group=self.tp_cpu_group)
+            if torch.distributed.is_initialized()
+            else 1
+        )
+        if world_size > 1:
+            results = [None] * world_size
+            torch.distributed.all_gather_object(
+                results, (success, message), group=self.tp_cpu_group
+            )
+            success = all(ok for ok, _ in results)
+            message = "; ".join(msg for ok, msg in results if not ok) or message
+        return PullWeightsReqOutput(success=success, message=message)
+
     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)

python/sglang/srt/managers/tokenizer_control_mixin.py

@@ -51,6 +51,8 @@
     ProfileReq,
     ProfileReqOutput,
     ProfileReqType,
+    PullWeightsReqInput,
+    PullWeightsReqOutput,
     ReleaseMemoryOccupationReqInput,
     ReleaseMemoryOccupationReqOutput,
     RemoveExternalCorpusReqInput,
@@ -105,6 +107,7 @@
     ("release_memory_occupation", ReleaseMemoryOccupationReqOutput),
     ("resume_memory_occupation", ResumeMemoryOccupationReqOutput),
     ("check_weights", CheckWeightsReqOutput),
+    ("pull_weights", PullWeightsReqOutput),
     ("slow_down", SlowDownReqOutput),
     ("flush_cache", FlushCacheReqOutput),
     ("add_external_corpus", AddExternalCorpusReqOutput),
@@ -770,6 +773,15 @@ async def resume_memory_occupation(
         self.auto_create_handle_loop()
         await self.resume_memory_occupation_communicator(obj)
 
+    async def pull_weights(
+        self: TokenizerManager,
+        obj: PullWeightsReqInput,
+        request: Optional[fastapi.Request] = None,
+    ) -> Tuple[bool, str]:
+        self.auto_create_handle_loop()
+        results = await self.pull_weights_communicator(obj)
+        return FanOutCommunicator.merge_results(results)
+
     async def check_weights(
         self: TokenizerManager,
         obj: CheckWeightsReqInput,

python/sglang/srt/server_args.py

@@ -2596,6 +2596,10 @@ class ServerArgs:
             nargs="*",
         ),
     ] = None
+    custom_pull_weights_pre_read_hook: A[
+        Optional[str],
+        "Import path of hook(source_dir, target_version) called before /pull_weights reads shared storage.",
+    ] = None
     weight_loader_disable_mmap: A[
         bool,
         "Disable mmap while loading weight using safetensors.",