test_weight_transfer.py · tduncan/Nvidia-Optimization-Stack-VLLM

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 2,369 lines 84.2 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3	"""Tests for weight transfer engine backends.
4
5	Unit tests for engine classes (parsing, validation, registry).
6	Integration tests for NCCL and IPC weight transfer between processes using Ray.
7	"""
8
9	import builtins
10	import pickle
11	import runpy
12	import threading
13	import time
14	from contextlib import nullcontext
15	from pathlib import Path
16	from types import SimpleNamespace
17	from unittest.mock import MagicMock
18
19	import pybase64 as base64
20	import pytest
21	import ray
22	import torch
23	from torch.multiprocessing.reductions import reduce_tensor
24
25	from vllm.config.parallel import ParallelConfig
26	from vllm.config.weight_transfer import WeightTransferConfig
27	from vllm.distributed.weight_transfer import (
28	HTTPVLLMWeightSyncClient,
29	ModuleSource,
30	ParamMeta,
31	RayVLLMWeightSyncClient,
32	TrainerWeightTransferEngine,
33	VLLMWeightSyncClient,
34	WeightSource,
35	WeightTransferEngineFactory,
36	WeightTransferTrainerFactory,
37	)
38	from vllm.distributed.weight_transfer.base import (
39	TrainerInitInfo,
40	WeightTransferEngine,
41	WeightTransferInitRequest,
42	WeightTransferUpdateRequest,
43	layerwise_groups,
44	)
45	from vllm.distributed.weight_transfer.ipc_engine import (
46	IPCTrainerInitInfo,
47	IPCTrainerWeightTransferEngine,
48	IPCWeightTransferEngine,
49	IPCWeightTransferInitInfo,
50	IPCWeightTransferUpdateInfo,
51	)
52	from vllm.distributed.weight_transfer.nccl_engine import (
53	NCCLTrainerInitInfo,
54	NCCLTrainerWeightTransferEngine,
55	NCCLWeightTransferEngine,
56	NCCLWeightTransferInitInfo,
57	NCCLWeightTransferUpdateInfo,
58	)
59	from vllm.distributed.weight_transfer.packed_tensor import (
60	DEFAULT_PACKED_BUFFER_SIZE_BYTES,
61	DEFAULT_PACKED_NUM_BUFFERS,
62	)
63	from vllm.distributed.weight_transfer.sparse_nccl_engine import (
64	SparseNCCLTrainerInitInfo,
65	SparseNCCLTrainerWeightTransferEngine,
66	SparseNCCLWeightTransferEngine,
67	SparseNCCLWeightTransferUpdateInfo,
68	SparseWeightPatch,
69	)
70	from vllm.platforms import current_platform
71	from vllm.utils.nccl import _nccl_has_no_cache
72	from vllm.utils.network_utils import get_open_port
73
74
75	def _init_ray_for_weight_transfer() -> None:
76	if ray.is_initialized():
77	return
78	ray.init(
79	ignore_reinit_error=True,
80	runtime_env={
81	"env_vars": {
82	"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES": "1",
83	"RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES": "1",
84	"RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES": "1",
85	}
86	},
87	)
88
89
90	def _get_ray_assigned_device() -> torch.device:
91	gpu_ids = ray.get_gpu_ids()
92	if not gpu_ids:
93	return torch.device("cuda:0")
94	return torch.device(f"cuda:{int(gpu_ids[0])}")
95
96
97	def _set_ray_assigned_device() -> torch.device:
98	device = _get_ray_assigned_device()
99	current_platform.set_device(device)
100	return device
101
102
103	def create_mock_parallel_config(
104	rank: int = 0,
105	world_size: int = 1,
106	dp_rank: int = 0,
107	) -> ParallelConfig:
108	"""Create a mock ParallelConfig for testing."""
109	config = MagicMock(spec=ParallelConfig)
110	config.rank = rank
111	config.world_size = world_size
112	config.data_parallel_rank = dp_rank
113	config.data_parallel_index = dp_rank
114	return config
115
116
117	def create_mock_vllm_config(
118	rank: int = 0,
119	world_size: int = 1,
120	dp_rank: int = 0,
121	) -> MagicMock:
122	"""Create a mock VllmConfig exposing parallel_config and model_config."""
123	vllm_config = MagicMock()
124	vllm_config.parallel_config = create_mock_parallel_config(rank, world_size, dp_rank)
125	vllm_config.model_config = MagicMock()
126	return vllm_config
127
128
129	# --- Unit Tests: NCCLWeightTransferUpdateInfo Validation ---
130
131
132	class TestNCCLWeightTransferUpdateInfoValidation:
133	"""Test NCCLWeightTransferUpdateInfo dataclass validation."""
134
135	def test_valid_update_info(self):
136	info = NCCLWeightTransferUpdateInfo(
137	names=["layer.weight", "layer.bias"],
138	dtype_names=["float32", "float32"],
139	shapes=[[10, 10], [10]],
140	)
141	assert info.names == ["layer.weight", "layer.bias"]
142	assert info.dtype_names == ["float32", "float32"]
143	assert info.shapes == [[10, 10], [10]]
144
145	def test_mismatched_dtype_names_raises(self):
146	with pytest.raises(ValueError, match="dtype_names"):
147	NCCLWeightTransferUpdateInfo(
148	names=["layer.weight", "layer.bias"],
149	dtype_names=["float32"], # Only one dtype
150	shapes=[[10, 10], [10]],
151	)
152
153	def test_mismatched_shapes_raises(self):
154	with pytest.raises(ValueError, match="shapes"):
155	NCCLWeightTransferUpdateInfo(
156	names=["layer.weight", "layer.bias"],
157	dtype_names=["float32", "float32"],
158	shapes=[[10, 10]], # Only one shape
159	)
160
161	def test_empty_lists_valid(self):
162	info = NCCLWeightTransferUpdateInfo(names=[], dtype_names=[], shapes=[])
163	assert len(info.names) == 0
164
165
166	# --- Unit Tests: SparseNCCLWeightTransferUpdateInfo Validation ---
167
168
169	class TestSparseNCCLWeightTransferUpdateInfoValidation:
170	"""Test SparseNCCLWeightTransferUpdateInfo dataclass validation."""
171
172	def test_valid_sparse_update_info(self):
173	info = SparseNCCLWeightTransferUpdateInfo(
174	names=["layer.weight", "layer.bias"],
175	dtype_names=["float32", "bfloat16"],
176	shapes=[[10, 10], [10]],
177	num_updates_list=[4, 2],
178	)
179	assert info.num_updates_list == [4, 2]
180
181	def test_mismatched_dtype_names_raises(self):
182	with pytest.raises(ValueError, match="dtype_names"):
183	SparseNCCLWeightTransferUpdateInfo(
184	names=["layer.weight", "layer.bias"],
185	dtype_names=["float32"],
186	shapes=[[10, 10], [10]],
187	num_updates_list=[4, 2],
188	)
189
190	def test_rejects_empty_num_updates_list(self):
191	with pytest.raises(ValueError, match="cannot be empty"):
192	SparseNCCLWeightTransferUpdateInfo(
193	names=[],
194	dtype_names=[],
195	shapes=[],
196	num_updates_list=[],
197	)
198
199	def test_rejects_mismatched_num_updates(self):
200	with pytest.raises(ValueError, match="`num_updates_list`"):
201	SparseNCCLWeightTransferUpdateInfo(
202	names=["layer.weight", "layer.bias"],
203	dtype_names=["float32", "float32"],
204	shapes=[[10, 10], [10]],
205	num_updates_list=[3],
206	)
207
208	def test_rejects_negative_num_updates(self):
209	with pytest.raises(ValueError, match="non-negative"):
210	SparseNCCLWeightTransferUpdateInfo(
211	names=["layer.weight"],
212	dtype_names=["float32"],
213	shapes=[[10, 10]],
214	num_updates_list=[-1],
215	)
216
217
218	# --- Unit Tests: Engine Parsing ---
219
220
221	class TestNCCLEngineParsing:
222	"""Test NCCLWeightTransferEngine parsing methods."""
223
224	def _make_engine(self):
225	config = WeightTransferConfig(backend="nccl")
226	return NCCLWeightTransferEngine(
227	config,
228	create_mock_vllm_config(),
229	torch.device("cuda"),
230	MagicMock(spec=torch.nn.Module),
231	)
232
233	def test_parse_init_info_valid(self):
234	engine = self._make_engine()
235	init_info = engine.parse_init_info(
236	{
237	"master_address": "127.0.0.1",
238	"master_port": 12345,
239	"rank_offset": 1,
240	"world_size": 3,
241	}
242	)
243	assert isinstance(init_info, NCCLWeightTransferInitInfo)
244	assert init_info.master_address == "127.0.0.1"
245	assert init_info.master_port == 12345
246	assert init_info.rank_offset == 1
247	assert init_info.world_size == 3
248
249	def test_parse_init_info_missing_field_raises(self):
250	engine = self._make_engine()
251	with pytest.raises(ValueError, match="Invalid init_info"):
252	engine.parse_init_info({"master_address": "127.0.0.1"})
253
254	def test_parse_update_info_valid(self):
255	engine = self._make_engine()
256	update_info = engine.parse_update_info(
257	{
258	"names": ["w1", "w2"],
259	"dtype_names": ["float32", "bfloat16"],
260	"shapes": [[100, 100], [50]],
261	}
262	)
263	assert isinstance(update_info, NCCLWeightTransferUpdateInfo)
264	assert update_info.names == ["w1", "w2"]
265	assert update_info.dtype_names == ["float32", "bfloat16"]
266	assert update_info.shapes == [[100, 100], [50]]
267
268
269	# --- Unit Tests: Engine Registry ---
270
271
272	class TestEngineRegistry:
273	"""Test weight transfer engine registry."""
274
275	def test_create_engine_nccl(self):
276	config = WeightTransferConfig(backend="nccl")
277	engine = WeightTransferEngineFactory.create_engine(
278	config,
279	create_mock_vllm_config(),
280	torch.device("cuda"),
281	MagicMock(spec=torch.nn.Module),
282	)
283	assert isinstance(engine, NCCLWeightTransferEngine)
284
285	def test_create_engine_ipc(self):
286	config = WeightTransferConfig(backend="ipc")
287	engine = WeightTransferEngineFactory.create_engine(
288	config,
289	create_mock_vllm_config(),
290	torch.device("cuda"),
291	MagicMock(spec=torch.nn.Module),
292	)
293	assert isinstance(engine, IPCWeightTransferEngine)
294
295	def test_create_engine_sparse_nccl(self):
296	config = WeightTransferConfig(backend="sparse_nccl")
297	engine = WeightTransferEngineFactory.create_engine(
298	config,
299	create_mock_vllm_config(),
300	torch.device("cuda"),
301	MagicMock(spec=torch.nn.Module),
302	)
303	assert isinstance(engine, SparseNCCLWeightTransferEngine)
304
305	def test_modelexpress_registration_is_native_and_lazy(self, monkeypatch):
306	from vllm.distributed.weight_transfer import factory
307
308	imported = []
309	module = MagicMock()
310	module.ModelExpressWeightTransferEngine.__name__ = (
311	"ModelExpressWeightTransferEngine"
312	)
313
314	def import_module(name):
315	imported.append(name)
316	return module
317
318	monkeypatch.setattr(factory.importlib, "import_module", import_module)
319	config = WeightTransferConfig(backend="modelexpress")
320	vllm_config = create_mock_vllm_config()
321	device = torch.device("cpu")
322	model = MagicMock(spec=torch.nn.Module)
323
324	assert imported == []
325	engine = WeightTransferEngineFactory.create_engine(
326	config, vllm_config, device, model
327	)
328
329	assert imported == ["vllm.distributed.weight_transfer.modelexpress_engine"]
330	module.ModelExpressWeightTransferEngine.assert_called_once_with(
331	config, vllm_config, device, model
332	)
333	assert engine is module.ModelExpressWeightTransferEngine.return_value
334
335	def test_create_engine_invalid_backend(self):
336	config = WeightTransferConfig(backend="invalid")
337	with pytest.raises(ValueError, match="Invalid weight transfer backend"):
338	WeightTransferEngineFactory.create_engine(
339	config,
340	create_mock_vllm_config(),
341	torch.device("cuda"),
342	MagicMock(spec=torch.nn.Module),
343	)
344
345	def test_register_duplicate_raises(self):
346	with pytest.raises(ValueError, match="already registered"):
347	WeightTransferEngineFactory.register_engine(
348	"nccl", NCCLWeightTransferEngine
349	)
350
351
352	# --- Test receive_weights without init raises ---
353
354
355	def test_nccl_receive_weights_without_init_raises():
356	"""Test that receive_weights raises if init_transfer_engine wasn't called."""
357	if torch.accelerator.device_count() < 1:
358	pytest.skip("Need at least 1 GPU for this test")
359
360	config = WeightTransferConfig(backend="nccl")
361	engine = NCCLWeightTransferEngine(
362	config,
363	create_mock_vllm_config(),
364	torch.device("cuda"),
365	MagicMock(spec=torch.nn.Module),
366	)
367
368	update_info = NCCLWeightTransferUpdateInfo(
369	names=["w"], dtype_names=["float32"], shapes=[[10]]
370	)
371
372	with pytest.raises(RuntimeError, match="not initialized"):
373	engine.receive_weights(update_info)
374
375
376	def test_sparse_nccl_receive_weights_without_init_raises():
377	"""Test that sparse receive raises if init_transfer_engine wasn't called."""
378	if torch.accelerator.device_count() < 1:
379	pytest.skip("Need at least 1 GPU for this test")
380
381	config = WeightTransferConfig(backend="sparse_nccl")
382	engine = SparseNCCLWeightTransferEngine(
383	config,
384	create_mock_vllm_config(),
385	torch.device("cuda"),
386	MagicMock(spec=torch.nn.Module),
387	)
388
389	update_info = SparseNCCLWeightTransferUpdateInfo(
390	names=["w"],
391	dtype_names=["float32"],
392	shapes=[[10]],
393	num_updates_list=[2],
394	)
395
396	with pytest.raises(RuntimeError, match="not initialized"):
397	engine.receive_weights(update_info)
398
399
400	# --- Integration Test: NCCL Weight Transfer Between Ray Tasks ---
401
402
403	@ray.remote(num_gpus=1)
404	def trainer_broadcast_tensor(
405	master_address: str,
406	master_port: int,
407	world_size: int,
408	tensor_shape: list[int],
409	tensor_dtype: str,
410	) -> bool:
411	"""Trainer task that broadcasts a tensor via NCCL."""
412	import torch
413
414	device = _set_ray_assigned_device()
415
416	from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
417	from vllm.distributed.utils import StatelessProcessGroup
418
419	# Create process group as rank 0 (trainer)
420	pg = StatelessProcessGroup.create(
421	host=master_address,
422	port=master_port,
423	rank=0,
424	world_size=world_size,
425	)
426	comm = PyNcclCommunicator(pg, device=device.index)
427
428	# Create and broadcast the tensor
429	dtype = getattr(torch, tensor_dtype)
430	tensor_to_send = torch.ones(tensor_shape, dtype=dtype, device=device)
431	comm.broadcast(tensor_to_send, src=0, stream=torch.cuda.current_stream())
432	torch.accelerator.synchronize()
433
434	return True
435
436
437	# max_calls=1: a batch-invariant run leaves NCCL pins in the worker process.
438	@ray.remote(num_gpus=1, max_calls=1)
439	def inference_receive_tensor(
440	master_address: str,
441	master_port: int,
442	world_size: int,
443	tensor_shape: list[int],
444	tensor_dtype: str,
445	batch_invariant: bool = False,
446	) -> dict:
447	"""Inference task that receives tensor via NCCLWeightTransferEngine."""
448	import contextlib
449	from unittest.mock import MagicMock
450
451	import torch
452
453	_set_ray_assigned_device()
454	if batch_invariant:
455	from vllm.model_executor.determinism.batch_invariant import (
456	override_envs_for_invariance,
457	)
458
459	override_envs_for_invariance()
460	# Like vLLM's own groups, a pinned communicator makes NCCL read the
461	# pins before the transfer group exists.
462	torch.distributed.init_process_group(
463	"nccl",
464	init_method=f"tcp://127.0.0.1:{get_open_port()}",
465	rank=0,
466	world_size=1,
467	)
468	torch.distributed.all_reduce(torch.ones(1, device="cuda"))
469
470	from vllm.config.parallel import ParallelConfig
471	from vllm.config.weight_transfer import WeightTransferConfig
472	from vllm.distributed.weight_transfer.nccl_engine import (
473	NCCLWeightTransferEngine,
474	NCCLWeightTransferInitInfo,
475	NCCLWeightTransferUpdateInfo,
476	)
477
478	class Recorder(torch.nn.Module):
479	def __init__(self):
480	super().__init__()
481	self.received = []
482
483	def load_weights(self, weights):
484	for name, tensor in weights:
485	self.received.append((name, tensor.clone()))
486
487	config = WeightTransferConfig(backend="nccl")
488	vllm_config = MagicMock()
489	parallel_config = MagicMock(spec=ParallelConfig)
490	parallel_config.rank = 0
491	parallel_config.world_size = 1
492	parallel_config.data_parallel_rank = 0
493	parallel_config.data_parallel_index = 0
494	vllm_config.parallel_config = parallel_config
495	vllm_config.model_config = MagicMock()
496
497	recorder = Recorder()
498	engine = NCCLWeightTransferEngine(
499	config, vllm_config, torch.device("cuda"), recorder
500	)
501	# Transport-only test: bypass the set_current_vllm_config context that
502	# receive_weights enters, since vllm_config here is a mock.
503	import vllm.config as _vllm_config_mod
504
505	_vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
506
507	# Initialize the engine (joins as rank 1)
508	# Trainer broadcasts a single tensor unpacked, so the worker must not
509	# expect the packed wire format (packed is a must-agree wire param shipped
510	# on the init info).
511	init_info = NCCLWeightTransferInitInfo(
512	master_address=master_address,
513	master_port=master_port,
514	rank_offset=1, # Trainer is rank 0, we become rank 1
515	world_size=world_size,
516	packed=False,
517	)
518	engine.init_transfer_engine(init_info)
519
520	update_info = NCCLWeightTransferUpdateInfo(
521	names=["test.weight"],
522	dtype_names=[tensor_dtype],
523	shapes=[tensor_shape],
524	)
525	engine.receive_weights(update_info)
526	torch.accelerator.synchronize()
527
528	# Verify we received the tensor
529	success = False
530	received_shape = None
531	received_sum = None
532
533	if len(recorder.received) == 1:
534	name, tensor = recorder.received[0]
535	received_shape = list(tensor.shape)
536	received_sum = tensor.sum().item()
537	if received_shape == tensor_shape:
538	expected_sum = 1.0 * torch.tensor(tensor_shape).prod().item()
539	if abs(received_sum - expected_sum) < 0.01:
540	success = True
541
542	engine.shutdown()
543
544	return {
545	"success": success,
546	"received_shape": received_shape,
547	"received_sum": received_sum,
548	}
549
550
551	@pytest.mark.skipif(
552	torch.accelerator.device_count() < 2,
553	reason="Need at least 2 GPUs to run NCCL weight transfer test.",
554	)
555	@pytest.mark.parametrize(
556	"batch_invariant",
557	[
558	pytest.param(False, id="default"),
559	pytest.param(
560	True,
561	id="batch-invariant-worker",
562	marks=pytest.mark.skipif(
563	not _nccl_has_no_cache(), reason="Needs CUDA NCCL >= 2.29.7."
564	),
565	),
566	],
567	)
568	def test_nccl_weight_transfer_between_processes(batch_invariant):
569	"""Test NCCL weight transfer from trainer to inference process using Ray.
570
571	This test verifies that the NCCLWeightTransferEngine can receive
572	tensors broadcast by a trainer process via NCCL, including when the
573	worker runs with batch-invariance NCCL pins the trainer does not have.
574	"""
575	_init_ray_for_weight_transfer()
576
577	master_address = "127.0.0.1"
578	master_port = get_open_port()
579	world_size = 2 # 1 trainer + 1 inference worker
580
581	tensor_shape = [100, 100]
582	tensor_dtype = "float32"
583
584	inference_future = inference_receive_tensor.remote(
585	master_address,
586	master_port,
587	world_size,
588	tensor_shape,
589	tensor_dtype,
590	batch_invariant,
591	)
592	trainer_future = trainer_broadcast_tensor.remote(
593	master_address, master_port, world_size, tensor_shape, tensor_dtype
594	)
595
596	# A mismatched NCCL config deadlocks instead of failing.
597	trainer_result, result = ray.get([trainer_future, inference_future], timeout=300)
598
599	assert trainer_result, "Trainer should complete successfully"
600	assert result["success"], (
601	f"Weight transfer failed. "
602	f"Received shape: {result['received_shape']}, "
603	f"Received sum: {result['received_sum']}"
604	)
605
606
607	def test_sparse_nccl_checkpoint_chunks_to_ep_local_experts_cpu(monkeypatch):
608	"""Replay global expert patches through two EP-local loaders on CPU."""
609
610	class CaptureClient:
611	def __init__(self):
612	self.order: list[str] = []
613	self.update_infos: list[dict] = []
614
615	def start_weight_update(self):
616	self.order.append("start")
617
618	def update_weights(self, update_info):
619	self.order.append("update")
620	self.update_infos.append(update_info)
621
622	def finish_weight_update(self):
623	self.order.append("finish")
624
625	monkeypatch.setattr(torch.cuda, "current_stream", lambda _, *__: None)
626	monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
627	monkeypatch.setattr(torch.accelerator, "device_index", lambda *_: nullcontext())
628
629	client = CaptureClient()
630	sender = SparseNCCLTrainerWeightTransferEngine(client=client)
631	sender.model_update_group = MagicMock()
632	sender.model_update_group.device = torch.device("cpu")
633	wire_payloads = []
634	sender.model_update_group.broadcast.side_effect = lambda tensor, **_: (
635	wire_payloads.append(tensor.clone())
636	)
637
638	expert_names = [
639	"model.layers.0.mlp.experts.0.gate_proj.weight",
640	"model.layers.0.mlp.experts.1.gate_proj.weight",
641	]
642	client.start_weight_update()
643	for name, index, value in zip(
644	expert_names,
645	(0, 3),
646	(5.0, 7.0),
647	strict=True,
648	):
649	sender.send_weight_chunk(
650	[
651	SparseWeightPatch(
652	name=name,
653	indices=torch.tensor([index], dtype=torch.int32),
654	values=torch.tensor([value]),
655	full_shape=(2, 2),
656	)
657	]
658	)
659	client.finish_weight_update()
660
661	assert client.order == ["start", "update", "update", "finish"]
662	assert [info["names"] for info in client.update_infos] == [
663	[expert_names[0]],
664	[expert_names[1]],
665	]
666	assert [info["shapes"] for info in client.update_infos] == [
667	[[2, 2]],
668	[[2, 2]],
669	]
670	assert [info["num_updates_list"] for info in client.update_infos] == [
671	[1],
672	[1],
673	]
674
675	expected_gates = (
676	[[5.0, -1.0], [-1.0, -1.0]],
677	[[-1.0, -1.0], [-1.0, 7.0]],
678	)
679	for ep_rank, expected_gate in enumerate(expected_gates):
680	model = torch.nn.Module()
681	model.register_parameter(
682	"w13_weight",
683	torch.nn.Parameter(torch.full((1, 4, 2), -1.0), requires_grad=False),
684	)
685	local_name = expert_names[ep_rank]
686	seen_names: list[str] = []
687	loaded_names: set[str] = set()
688	copy_count = [0]
689
690	def load_weights(
691	weights,
692	local_name=local_name,
693	target=model,
694	seen=seen_names,
695	loaded=loaded_names,
696	count=copy_count,
697	):
698	for name, checkpoint_weight in weights:
699	seen.append(name)
700	if name != local_name:
701	continue
702	target.w13_weight.data[0, :2].copy_(checkpoint_weight)
703	loaded.add(name)
704	count[0] += 1
705	return loaded
706
707	model.load_weights = load_weights
708	receiver = SparseNCCLWeightTransferEngine(
709	WeightTransferConfig(backend="sparse_nccl"),
710	create_mock_vllm_config(rank=ep_rank, world_size=2),
711	torch.device("cpu"),
712	model,
713	)
714	receiver.model_update_group = MagicMock()
715	receiver.model_update_group.device = torch.device("cpu")
716	payloads = iter(wire_payloads)
717	receiver.model_update_group.broadcast.side_effect = (
718	lambda tensor, payloads=payloads, **_: tensor.copy_(next(payloads))
719	)
720
721	receiver.start_weight_update()
722	for update_info in client.update_infos:
723	receiver.receive_weights(SparseNCCLWeightTransferUpdateInfo(**update_info))
724	receiver.finish_weight_update()
725
726	assert seen_names == expert_names
727	assert loaded_names == {local_name}
728	assert copy_count == [1]
729	assert model.w13_weight[0, :2].tolist() == expected_gate
730	assert model.w13_weight[0, 2:].eq(-1).all()
731	foreign_value = 7.0 if ep_rank == 0 else 5.0
732	assert not model.w13_weight.eq(foreign_value).any()
733	receiver.shutdown()
734
735	sender.shutdown()
736
737
738	# --- Unit Tests: IPCWeightTransferUpdateInfo Validation ---
739
740
741	class TestIPCWeightTransferUpdateInfoValidation:
742	"""Test IPCWeightTransferUpdateInfo dataclass validation."""
743
744	def test_valid_update_info(self):
745	if torch.accelerator.device_count() < 1:
746	pytest.skip("Need at least 1 GPU for this test")
747
748	dummy_tensor = torch.ones(10, 10, device="cuda:0")
749	_, ipc_handle = reduce_tensor(dummy_tensor)
750	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
751	ipc_handles = [{gpu_uuid: ipc_handle}]
752
753	info = IPCWeightTransferUpdateInfo(
754	names=["layer.weight"],
755	dtype_names=["float32"],
756	shapes=[[10, 10]],
757	ipc_handles=ipc_handles,
758	)
759	assert info.names == ["layer.weight"]
760	assert info.dtype_names == ["float32"]
761	assert info.shapes == [[10, 10]]
762	assert len(info.ipc_handles) == 1
763
764	def test_mismatched_dtype_names_raises(self):
765	if torch.accelerator.device_count() < 1:
766	pytest.skip("Need at least 1 GPU for this test")
767
768	dummy_tensor = torch.ones(10, 10, device="cuda:0")
769	_, ipc_handle = reduce_tensor(dummy_tensor)
770	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
771	ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
772
773	with pytest.raises(ValueError, match="dtype_names"):
774	IPCWeightTransferUpdateInfo(
775	names=["layer.weight", "layer.bias"],
776	dtype_names=["float32"], # Only one dtype
777	shapes=[[10, 10], [10]],
778	ipc_handles=ipc_handles,
779	)
780
781	def test_mismatched_shapes_raises(self):
782	if torch.accelerator.device_count() < 1:
783	pytest.skip("Need at least 1 GPU for this test")
784
785	dummy_tensor = torch.ones(10, 10, device="cuda:0")
786	_, ipc_handle = reduce_tensor(dummy_tensor)
787	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
788	ipc_handles = [{gpu_uuid: ipc_handle}, {gpu_uuid: ipc_handle}]
789
790	with pytest.raises(ValueError, match="shapes"):
791	IPCWeightTransferUpdateInfo(
792	names=["layer.weight", "layer.bias"],
793	dtype_names=["float32", "float32"],
794	shapes=[[10, 10]], # Only one shape
795	ipc_handles=ipc_handles,
796	)
797
798	def test_mismatched_ipc_handles_raises(self):
799	if torch.accelerator.device_count() < 1:
800	pytest.skip("Need at least 1 GPU for this test")
801
802	dummy_tensor = torch.ones(10, 10, device="cuda:0")
803	_, ipc_handle = reduce_tensor(dummy_tensor)
804	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
805	ipc_handles = [{gpu_uuid: ipc_handle}] # Only one handle
806
807	with pytest.raises(ValueError, match="ipc_handles"):
808	IPCWeightTransferUpdateInfo(
809	names=["layer.weight", "layer.bias"],
810	dtype_names=["float32", "float32"],
811	shapes=[[10, 10], [10]],
812	ipc_handles=ipc_handles,
813	)
814
815	def test_valid_update_info_from_pickled(self, monkeypatch):
816	if torch.accelerator.device_count() < 1:
817	pytest.skip("Need at least 1 GPU for this test")
818
819	monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
820
821	dummy_tensor = torch.ones(10, 10, device="cuda:0")
822	ipc_handle = reduce_tensor(dummy_tensor)
823	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
824	ipc_handles = [{gpu_uuid: ipc_handle}]
825
826	pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
827
828	info = IPCWeightTransferUpdateInfo(
829	names=["layer.weight"],
830	dtype_names=["float32"],
831	shapes=[[10, 10]],
832	ipc_handles_pickled=pickled,
833	)
834	assert info.ipc_handles == ipc_handles
835	assert info.ipc_handles_pickled is None
836
837	def test_pickled_requires_insecure_serialization_flag(self, monkeypatch):
838	monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "0")
839
840	with pytest.raises(ValueError, match="VLLM_ALLOW_INSECURE_SERIALIZATION=1"):
841	IPCWeightTransferUpdateInfo(
842	names=[],
843	dtype_names=[],
844	shapes=[],
845	ipc_handles_pickled=base64.b64encode(pickle.dumps([])).decode("utf-8"),
846	)
847
848	def test_both_handles_and_pickled_raises(self):
849	if torch.accelerator.device_count() < 1:
850	pytest.skip("Need at least 1 GPU for this test")
851
852	dummy_tensor = torch.ones(10, 10, device="cuda:0")
853	ipc_handle = reduce_tensor(dummy_tensor)
854	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
855	ipc_handles = [{gpu_uuid: ipc_handle}]
856
857	pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
858
859	with pytest.raises(ValueError, match="Cannot specify both"):
860	IPCWeightTransferUpdateInfo(
861	names=["layer.weight"],
862	dtype_names=["float32"],
863	shapes=[[10, 10]],
864	ipc_handles=ipc_handles,
865	ipc_handles_pickled=pickled,
866	)
867
868	def test_neither_handles_nor_pickled_raises(self):
869	with pytest.raises(ValueError, match="must be provided"):
870	IPCWeightTransferUpdateInfo(
871	names=["layer.weight"],
872	dtype_names=["float32"],
873	shapes=[[10, 10]],
874	)
875
876	def test_empty_lists_valid(self):
877	info = IPCWeightTransferUpdateInfo(
878	names=[],
879	dtype_names=[],
880	shapes=[],
881	ipc_handles=[],
882	)
883	assert len(info.names) == 0
884
885
886	# --- Unit Tests: IPC Engine Parsing ---
887
888
889	class TestIPCEngineParsing:
890	"""Test IPCWeightTransferEngine parsing methods."""
891
892	def _make_engine(self):
893	config = WeightTransferConfig(backend="ipc")
894	return IPCWeightTransferEngine(
895	config,
896	create_mock_vllm_config(),
897	torch.device("cuda"),
898	MagicMock(spec=torch.nn.Module),
899	)
900
901	def test_parse_update_info_valid(self):
902	if torch.accelerator.device_count() < 1:
903	pytest.skip("Need at least 1 GPU for this test")
904
905	engine = self._make_engine()
906
907	dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
908	dummy_tensor2 = torch.ones(50, device="cuda:0")
909	_, ipc_args1 = reduce_tensor(dummy_tensor1)
910	_, ipc_args2 = reduce_tensor(dummy_tensor2)
911	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
912	ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]
913
914	update_info = engine.parse_update_info(
915	{
916	"names": ["w1", "w2"],
917	"dtype_names": ["float32", "bfloat16"],
918	"shapes": [[100, 100], [50]],
919	"ipc_handles": ipc_handles,
920	}
921	)
922
923	assert isinstance(update_info, IPCWeightTransferUpdateInfo)
924	assert update_info.names == ["w1", "w2"]
925	assert update_info.dtype_names == ["float32", "bfloat16"]
926	assert update_info.shapes == [[100, 100], [50]]
927	assert len(update_info.ipc_handles) == 2
928
929	def test_parse_update_info_pickled(self, monkeypatch):
930	if torch.accelerator.device_count() < 1:
931	pytest.skip("Need at least 1 GPU for this test")
932
933	monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
934
935	engine = self._make_engine()
936
937	dummy_tensor1 = torch.ones(100, 100, device="cuda:0")
938	dummy_tensor2 = torch.ones(50, device="cuda:0")
939	_, ipc_args1 = reduce_tensor(dummy_tensor1)
940	_, ipc_args2 = reduce_tensor(dummy_tensor2)
941	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
942	ipc_handles = [{gpu_uuid: ipc_args1}, {gpu_uuid: ipc_args2}]
943
944	pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
945
946	update_info = engine.parse_update_info(
947	{
948	"names": ["w1", "w2"],
949	"dtype_names": ["float32", "bfloat16"],
950	"shapes": [[100, 100], [50]],
951	"ipc_handles_pickled": pickled,
952	}
953	)
954
955	assert isinstance(update_info, IPCWeightTransferUpdateInfo)
956	assert update_info.names == ["w1", "w2"]
957	assert len(update_info.ipc_handles) == 2
958	assert gpu_uuid in update_info.ipc_handles[0]
959	assert gpu_uuid in update_info.ipc_handles[1]
960
961	def test_parse_update_info_ignores_none_pickled_handles(self):
962	engine = self._make_engine()
963	ipc_handles = [{"gpu-uuid": ("ipc-args",)}]
964
965	update_info = engine.parse_update_info(
966	{
967	"names": ["w1"],
968	"dtype_names": ["float32"],
969	"shapes": [[1]],
970	"ipc_handles": ipc_handles,
971	"ipc_handles_pickled": None,
972	}
973	)
974
975	assert isinstance(update_info, IPCWeightTransferUpdateInfo)
976	assert update_info.ipc_handles == ipc_handles
977
978	def test_parse_update_info_both_handles_and_pickled_raises(self):
979	if torch.accelerator.device_count() < 1:
980	pytest.skip("Need at least 1 GPU for this test")
981
982	engine = self._make_engine()
983
984	dummy_tensor = torch.ones(10, 10, device="cuda:0")
985	_, ipc_handle = reduce_tensor(dummy_tensor)
986	gpu_uuid = str(torch.cuda.get_device_properties(0).uuid)
987	ipc_handles = [{gpu_uuid: ipc_handle}]
988
989	pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
990
991	with pytest.raises(ValueError, match="Cannot specify both"):
992	engine.parse_update_info(
993	{
994	"names": ["layer.weight"],
995	"dtype_names": ["float32"],
996	"shapes": [[10, 10]],
997	"ipc_handles": ipc_handles,
998	"ipc_handles_pickled": pickled,
999	}
1000	)
1001
1002
1003	# --- Integration Test: IPC Weight Transfer Between Ray Tasks ---
1004
1005
1006	def get_physical_gpu_id(device_index: int = 0) -> str:
1007	"""Get physical GPU UUID for a device."""
1008	props = torch.cuda.get_device_properties(device_index)
1009	return str(props.uuid)
1010
1011
1012	@ray.remote(num_gpus=0.5)
1013	class TrainerActor:
1014	"""Trainer actor that creates and holds CUDA IPC handles."""
1015
1016	def __init__(self, tensor_shape: list[int], tensor_dtype: str):
1017	device = _set_ray_assigned_device()
1018
1019	# Create tensor on GPU and keep it alive
1020	dtype = getattr(torch, tensor_dtype)
1021	self.tensor = torch.ones(tensor_shape, dtype=dtype, device=device)
1022	self.tensor.fill_(42.0) # Fill with 42 to verify correct transfer
1023
1024	_, ipc_args = reduce_tensor(self.tensor)
1025	gpu_uuid = get_physical_gpu_id(device.index)
1026
1027	torch.accelerator.synchronize()
1028
1029	self.ipc_handle_dict = {
1030	"ipc_handle": ipc_args,
1031	"gpu_uuid": gpu_uuid,
1032	"shape": tensor_shape,
1033	"dtype": tensor_dtype,
1034	}
1035
1036	def get_ipc_handle_dict(self) -> dict:
1037	"""Return IPC handle dict. Tensor stays alive in this actor."""
1038	return self.ipc_handle_dict
1039
1040
1041	@ray.remote(num_gpus=0.5)
1042	def inference_receive_ipc_tensor(
1043	ipc_handle_dict: dict,
1044	mode: str = "ray",
1045	) -> dict:
1046	"""Inference task that receives tensor via IPCWeightTransferEngine."""
1047	import contextlib
1048	import os
1049
1050	# Worker-side: ipc_handles_pickled is deserialized via pickle.
1051	if mode == "http":
1052	os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
1053
1054	from unittest.mock import MagicMock
1055
1056	import torch
1057
1058	device = _set_ray_assigned_device()
1059
1060	from vllm.config.parallel import ParallelConfig
1061	from vllm.config.weight_transfer import WeightTransferConfig
1062	from vllm.distributed.weight_transfer.ipc_engine import (
1063	IPCWeightTransferEngine,
1064	)
1065
1066	class Recorder(torch.nn.Module):
1067	def __init__(self):
1068	super().__init__()
1069	self.received = []
1070
1071	def load_weights(self, weights):
1072	for name, tensor in weights:
1073	self.received.append((name, tensor.clone()))
1074
1075	# Trainer sends unpacked IPC handles; the worker learns packed=False from
1076	# the init handshake below (IPCWeightTransferInitInfo defaults to False).
1077	config = WeightTransferConfig(backend="ipc")
1078	vllm_config = MagicMock()
1079	parallel_config = MagicMock(spec=ParallelConfig)
1080	parallel_config.rank = 0
1081	parallel_config.world_size = 1
1082	parallel_config.data_parallel_rank = 0
1083	parallel_config.data_parallel_index = 0
1084	vllm_config.parallel_config = parallel_config
1085	vllm_config.model_config = MagicMock()
1086
1087	recorder = Recorder()
1088	engine = IPCWeightTransferEngine(config, vllm_config, device, recorder)
1089	# Transport-only test: bypass the set_current_vllm_config context that
1090	# receive_weights enters, since vllm_config here is a mock.
1091	import vllm.config as _vllm_config_mod
1092
1093	_vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext()
1094
1095	init_info = IPCWeightTransferInitInfo()
1096	engine.init_transfer_engine(init_info)
1097
1098	ipc_handles = [{ipc_handle_dict["gpu_uuid"]: ipc_handle_dict["ipc_handle"]}]
1099
1100	if mode == "ray":
1101	update_dict: dict = {
1102	"names": ["test.weight"],
1103	"dtype_names": [ipc_handle_dict["dtype"]],
1104	"shapes": [ipc_handle_dict["shape"]],
1105	"ipc_handles": ipc_handles,
1106	}
1107	elif mode == "http":
1108	pickled = base64.b64encode(pickle.dumps(ipc_handles)).decode("utf-8")
1109	update_dict = {
1110	"names": ["test.weight"],
1111	"dtype_names": [ipc_handle_dict["dtype"]],
1112	"shapes": [ipc_handle_dict["shape"]],
1113	"ipc_handles_pickled": pickled,
1114	}
1115	else:
1116	raise ValueError(f"Unknown mode: {mode}")
1117
1118	update_info = engine.parse_update_info(update_dict)
1119	engine.receive_weights(update_info)
1120	torch.accelerator.synchronize()
1121
1122	success = False
1123	received_shape = None
1124	received_sum = None
1125
1126	if len(recorder.received) == 1:
1127	name, tensor = recorder.received[0]
1128	received_shape = list(tensor.shape)
1129	received_sum = tensor.sum().item()
1130	if received_shape == ipc_handle_dict["shape"]:
1131	expected_sum = 42.0 * torch.tensor(ipc_handle_dict["shape"]).prod().item()
1132	if abs(received_sum - expected_sum) < 0.01:
1133	success = True
1134
1135	engine.shutdown()
1136
1137	return {
1138	"success": success,
1139	"received_shape": received_shape,
1140	"received_sum": received_sum,
1141	}
1142
1143
1144	@pytest.mark.skipif(
1145	torch.accelerator.device_count() < 1,
1146	reason="Need at least 1 GPU to run IPC weight transfer test.",
1147	)
1148	@pytest.mark.parametrize("mode", ["ray", "http"])
1149	def test_ipc_weight_transfer_between_processes(mode: str):
1150	"""Test IPC weight transfer from trainer to inference process using Ray."""
1151	from ray.util.placement_group import placement_group
1152	from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
1153
1154	_init_ray_for_weight_transfer()
1155
1156	pg = placement_group([{"GPU": 1, "CPU": 2}])
1157	ray.get(pg.ready())
1158
1159	scheduling_strategy = PlacementGroupSchedulingStrategy(
1160	placement_group=pg,
1161	placement_group_capture_child_tasks=True,
1162	)
1163
1164	tensor_shape = [100, 100]
1165	tensor_dtype = "float32"
1166
1167	trainer_actor = TrainerActor.options( # type: ignore[attr-defined]
1168	scheduling_strategy=scheduling_strategy
1169	).remote(tensor_shape, tensor_dtype)
1170
1171	ipc_handle_dict = ray.get(trainer_actor.get_ipc_handle_dict.remote())
1172
1173	inference_result = ray.get(
1174	inference_receive_ipc_tensor.options(
1175	scheduling_strategy=scheduling_strategy
1176	).remote(ipc_handle_dict, mode=mode)
1177	)
1178
1179	assert inference_result["success"], (
1180	f"IPC weight transfer failed (mode={mode}). "
1181	f"Received shape: {inference_result['received_shape']}, "
1182	f"Received sum: {inference_result['received_sum']}"
1183	)
1184
1185
1186	def test_ipc_receive_weights_missing_gpu_uuid_raises():
1187	"""Test that receive_weights raises if GPU UUID not found in IPC handles."""
1188	if torch.accelerator.device_count() < 1:
1189	pytest.skip("Need at least 1 GPU for this test")
1190
1191	config = WeightTransferConfig(backend="ipc")
1192	engine = IPCWeightTransferEngine(
1193	config,
1194	create_mock_vllm_config(),
1195	torch.device("cuda:0"),
1196	MagicMock(spec=torch.nn.Module),
1197	)
1198	# No init handshake here, so the engine keeps its default packed=False.
1199
1200	dummy_tensor = torch.ones(10, 10, device="cuda:0")
1201	_, ipc_handle = reduce_tensor(dummy_tensor)
1202	wrong_uuid = "wrong-uuid-12345"
1203	ipc_handles = [{wrong_uuid: ipc_handle}]
1204
1205	update_info = IPCWeightTransferUpdateInfo(
1206	names=["w"],
1207	dtype_names=["float32"],
1208	shapes=[[10, 10]],
1209	ipc_handles=ipc_handles,
1210	)
1211
1212	with pytest.raises(ValueError, match="IPC handle not found"):
1213	engine.receive_weights(update_info)
1214
1215
1216	class RecordingClient:
1217	"""A fake VLLMWeightSyncClient that records the order of calls."""
1218
1219	def __init__(self):
1220	self.order: list[str] = []
1221	self.last_init_info: dict \| None = None
1222	self.last_update_info: dict \| None = None
1223
1224	def init_weight_transfer_engine(self, init_info: dict) -> None:
1225	self.order.append("init")
1226	self.last_init_info = init_info
1227
1228	def start_weight_update(self) -> None:
1229	self.order.append("start")
1230
1231	def update_weights(self, update_info: dict) -> None:
1232	self.order.append("update")
1233	self.last_update_info = update_info
1234
1235	def finish_weight_update(self, weight_version: str \| None = None) -> None:
1236	self.order.append("finish")
1237
1238
1239	def _module_with(*pairs):
1240	"""A tiny nn.Module exposing the given (name, tensor) pairs as parameters,
1241	so trainer tests can build a ModuleSource without a real model."""
1242	module = torch.nn.Module()
1243	for name, tensor in pairs:
1244	module.register_parameter(name, torch.nn.Parameter(tensor, requires_grad=False))
1245	return module
1246
1247
1248	class _DummyTrainerEngine(TrainerWeightTransferEngine):
1249	"""Minimal concrete trainer engine to exercise base-class + factory."""
1250
1251	@classmethod
1252	def trainer_init(cls, init_info, *, client, source):
1253	return cls(client=client, source=source)
1254
1255	def send_weights(self):
1256	pass
1257
1258
1259	class TestTrainerClients:
1260	"""Structural protocol conformance for the built-in clients."""
1261
1262	def test_recording_client_is_protocol(self):
1263	assert isinstance(RecordingClient(), VLLMWeightSyncClient)
1264
1265	def test_http_client_is_protocol(self):
1266	assert isinstance(
1267	HTTPVLLMWeightSyncClient("http://localhost:8000"), VLLMWeightSyncClient
1268	)
1269
1270	def test_ray_client_is_protocol(self):
1271	assert isinstance(RayVLLMWeightSyncClient(MagicMock()), VLLMWeightSyncClient)
1272
1273	def test_ray_client_sends_typed_requests(self, monkeypatch):
1274	"""Ray client must hand the actor typed Request objects, not raw dicts."""
1275	import ray
1276
1277	monkeypatch.setattr(ray, "get", lambda refs: None)
1278	handle = MagicMock()
1279	client = RayVLLMWeightSyncClient(handle)
1280
1281	client.init_weight_transfer_engine({"master_addr": "x"})
1282	(init_req,), _ = handle.init_weight_transfer_engine.remote.call_args
1283	assert isinstance(init_req, WeightTransferInitRequest)
1284	assert init_req.init_info == {"master_addr": "x"}
1285
1286	client.update_weights({"names": ["w"]})
1287	(update_req,), _ = handle.update_weights.remote.call_args
1288	assert isinstance(update_req, WeightTransferUpdateRequest)
1289	assert update_req.update_info == {"names": ["w"]}
1290
1291	client.finish_weight_update("step-42")
1292	handle.finish_weight_update.remote.assert_called_once_with()
1293	handle.update_weight_version.remote.assert_called_once_with("step-42")
1294
1295	def test_http_client_pickles_ipc_handles_for_json(self, monkeypatch):
1296	"""HTTP update_weights must encode raw ipc_handles as a base64 pickle."""
1297	captured = {}
1298
1299	def fake_post(self, path, json=None):
1300	captured["path"] = path
1301	captured["json"] = json
1302
1303	monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post)
1304	client = HTTPVLLMWeightSyncClient("http://localhost:8000")
1305	client.update_weights({"names": ["w"], "ipc_handles": [{"gpu": ("args",)}]})
1306	sent = captured["json"]["update_info"]
1307	assert "ipc_handles" not in sent
1308	assert "ipc_handles_pickled" in sent
1309	assert pickle.loads(base64.b64decode(sent["ipc_handles_pickled"])) == [
1310	{"gpu": ("args",)}
1311	]
1312
1313	client.update_weights([{"ipc_handles": {"gpu": ("args",)}}, {"names": []}])
1314	sent = captured["json"]["update_info"]
1315	assert sent[1] == {"names": []}
1316	assert pickle.loads(base64.b64decode(sent[0]["ipc_handles_pickled"])) == {
1317	"gpu": ("args",)
1318	}
1319
1320	def test_http_client_passes_through_nccl_update_info(self, monkeypatch):
1321	"""NCCL update_info has only JSON-native fields and passes unchanged."""
1322	captured = {}
1323
1324	def fake_post(self, path, json=None):
1325	captured["json"] = json
1326
1327	monkeypatch.setattr(HTTPVLLMWeightSyncClient, "_post", fake_post)
1328	client = HTTPVLLMWeightSyncClient("http://localhost:8000")
1329	update_info = {"names": ["w"], "dtype_names": ["float32"], "shapes": [[4]]}
1330	client.update_weights(update_info)
1331	assert captured["json"]["update_info"] == update_info
1332
1333	client.finish_weight_update("step-42")
1334	assert captured["json"] == {"weight_version": "step-42"}
1335
1336
1337	class TestModuleSource:
1338	"""`ModuleSource` metadata vs. materialized iteration (dense, no GPU)."""
1339
1340	def test_metadata_reads_shape_and_dtype(self):
1341	source = ModuleSource(
1342	_module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3)))
1343	)
1344	meta = source.metadata()
1345	assert [m.name for m in meta] == ["w", "b"]
1346	assert [m.shape for m in meta] == [(2, 3), (3,)]
1347	assert all(m.dtype == torch.float32 for m in meta)
1348
1349	def test_iteration_yields_materialized_tensors(self):
1350	w = torch.arange(6, dtype=torch.float32).reshape(2, 3)
1351	source = ModuleSource(_module_with(("w", w)))
1352	pairs = list(source)
1353	assert [name for name, _ in pairs] == ["w"]
1354	assert torch.equal(pairs[0][1], w)
1355
1356	def test_source_is_reiterable(self):
1357	source = ModuleSource(_module_with(("w", torch.zeros(2))))
1358	assert [n for n, _ in source] == [n for n, _ in source] == ["w"]
1359
1360	def test_metadata_agrees_with_iteration(self):
1361	"""The two channels must line up element-for-element: engines declare
1362	the round from `metadata()` and then send what iteration yields."""
1363	source = ModuleSource(
1364	_module_with(("w", torch.zeros(2, 3)), ("b", torch.zeros(3)))
1365	)
1366	meta = source.metadata()
1367	pairs = list(source)
1368	assert [m.name for m in meta] == [name for name, _ in pairs]
1369	assert [m.dtype for m in meta] == [t.dtype for _, t in pairs]
1370	assert [m.shape for m in meta] == [tuple(t.shape) for _, t in pairs]
1371
1372
1373	class TestWeightSourceGroupContract:
1374	"""`groups()` / `iter_groups()` on the WeightSource ABC. Groups are what
1375	backends gather and free by, so the default must agree with
1376	`layerwise_groups` over `metadata()`, restricted to what this rank holds."""
1377
1378	class _Source(WeightSource):
1379	"""Minimal source over an ordered (name, tensor) list, optionally holding
1380	only some names (in which case it iterates only their groups)."""
1381
1382	def __init__(self, names, held=None, reverse=False):
1383	self._pairs = [(n, torch.full((2,), float(i))) for i, n in enumerate(names)]
1384	self._held = held
1385	self._reverse = reverse
1386
1387	def metadata(self):
1388	return [ParamMeta(n, t.dtype, tuple(t.shape)) for n, t in self._pairs]
1389
1390	def held_names(self):
1391	return self._held
1392
1393	def __iter__(self):
1394	pairs = self._pairs
1395	if self._held is not None:
1396	keep = set(self._held)
1397	pairs = [(n, t) for n, t in pairs if n in keep]
1398	return iter(list(reversed(pairs)) if self._reverse else pairs)
1399
1400	def _source(self, names, held=None, reverse=False):
1401	return self._Source(names, held, reverse)
1402
1403	def test_groups_defaults_to_the_layerwise_partition(self):
1404	names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"]
1405	assert self._source(names).groups() == layerwise_groups(names)
1406
1407	def test_groups_keeps_only_groups_holding_a_held_name(self):
1408	names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"]
1409	held = ["model.layers.0.a", "model.layers.1.a"]
1410	assert self._source(names, held=held).groups() == [
1411	["model.layers.0.a"],
1412	["model.layers.1.a"],
1413	]
1414
1415	def test_groups_keeps_a_partially_held_group_whole(self):
1416	"""One held name selects the WHOLE group, unheld members included.
1417
1418	A group is the collective unit, so its membership has to be
1419	rank-uniform: a rank that narrowed it to what it holds would disagree
1420	with its peers about which names group g covers, and the gather would
1421	desynchronize. Selection is per group, but only whole groups; the
1422	per-name split is the backend's, via owner sets. This is the
1423	foreign-expert shape under expert parallelism — a rank holds some
1424	experts of a layer, not all of them.
1425	"""
1426	names = ["model.layers.0.a", "model.layers.0.b", "model.layers.1.a"]
1427	held = ["model.layers.0.a"]
1428	assert self._source(names, held=held).groups() == [
1429	["model.layers.0.a", "model.layers.0.b"],
1430	]
1431
1432	def test_groups_order_follows_metadata_not_the_declaration(self):
1433	"""The declaration is a SET of names; the groups it selects still come
1434	out in metadata order, so each pairs with the right ``iter_groups()``
1435	batch however the source listed them."""
1436	names = ["embed.w", "model.layers.0.a", "model.layers.1.a", "norm.w"]
1437	source = self._source(
1438	names, held=["model.layers.1.a", "model.layers.0.a", "model.layers.1.a"]
1439	)
1440	assert source.groups() == [["model.layers.0.a"], ["model.layers.1.a"]]
1441	assert [ns for ns, _ in source.iter_groups()] == [
1442	["model.layers.0.a"],
1443	["model.layers.1.a"],
1444	]
1445
1446	def test_iter_groups_batches_the_stream_per_group(self):
1447	names = ["embed.w", "model.layers.0.a", "model.layers.0.b", "norm.w"]
1448	batches = list(self._source(names).iter_groups())
1449	assert [ns for ns, _ in batches] == [
1450	["embed.w"],
1451	["model.layers.0.a", "model.layers.0.b"],
1452	["norm.w"],
1453	]
1454	assert all(len(ns) == len(ts) for ns, ts in batches)
1455
1456	def test_iter_groups_yields_the_tensors_iteration_produced(self):
1457	names = ["model.layers.0.a", "model.layers.0.b"]
1458	(batch,) = list(self._source(names).iter_groups())
1459	_names, tensors = batch
1460	assert [float(t[0]) for t in tensors] == [0.0, 1.0]
1461
1462	def test_iter_groups_yields_only_held_groups(self):
1463	names = ["embed.w", "model.layers.0.a", "model.layers.1.a"]
1464	batches = list(self._source(names, held=["model.layers.1.a"]).iter_groups())
1465	assert [ns for ns, _ in batches] == [["model.layers.1.a"]]
1466
1467	def test_out_of_order_iteration_raises(self):
1468	"""Materializing is usually a collective, so a rank that iterates out of
1469	order deadlocks its peers -- fail loudly instead."""
1470	source = self._source(["model.layers.0.a", "model.layers.0.b"], reverse=True)
1471	with pytest.raises(RuntimeError, match="iteration order must match"):
1472	list(source.iter_groups())
1473
1474	def test_a_source_may_override_iter_groups(self):
1475	"""The extension point: materialize a whole group in one step instead of
1476	one generator resume per tensor."""
1477	calls = []
1478
1479	class _Batched(TestWeightSourceGroupContract._Source):
1480	def iter_groups(self):
1481	for group in self.groups():
1482	calls.append(len(group))
1483	yield group, [torch.zeros(2) for _ in group]
1484
1485	source = _Batched(["model.layers.0.a", "model.layers.0.b"])
1486	assert [ns for ns, _ in source.iter_groups()] == [
1487	["model.layers.0.a", "model.layers.0.b"]
1488	]
1489	assert calls == [2]
1490
1491
1492	class TestDeferredProcessingContract:
1493	"""`defers_processing` and `drain_pending` are two halves of one contract: a
1494	caller that takes over the update tail (running its own
1495	`finalize_layerwise_reload` instead of going through `finish_weight_update`)
1496	reads the flag and calls the method. Both must be answerable on any engine, or
1497	that caller ends up reaching through a getattr."""
1498
1499	def _engines(self):
1500	return {
1501	name: loader()
1502	for name, loader in WeightTransferEngineFactory._registry.items()
1503	}
1504
1505	def test_every_engine_declares_whether_it_defers(self):
1506	for name, cls in self._engines().items():
1507	assert isinstance(cls.defers_processing, bool), name
1508
1509	def test_every_engine_can_be_drained(self):
1510	"""The default is a no-op, so a caller never has to check whether the
1511	method exists before calling it."""
1512	for name, cls in self._engines().items():
1513	assert callable(cls.drain_pending), name
1514
1515	def test_the_default_is_not_to_defer(self):
1516	assert WeightTransferEngine.defers_processing is False
1517
1518	def test_a_synchronous_engine_drains_as_a_no_op(self):
1519	engine = object.__new__(WeightTransferEngineFactory._registry["nccl"]())
1520	engine.drain_pending() # must not raise, and must not need any state
1521
1522	def test_the_rdt_engine_defers_and_overrides_the_drain(self):
1523	"""The one engine the contract exists for."""
1524	cls = WeightTransferEngineFactory._registry["sharded_rdt"]()
1525	assert cls.defers_processing is True
1526	assert cls.drain_pending is not WeightTransferEngine.drain_pending
1527
1528
1529	class TestTrainerFactory:
1530	"""WeightTransferTrainerFactory registry mechanics."""
1531
1532	def test_registry_has_all_backends(self):
1533	assert "nccl" in WeightTransferTrainerFactory._registry
1534	assert "ipc" in WeightTransferTrainerFactory._registry
1535	assert "sparse_nccl" in WeightTransferTrainerFactory._registry
1536
1537	def test_register_and_dispatch(self):
1538	saved = dict(WeightTransferTrainerFactory._registry)
1539	try:
1540	WeightTransferTrainerFactory.register_engine("dummy", _DummyTrainerEngine)
1541	engine = WeightTransferTrainerFactory.trainer_init(
1542	MagicMock(backend="dummy"), # backend read from the init info
1543	client=RecordingClient(),
1544	source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1545	)
1546	assert isinstance(engine, _DummyTrainerEngine)
1547	with pytest.raises(ValueError, match="already registered"):
1548	WeightTransferTrainerFactory.register_engine(
1549	"dummy", _DummyTrainerEngine
1550	)
1551	finally:
1552	WeightTransferTrainerFactory._registry = saved
1553
1554	def test_unknown_backend_raises(self):
1555	with pytest.raises(ValueError, match="Invalid weight transfer backend"):
1556	WeightTransferTrainerFactory.trainer_init(
1557	MagicMock(backend="nope"),
1558	client=RecordingClient(),
1559	source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1560	)
1561
1562	def test_ipc_init_info_declares_backend(self):
1563	assert IPCTrainerInitInfo.backend == "ipc"
1564
1565	def test_nccl_init_info_declares_backend(self):
1566	assert NCCLTrainerInitInfo.backend == "nccl"
1567
1568	def test_sparse_nccl_init_info_declares_backend(self):
1569	assert SparseNCCLTrainerInitInfo.backend == "sparse_nccl"
1570
1571	def test_trainer_init_info_subclass_must_set_backend(self):
1572	with pytest.raises(TypeError, match="class-level `backend`"):
1573
1574	class _NoBackend(TrainerInitInfo):
1575	pass
1576
1577
1578	class TestTrainerEngineBase:
1579	"""Base-class construction (no GPU)."""
1580
1581	def test_source_stored_and_sender_by_default(self):
1582	engine = _DummyTrainerEngine(
1583	client=RecordingClient(),
1584	source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1585	)
1586	assert engine.is_sender is True
1587	assert [name for name, _ in engine.source] == ["w"]
1588
1589	def test_shutdown_default_is_noop(self):
1590	engine = _DummyTrainerEngine(
1591	client=RecordingClient(),
1592	source=ModuleSource(_module_with(("w", torch.zeros(2)))),
1593	is_sender=False,
1594	)
1595	assert engine.is_sender is False
1596	engine.shutdown() # must not raise
1597
1598
1599	@pytest.mark.skipif(
1600	torch.accelerator.device_count() < 1,
1601	reason="Need at least 1 GPU (CUDA IPC handles).",
1602	)
1603	def test_ipc_trainer_send_weights_drives_client_in_order():
1604	"""send_weights issues start -> update -> finish and ships per-round metadata;
1605	the packed wire param rides the init info, not the per-round update_info."""
1606	client = RecordingClient()
1607	engine = IPCTrainerWeightTransferEngine(
1608	client=client,
1609	source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))),
1610	packed=False,
1611	)
1612
1613	engine.send_weights()
1614
1615	assert client.order == ["start", "update", "finish"]
1616	assert client.last_update_info is not None
1617	assert client.last_update_info["names"] == ["w"]
1618	assert client.last_update_info["shapes"] == [[4]]
1619	assert "packed" not in client.last_update_info
1620
1621
1622	def test_ipc_trainer_init_ships_packed_to_worker():
1623	"""trainer_init drives the inference-side init handshake and propagates the
1624	must-agree `packed` flag to the worker."""
1625	if torch.accelerator.device_count() < 1:
1626	pytest.skip("Need at least 1 GPU (CUDA IPC handles).")
1627
1628	client = RecordingClient()
1629	engine = WeightTransferTrainerFactory.trainer_init(
1630	init_info=IPCTrainerInitInfo(rank=0, packed=True), # backend from init info
1631	client=client,
1632	source=ModuleSource(_module_with(("w", torch.ones(4, device="cuda")))),
1633	)
1634
1635	assert isinstance(engine, IPCTrainerWeightTransferEngine)
1636	assert engine.is_sender is True
1637	assert engine.packed is True
1638	assert client.order == ["init"]
1639	assert client.last_init_info == {"packed": True}
1640
1641
1642	def test_nccl_trainer_init_ships_worker_init_info(monkeypatch):
1643	"""The sender's trainer_init drives the inference-side init handshake with
1644	the worker-shaped init info (rank_offset=1) while opening its own endpoint,
1645	and propagates the must-agree wire params to the worker."""
1646	import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod
1647
1648	# Bypass the real NCCL rendezvous.
1649	monkeypatch.setattr(
1650	nccl_engine_mod, "open_trainer_endpoint", lambda info: MagicMock()
1651	)
1652
1653	client = RecordingClient()
1654	engine = WeightTransferTrainerFactory.trainer_init(
1655	init_info=NCCLTrainerInitInfo(
1656	master_address="127.0.0.1",
1657	master_port=29500,
1658	world_size=3,
1659	rank=0,
1660	packed=True,
1661	packed_buffer_size_bytes=1024,
1662	packed_num_buffers=3,
1663	),
1664	client=client,
1665	source=ModuleSource(_module_with(("w", torch.zeros(4)))),
1666	)
1667
1668	assert isinstance(engine, NCCLTrainerWeightTransferEngine)
1669	assert engine.is_sender is True
1670	assert engine.packed is True
1671	assert client.order == ["init"]
1672	assert client.last_init_info == {
1673	"master_address": "127.0.0.1",
1674	"master_port": 29500,
1675	"rank_offset": 1,
1676	"world_size": 3,
1677	"packed": True,
1678	"packed_buffer_size_bytes": 1024,
1679	"packed_num_buffers": 3,
1680	}
1681
1682
1683	def test_nccl_worker_learns_wire_params_from_init_handshake(monkeypatch):
1684	"""The worker engine reads packed + buffer geometry from the
1685	trainer-supplied init info at the handshake, not from the config or the
1686	per-round update info."""
1687	import vllm.distributed.weight_transfer.nccl_engine as nccl_engine_mod
1688
1689	monkeypatch.setattr(
1690	nccl_engine_mod, "worker_init_process_group", lambda info, pc: MagicMock()
1691	)
1692
1693	engine = NCCLWeightTransferEngine(
1694	WeightTransferConfig(backend="nccl"),
1695	create_mock_vllm_config(),
1696	torch.device("cuda:0"),
1697	MagicMock(spec=torch.nn.Module),
1698	)
1699	assert engine.packed is False # pre-handshake default (legacy unpacked)
1700	engine.init_transfer_engine(
1701	NCCLWeightTransferInitInfo(
1702	master_address="127.0.0.1",
1703	master_port=29500,
1704	rank_offset=1,
1705	world_size=2,
1706	packed=True,
1707	packed_buffer_size_bytes=2048,
1708	packed_num_buffers=4,
1709	)
1710	)
1711
1712	assert engine.packed is True
1713	assert engine.packed_buffer_size_bytes == 2048
1714	assert engine.packed_num_buffers == 4
1715
1716
1717	def test_nccl_trainer_init_non_sender_skips_rendezvous_and_client():
1718	"""Non-sender trainer ranks build an engine without opening an endpoint or
1719	touching the client; they only join the collectives in send_weights."""
1720	client = RecordingClient()
1721	engine = WeightTransferTrainerFactory.trainer_init(
1722	init_info=NCCLTrainerInitInfo(
1723	master_address="127.0.0.1",
1724	master_port=29500,
1725	world_size=3,
1726	rank=1,
1727	),
1728	client=client,
1729	source=ModuleSource(_module_with(("w", torch.zeros(4)))),
1730	)
1731
1732	assert engine.is_sender is False
1733	assert engine.model_update_group is None
1734	assert client.order == []
1735
1736	# send_weights on a non-sender only iterates the source (packed mode needs
1737	# no CUDA stream on non-senders), never the client.
1738	engine.send_weights()
1739	assert client.order == []
1740
1741
1742	@pytest.mark.skipif(
1743	torch.accelerator.device_count() < 1,
1744	reason="Need at least 1 GPU (NCCL broadcast / CUDA stream).",
1745	)
1746	def test_nccl_trainer_send_weights_drives_client_in_order():
1747	"""send_weights issues start -> update -> finish and ships per-round
1748	metadata; the packed wire params ride the init handshake, not the
1749	per-round update_info."""
1750	client = RecordingClient()
1751	engine = NCCLTrainerWeightTransferEngine(
1752	client=client,
1753	source=ModuleSource(_module_with(("w", torch.zeros(4, device="cuda")))),
1754	packed=False,
1755	)
1756	# Bypass the real NCCL rendezvous; broadcast is a no-op.
1757	engine.model_update_group = MagicMock()
1758
1759	engine.send_weights()
1760
1761	assert client.order == ["start", "update", "finish"]
1762	assert client.last_update_info is not None
1763	assert client.last_update_info["names"] == ["w"]
1764	assert client.last_update_info["shapes"] == [[4]]
1765	assert "packed" not in client.last_update_info
1766
1767
1768	class _ScriptedSource(WeightSource):
1769	"""Declares `meta` but yields whatever `pairs` says — used to drive the
1770	metadata/iteration agreement checks."""
1771
1772	def __init__(self, meta, pairs):
1773	self._meta = meta
1774	self._pairs = pairs
1775
1776	def metadata(self):
1777	return list(self._meta)
1778
1779	def __iter__(self):
1780	yield from self._pairs
1781
1782
1783	def _mock_group_engine(source, monkeypatch, **kwargs):
1784	"""Unpacked trainer engine with a mocked group and stream (no GPU needed)."""
1785	engine = NCCLTrainerWeightTransferEngine(
1786	client=RecordingClient(), source=source, packed=False, **kwargs
1787	)
1788	engine.model_update_group = MagicMock()
1789	monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
1790	return engine
1791
1792
1793	def test_nccl_trainer_init_requires_source():
1794	"""NCCL is a full-resync backend: it cannot run without a WeightSource."""
1795	with pytest.raises(ValueError, match="requires a WeightSource"):
1796	NCCLTrainerWeightTransferEngine.trainer_init(
1797	NCCLTrainerInitInfo(
1798	master_address="127.0.0.1", master_port=29500, world_size=2, rank=0
1799	),
1800	client=RecordingClient(),
1801	)
1802
1803
1804	def test_nccl_trainer_send_weights_rejects_reordered_source(monkeypatch):
1805	"""The worker sizes its buffers (and cuts packed chunks) from metadata(), so
1806	iteration disagreeing with it must raise rather than corrupt the stream."""
1807	meta = [
1808	ParamMeta("w", torch.float32, (4,)),
1809	ParamMeta("b", torch.float32, (2,)),
1810	]
1811	reordered = [("b", torch.zeros(2)), ("w", torch.zeros(4))]
1812	engine = _mock_group_engine(_ScriptedSource(meta, reordered), monkeypatch)
1813
1814	with pytest.raises(ValueError, match="disagrees with iteration at index 0"):
1815	engine.send_weights()
1816
1817
1818	def test_nccl_trainer_send_weights_rejects_dtype_disagreement(monkeypatch):
1819	"""A source that declares one wire dtype and materializes another would make
1820	the two sides disagree on every byte offset."""
1821	meta = [ParamMeta("w", torch.float32, (4,))]
1822	engine = _mock_group_engine(
1823	_ScriptedSource(meta, [("w", torch.zeros(4, dtype=torch.bfloat16))]),
1824	monkeypatch,
1825	)
1826
1827	with pytest.raises(ValueError, match="disagrees with iteration"):
1828	engine.send_weights()
1829
1830
1831	def test_nccl_trainer_send_weights_rejects_truncated_source(monkeypatch):
1832	"""Yielding fewer parameters than declared leaves the worker waiting."""
1833	meta = [
1834	ParamMeta("w", torch.float32, (4,)),
1835	ParamMeta("b", torch.float32, (2,)),
1836	]
1837	engine = _mock_group_engine(
1838	_ScriptedSource(meta, [("w", torch.zeros(4))]), monkeypatch
1839	)
1840
1841	with pytest.raises(ValueError, match="yielded 1 parameters"):
1842	engine.send_weights()
1843
1844
1845	def test_nccl_trainer_send_weights_broadcasts_contiguous(monkeypatch):
1846	"""NCCL sends numel elements from data_ptr(), so a non-contiguous view must
1847	be linearized first or the worker receives unrelated memory."""
1848	base = torch.arange(6, dtype=torch.float32).reshape(2, 3)
1849	view = base.t() # non-contiguous
1850	meta = [ParamMeta("w", torch.float32, tuple(view.shape))]
1851	engine = _mock_group_engine(_ScriptedSource(meta, [("w", view)]), monkeypatch)
1852
1853	engine.send_weights()
1854
1855	sent = engine.model_update_group.broadcast.call_args.args[0]
1856	assert sent.is_contiguous()
1857	assert torch.equal(sent, view)
1858
1859
1860	def test_nccl_trainer_send_weights_raises_instead_of_hanging(monkeypatch):
1861	"""A failed broadcast must surface even while the inference-side
1862	update_weights is still blocked in its matching NCCL call.
1863
1864	Joining the RPC thread there would deadlock: the worker only returns once
1865	the broadcast it is waiting for arrives, which never happens.
1866	"""
1867	rpc_entered = threading.Event()
1868	release_rpc = threading.Event()
1869
1870	class _BlockingClient(RecordingClient):
1871	def update_weights(self, update_info):
1872	rpc_entered.set()
1873	release_rpc.wait(timeout=30) # stands in for a wedged NCCL recv
1874	super().update_weights(update_info)
1875
1876	class _FailingSource(WeightSource):
1877	def metadata(self):
1878	return [ParamMeta("w", torch.float32, (4,))]
1879
1880	def __iter__(self):
1881	# Fail only once the RPC is provably in flight.
1882	rpc_entered.wait(timeout=60)
1883	raise RuntimeError("broadcast blew up")
1884
1885	engine = NCCLTrainerWeightTransferEngine(
1886	client=_BlockingClient(), source=_FailingSource(), packed=False
1887	)
1888	engine.model_update_group = MagicMock()
1889	monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
1890
1891	started = time.perf_counter()
1892	try:
1893	with pytest.raises(RuntimeError, match="broadcast blew up"):
1894	engine.send_weights()
1895	elapsed = time.perf_counter() - started
1896	assert rpc_entered.is_set(), "the RPC was never in flight"
1897	assert not release_rpc.is_set(), "send_weights waited for the wedged RPC"
1898	# Regression guard: joining the RPC thread would park here until the
1899	# client's own timeout expires instead of raising immediately.
1900	assert elapsed < 5.0, f"send_weights blocked {elapsed:.1f}s on the RPC"
1901	finally:
1902	release_rpc.set()
1903
1904
1905	def _sparse_patch(device: str = "cpu") -> SparseWeightPatch:
1906	return SparseWeightPatch(
1907	name="w",
1908	indices=torch.tensor([1, 3], dtype=torch.int32, device=device),
1909	values=torch.tensor([1.0, 2.0], dtype=torch.float32, device=device),
1910	full_shape=(4, 4),
1911	)
1912
1913
1914	def test_sparse_nccl_trainer_init_ships_worker_init_info(monkeypatch):
1915	"""The sender's trainer_init drives the init handshake with the
1916	worker-shaped init info; sparse ships no packed wire params, so the worker
1917	keeps its unpacked defaults. Sparse takes no `source`."""
1918	import vllm.distributed.weight_transfer.sparse_nccl_engine as sparse_mod
1919
1920	monkeypatch.setattr(sparse_mod, "open_trainer_endpoint", lambda info: MagicMock())
1921
1922	client = RecordingClient()
1923	engine = WeightTransferTrainerFactory.trainer_init(
1924	init_info=SparseNCCLTrainerInitInfo(
1925	master_address="127.0.0.1",
1926	master_port=29500,
1927	world_size=2,
1928	rank=0,
1929	),
1930	client=client,
1931	)
1932
1933	assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine)
1934	assert client.order == ["init"]
1935	assert client.last_init_info == {
1936	"master_address": "127.0.0.1",
1937	"master_port": 29500,
1938	"rank_offset": 1,
1939	"world_size": 2,
1940	"packed": False,
1941	"packed_buffer_size_bytes": DEFAULT_PACKED_BUFFER_SIZE_BYTES,
1942	"packed_num_buffers": DEFAULT_PACKED_NUM_BUFFERS,
1943	}
1944
1945
1946	def test_sparse_nccl_trainer_send_weights_drives_client_in_order(monkeypatch):
1947	"""One-shot send drives start, update, and finish in order."""
1948	client = RecordingClient()
1949	engine = SparseNCCLTrainerWeightTransferEngine(client=client)
1950	engine.model_update_group = MagicMock()
1951	engine.model_update_group.device = torch.device("cpu")
1952	# The group is a mock, so the stream is just a handle it is handed (and a
1953	# handle _post_send_sync can synchronize).
1954	monkeypatch.setattr(torch.cuda, "current_stream", MagicMock())
1955
1956	engine.send_weights([_sparse_patch()])
1957
1958	assert client.order == ["start", "update", "finish"]
1959	assert client.last_update_info is not None
1960	assert client.last_update_info["names"] == ["w"]
1961	assert client.last_update_info["shapes"] == [[4, 4]]
1962	assert client.last_update_info["num_updates_list"] == [2]
1963	# One broadcast for indices + one for values per patch.
1964	assert engine.model_update_group.broadcast.call_count == 2
1965	engine.shutdown()
1966
1967
1968	def test_sparse_nccl_trainer_send_weights_empty_round_is_noop():
1969	"""A round with no patches must not touch the client (an empty sparse
1970	update info is invalid by construction)."""
1971	client = RecordingClient()
1972	engine = SparseNCCLTrainerWeightTransferEngine(client=client)
1973	engine.model_update_group = MagicMock()
1974
1975	engine.send_weights([])
1976	engine.send_weights() # no argument is also a no-op round
1977	engine.send_weight_chunk([])
1978	engine.send_weight_chunk()
1979
1980	assert client.order == []
1981
1982
1983	def test_sparse_nccl_trainer_send_weights_requires_full_shape():
1984	patch = _sparse_patch()
1985	patch.full_shape = None
1986	engine = SparseNCCLTrainerWeightTransferEngine(client=RecordingClient())
1987	engine.model_update_group = MagicMock()
1988	engine.model_update_group.device = torch.device("cpu")
1989
1990	with pytest.raises(ValueError, match="full_shape"):
1991	engine.send_weights([patch])
1992
1993
1994	def test_sparse_nccl_trainer_rejects_source():
1995	"""Sparse is a delta backend; a WeightSource would silently never be sent."""
1996	with pytest.raises(ValueError, match="takes no WeightSource"):
1997	SparseNCCLTrainerWeightTransferEngine(
1998	client=RecordingClient(),
1999	source=ModuleSource(_module_with(("w", torch.zeros(2)))),
2000	)
2001
2002
2003	def test_sparse_nccl_trainer_validates_patch_before_any_rpc():
2004	"""Malformed patches must fail on the trainer, before start_weight_update:
2005	the worker's own checks only run once the broadcasts are already in flight,
2006	where a size mismatch wedges both sides instead of raising."""
2007	client = RecordingClient()
2008	engine = SparseNCCLTrainerWeightTransferEngine(client=client)
2009	engine.model_update_group = MagicMock()
2010	engine.model_update_group.device = torch.device("cpu")
2011
2012	mismatched = SparseWeightPatch(
2013	name="w",
2014	indices=torch.tensor([1, 3], dtype=torch.int32),
2015	values=torch.tensor([1.0], dtype=torch.float32),
2016	full_shape=(4, 4),
2017	)
2018	with pytest.raises(ValueError, match="matching lengths"):
2019	engine.send_weights([mismatched])
2020
2021	wrong_index_dtype = SparseWeightPatch(
2022	name="w",
2023	indices=torch.tensor([1, 3], dtype=torch.int64),
2024	values=torch.tensor([1.0, 2.0], dtype=torch.float32),
2025	full_shape=(4, 4),
2026	)
2027	with pytest.raises(ValueError, match="int32 indices"):
2028	engine.send_weights([wrong_index_dtype])
2029
2030	assert client.order == []
2031
2032
2033	def test_sparse_nccl_trainer_non_sender_skips_client():
2034	client = RecordingClient()
2035	engine = WeightTransferTrainerFactory.trainer_init(
2036	init_info=SparseNCCLTrainerInitInfo(
2037	master_address="127.0.0.1",
2038	master_port=29500,
2039	world_size=2,
2040	rank=1,
2041	),
2042	client=client,
2043	)
2044
2045	assert engine.is_sender is False
2046	assert engine.model_update_group is None
2047	assert isinstance(engine, SparseNCCLTrainerWeightTransferEngine)
2048	engine.send_weights([_sparse_patch()])
2049	engine.send_weight_chunk([_sparse_patch()])
2050	assert client.order == []
2051
2052
2053	# --- Optional ModelExpress Client Lifecycle ---
2054
2055
2056	def _import_modelexpress_shim():
2057	path = (
2058	Path(__file__).parents[2]
2059	/ "vllm/distributed/weight_transfer/modelexpress_engine.py"
2060	)
2061	runpy.run_path(str(path), run_name="_test_modelexpress_shim")
2062
2063
2064	@pytest.mark.parametrize(
2065	"missing_module",
2066	[
2067	"modelexpress",
2068	"modelexpress_rl",
2069	"modelexpress_rl.inference",
2070	"modelexpress_rl.inference.engines",
2071	"modelexpress_rl.inference.engines.vllm",
2072	"modelexpress_rl.inference.engines.vllm.weight_transfer_engine",
2073	],
2074	)
2075	def test_modelexpress_missing_backend_explains_source_install(
2076	monkeypatch, missing_module
2077	):
2078	original_import = builtins.__import__
2079	error = ModuleNotFoundError(name=missing_module)
2080
2081	def import_without_backend(name, args, *kwargs):
2082	if name == "modelexpress":
2083	raise error
2084	return original_import(name, args, *kwargs)
2085
2086	monkeypatch.setattr(builtins, "__import__", import_without_backend)
2087	with pytest.raises(ImportError, match="modelexpress_client/python") as exc_info:
2088	_import_modelexpress_shim()
2089	assert exc_info.value.__cause__ is error
2090	assert "uv pip install" in str(exc_info.value)
2091
2092
2093	@pytest.mark.parametrize("missing_module", ["grpc", "modelexpress.internal_dependency"])
2094	def test_modelexpress_preserves_transitive_import_errors(monkeypatch, missing_module):
2095	original_import = builtins.__import__
2096	error = ModuleNotFoundError(name=missing_module)
2097
2098	def import_without_dependency(name, args, *kwargs):
2099	if name == "modelexpress":
2100	raise error
2101	return original_import(name, args, *kwargs)
2102
2103	monkeypatch.setattr(builtins, "__import__", import_without_dependency)
2104	with pytest.raises(ModuleNotFoundError) as exc_info:
2105	_import_modelexpress_shim()
2106	assert exc_info.value is error
2107
2108
2109	@pytest.fixture
2110	def mx_client(monkeypatch):
2111	client_module = pytest.importorskip("modelexpress_rl.inference.client")
2112	client = MagicMock()
2113	monkeypatch.setattr(
2114	client_module.ModelExpressGeneratorClient,
2115	"initialize",
2116	MagicMock(return_value=client),
2117	)
2118	monkeypatch.setattr(torch.accelerator, "synchronize", MagicMock())
2119	return client
2120
2121
2122	def _make_modelexpress_engine():
2123	return WeightTransferEngineFactory.create_engine(
2124	WeightTransferConfig(backend="modelexpress"),
2125	SimpleNamespace(
2126	parallel_config=SimpleNamespace(),
2127	model_config=SimpleNamespace(model="test/model"),
2128	),
2129	torch.device("cpu"),
2130	torch.nn.Linear(2, 2),
2131	)
2132
2133
2134	@pytest.fixture
2135	def mx_worker(mx_client, monkeypatch):
2136	from vllm.v1.worker import gpu_worker
2137
2138	monkeypatch.setattr(gpu_worker, "set_current_vllm_config", lambda _: nullcontext())
2139	worker = object.__new__(gpu_worker.Worker)
2140	worker.weight_transfer_engine = _make_modelexpress_engine()
2141	worker.vllm_config = worker.weight_transfer_engine.vllm_config
2142	worker._weight_update_active = False
2143	worker._weight_update_is_draft = False
2144	worker.model_runner = MagicMock()
2145	return worker
2146
2147
2148	def test_modelexpress_worker_rejects_updates_before_initialization(
2149	mx_worker, mx_client
2150	):
2151	with pytest.raises(RuntimeError, match="not initialized"):
2152	mx_worker.start_weight_update()
2153	with pytest.raises(RuntimeError, match="start_weight_update must be called"):
2154	mx_worker.update_weights({"version_id": "version-a"})
2155	with pytest.raises(RuntimeError, match="without a matching"):
2156	mx_worker.finish_weight_update()
2157	mx_client.stage_weight.assert_not_called()
2158
2159
2160	@pytest.mark.parametrize(
2161	"second_update, error_type",
2162	[
2163	({"version_id": "version-a"}, RuntimeError),
2164	({"version_id": "version-b"}, RuntimeError),
2165	({"version_id": ""}, ValueError),
2166	({"version_id": None}, ValueError),
2167	({}, ValueError),
2168	],
2169	)
2170	def test_modelexpress_worker_recovers_after_rejected_update(
2171	mx_worker, mx_client, second_update, error_type
2172	):
2173	"""A rejected payload must release A so a fresh session can actually install B."""
2174	mx_worker.weight_transfer_engine.init_transfer_engine(
2175	mx_worker.weight_transfer_engine.parse_init_info({})
2176	)
2177	staged_a = SimpleNamespace(version_id="version-a", metrics={}, release=MagicMock())
2178	staged_b = SimpleNamespace(version_id="version-b", metrics={}, release=MagicMock())
2179	mx_client.stage_weight.side_effect = [staged_a, staged_b]
2180	mx_client.apply_weight.return_value = {}
2181	mx_worker.start_weight_update()
2182	mx_worker.update_weights({"version_id": "version-a"})
2183
2184	with pytest.raises(error_type):
2185	mx_worker.update_weights(second_update)
2186	staged_a.release.assert_called_once_with()
2187	mx_client.apply_weight.assert_called_once_with(staged_a)
2188
2189	mx_worker.start_weight_update()
2190	mx_worker.update_weights({"version_id": "version-b"})
2191	mx_worker.finish_weight_update()
2192
2193	assert [
2194	call.kwargs["version"].version_id
2195	for call in mx_client.stage_weight.call_args_list
2196	] == ["version-a", "version-b"]
2197	mx_client.apply_weight.assert_called_with(staged_b)
2198	staged_b.release.assert_called_once_with()
2199
2200
2201	def test_modelexpress_worker_can_retry_failed_finish(mx_worker, mx_client):
2202	mx_worker.weight_transfer_engine.init_transfer_engine(
2203	mx_worker.weight_transfer_engine.parse_init_info({})
2204	)
2205	staged = mx_client.stage_weight.return_value
2206	staged.release.side_effect = [RuntimeError("release failed"), None]
2207	mx_worker.start_weight_update()
2208	mx_worker.update_weights({"version_id": "version-a"})
2209
2210	with pytest.raises(RuntimeError, match="release failed"):
2211	mx_worker.finish_weight_update()
2212	mx_worker.finish_weight_update()
2213
2214	assert staged.release.call_count == 2
2215	assert not mx_worker._weight_update_active
2216	mx_worker.start_weight_update()
2217
2218
2219	def test_native_engine_applies_and_releases_exact_versions(mx_client):
2220	from modelexpress_rl.inference.engines.vllm import (
2221	weight_transfer_engine as mx_engine,
2222	)
2223
2224	from vllm.distributed.weight_transfer import modelexpress_engine
2225
2226	engine = _make_modelexpress_engine()
2227	assert type(engine) is mx_engine.ModelExpressWeightTransferEngine
2228	assert (
2229	modelexpress_engine.ModelExpressWeightTransferInitInfo
2230	is mx_engine.ModelExpressWeightTransferInitInfo
2231	)
2232	assert (
2233	modelexpress_engine.ModelExpressWeightTransferUpdateInfo
2234	is mx_engine.ModelExpressWeightTransferUpdateInfo
2235	)
2236	assert not engine.supports_draft_weight_update
2237	engine.init_transfer_engine(engine.parse_init_info({}))
2238
2239	for version_id in ("version-a", "version-b"):
2240	staged = SimpleNamespace(version_id=version_id, metrics={}, release=MagicMock())
2241	mx_client.stage_weight.return_value = staged
2242	mx_client.apply_weight.return_value = {}
2243	engine.start_weight_update()
2244	engine.update_weights({"version_id": version_id})
2245
2246	assert (
2247	mx_client.stage_weight.call_args.kwargs["version"].version_id == version_id
2248	)
2249	mx_client.apply_weight.assert_called_with(staged)
2250	torch.accelerator.synchronize.assert_called()
2251	staged.release.assert_not_called()
2252	engine.finish_weight_update()
2253	staged.release.assert_called_once_with()
2254
2255	engine.shutdown()
2256	engine.shutdown()
2257	mx_client.close.assert_called_once_with()
2258
2259
2260	def test_vime_init_config_reaches_mx_client(mx_client):
2261	from modelexpress_rl.inference.client import ModelExpressGeneratorClient
2262
2263	engine = _make_modelexpress_engine()
2264	ModelExpressGeneratorClient.initialize.assert_not_called()
2265	engine.init_transfer_engine(
2266	engine.parse_init_info(
2267	{
2268	"model_name": "policy",
2269	"server_url": "mx:8001",
2270	"initial_serving_version_id": "serving-a",
2271	"object_storage_type": "S3",
2272	"initial_base_version_id": "base-a",
2273	"seed_checkpoint_path": "/models/launch",
2274	"refit_checkpoint_dir": "/mxdelta/receiver",
2275	"refit_checkpoint_max_size_gb": 200,
2276	"object_storage_endpoint_url": "http://minio:9000",
2277	"object_storage_region_name": "us-west-2",
2278	"registration_ttl_seconds": 90,
2279	"lease_ttl_seconds": 60,
2280	"max_transfer_attempts": 4,
2281	"max_replay_chain_length": 17,
2282	"rpc_timeout_seconds": 12.5,
2283	}
2284	)
2285	)
2286	config = ModelExpressGeneratorClient.initialize.call_args.args[0]
2287	assert config.engine_context.model is engine.model
2288	assert config.engine_context.vllm_config is engine.vllm_config
2289	assert config.model_name == "policy"
2290	assert config.server_url == "mx:8001"
2291	assert config.initial_serving_version_id == "serving-a"
2292	assert config.registration_ttl_seconds == 90
2293	assert config.lease_ttl_seconds == 60
2294	assert config.max_transfer_attempts == 4
2295	assert config.max_replay_chain_length == 17
2296	assert config.rpc_timeout_seconds == 12.5
2297	storage = config.object_storage
2298	assert storage.storage_type.value == "S3"
2299	assert storage.initial_base_version_id == "base-a"
2300	assert storage.seed_checkpoint_path == "/models/launch"
2301	assert storage.refit_checkpoint_dir == "/mxdelta/receiver"
2302	assert storage.refit_checkpoint_max_size_gb == 200
2303	assert storage.endpoint_url == "http://minio:9000"
2304	assert storage.region_name == "us-west-2"
2305
2306
2307	@pytest.mark.parametrize(
2308	"key",
2309	[
2310	"object_storage_type",
2311	"initial_base_version_id",
2312	"seed_checkpoint_path",
2313	"refit_checkpoint_dir",
2314	"object_storage_endpoint_url",
2315	"object_storage_region_name",
2316	],
2317	)
2318	def test_incomplete_storage_config_rejected_before_init(mx_client, key):
2319	from modelexpress_rl.inference.client import ModelExpressGeneratorClient
2320
2321	engine = _make_modelexpress_engine()
2322	with pytest.raises(ValueError, match="object storage requires"):
2323	engine.init_transfer_engine(engine.parse_init_info({key: "value"}))
2324	ModelExpressGeneratorClient.initialize.assert_not_called()
2325
2326
2327	@pytest.mark.parametrize("version_id", ["", " ", None, 1])
2328	def test_invalid_version_rejected_before_staging(mx_client, version_id):
2329	engine = _make_modelexpress_engine()
2330	engine.init_transfer_engine(engine.parse_init_info({}))
2331	engine.start_weight_update()
2332	with pytest.raises(ValueError, match="version_id is required"):
2333	engine.update_weights({"version_id": version_id})
2334	mx_client.stage_weight.assert_not_called()
2335
2336
2337	@pytest.mark.parametrize("failure", ["stage", "apply", "release"])
2338	def test_update_failure_preserves_error_and_releases_handle(mx_client, failure):
2339	engine = _make_modelexpress_engine()
2340	engine.init_transfer_engine(engine.parse_init_info({}))
2341	staged = SimpleNamespace(version_id="version-a", metrics={}, release=MagicMock())
2342	mx_client.stage_weight.return_value = staged
2343	if failure == "stage":
2344	mx_client.stage_weight.side_effect = RuntimeError("stage failed")
2345	else:
2346	mx_client.apply_weight.side_effect = RuntimeError("apply failed")
2347	if failure == "release":
2348	staged.release.side_effect = RuntimeError("release failed")
2349	engine.start_weight_update()
2350	with pytest.raises(RuntimeError, match="stage failed\|apply failed"):
2351	engine.update_weights({"version_id": "version-a"})
2352	assert staged.release.call_count == (0 if failure == "stage" else 1)
2353	torch.accelerator.synchronize.assert_not_called()
2354	engine.shutdown()
2355	mx_client.close.assert_called_once_with()
2356
2357
2358	def test_plugin_does_not_replace_native_backend(mx_client):
2359	from modelexpress.engines.vllm.registration import (
2360	register_plugin_weight_transfer_engine,
2361	)
2362	from modelexpress_rl.inference.engines.vllm.weight_transfer_engine import (
2363	ModelExpressWeightTransferEngine,
2364	)
2365
2366	native_loader = WeightTransferEngineFactory._registry["modelexpress"]
2367	register_plugin_weight_transfer_engine()
2368	assert WeightTransferEngineFactory._registry["modelexpress"] is native_loader
2369	assert type(_make_modelexpress_engine()) is ModelExpressWeightTransferEngine