From 9459fc647105f10f754697b3bf136d194564d603 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Mon, 20 Jul 2026 15:02:56 +0800 Subject: [PATCH] [Bugfix][RL] Set vLLM config during weight reload (#45989) Signed-off-by: aoshen02 --- .../worker/test_gpu_worker_weight_transfer.py | 28 ++++++++++++++++++ vllm/v1/worker/gpu_worker.py | 29 +++++++++++-------- 2 files changed, 45 insertions(+), 12 deletions(-) diff --git a/tests/v1/worker/test_gpu_worker_weight_transfer.py b/tests/v1/worker/test_gpu_worker_weight_transfer.py index aeb727d9ce3..6a97d64c6be 100644 --- a/tests/v1/worker/test_gpu_worker_weight_transfer.py +++ b/tests/v1/worker/test_gpu_worker_weight_transfer.py @@ -9,6 +9,7 @@ session is active. These tests verify that delegation and the session guard. import pytest +from vllm.config import VllmConfig, get_current_vllm_config from vllm.v1.worker.gpu_worker import Worker @@ -21,29 +22,55 @@ class _RecordingEngine: self.finished = False self.reset_count = 0 self.update_calls: list[dict] = [] + self.seen_configs: list[VllmConfig] = [] + + def _record_config(self) -> None: + self.seen_configs.append(get_current_vllm_config()) def start_weight_update(self) -> None: + self._record_config() self.started = True def update_weights(self, update_info: dict) -> None: + self._record_config() self.update_calls.append(update_info) if self.raise_on_update: raise ValueError("boom") def finish_weight_update(self) -> None: + self._record_config() self.finished = True def reset_weight_update_target(self) -> None: self.reset_count += 1 +class _RecordingModelRunner: + def __init__(self) -> None: + self.seen_config: VllmConfig | None = None + + def reload_weights(self) -> None: + self.seen_config = get_current_vllm_config() + + def _make_worker(engine: _RecordingEngine | None) -> Worker: worker = object.__new__(Worker) + worker.vllm_config = VllmConfig() worker.weight_transfer_engine = engine worker._weight_update_active = False return worker +def test_reload_weights_sets_current_config(): + worker = _make_worker(None) + model_runner = _RecordingModelRunner() + worker.model_runner = model_runner # type: ignore[assignment] + + Worker.reload_weights(worker) + + assert model_runner.seen_config is worker.vllm_config + + def test_start_update_finish_delegates_to_engine(): engine = _RecordingEngine() worker = _make_worker(engine) @@ -60,6 +87,7 @@ def test_start_update_finish_delegates_to_engine(): assert engine.finished is True assert engine.reset_count == 1 assert worker._weight_update_active is False + assert engine.seen_configs == [worker.vllm_config] * 3 def test_double_start_raises(): diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 9c20df0d18c..7ac35ad2329 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -442,7 +442,8 @@ class Worker(WorkerBase): self.model_runner.update_config(overrides) def reload_weights(self, *args, **kwargs) -> None: - self.model_runner.reload_weights(*args, **kwargs) + with set_current_vllm_config(self.vllm_config): + self.model_runner.reload_weights(*args, **kwargs) @torch.inference_mode() def determine_available_memory(self) -> int: @@ -1301,14 +1302,16 @@ class Worker(WorkerBase): the configured weight transfer engine. The worker only tracks that a session is active. """ - self._start_weight_update() + with set_current_vllm_config(self.vllm_config): + self._start_weight_update() def start_draft_weight_update(self) -> None: """ Like start_weight_update, but retargets the engine at the speculative draft model for this session. """ - self._start_weight_update(is_draft=True) + with set_current_vllm_config(self.vllm_config): + self._start_weight_update(is_draft=True) def _start_weight_update(self, is_draft: bool = False) -> None: self._check_weight_transfer_engine() @@ -1355,12 +1358,13 @@ class Worker(WorkerBase): "start_weight_update must be called before update_weights." ) - try: - self.weight_transfer_engine.update_weights(update_info) - except BaseException: - self._weight_update_active = False - self.weight_transfer_engine.reset_weight_update_target() - raise + with set_current_vllm_config(self.vllm_config): + try: + self.weight_transfer_engine.update_weights(update_info) + except BaseException: + self._weight_update_active = False + self.weight_transfer_engine.reset_weight_update_target() + raise def finish_weight_update(self) -> None: """Finish the current weight update session.""" @@ -1372,9 +1376,10 @@ class Worker(WorkerBase): "finish_weight_update called without a matching start_weight_update." ) - self.weight_transfer_engine.finish_weight_update() - self.weight_transfer_engine.reset_weight_update_target() - self._weight_update_active = False + with set_current_vllm_config(self.vllm_config): + self.weight_transfer_engine.finish_weight_update() + self.weight_transfer_engine.reset_weight_update_target() + self._weight_update_active = False def shutdown(self) -> None: gc.unfreeze()