[Bugfix][RL] Set vLLM config during weight reload (#45989)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Jul 20, 2026 at 15:02 UTC
9459fc647105f10f754697b3bf136d194564d603
2 files changed
+45
-12
tests/v1/worker/test_gpu_worker_weight_transfer.py
+28
@@ -9,6 +9,7 @@ session is active. These tests verify that delegation and the session guard.
9
10
import pytest
11
12
+from vllm.config import VllmConfig, get_current_vllm_config
13
from vllm.v1.worker.gpu_worker import Worker
14
15
@@ -21,29 +22,55 @@ class _RecordingEngine:
22
self.finished = False
23
self.reset_count = 0
24
self.update_calls: list[dict] = []
25
+ self.seen_configs: list[VllmConfig] = []
26
+
27
+ def _record_config(self) -> None:
28
+ self.seen_configs.append(get_current_vllm_config())
29
30
def start_weight_update(self) -> None:
31
+ self._record_config()
32
self.started = True
33
34
def update_weights(self, update_info: dict) -> None:
35
+ self._record_config()
36
self.update_calls.append(update_info)
37
if self.raise_on_update:
38
raise ValueError("boom")
39
40
def finish_weight_update(self) -> None:
41
+ self._record_config()
42
self.finished = True
43
44
def reset_weight_update_target(self) -> None:
45
self.reset_count += 1
46
47
48
+class _RecordingModelRunner:
49
+ def __init__(self) -> None:
50
+ self.seen_config: VllmConfig | None = None
51
+
52
+ def reload_weights(self) -> None:
53
+ self.seen_config = get_current_vllm_config()
54
+
55
+
56
def _make_worker(engine: _RecordingEngine | None) -> Worker:
57
worker = object.__new__(Worker)
58
+ worker.vllm_config = VllmConfig()
59
worker.weight_transfer_engine = engine
60
worker._weight_update_active = False
61
return worker
62
63
64
+def test_reload_weights_sets_current_config():
65
+ worker = _make_worker(None)
66
+ model_runner = _RecordingModelRunner()
67
+ worker.model_runner = model_runner # type: ignore[assignment]
68
+
69
+ Worker.reload_weights(worker)
70
+
71
+ assert model_runner.seen_config is worker.vllm_config
72
+
73
+
74
def test_start_update_finish_delegates_to_engine():
75
engine = _RecordingEngine()
76
worker = _make_worker(engine)
@@ -60,6 +87,7 @@ def test_start_update_finish_delegates_to_engine():
87
assert engine.finished is True
88
assert engine.reset_count == 1
89
assert worker._weight_update_active is False
90
+ assert engine.seen_configs == [worker.vllm_config] * 3
91
92
93
def test_double_start_raises():
vllm/v1/worker/gpu_worker.py
+17
-12
@@ -442,7 +442,8 @@ class Worker(WorkerBase):
442
self.model_runner.update_config(overrides)
443
444
def reload_weights(self, *args, **kwargs) -> None:
445
- self.model_runner.reload_weights(*args, **kwargs)
445
+ with set_current_vllm_config(self.vllm_config):
446
+ self.model_runner.reload_weights(*args, **kwargs)
447
448
@torch.inference_mode()
449
def determine_available_memory(self) -> int:
@@ -1301,14 +1302,16 @@ class Worker(WorkerBase):
1302
the configured weight transfer engine. The worker only tracks that a
1303
session is active.
1304
"""
1304
- self._start_weight_update()
1305
+ with set_current_vllm_config(self.vllm_config):
1306
+ self._start_weight_update()
1307
1308
def start_draft_weight_update(self) -> None:
1309
"""
1310
Like start_weight_update, but retargets the engine at the speculative
1311
draft model for this session.
1312
"""
1311
- self._start_weight_update(is_draft=True)
1313
+ with set_current_vllm_config(self.vllm_config):
1314
+ self._start_weight_update(is_draft=True)
1315
1316
def _start_weight_update(self, is_draft: bool = False) -> None:
1317
self._check_weight_transfer_engine()
@@ -1355,12 +1358,13 @@ class Worker(WorkerBase):
1358
"start_weight_update must be called before update_weights."
1359
)
1360
1358
- try:
1359
- self.weight_transfer_engine.update_weights(update_info)
1360
- except BaseException:
1361
- self._weight_update_active = False
1362
- self.weight_transfer_engine.reset_weight_update_target()
1363
- raise
1361
+ with set_current_vllm_config(self.vllm_config):
1362
+ try:
1363
+ self.weight_transfer_engine.update_weights(update_info)
1364
+ except BaseException:
1365
+ self._weight_update_active = False
1366
+ self.weight_transfer_engine.reset_weight_update_target()
1367
+ raise
1368
1369
def finish_weight_update(self) -> None:
1370
"""Finish the current weight update session."""
@@ -1372,9 +1376,10 @@ class Worker(WorkerBase):
1376
"finish_weight_update called without a matching start_weight_update."
1377
)
1378
1375
- self.weight_transfer_engine.finish_weight_update()
1376
- self.weight_transfer_engine.reset_weight_update_target()
1377
- self._weight_update_active = False
1379
+ with set_current_vllm_config(self.vllm_config):
1380
+ self.weight_transfer_engine.finish_weight_update()
1381
+ self.weight_transfer_engine.reset_weight_update_target()
1382
+ self._weight_update_active = False
1383
1384
def shutdown(self) -> None:
1385
gc.unfreeze()