diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index d8b436c646..49910d80c7 100644 --- a/rust/src/engine-core-client/src/client.rs +++ 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::::new(), + BTreeMap::::new(), + ) + .await?; + Ok(()) + } + /// Start a weight update for the base model. pub async fn start_weight_update(&self) -> Result<()> { self.collective_rpc( diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 251398611b..23c4244c82 100644 --- a/rust/src/server/src/routes.rs +++ b/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), diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index d574e02a92..9b8378f899 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/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); diff --git a/rust/src/server/src/routes/weight_transfer.rs b/rust/src/server/src/routes/weight_transfer.rs index 5ccd66e39b..7e1d6af61c 100644 --- a/rust/src/server/src/routes/weight_transfer.rs +++ b/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>, +) -> Result, 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>, diff --git a/tests/distributed/test_weight_transfer.py b/tests/distributed/test_weight_transfer.py index d044ef3bb7..642193ce88 100644 --- a/tests/distributed/test_weight_transfer.py +++ b/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).""" diff --git a/tests/distributed/test_weight_transfer_m2n.py b/tests/distributed/test_weight_transfer_m2n.py index 0591adff4d..588ffeeb26 100644 --- a/tests/distributed/test_weight_transfer_m2n.py +++ b/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 diff --git a/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py b/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py index b51d08db07..9e4df28d03 100644 --- a/tests/entrypoints/rl/return_routed_experts/test_route_request_bodies.py +++ b/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): diff --git a/tests/entrypoints/rl/weight_transfer/test_offline.py b/tests/entrypoints/rl/weight_transfer/test_offline.py index e4a30f17e6..0af71a1dee 100644 --- a/tests/entrypoints/rl/weight_transfer/test_offline.py +++ b/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): diff --git a/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py b/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py index 22ec5e6112..a3bb67ed2b 100644 --- a/tests/entrypoints/rl/weight_transfer_metrics/test_weight_operation_metrics.py +++ b/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": {}}), diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index f1965cb1fa..ec63c08c70 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/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)) diff --git a/vllm/distributed/weight_transfer/base.py b/vllm/distributed/weight_transfer/base.py index 3140b70e96..6eba3fa96f 100644 --- a/vllm/distributed/weight_transfer/base.py +++ b/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 diff --git a/vllm/distributed/weight_transfer/clients.py b/vllm/distributed/weight_transfer/clients.py index 745e67779e..0ccce60f97 100644 --- a/vllm/distributed/weight_transfer/clients.py +++ b/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 diff --git a/vllm/distributed/weight_transfer/m2n_engine.py b/vllm/distributed/weight_transfer/m2n_engine.py index c61bac4f2d..e480e60eaf 100644 --- a/vllm/distributed/weight_transfer/m2n_engine.py +++ b/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 = [] diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py index 2d1d9f47ec..8b9b73f9f9 100644 --- a/vllm/distributed/weight_transfer/m2n_trainer.py +++ b/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 diff --git a/vllm/distributed/weight_transfer/nccl_engine.py b/vllm/distributed/weight_transfer/nccl_engine.py index 77d8b60872..331cb5eacc 100644 --- a/vllm/distributed/weight_transfer/nccl_engine.py +++ b/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 diff --git a/vllm/distributed/weight_transfer/sparse_nccl_engine.py b/vllm/distributed/weight_transfer/sparse_nccl_engine.py index 01b29529de..f9233516a5 100644 --- a/vllm/distributed/weight_transfer/sparse_nccl_engine.py +++ b/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 diff --git a/vllm/engine/protocol.py b/vllm/engine/protocol.py index 15f6b801bc..15f8f8b744 100644 --- a/vllm/engine/protocol.py +++ b/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 diff --git a/vllm/entrypoints/rl/offline.py b/vllm/entrypoints/rl/offline.py index 1db6a82533..be72a9a554 100644 --- a/vllm/entrypoints/rl/offline.py +++ b/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") diff --git a/vllm/entrypoints/rl/online/api_router.py b/vllm/entrypoints/rl/online/api_router.py index 51182fd996..847287336e 100644 --- a/vllm/entrypoints/rl/online/api_router.py +++ b/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"): diff --git a/vllm/entrypoints/rl/online/metrics.py b/vllm/entrypoints/rl/online/metrics.py index ed16f9a820..b67f479f6d 100644 --- a/vllm/entrypoints/rl/online/metrics.py +++ b/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: diff --git a/vllm/v1/engine/async_llm.py b/vllm/v1/engine/async_llm.py index 71e30990c5..60588e0760 100644 --- a/vllm/v1/engine/async_llm.py +++ b/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") diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 152def6222..3efe77a3a4 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/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.