[RL] Release weight-transfer communicators without stopping inference (#61060)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Signed-off-by: fh3 <a@b>

rust/src/engine-core-client/src/client.rs

@@ -828,6 +828,18 @@ impl EngineCoreClient {
         Ok(())
     }
 
+    /// Release the trainer connection without shutting down inference.
+    pub async fn shutdown_weight_transfer_engine(&self) -> Result<()> {
+        self.collective_rpc(
+            "shutdown_weight_transfer_engine",
+            None,
+            Vec::<JsonValue>::new(),
+            BTreeMap::<String, JsonValue>::new(),
+        )
+        .await?;
+        Ok(())
+    }
+
     /// Start a weight update for the base model.
     pub async fn start_weight_update(&self) -> Result<()> {
         self.collective_rpc(

rust/src/server/src/routes.rs

@@ -145,6 +145,10 @@ fn build_router_with_options(
                 "/init_weight_transfer_engine",
                 post(weight_transfer::init_weight_transfer_engine),
             )
+            .route(
+                "/shutdown_weight_transfer_engine",
+                post(weight_transfer::shutdown_weight_transfer_engine),
+            )
             .route(
                 "/start_weight_update",
                 post(weight_transfer::start_weight_update),

rust/src/server/src/routes/tests.rs

@@ -6708,6 +6708,7 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
         ("POST", "/reset_mm_cache"),
         ("POST", "/reset_encoder_cache"),
         ("POST", "/init_weight_transfer_engine"),
+        ("POST", "/shutdown_weight_transfer_engine"),
         ("POST", "/start_weight_update"),
         ("POST", "/start_draft_weight_update"),
         ("POST", "/update_weights"),
@@ -6739,7 +6740,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {
     let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
         boxed_test_future(async move {
             let mut calls = Vec::new();
-            for _ in 0..10 {
+            for _ in 0..11 {
                 let utility = recv_engine_message(dealer).await;
                 assert_eq!(utility[0].as_ref(), &[0x03]);
                 let payload = decode_value(&utility[1]).expect("decode utility");
@@ -6769,6 +6770,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {
                     "collective_rpc [\"finish_weight_update\",null,[],{}]",
                     "set_weight_version [\"step-2\"]",
                     "get_weight_version []",
+                    "collective_rpc [\"shutdown_weight_transfer_engine\",null,[],{}]",
                 ]
             "#]].assert_debug_eq(&calls);
         })
@@ -6817,6 +6819,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {
             Some(json!({"new_version": "step-2"})),
         ),
         ("GET", "/weight_info", None),
+        ("POST", "/shutdown_weight_transfer_engine", None),
     ] {
         let mut request = Request::builder().method(method).uri(path);
         let body = match body {
@@ -6838,7 +6841,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {
     }
     engine_task.await.expect("mock engine task");
     let metrics_after = METRICS.render().unwrap();
-    for operation in ["update", "finish"] {
+    for (operation, count) in [("update", 2.0), ("finish", 2.0), ("shutdown", 1.0)] {
         assert_eq!(
             metric_delta(
                 &metrics_before,
@@ -6846,7 +6849,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {
                 "vllm:rl_weight_update_operation_duration_seconds_count",
                 Some(&format!("operation=\"{operation}\"")),
             ),
-            2.0
+            count
         );
         assert_eq!(
             metric_value(
@@ -6868,6 +6871,7 @@ async fn weight_transfer_routes_support_the_http_training_lifecycle() {
             "{\"message\":\"Weight update finished\"}",
             "{\"success\":true,\"new_version\":\"step-2\"}",
             "{\"weight_version\":\"step-2\"}",
+            "{\"message\":\"Weight transfer shut down\"}",
         ]
     "#]]
     .assert_debug_eq(&responses);

rust/src/server/src/routes/weight_transfer.rs

@@ -75,6 +75,22 @@ pub async fn init_weight_transfer_engine(
     }))
 }
 
+/// Release the trainer connection without shutting down inference.
+pub async fn shutdown_weight_transfer_engine(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    let _recorder = METRICS.api_server.record_weight_operation("shutdown");
+    state
+        .engine_core_client()
+        .shutdown_weight_transfer_engine()
+        .await
+        .map_err(|error| utility_call_error("shutdown_weight_transfer_engine", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Weight transfer shut down",
+    }))
+}
+
 /// Start a weight update transaction.
 pub async fn start_weight_update(
     State(state): State<Arc<AppState>>,

tests/distributed/test_weight_transfer.py

@@ -218,6 +218,27 @@ class TestSparseNCCLWeightTransferUpdateInfoValidation:
 # --- Unit Tests: Engine Parsing ---
 
 
+@pytest.mark.parametrize(
+    "engine_type",
+    [
+        NCCLWeightTransferEngine,
+        NCCLTrainerWeightTransferEngine,
+        SparseNCCLWeightTransferEngine,
+        SparseNCCLTrainerWeightTransferEngine,
+    ],
+)
+def test_nccl_shutdown_destroys_communicator_once(engine_type):
+    engine = object.__new__(engine_type)
+    group = MagicMock()
+    engine.model_update_group = group
+
+    engine.shutdown()
+    engine.shutdown()
+
+    group.destroy.assert_called_once_with()
+    assert engine.model_update_group is None
+
+
 class TestNCCLEngineParsing:
     """Test NCCLWeightTransferEngine parsing methods."""
 
@@ -1292,6 +1313,9 @@ class TestTrainerClients:
         handle.finish_weight_update.remote.assert_called_once_with()
         handle.update_weight_version.remote.assert_called_once_with("step-42")
 
+        client.shutdown_weight_transfer_engine()
+        handle.shutdown_weight_transfer_engine.remote.assert_called_once_with()
+
     def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch):
         """HTTP update_weights must encode raw ipc_handles as a base64 pickle."""
         captured = {}
@@ -1322,6 +1346,7 @@ class TestTrainerClients:
         captured = {}
 
         def fake_post(self, path, json=None):
+            captured["path"] = path
             captured["json"] = json
 
         monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post)
@@ -1333,6 +1358,9 @@ class TestTrainerClients:
         client.finish_weight_update("step-42")
         assert captured["json"] == {"weight_version": "step-42"}
 
+        client.shutdown_weight_transfer_engine()
+        assert captured == {"path": "shutdown_weight_transfer_engine", "json": None}
+
 
 class TestModuleSource:
     """`ModuleSource` metadata vs. materialized iteration (dense, no GPU)."""

tests/distributed/test_weight_transfer_m2n.py

@@ -46,7 +46,10 @@ from vllm.distributed.weight_transfer.m2n_source import (
     mesh_from_tensor,
     placements_from_tensor,
 )
-from vllm.distributed.weight_transfer.m2n_trainer import M2NTrainerInitInfo
+from vllm.distributed.weight_transfer.m2n_trainer import (
+    M2NTrainerInitInfo,
+    M2NTrainerWeightTransferEngine,
+)
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     RowParallelLinear,
@@ -57,6 +60,27 @@ from vllm.platforms import current_platform
 VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode()
 
 
+@pytest.mark.parametrize(
+    "engine_type,group_attribute",
+    [
+        (M2NWeightTransferEngine, "model_update_group"),
+        (M2NTrainerWeightTransferEngine, "group"),
+    ],
+)
+def test_shutdown_destroys_communicator_once(engine_type, group_attribute):
+    engine = object.__new__(engine_type)
+    engine._handle = None
+    engine._executor = None
+    group = Mock()
+    setattr(engine, group_attribute, group)
+
+    engine.shutdown()
+    engine.shutdown()
+
+    group.destroy.assert_called_once_with()
+    assert getattr(engine, group_attribute) is None
+
+
 class TestLayout:
     def test_replicated_splits_nothing(self):
         """A replicated tensor imposes no divisibility constraint. The shape

tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py

@@ -89,6 +89,11 @@ def test_invalid_field_is_rejected_before_the_engine(client, path, body):
         ("/update_weights", {"update_info": {"names": []}}, "update_weights"),
         ("/update_weights", {"update_info": [{}, {}]}, "update_weights"),
         ("/collective_rpc", {"method": "m"}, "collective_rpc"),
+        (
+            "/shutdown_weight_transfer_engine",
+            {},
+            "shutdown_weight_transfer_engine",
+        ),
     ],
 )
 def test_valid_body_reaches_the_engine(client, path, body, engine_call):

tests/entrypoints/rl/weight_transfer/test_offline.py

@@ -169,6 +169,13 @@ def test_init_weight_transfer_engine_calls_engine(vllm_runner):
             assert called, "init_transfer_engine should have been called"
             assert param == "hello", f"Expected 'hello', got {param}"
 
+        llm.shutdown_weight_transfer_engine()
+        assert all(
+            llm.collective_rpc(
+                lambda worker: worker.weight_transfer_engine.shutdown_called
+            )
+        )
+
 
 @create_new_process_for_each_test()
 def test_update_weights_calls_engine(vllm_runner):

tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py

@@ -75,6 +75,7 @@ def test_record_releases_in_flight_and_observes_duration(registry, fail):
     "operation,route,body",
     [
         ("init", api_router.init_weight_transfer_engine, {"init_info": {}}),
+        ("shutdown", api_router.shutdown_weight_transfer_engine, None),
         ("start", api_router.start_weight_update, None),
         ("start_draft", api_router.start_draft_weight_update, None),
         ("update", api_router.update_weights, {"update_info": {}}),

tests/v1/worker/test_gpu_worker_weight_transfer.py

@@ -24,6 +24,7 @@ class _RecordingEngine:
         self.raise_on_update = raise_on_update
         self.started = False
         self.finished = False
+        self.closed = False
         self.reset_count = 0
         self.supports_draft_weight_update = False
         self.update_calls: list[dict] = []
@@ -49,6 +50,9 @@ class _RecordingEngine:
     def reset_weight_update_target(self) -> None:
         self.reset_count += 1
 
+    def shutdown(self) -> None:
+        self.closed = True
+
 
 class _RecordingModelRunner:
     def __init__(self) -> None:
@@ -82,6 +86,25 @@ def test_reload_weights_sets_current_config():
     assert model_runner.seen_config is worker.vllm_config
 
 
+@pytest.mark.parametrize(
+    "enabled,active", [(True, False), (True, True), (False, False)]
+)
+def test_shutdown_weight_transfer_engine_preserves_inference(enabled, active):
+    engine = _RecordingEngine() if enabled else None
+    worker = _make_worker(engine)
+    worker._weight_update_active = active
+    model_runner = worker.model_runner
+    if active:
+        with pytest.raises(RuntimeError, match="during a weight update"):
+            worker.shutdown_weight_transfer_engine()
+    else:
+        worker.shutdown_weight_transfer_engine()
+    if engine is not None:
+        assert engine.closed is not active
+    assert worker.weight_transfer_engine is engine
+    assert worker.model_runner is model_runner
+
+
 def test_reload_parameter_lookup_preserves_lora_module_names():
     base_layer = nn.Module()
     qweight = nn.Parameter(torch.ones(1))

vllm/distributed/weight_transfer/base.py

@@ -514,9 +514,7 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
 
     @abstractmethod
     def shutdown(self) -> None:
-        """Shutdown the weight transfer engine.
-        This should be called when the worker is shutting down.
-        """
+        """Release transfer resources without shutting down the inference worker."""
         raise NotImplementedError
 
 

vllm/distributed/weight_transfer/clients.py

@@ -60,7 +60,7 @@ class HTTPVLLMWeightSyncClient:
 
     Mirrors `vllm/entrypoints/rl/online/api_router.py`:
     `/init_weight_transfer_engine`, `/start_weight_update`, `/update_weights`,
-    `/finish_weight_update`.
+    `/finish_weight_update`, `/shutdown_weight_transfer_engine`.
     """
 
     def __init__(self, base_url: str, timeout: float = 300) -> None:
@@ -78,6 +78,9 @@ class HTTPVLLMWeightSyncClient:
     def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None:
         self._post("init_weight_transfer_engine", {"init_info": init_info})
 
+    def shutdown_weight_transfer_engine(self) -> None:
+        self._post("shutdown_weight_transfer_engine")
+
     def start_weight_update(self) -> None:
         self._post("start_weight_update")
 
@@ -109,6 +112,11 @@ class RayVLLMWeightSyncClient:
         request = WeightTransferInitRequest(init_info=init_info)
         ray.get([h.init_weight_transfer_engine.remote(request) for h in self.handles])
 
+    def shutdown_weight_transfer_engine(self) -> None:
+        import ray
+
+        ray.get([h.shutdown_weight_transfer_engine.remote() for h in self.handles])
+
     def start_weight_update(self) -> None:
         import ray
 

vllm/distributed/weight_transfer/m2n_engine.py

@@ -449,5 +449,7 @@ class M2NWeightTransferEngine(
             torch.accelerator.synchronize()
             self._handle.destroy()
             self._handle = None
-        self.model_update_group = None
+        if self.model_update_group is not None:
+            self.model_update_group.destroy()
+            self.model_update_group = None
         self._parameter_destinations = []

vllm/distributed/weight_transfer/m2n_trainer.py

@@ -410,4 +410,6 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI
         if self._executor is not None:
             self._executor.shutdown()
             self._executor = None
-        self.group = None
+        if self.group is not None:
+            self.group.destroy()
+            self.group = None

vllm/distributed/weight_transfer/nccl_engine.py

@@ -229,7 +229,7 @@ class NCCLWeightTransferEngine(
 
     def shutdown(self) -> None:
         if self.model_update_group is not None:
-            # Clean up the communicator by removing the reference
+            self.model_update_group.destroy()
             self.model_update_group = None
 
 
@@ -456,4 +456,6 @@ class NCCLTrainerWeightTransferEngine(TrainerWeightTransferEngine[NCCLTrainerIni
             torch.cuda.current_stream().synchronize()
 
     def shutdown(self) -> None:
-        self.model_update_group = None
+        if self.model_update_group is not None:
+            self.model_update_group.destroy()
+            self.model_update_group = None

vllm/distributed/weight_transfer/sparse_nccl_engine.py

@@ -186,6 +186,7 @@ class SparseNCCLWeightTransferEngine(
 
     def shutdown(self) -> None:
         if self.model_update_group is not None:
+            self.model_update_group.destroy()
             self.model_update_group = None
 
 
@@ -371,4 +372,6 @@ class SparseNCCLTrainerWeightTransferEngine(
             torch.cuda.current_stream().synchronize()
 
     def shutdown(self) -> None:
-        self.model_update_group = None
+        if self.model_update_group is not None:
+            self.model_update_group.destroy()
+            self.model_update_group = None

vllm/engine/protocol.py

@@ -296,6 +296,10 @@ class EngineClient(ABC):
         """Initialize weight transfer for RL training."""
         raise NotImplementedError
 
+    async def shutdown_weight_transfer_engine(self) -> None:
+        """Release the trainer connection without shutting down inference."""
+        raise NotImplementedError
+
     async def start_weight_update(self) -> None:
         """Start a new weight update."""
         raise NotImplementedError

vllm/entrypoints/rl/offline.py

@@ -26,6 +26,10 @@ class RLOfflineMixin(OfflineInferenceMixin):
             "init_weight_transfer_engine", kwargs={"init_info": init_info_dict}
         )
 
+    def shutdown_weight_transfer_engine(self) -> None:
+        """Release the trainer connection without shutting down inference."""
+        self.llm_engine.collective_rpc("shutdown_weight_transfer_engine")
+
     def start_weight_update(self) -> None:
         """Start a new weight update."""
         self.llm_engine.collective_rpc("start_weight_update")

vllm/entrypoints/rl/online/api_router.py

@@ -168,6 +168,13 @@ async def init_weight_transfer_engine(
     return JSONResponse(content={"message": "Weight transfer initialized"})
 
 
+@router.post("/shutdown_weight_transfer_engine")
+async def shutdown_weight_transfer_engine(raw_request: Request):
+    with weight_operation_metrics().record("shutdown"):
+        await engine_client(raw_request).shutdown_weight_transfer_engine()
+    return JSONResponse(content={"message": "Weight transfer shut down"})
+
+
 @router.post("/start_weight_update")
 async def start_weight_update(raw_request: Request):
     with weight_operation_metrics().record("start"):

vllm/entrypoints/rl/online/metrics.py

@@ -10,7 +10,7 @@ from typing import Literal
 
 from prometheus_client import REGISTRY, CollectorRegistry, Gauge, Histogram
 
-Operation = Literal["init", "start", "start_draft", "update", "finish"]
+Operation = Literal["init", "shutdown", "start", "start_draft", "update", "finish"]
 
 
 class WeightOperationMetrics:

vllm/v1/engine/async_llm.py

@@ -1318,6 +1318,10 @@ class AsyncLLM(EngineClient):
             "init_weight_transfer_engine", kwargs={"init_info": request.init_info}
         )
 
+    async def shutdown_weight_transfer_engine(self) -> None:
+        """Release the trainer connection without shutting down inference."""
+        await self.collective_rpc("shutdown_weight_transfer_engine")
+
     async def start_weight_update(self) -> None:
         """Start a new weight update."""
         await self.collective_rpc("start_weight_update")

vllm/v1/worker/gpu_worker.py

@@ -1482,6 +1482,15 @@ class Worker(WorkerBase):
         typed_init_info = self.weight_transfer_engine.parse_init_info(init_info)
         self.weight_transfer_engine.init_transfer_engine(typed_init_info)
 
+    def shutdown_weight_transfer_engine(self) -> None:
+        """Release a trainer connection without shutting down inference."""
+        if self._weight_update_active:
+            raise RuntimeError(
+                "Cannot release a trainer connection during a weight update"
+            )
+        if self.weight_transfer_engine is not None:
+            self.weight_transfer_engine.shutdown()
+
     def start_weight_update(self) -> None:
         """Start a new weight update session.