A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 06:30 UTC
[RL] [Weight Sync] Guard IPC update-info pickle deserialization behind insecure serialization flag (#35928)
Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Simon Mo committed
Mar 4, 2026 at 14:05 UTC
f678c3f61a2f3f224f29d3574225a6660e818e7e
2 files changed
+26
-2
tests/distributed/test_weight_transfer.py
+18
-2
@@ -456,11 +456,13 @@ class TestIPCWeightTransferUpdateInfoValidation:
456
ipc_handles=ipc_handles,
457
)
458
459
- def test_valid_update_info_from_pickled(self):
459
+ def test_valid_update_info_from_pickled(self, monkeypatch):
460
"""Test creating IPCWeightTransferUpdateInfo from pickled handles."""
461
if torch.cuda.device_count() < 1:
462
pytest.skip("Need at least 1 GPU for this test")
463
464
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
465
+
466
dummy_tensor = torch.ones(10, 10, device="cuda:0")
467
ipc_handle = reduce_tensor(dummy_tensor)
468
gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
@@ -477,6 +479,18 @@ class TestIPCWeightTransferUpdateInfoValidation:
479
assert info.ipc_handles == ipc_handles
480
assert info.ipc_handles_pickled is None
481
482
+ def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
483
+ """Test that pickled handles are rejected unless env flag is enabled."""
484
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
485
+
486
+ with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
487
+ IPCWeightTransferUpdateInfo(
488
+ names=[],
489
+ dtype_names=[],
490
+ shapes=[],
491
+ ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
492
+ )
493
+
494
def test_both_handles_and_pickled_raises(self):
495
"""Test that providing both ipc_handles and ipc_handles_pickled raises."""
496
if torch.cuda.device_count() < 1:
@@ -556,11 +570,13 @@ class TestIPCEngineParsing:
570
assert update_info.shapes == [[100, 100], [50]]
571
assert len(update_info.ipc_handles) == 2
572
559
- def test_parse_update_info_pickled(self):
573
+ def test_parse_update_info_pickled(self, monkeypatch):
574
"""Test parsing update info with pickled IPC handles (HTTP path)."""
575
if torch.cuda.device_count() < 1:
576
pytest.skip("Need at least 1 GPU for this test")
577
578
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
579
+
580
config = WeightTransferConfig(backend="ipc")
581
parallel_config = create_mock_parallel_config()
582
engine = IPCWeightTransferEngine(config, parallel_config)
vllm/distributed/weight_transfer/ipc_engine.py
+8
@@ -12,6 +12,7 @@ import requests
12
import torch
13
from torch.multiprocessing.reductions import reduce_tensor
14
15
+from vllm import envs
16
from vllm.config.parallel import ParallelConfig
17
from vllm.config.weight_transfer import WeightTransferConfig
18
from vllm.distributed.weight_transfer.base import (
@@ -74,6 +75,13 @@ class IPCWeightTransferUpdateInfo(WeightTransferUpdateInfo):
75
raise ValueError(
76
"Cannot specify both `ipc_handles` and `ipc_handles_pickled`"
77
)
78
+
79
+ if not envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
80
+ raise ValueError(
81
+ "Refusing to deserialize `ipc_handles_pickled` without "
82
+ "VLLM_ALLOW_INSECURE_SERIALIZATION=1"
83
+ )
84
+
85
self.ipc_handles = pickle.loads(base64.b64decode(self.ipc_handles_pickled))
86
self.ipc_handles_pickled = None
87