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

@tduncan / Nvidia-Optimization-Stack-VLLM

Code Commits Issues Pulls Cloud Agent CI/CD

main

py 1,582 lines 54.5 KB

Raw

1	# SPDX-License-Identifier: Apache-2.0
2	# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4	import asyncio
5	import importlib
6	import inspect
7	import os
8	import signal
9	import time
10	import uuid
11	from collections import Counter
12	from concurrent.futures import Future
13	from dataclasses import dataclass
14	from threading import Thread
15	from types import SimpleNamespace
16	from typing import Any
17	from unittest.mock import MagicMock
18
19	import pytest
20	import torch
21	from transformers import AutoTokenizer
22
23	from tests.utils import multi_gpu_test
24	from vllm import SamplingParams
25	from vllm.distributed.kv_events import BlockStored, KVEventBatch, ZmqEventPublisher
26	from vllm.engine.arg_utils import EngineArgs
27	from vllm.platforms import current_platform
28	from vllm.pooling_params import LateInteractionParams, PoolingParams
29	from vllm.usage.usage_lib import UsageContext
30	from vllm.utils.torch_utils import set_default_torch_num_threads
31	from vllm.v1.engine import (
32	EngineCoreOutputs,
33	EngineCoreReadyResponse,
34	EngineCoreRequest,
35	)
36	from vllm.v1.engine.core import EngineCore
37	from vllm.v1.engine.core_client import (
38	AsyncMPClient,
39	DPLBAsyncMPClient,
40	EngineCoreClient,
41	MPClient,
42	SyncMPClient,
43	)
44	from vllm.v1.engine.utils import CoreEngineProcManager
45	from vllm.v1.executor.abstract import Executor
46	from vllm.v1.pool.late_interaction import (
47	LATE_INTERACTION_MODE_CACHE_QUERY,
48	LATE_INTERACTION_MODE_SCORE_DOC,
49	)
50
51	from ...distributed.conftest import MockSubscriber
52	from ...utils import create_new_process_for_each_test
53
54	if not current_platform.is_cuda_alike():
55	pytest.skip(
56	reason="V1 currently only supported on CUDA-alike platforms.",
57	allow_module_level=True,
58	)
59
60	MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
61	TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
62	PROMPT = "Hello my name is Robert and I love quantization kernels"
63	PROMPT_TOKENS = TOKENIZER(PROMPT).input_ids
64	TEST_MODULE = "tests.v1.engine.test_engine_core_client"
65
66	_REQUEST_COUNTER = 0
67
68
69	def make_request(
70	params: SamplingParams, prompt_tokens_ids: list[int] \| None = None
71	) -> EngineCoreRequest:
72	if not prompt_tokens_ids:
73	prompt_tokens_ids = PROMPT_TOKENS
74
75	global _REQUEST_COUNTER
76	_REQUEST_COUNTER += 1
77	request_id = f"request-{_REQUEST_COUNTER}"
78	return EngineCoreRequest(
79	request_id=request_id,
80	external_req_id=f"{request_id}-{uuid.uuid4()}",
81	prompt_token_ids=prompt_tokens_ids,
82	mm_features=None,
83	sampling_params=params,
84	pooling_params=None,
85	arrival_time=time.time(),
86	lora_request=None,
87	cache_salt=None,
88	data_parallel_rank=None,
89	)
90
91
92	def _reload_envs_module():
93	import vllm.envs as envs_mod
94
95	cache_clear = getattr(getattr(envs_mod, "__getattr__", None), "cache_clear", None)
96	if cache_clear is not None:
97	cache_clear()
98	return importlib.reload(envs_mod)
99
100
101	def _reload_core_client_module():
102	module = importlib.import_module("vllm.v1.engine.core_client")
103	return importlib.reload(module)
104
105
106	def test_mp_client_uses_env_timeout(monkeypatch: pytest.MonkeyPatch):
107	timeout_value = 654
108	monkeypatch.setenv("VLLM_ENGINE_READY_TIMEOUT_S", str(timeout_value))
109
110	# Ensure that the environment variable is loaded if caching is enabled
111	_reload_envs_module()
112	core_client_mod = _reload_core_client_module()
113
114	poll_timeouts: list[int] = []
115
116	class ShadowSocket:
117	def poll(self, timeout: int) -> int:
118	# Capture the timeout value for each poll call
119	poll_timeouts.append(timeout)
120	return 1
121
122	def recv_multipart(self):
123	return (b"\x00\x00", b"")
124
125	class DummySocket:
126	def send_multipart(self, _msg, *, copy: bool = False, track: bool = False):
127	if track:
128	return SimpleNamespace(done=True)
129
130	def recv_multipart(self, *, copy: bool = False):
131	return (b"", b"")
132
133	def close(self, *, linger: int = 0):
134	pass
135
136	def bind(self, _address):
137	pass
138
139	def connect(self, _address):
140	pass
141
142	def setsockopt(self, _args, *_kwargs):
143	pass
144
145	monkeypatch.setattr(core_client_mod.zmq.Socket, "shadow", lambda *_: ShadowSocket())
146	monkeypatch.setattr(
147	core_client_mod, "make_zmq_socket", lambda _, *__: DummySocket()
148	)
149
150	parallel_config = SimpleNamespace(
151	data_parallel_size=1,
152	data_parallel_rank=0,
153	data_parallel_index=0,
154	data_parallel_size_local=1,
155	data_parallel_rank_local=None,
156	data_parallel_hybrid_lb=False,
157	data_parallel_external_lb=False,
158	local_engines_only=False,
159	enable_elastic_ep=False,
160	)
161	vllm_config = SimpleNamespace(
162	parallel_config=parallel_config,
163	cache_config=SimpleNamespace(effective_attention_block_size=None),
164	)
165
166	client = core_client_mod.MPClient(
167	asyncio_mode=False,
168	vllm_config=vllm_config,
169	executor_class=object,
170	log_stats=False,
171	client_addresses={
172	"input_address": "inproc://input",
173	"output_address": "inproc://output",
174	},
175	)
176	try:
177	# timeout_value is in seconds, but poll receives milliseconds
178	assert poll_timeouts == [timeout_value * 1000]
179	finally:
180	client.shutdown()
181
182
183	def _make_pooling_request(
184	request_id: str, *, mode: str \| None = None, query_key: str \| None = None
185	) -> EngineCoreRequest:
186	late_interaction_params = None
187	if mode is not None and query_key is not None:
188	late_interaction_params = LateInteractionParams(
189	mode=mode,
190	query_key=query_key,
191	)
192
193	return EngineCoreRequest(
194	request_id=request_id,
195	prompt_token_ids=[1, 2, 3],
196	mm_features=None,
197	sampling_params=None,
198	pooling_params=PoolingParams(
199	task="token_embed",
200	late_interaction_params=late_interaction_params,
201	),
202	arrival_time=time.time(),
203	lora_request=None,
204	cache_salt=None,
205	data_parallel_rank=None,
206	)
207
208
209	def _make_dplb_client(num_engines: int = 3, client_count: int = 1) -> DPLBAsyncMPClient:
210	client = object.__new__(DPLBAsyncMPClient)
211	client.client_count = client_count
212	client.reqs_in_flight = {}
213	client.engine_inflight = Counter()
214	client.core_engines = [bytes([i, 0]) for i in range(num_engines)]
215	client.lb_engines = [[0, 0, 0.0] for _ in range(num_engines)]
216	client.eng_start_index = 0
217	client._kv_event_sources = {}
218	client.vllm_config = SimpleNamespace(
219	use_v2_model_runner=True,
220	parallel_config=SimpleNamespace(all2all_backend="allgather_reducescatter"),
221	)
222	return client
223
224
225	def test_dplb_weight_checksums_gather_every_engine():
226	"""call_utility_async on a DPLB client would keep only the first engine."""
227	client = _make_dplb_client(num_engines=2)
228
229	async def call_utility(method, *args, engine):
230	assert (method, args) == ("compute_weight_checksums", ())
231	return [{f"dp{engine[0]}:tp{tp}:w": "x"} for tp in (0, 1)]
232
233	client._call_utility_async = call_utility
234	workers = asyncio.run(client.compute_weight_checksums_async())
235	assert [key for worker in workers for key in worker] == [
236	"dp0:tp0:w",
237	"dp0:tp1:w",
238	"dp1:tp0:w",
239	"dp1:tp1:w",
240	]
241
242
243	def test_dplb_late_interaction_sticky_routing():
244	client = _make_dplb_client()
245
246	query_key = "rerank-abc-query-0"
247	query_request = _make_pooling_request(
248	"query-req", mode=LATE_INTERACTION_MODE_CACHE_QUERY, query_key=query_key
249	)
250	doc_request = _make_pooling_request(
251	"doc-req", mode=LATE_INTERACTION_MODE_SCORE_DOC, query_key=query_key
252	)
253
254	query_engine = client.get_core_engine_for_request(query_request)
255	doc_engine = client.get_core_engine_for_request(doc_request)
256
257	assert query_engine == doc_engine
258	assert client.reqs_in_flight["query-req"] == query_engine
259	assert client.reqs_in_flight["doc-req"] == doc_engine
260
261
262	def test_dplb_non_late_interaction_still_uses_lb():
263	client = _make_dplb_client()
264	client.lb_engines = [[2, 1, 0.0], [0, 0, 0.0], [1, 0, 0.0]]
265
266	request = make_request(SamplingParams(max_tokens=1))
267	chosen_engine = client.get_core_engine_for_request(request)
268
269	assert chosen_engine == client.core_engines[1]
270	assert client.lb_engines[1][0] == 1
271
272
273	def test_dplb_in_flight_request_stays_on_same_engine():
274	client = _make_dplb_client()
275	request = make_request(SamplingParams(max_tokens=1))
276	engine = client.core_engines[2]
277	client.reqs_in_flight[request.request_id] = engine
278	client.engine_inflight[engine] = 1
279
280	chosen_engine = client.get_core_engine_for_request(request)
281
282	assert chosen_engine == engine
283	assert client.engine_inflight[engine] == 1
284
285
286	def test_dplb_burst_round_robins_despite_snapshot_rebinds():
287	"""A stats snapshot rebind wipes the optimistic lb_engines increments;
288	the exact in-flight floor must keep a burst spreading round-robin."""
289	client = _make_dplb_client(num_engines=4)
290
291	for _ in range(4):
292	client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1)))
293	# Coordinator snapshot arrives, not yet reflecting the 4 routed requests.
294	client.lb_engines = [[0, 0, 0.0] for _ in range(4)]
295	for _ in range(4):
296	client.get_core_engine_for_request(make_request(SamplingParams(max_tokens=1)))
297
298	assert sorted(client.engine_inflight.values()) == [2, 2, 2, 2]
299
300
301	@pytest.mark.asyncio
302	async def test_dplb_scale_down_routes_after_stale_stats_snapshot():
303	"""A coordinator snapshot during scale-down must only route to survivors."""
304	import msgspec
305	import zmq.asyncio
306
307	client = _make_dplb_client(num_engines=4)
308	client.engine_ranks_managed = list(range(4))
309	client.engines_running = True
310	client.current_wave = 0
311	client.resources = SimpleNamespace(stats_update_task=None)
312	client.stats_update_address = "inproc://scale-down-stats"
313	client.first_req_sock_addr = "inproc://scale-down-first-request"
314	pause_started = asyncio.Event()
315
316	async def pause_scheduler(args, *kwargs):
317	pause_started.set()
318	await asyncio.Future()
319
320	client._call_utility_async = pause_scheduler
321
322	with (
323	zmq.asyncio.Context() as ctx,
324	ctx.socket(zmq.XPUB) as coordinator,
325	ctx.socket(zmq.PAIR) as first_req_socket,
326	):
327	client.ctx = ctx
328	coordinator.bind(client.stats_update_address)
329	first_req_socket.bind(client.first_req_sock_addr)
330	client._ensure_stats_update_task()
331	scale_task = None
332	try:
333	# Wait for subscription so the snapshot cannot be dropped.
334	await asyncio.wait_for(coordinator.recv(), timeout=5)
335	scale_task = asyncio.create_task(client._commit_scale_down_elastic_ep(2))
336	await asyncio.wait_for(pause_started.wait(), timeout=5)
337
338	# The coordinator still knows about all four engines while the
339	# client has already stopped routing to the two removed engines.
340	counts = [[0, 0, 0.0], [0, 0, 0.0], [1, 0, 0.0], [1, 0, 0.0]]
341	await coordinator.send(msgspec.msgpack.encode((counts, 1, True)))
342
343	async def wait_for_snapshot():
344	while client.current_wave != 1:
345	await asyncio.sleep(0)
346
347	await asyncio.wait_for(wait_for_snapshot(), timeout=5)
348	chosen = client.get_core_engine_for_request(
349	_make_pooling_request("scale-down-request")
350	)
351	assert chosen == client.core_engines[0]
352	finally:
353	tasks = [client.resources.stats_update_task]
354	if scale_task is not None:
355	tasks.append(scale_task)
356	for task in tasks:
357	task.cancel()
358	await asyncio.gather(*tasks, return_exceptions=True)
359
360
361	def test_dplb_snapshot_backpressure_overrides_inflight():
362	"""An engine reported heavily loaded by the coordinator is avoided even
363	when this client has routed nothing to it."""
364	client = _make_dplb_client(num_engines=2)
365	client.lb_engines = [[5, 10, 0.0], [0, 0, 0.0]]
366
367	chosen = client.get_core_engine_for_request(
368	make_request(SamplingParams(max_tokens=1))
369	)
370
371	assert chosen == client.core_engines[1]
372
373
374	def test_dplb_kv_pressure_amplifies_waiting_penalty():
375	"""A waiting queue on a KV-bound engine (slow drain) is penalized, while
376	the same queue with low KV usage is not (e.g. transient burst)."""
377	client = _make_dplb_client(num_engines=2)
378	# Engine 0 has a smaller total but is KV-bound with a queue.
379	client.lb_engines = [[5, 10, 1.0], [0, 20, 0.2]]
380
381	chosen = client.get_core_engine_for_request(
382	make_request(SamplingParams(max_tokens=1))
383	)
384	assert chosen == client.core_engines[1]
385
386	# Same counts without KV pressure: the smaller total wins.
387	client = _make_dplb_client(num_engines=2)
388	client.lb_engines = [[5, 10, 0.2], [0, 20, 0.2]]
389
390	chosen = client.get_core_engine_for_request(
391	make_request(SamplingParams(max_tokens=1))
392	)
393	assert chosen == client.core_engines[0]
394
395
396	def test_dplb_finished_requests_release_inflight():
397	client = _make_dplb_client(num_engines=2)
398
399	req = make_request(SamplingParams(max_tokens=1))
400	engine = client.get_core_engine_for_request(req)
401	assert client.engine_inflight[engine] == 1
402
403	outputs = EngineCoreOutputs(finished_requests={req.request_id})
404	asyncio.run(DPLBAsyncMPClient.process_engine_outputs(client, outputs))
405
406	assert client.engine_inflight[engine] == 0
407	assert req.request_id not in client.reqs_in_flight
408
409
410	@pytest.mark.parametrize(
411	("effective_size", "other_size"),
412	[(None, None), (4224, 4224), (4224, 1056), (4224, None)],
413	)
414	def test_apply_ready_response_syncs_block_size(effective_size, other_size):
415	import msgspec
416
417	client = object.__new__(MPClient)
418	client._effective_attention_block_sizes = set()
419	client.vllm_config = SimpleNamespace(
420	cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
421	model_config=SimpleNamespace(max_model_len=8192),
422	)
423	client.stats_update_address = None
424
425	payload = msgspec.msgpack.encode(
426	EngineCoreReadyResponse(
427	max_model_len=8192,
428	num_gpu_blocks=100,
429	block_size=1056,
430	dp_stats_address=None,
431	dtype="bfloat16",
432	vllm_version="test",
433	world_size=1,
434	data_parallel_size=1,
435	tensor_parallel_size=1,
436	pipeline_parallel_size=1,
437	decode_context_parallel_size=1,
438	data_parallel_rank=0,
439	max_num_seqs=256,
440	max_num_batched_tokens=8192,
441	instance_id="test-instance",
442	supports_lora=False,
443	max_loras=0,
444	)
445	)
446	fields = msgspec.msgpack.decode(payload)
447	if effective_size is None:
448	del fields["effective_attention_block_size"]
449	else:
450	fields["effective_attention_block_size"] = effective_size
451	client._apply_ready_response(msgspec.msgpack.encode(fields))
452	assert client.vllm_config.cache_config.block_size == 1056
453	cache_config = client.vllm_config.cache_config
454	assert cache_config.effective_attention_block_size == effective_size
455
456	fields["effective_attention_block_size"] = other_size
457	client._apply_ready_response(msgspec.msgpack.encode(fields))
458	assert cache_config.effective_attention_block_size == (
459	effective_size if effective_size == other_size else None
460	)
461
462	client._apply_ready_response(b"")
463	assert cache_config.effective_attention_block_size is None
464
465
466	def test_apply_ready_response_syncs_mamba_block_size():
467	import msgspec
468
469	client = object.__new__(MPClient)
470	client._effective_attention_block_sizes = set()
471	client.vllm_config = SimpleNamespace(
472	cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
473	model_config=SimpleNamespace(max_model_len=8192),
474	)
475	client.stats_update_address = None
476
477	payload = msgspec.msgpack.encode(
478	EngineCoreReadyResponse(
479	max_model_len=8192,
480	num_gpu_blocks=100,
481	block_size=1056,
482	dp_stats_address=None,
483	dtype="bfloat16",
484	vllm_version="test",
485	world_size=1,
486	data_parallel_size=1,
487	tensor_parallel_size=1,
488	pipeline_parallel_size=1,
489	decode_context_parallel_size=1,
490	data_parallel_rank=0,
491	max_num_seqs=256,
492	max_num_batched_tokens=8192,
493	instance_id="test-instance",
494	supports_lora=False,
495	max_loras=0,
496	mamba_block_size=1056,
497	)
498	)
499	client._apply_ready_response(payload)
500	assert client.vllm_config.cache_config.mamba_block_size == 1056
501
502
503	def loop_until_done(client: EngineCoreClient, outputs: dict):
504	while True:
505	engine_core_outputs = client.get_output().outputs
506
507	if len(engine_core_outputs) == 0:
508	continue
509
510	all_finished = True
511	for out in engine_core_outputs:
512	outputs[out.request_id].append(out)
513	if not out.finished:
514	all_finished = False
515
516	if all_finished:
517	break
518
519
520	async def loop_until_done_async(client: EngineCoreClient, outputs: dict):
521	while True:
522	engine_core_outputs = (await client.get_output_async()).outputs
523
524	if len(engine_core_outputs) == 0:
525	continue
526
527	all_finished = True
528	for out in engine_core_outputs:
529	outputs[out.request_id].append(out)
530	if not out.finished:
531	all_finished = False
532
533	if all_finished:
534	break
535
536
537	async def loop_until_fully_done_async(client: EngineCoreClient, outputs: dict):
538	while True:
539	engine_core_outputs = (await client.get_output_async()).outputs
540
541	if len(engine_core_outputs) == 0:
542	continue
543
544	# Add outputs to the dict
545	for out in engine_core_outputs:
546	outputs[out.request_id].append(out)
547
548	# Check if all request IDs in outputs have finished
549	if all(outs and outs[-1].finished for outs in outputs.values()):
550	break
551
552	await asyncio.sleep(0.1)
553
554
555	# Dummy utility function to monkey-patch into engine core.
556	def echo(self, msg: str, err_msg: str \| None = None, sleep: float \| None = None) -> str:
557	print(f"echo util function called: {msg}, {err_msg}")
558	if sleep is not None:
559	time.sleep(sleep)
560	if err_msg is not None:
561	raise ValueError(err_msg)
562	return msg
563
564
565	@dataclass
566	class TestMessage:
567	"""Test dataclass for verifying custom type serialization."""
568
569	message: str
570
571
572	# Dummy utility function to monkey-patch into engine core.
573	def echo_dc(
574	self,
575	msg: str,
576	return_list: bool = False,
577	) -> TestMessage \| list[TestMessage]:
578	print(f"echo dc util function called: {msg}")
579	val = None if msg is None else TestMessage(msg)
580	# Return dataclass to verify support for returning custom types
581	# (for which there is special handling to make it work with msgspec).
582	return [val for _ in range(3)] if return_list else val
583
584
585	# Dummy utility function to test dict serialization with custom types.
586	def echo_dc_dict(
587	self,
588	msg: str,
589	return_dict: bool = False,
590	) -> TestMessage \| dict[str, TestMessage]:
591	print(f"echo dc dict util function called: {msg}")
592	val = None if msg is None else TestMessage(msg)
593	# Return dict of dataclasses to verify support for returning dicts
594	# with custom value types.
595	if return_dict:
596	return {"key1": val, "key2": val, "key3": val}
597	else:
598	return val
599
600
601	# Dummy utility function to test nested structures with custom types.
602	def echo_dc_nested(
603	self,
604	msg: str,
605	structure_type: str = "list_of_dicts",
606	) -> Any:
607	print(f"echo dc nested util function called: {msg}, structure: {structure_type}")
608	val = None if msg is None else TestMessage(msg)
609
610	structures = {
611	"list_of_dicts": [{"a": val, "b": val}, {"c": val, "d": val}],
612	"dict_of_lists": {"list1": [val, val], "list2": [val, val]},
613	"deep_nested": {"outer": [{"inner": [val, val]}, {"inner": [val]}]},
614	}
615	return structures.get(structure_type, val)
616
617
618	def future_echo(self, value: Any, num_wait_loops: int = 2) -> Future:
619	"""Utility that returns a Future completed once the engine is idle
620	(tests deferred utility path).
621	"""
622	future: Future = Future()
623
624	def idle(engine: EngineCore):
625	future.set_result(value)
626
627	self._idle_state_callbacks.append(idle)
628	return future
629
630
631	# --- Fixtures for subprocess patching ---
632	# These create sitecustomize.py files that patch EngineCore in spawned
633	# subprocesses. This is necessary because ROCm requires 'spawn' multiprocessing
634	# start method, which creates fresh Python interpreters that don't inherit
635	# monkey-patches from the parent process.
636
637
638	@pytest.fixture
639	def subprocess_echo_patch(monkeypatch, tmp_path):
640	"""Create sitecustomize.py so spawned subprocesses have echo method.
641
642	This is needed because ROCm uses 'spawn' multiprocessing start method,
643	which creates a fresh Python interpreter that doesn't inherit monkey-patches.
644	By using sitecustomize.py, we ensure the patch is applied when Python starts.
645	"""
646	sc = tmp_path / "sitecustomize.py"
647	sc.write_text(
648	"\n".join(
649	[
650	"import time",
651	"from vllm.v1.engine.core import EngineCore",
652	inspect.getsource(echo),
653	"EngineCore.echo = echo",
654	]
655	)
656	)
657	monkeypatch.setenv(
658	"PYTHONPATH",
659	os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
660	)
661
662
663	@pytest.fixture
664	def subprocess_echo_dc_patch(monkeypatch, tmp_path):
665	"""Create sitecustomize.py so spawned subprocesses have echo_dc method."""
666	sc = tmp_path / "sitecustomize.py"
667	sc.write_text(
668	"\n".join(
669	[
670	"from dataclasses import dataclass",
671	"",
672	inspect.getsource(TestMessage),
673	f"TestMessage.__module__ = '{TEST_MODULE}'",
674	"",
675	"from vllm.v1.engine.core import EngineCore",
676	inspect.getsource(echo_dc),
677	"EngineCore.echo_dc = echo_dc",
678	]
679	)
680	)
681	monkeypatch.setenv(
682	"PYTHONPATH",
683	os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
684	)
685
686
687	@pytest.fixture
688	def subprocess_echo_dc_dict_patch(monkeypatch, tmp_path):
689	"""Create sitecustomize.py so spawned subprocesses have echo_dc_dict method."""
690	sc = tmp_path / "sitecustomize.py"
691	sc.write_text(
692	"\n".join(
693	[
694	"from dataclasses import dataclass",
695	"",
696	inspect.getsource(TestMessage),
697	f"TestMessage.__module__ = '{TEST_MODULE}'",
698	"",
699	"from vllm.v1.engine.core import EngineCore",
700	inspect.getsource(echo_dc_dict),
701	"EngineCore.echo_dc_dict = echo_dc_dict",
702	]
703	)
704	)
705	monkeypatch.setenv(
706	"PYTHONPATH",
707	os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
708	)
709
710
711	@pytest.fixture
712	def subprocess_echo_dc_nested_patch(monkeypatch, tmp_path):
713	"""Create sitecustomize.py so spawned subprocesses have echo_dc_nested method."""
714	sc = tmp_path / "sitecustomize.py"
715	sc.write_text(
716	"\n".join(
717	[
718	"from dataclasses import dataclass",
719	"from typing import Any",
720	"",
721	inspect.getsource(TestMessage),
722	f"TestMessage.__module__ = '{TEST_MODULE}'",
723	"",
724	"from vllm.v1.engine.core import EngineCore",
725	inspect.getsource(echo_dc_nested),
726	"EngineCore.echo_dc_nested = echo_dc_nested",
727	]
728	)
729	)
730	monkeypatch.setenv(
731	"PYTHONPATH",
732	os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
733	)
734
735
736	@pytest.fixture
737	def subprocess_future_echo_patch(monkeypatch, tmp_path):
738	"""Create sitecustomize.py so spawned subprocesses have future_echo method."""
739	sc = tmp_path / "sitecustomize.py"
740	sc.write_text(
741	"\n".join(
742	[
743	"from concurrent.futures import Future",
744	"from typing import Any",
745	"",
746	"from vllm.v1.engine.core import EngineCore",
747	inspect.getsource(future_echo),
748	"EngineCore.future_echo = future_echo",
749	]
750	)
751	)
752	monkeypatch.setenv(
753	"PYTHONPATH",
754	os.pathsep.join(filter(None, [str(tmp_path), os.getenv("PYTHONPATH")])),
755	)
756
757
758	@create_new_process_for_each_test()
759	@pytest.mark.parametrize("multiprocessing_mode", [True, False])
760	def test_engine_core_client(
761	monkeypatch: pytest.MonkeyPatch,
762	multiprocessing_mode: bool,
763	subprocess_echo_patch,
764	):
765	with monkeypatch.context() as m:
766	# Monkey-patch core engine utility function to test.
767	m.setattr(EngineCore, "echo", echo, raising=False)
768
769	engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)
770	vllm_config = engine_args.create_engine_config(UsageContext.UNKNOWN_CONTEXT)
771	executor_class = Executor.get_class(vllm_config)
772
773	with set_default_torch_num_threads(1):
774	client = EngineCoreClient.make_client(
775	multiprocess_mode=multiprocessing_mode,
776	asyncio_mode=False,
777	vllm_config=vllm_config,
778	executor_class=executor_class,
779	log_stats=False,
780	)
781
782	MAX_TOKENS = 20
783	params = SamplingParams(max_tokens=MAX_TOKENS)
784	"""Normal Request Cycle."""
785	requests = [make_request(params) for _ in range(10)]
786	request_ids = [req.request_id for req in requests]
787
788	# Add requests to the engine.
789	for request in requests:
790	client.add_request(request)
791	time.sleep(0.01)
792
793	outputs: dict[str, list] = {req_id: [] for req_id in request_ids}
794	loop_until_done(client, outputs)
795
796	for req_id in request_ids:
797	assert len(outputs[req_id]) == MAX_TOKENS, (
798	f"{outputs[req_id]=}, {MAX_TOKENS=}"
799	)
800	"""Abort Request Cycle."""
801
802	# Note: this code pathway will only work for multiprocessing
803	# since we have to call get_output() explicitly
804
805	# Add requests to the engine.
806	for idx, request in enumerate(requests):
807	client.add_request(request)
808	time.sleep(0.01)
809	if idx % 2 == 0:
810	client.abort_requests([request.request_id])
811
812	outputs = {req_id: [] for req_id in request_ids}
813	loop_until_done(client, outputs)
814
815	for idx, req_id in enumerate(request_ids):
816	if idx % 2 == 0:
817	assert len(outputs[req_id]) < MAX_TOKENS, (
818	f"{len(outputs[req_id])=}, {MAX_TOKENS=}"
819	)
820	else:
821	assert len(outputs[req_id]) == MAX_TOKENS, (
822	f"{len(outputs[req_id])=}, {MAX_TOKENS=}"
823	)
824	"""Abort after request is finished."""
825
826	# Note: this code pathway will only work for multiprocessing
827	# since we have to call get_output() explicitly
828
829	request = requests[0]
830	client.add_request(request)
831	time.sleep(10.0)
832
833	client.abort_requests([request.request_id])
834
835	if multiprocessing_mode:
836	"""Utility method invocation"""
837
838	core_client: SyncMPClient = client
839
840	result = core_client.call_utility("echo", "testarg")
841	assert result == "testarg"
842
843	with pytest.raises(Exception) as e_info:
844	core_client.call_utility("echo", None, "help!")
845
846	assert str(e_info.value) == "Call to echo method failed: help!"
847
848
849	@pytest.mark.asyncio(loop_scope="function")
850	async def test_engine_core_client_asyncio(
851	monkeypatch: pytest.MonkeyPatch,
852	subprocess_echo_patch,
853	):
854	with monkeypatch.context() as m:
855	# Monkey-patch core engine utility function to test.
856	m.setattr(EngineCore, "echo", echo, raising=False)
857
858	engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)
859	vllm_config = engine_args.create_engine_config(
860	usage_context=UsageContext.UNKNOWN_CONTEXT
861	)
862	executor_class = Executor.get_class(vllm_config)
863
864	with set_default_torch_num_threads(1):
865	client = EngineCoreClient.make_client(
866	multiprocess_mode=True,
867	asyncio_mode=True,
868	vllm_config=vllm_config,
869	executor_class=executor_class,
870	log_stats=True,
871	)
872
873	try:
874	MAX_TOKENS = 20
875	params = SamplingParams(max_tokens=MAX_TOKENS)
876	"""Normal Request Cycle."""
877
878	requests = [make_request(params) for _ in range(10)]
879	request_ids = [req.request_id for req in requests]
880
881	# Add requests to the engine.
882	for request in requests:
883	await client.add_request_async(request)
884	await asyncio.sleep(0.01)
885
886	outputs: dict[str, list] = {req_id: [] for req_id in request_ids}
887	await loop_until_done_async(client, outputs)
888
889	for req_id in request_ids:
890	assert len(outputs[req_id]) == MAX_TOKENS, (
891	f"{outputs[req_id]=}, {MAX_TOKENS=}"
892	)
893	"""Abort Request Cycle."""
894
895	# Add requests to the engine.
896	for idx, request in enumerate(requests):
897	await client.add_request_async(request)
898	await asyncio.sleep(0.01)
899	if idx % 2 == 0:
900	await client.abort_requests_async([request.request_id])
901
902	outputs = {req_id: [] for req_id in request_ids}
903	await loop_until_done_async(client, outputs)
904
905	for idx, req_id in enumerate(request_ids):
906	if idx % 2 == 0:
907	assert len(outputs[req_id]) < MAX_TOKENS, (
908	f"{len(outputs[req_id])=}, {MAX_TOKENS=}"
909	)
910	else:
911	assert len(outputs[req_id]) == MAX_TOKENS, (
912	f"{len(outputs[req_id])=}, {MAX_TOKENS=}"
913	)
914	"""Utility method invocation"""
915
916	core_client: AsyncMPClient = client
917
918	result = await core_client.call_utility_async("echo", "testarg")
919	assert result == "testarg"
920
921	with pytest.raises(Exception) as e_info:
922	await core_client.call_utility_async("echo", None, "help!")
923
924	assert str(e_info.value) == "Call to echo method failed: help!"
925
926	# Test that cancelling the utility call doesn't destabilize the
927	# engine.
928	util_task = asyncio.create_task(
929	core_client.call_utility_async("echo", "testarg2", None, 0.5)
930	) # sleep for 0.5 sec
931	await asyncio.sleep(0.05)
932	cancelled = util_task.cancel()
933	assert cancelled
934
935	# Ensure client is still functional. The engine runs utility
936	# methods in a single thread so this request won't be processed
937	# until the cancelled sleeping one is complete.
938	result = await asyncio.wait_for(
939	core_client.call_utility_async("echo", "testarg3"), timeout=1.0
940	)
941	assert result == "testarg3"
942	finally:
943	client.shutdown()
944
945
946	@pytest.mark.asyncio(loop_scope="function")
947	async def test_engine_core_client_util_method_custom_return(
948	monkeypatch: pytest.MonkeyPatch,
949	subprocess_echo_dc_patch,
950	):
951	with monkeypatch.context() as m:
952	# Must set insecure serialization to allow returning custom types.
953	m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
954
955	# Monkey-patch core engine utility function to test.
956	m.setattr(EngineCore, "echo_dc", echo_dc, raising=False)
957
958	engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)
959	vllm_config = engine_args.create_engine_config(
960	usage_context=UsageContext.UNKNOWN_CONTEXT
961	)
962	executor_class = Executor.get_class(vllm_config)
963
964	with set_default_torch_num_threads(1):
965	client = EngineCoreClient.make_client(
966	multiprocess_mode=True,
967	asyncio_mode=True,
968	vllm_config=vllm_config,
969	executor_class=executor_class,
970	log_stats=True,
971	)
972
973	try:
974	# Test utility method returning custom / non-native data type.
975	core_client: AsyncMPClient = client
976
977	result = await core_client.call_utility_async("echo_dc", "testarg2", False)
978	assert isinstance(result, TestMessage) and result.message == "testarg2"
979	result = await core_client.call_utility_async("echo_dc", "testarg2", True)
980	assert isinstance(result, list) and all(
981	isinstance(r, TestMessage) and r.message == "testarg2" for r in result
982	)
983
984	# Test returning None and list of Nones
985	result = await core_client.call_utility_async("echo_dc", None, False)
986	assert result is None
987	result = await core_client.call_utility_async("echo_dc", None, True)
988	assert isinstance(result, list) and all(r is None for r in result)
989
990	finally:
991	client.shutdown()
992
993
994	@pytest.mark.asyncio(loop_scope="function")
995	async def test_engine_core_client_util_method_custom_dict_return(
996	monkeypatch: pytest.MonkeyPatch,
997	subprocess_echo_dc_dict_patch,
998	):
999	with monkeypatch.context() as m:
1000	# Must set insecure serialization to allow returning custom types.
1001	m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
1002
1003	# Monkey-patch core engine utility function to test.
1004	m.setattr(EngineCore, "echo_dc_dict", echo_dc_dict, raising=False)
1005
1006	engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)
1007	vllm_config = engine_args.create_engine_config(
1008	usage_context=UsageContext.UNKNOWN_CONTEXT
1009	)
1010	executor_class = Executor.get_class(vllm_config)
1011
1012	with set_default_torch_num_threads(1):
1013	client = EngineCoreClient.make_client(
1014	multiprocess_mode=True,
1015	asyncio_mode=True,
1016	vllm_config=vllm_config,
1017	executor_class=executor_class,
1018	log_stats=True,
1019	)
1020
1021	try:
1022	# Test utility method returning custom / non-native data type.
1023	core_client: AsyncMPClient = client
1024
1025	# Test single object return
1026	result = await core_client.call_utility_async(
1027	"echo_dc_dict", "testarg3", False
1028	)
1029	assert isinstance(result, TestMessage) and result.message == "testarg3"
1030
1031	# Test dict return with custom value types
1032	result = await core_client.call_utility_async(
1033	"echo_dc_dict", "testarg3", True
1034	)
1035	assert isinstance(result, dict) and len(result) == 3
1036	for key, val in result.items():
1037	assert key in ["key1", "key2", "key3"]
1038	assert isinstance(val, TestMessage) and val.message == "testarg3"
1039
1040	# Test returning dict with None values
1041	result = await core_client.call_utility_async("echo_dc_dict", None, True)
1042	assert isinstance(result, dict) and len(result) == 3
1043	for key, val in result.items():
1044	assert key in ["key1", "key2", "key3"]
1045	assert val is None
1046
1047	finally:
1048	client.shutdown()
1049
1050
1051	@pytest.mark.asyncio(loop_scope="function")
1052	async def test_engine_core_client_util_method_nested_structures(
1053	monkeypatch: pytest.MonkeyPatch,
1054	subprocess_echo_dc_nested_patch,
1055	):
1056	with monkeypatch.context() as m:
1057	# Must set insecure serialization to allow returning custom types.
1058	m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
1059
1060	# Monkey-patch core engine utility function to test.
1061	m.setattr(EngineCore, "echo_dc_nested", echo_dc_nested, raising=False)
1062
1063	engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)
1064	vllm_config = engine_args.create_engine_config(
1065	usage_context=UsageContext.UNKNOWN_CONTEXT
1066	)
1067	executor_class = Executor.get_class(vllm_config)
1068
1069	with set_default_torch_num_threads(1):
1070	client = EngineCoreClient.make_client(
1071	multiprocess_mode=True,
1072	asyncio_mode=True,
1073	vllm_config=vllm_config,
1074	executor_class=executor_class,
1075	log_stats=True,
1076	)
1077
1078	try:
1079	core_client: AsyncMPClient = client
1080
1081	# Test list of dicts: [{"a": val, "b": val}, {"c": val, "d": val}]
1082	result = await core_client.call_utility_async(
1083	"echo_dc_nested", "nested1", "list_of_dicts"
1084	)
1085	assert isinstance(result, list) and len(result) == 2
1086	for i, item in enumerate(result):
1087	assert isinstance(item, dict)
1088	if i == 0:
1089	assert "a" in item and "b" in item
1090	assert (
1091	isinstance(item["a"], TestMessage)
1092	and item["a"].message == "nested1"
1093	)
1094	assert (
1095	isinstance(item["b"], TestMessage)
1096	and item["b"].message == "nested1"
1097	)
1098	else:
1099	assert "c" in item and "d" in item
1100	assert (
1101	isinstance(item["c"], TestMessage)
1102	and item["c"].message == "nested1"
1103	)
1104	assert (
1105	isinstance(item["d"], TestMessage)
1106	and item["d"].message == "nested1"
1107	)
1108
1109	# Test dict of lists: {"list1": [val, val], "list2": [val, val]}
1110	result = await core_client.call_utility_async(
1111	"echo_dc_nested", "nested2", "dict_of_lists"
1112	)
1113	assert isinstance(result, dict) and len(result) == 2
1114	assert "list1" in result and "list2" in result
1115	for key, lst in result.items():
1116	assert isinstance(lst, list) and len(lst) == 2
1117	for item in lst:
1118	assert isinstance(item, TestMessage) and item.message == "nested2"
1119
1120	# Test deeply nested: {"outer": [{"inner": [val, val]},
1121	# {"inner": [val]}]}
1122	result = await core_client.call_utility_async(
1123	"echo_dc_nested", "nested3", "deep_nested"
1124	)
1125	assert isinstance(result, dict) and "outer" in result
1126	outer_list = result["outer"]
1127	assert isinstance(outer_list, list) and len(outer_list) == 2
1128
1129	# First dict in outer list should have "inner" with 2 items
1130	inner_dict1 = outer_list[0]
1131	assert isinstance(inner_dict1, dict) and "inner" in inner_dict1
1132	inner_list1 = inner_dict1["inner"]
1133	assert isinstance(inner_list1, list) and len(inner_list1) == 2
1134	for item in inner_list1:
1135	assert isinstance(item, TestMessage) and item.message == "nested3"
1136
1137	# Second dict in outer list should have "inner" with 1 item
1138	inner_dict2 = outer_list[1]
1139	assert isinstance(inner_dict2, dict) and "inner" in inner_dict2
1140	inner_list2 = inner_dict2["inner"]
1141	assert isinstance(inner_list2, list) and len(inner_list2) == 1
1142	assert (
1143	isinstance(inner_list2[0], TestMessage)
1144	and inner_list2[0].message == "nested3"
1145	)
1146
1147	# Test with None values in nested structures
1148	result = await core_client.call_utility_async(
1149	"echo_dc_nested", None, "list_of_dicts"
1150	)
1151	assert isinstance(result, list) and len(result) == 2
1152	for item in result:
1153	assert isinstance(item, dict)
1154	for val in item.values():
1155	assert val is None
1156
1157	finally:
1158	client.shutdown()
1159
1160
1161	@pytest.mark.asyncio(loop_scope="function")
1162	async def test_engine_core_client_future_utility_async(
1163	monkeypatch: pytest.MonkeyPatch,
1164	subprocess_future_echo_patch,
1165	):
1166	"""Test that a utility returning a Future completes when the future is done
1167	(engine uses add_done_callback).
1168	"""
1169	with monkeypatch.context() as m:
1170	m.setattr(EngineCore, "future_echo", future_echo, raising=False)
1171
1172	engine_args = EngineArgs(model=MODEL_NAME, enforce_eager=True)
1173	vllm_config = engine_args.create_engine_config(
1174	usage_context=UsageContext.UNKNOWN_CONTEXT
1175	)
1176	executor_class = Executor.get_class(vllm_config)
1177
1178	with set_default_torch_num_threads(1):
1179	client = EngineCoreClient.make_client(
1180	multiprocess_mode=True,
1181	asyncio_mode=True,
1182	vllm_config=vllm_config,
1183	executor_class=executor_class,
1184	log_stats=True,
1185	)
1186
1187	try:
1188	core_client: AsyncMPClient = client
1189
1190	# Completes after 2 engine steps (num_wait_loops=2)
1191	result = await core_client.call_utility_async(
1192	"future_echo", "future_result", 2
1193	)
1194	assert result == "future_result"
1195
1196	# None is a valid result (num_wait_loops=0 → completes on first step)
1197	result = await core_client.call_utility_async("future_echo", None, 0)
1198	assert result is None
1199	finally:
1200	client.shutdown()
1201
1202
1203	@pytest.mark.parametrize(
1204	"model_name,num_groups",
1205	[
1206	("meta-llama/Llama-3.2-1B-Instruct", 1),
1207	("google/gemma-3-1b-it", 7),
1208	],
1209	)
1210	@pytest.mark.parametrize(
1211	"multiprocessing_mode,publisher_config",
1212	[(True, "tcp"), (False, "inproc")],
1213	indirect=["publisher_config"],
1214	)
1215	def test_kv_cache_events(
1216	multiprocessing_mode: bool,
1217	publisher_config,
1218	model_name: str,
1219	num_groups: int,
1220	):
1221	block_size = 16
1222	num_blocks = 2
1223
1224	engine_args = EngineArgs(
1225	model=model_name,
1226	enforce_eager=True,
1227	enable_prefix_caching=True,
1228	prefix_cache_retention_interval=None,
1229	block_size=block_size,
1230	)
1231	engine_args.kv_events_config = publisher_config
1232
1233	vllm_config = engine_args.create_engine_config(UsageContext.UNKNOWN_CONTEXT)
1234
1235	executor_class = Executor.get_class(vllm_config)
1236	with set_default_torch_num_threads(1):
1237	client = EngineCoreClient.make_client(
1238	multiprocess_mode=multiprocessing_mode,
1239	asyncio_mode=False,
1240	vllm_config=vllm_config,
1241	executor_class=executor_class,
1242	log_stats=False,
1243	)
1244	endpoint = publisher_config.endpoint.replace("*", "127.0.0.1")
1245	subscriber = MockSubscriber(
1246	endpoint, topic=publisher_config.topic, decode_type=KVEventBatch
1247	)
1248
1249	try:
1250	custom_tokens = list(range(num_blocks * block_size))
1251	sampling_params = SamplingParams(max_tokens=1)
1252	request = make_request(sampling_params, custom_tokens)
1253	client.add_request(request)
1254
1255	outputs: dict[str, list] = {request.request_id: []}
1256	loop_until_done(client, outputs)
1257
1258	result = subscriber.receive_one(timeout=1000)
1259	assert result is not None, "No message received"
1260
1261	seq, received = result
1262	assert seq == 0, "Sequence number mismatch"
1263	assert len(received.events) == num_groups, (
1264	f"Expected {num_groups} BlockStored event(s), got {len(received.events)}"
1265	)
1266
1267	for index, event in enumerate(received.events):
1268	assert isinstance(event, BlockStored), "We should have a BlockStored event"
1269	assert len(event.block_hashes) == num_blocks, (
1270	"We should have a BlockStored event with 2 block_hashes"
1271	)
1272	assert event.block_size == block_size, (
1273	"Block size should be the same as the block size"
1274	)
1275	assert event.parent_block_hash is None, "Parent block hash should be None"
1276	assert event.lora_id is None, "Lora id should be None"
1277	assert event.lora_name is None, "Lora name should be None"
1278	assert len(event.token_ids) == num_blocks * block_size, (
1279	"Token ids should be the same as the custom tokens"
1280	)
1281	assert event.token_ids == custom_tokens, (
1282	"Token ids should be the same as the custom tokens"
1283	)
1284	assert event.group_idx == index
1285	finally:
1286	client.shutdown()
1287	subscriber.close()
1288
1289
1290	@pytest.mark.asyncio
1291	@pytest.mark.parametrize(
1292	"multiprocessing_mode,publisher_config",
1293	[(True, "tcp")],
1294	indirect=["publisher_config"],
1295	)
1296	@multi_gpu_test(num_gpus=4)
1297	async def test_kv_cache_events_dp(
1298	multiprocessing_mode: bool,
1299	publisher_config,
1300	):
1301	block_size = 16
1302	num_blocks = 2
1303	dp_size = 2
1304	tp_size = 2
1305
1306	engine_args = EngineArgs(
1307	model=MODEL_NAME,
1308	enforce_eager=True,
1309	enable_prefix_caching=True,
1310	data_parallel_size=dp_size,
1311	tensor_parallel_size=tp_size,
1312	block_size=block_size,
1313	)
1314	engine_args.kv_events_config = publisher_config
1315
1316	vllm_config = engine_args.create_engine_config(UsageContext.UNKNOWN_CONTEXT)
1317
1318	executor_class = Executor.get_class(vllm_config)
1319	with set_default_torch_num_threads(1):
1320	client = EngineCoreClient.make_client(
1321	multiprocess_mode=multiprocessing_mode,
1322	asyncio_mode=True,
1323	vllm_config=vllm_config,
1324	executor_class=executor_class,
1325	log_stats=False,
1326	)
1327	await asyncio.sleep(1)
1328
1329	# Build endpoints for all DP ranks
1330	base_endpoint = publisher_config.endpoint.replace("*", "127.0.0.1")
1331	endpoints = []
1332	for i in range(dp_size):
1333	offset_endpoint = ZmqEventPublisher.offset_endpoint_port(base_endpoint, i)
1334	endpoints.append(offset_endpoint)
1335
1336	subscriber = MockSubscriber(
1337	endpoints, topic=publisher_config.topic, decode_type=KVEventBatch
1338	)
1339
1340	try:
1341	custom_tokens = list(range(num_blocks * block_size))
1342	sampling_params = SamplingParams(max_tokens=1)
1343	all_request_ids = []
1344
1345	# Create and add 25 requests
1346	# NOTE: attempts to force routing to both dp groups but can be flaky
1347	for i in range(25):
1348	await asyncio.sleep(0.01)
1349	request = make_request(sampling_params, custom_tokens)
1350	await client.add_request_async(request)
1351	all_request_ids.append(request.request_id)
1352
1353	await asyncio.sleep(0.1)
1354
1355	# Initialize outputs dict for all requests
1356	outputs: dict[str, list] = {req_id: [] for req_id in all_request_ids}
1357
1358	print("processing requests...")
1359	await asyncio.wait_for(
1360	loop_until_fully_done_async(client, outputs), timeout=20.0
1361	)
1362
1363	# Receive from subscriber until no more messages
1364	print("collecting results...")
1365	results = []
1366	while True:
1367	result = subscriber.receive_one(timeout=1)
1368	print(result)
1369	if result is None:
1370	break
1371	results.append(result)
1372
1373	# Collect all events and data_parallel_ranks from all results
1374	all_dp_ranks = [received.data_parallel_rank for (_, received) in results]
1375	unique_dps = set(all_dp_ranks)
1376	assert len(unique_dps) == 2, (
1377	f"Expected 2 unique data_parallel_ranks, got {len(unique_dps)}"
1378	)
1379
1380	finally:
1381	client.shutdown()
1382	subscriber.close()
1383
1384
1385	@pytest.mark.timeout(20)
1386	def test_startup_failure(monkeypatch: pytest.MonkeyPatch):
1387	with monkeypatch.context() as m, pytest.raises(Exception) as e_info:
1388	# Monkey-patch to extract core process pid while it's starting.
1389	core_proc_pid = [None]
1390	cepm_ctor = CoreEngineProcManager.__init__
1391
1392	def patched_cepm_ctor(self: CoreEngineProcManager, args, *kwargs):
1393	cepm_ctor(self, args, *kwargs)
1394	core_proc_pid[0] = self.processes[0].pid
1395
1396	m.setattr(CoreEngineProcManager, "__init__", patched_cepm_ctor)
1397
1398	t = time.time()
1399	engine_args = EngineArgs(model=MODEL_NAME)
1400	vllm_config = engine_args.create_engine_config(
1401	usage_context=UsageContext.UNKNOWN_CONTEXT
1402	)
1403	executor_class = Executor.get_class(vllm_config)
1404	print(f"VllmConfig creation took {time.time() - t:.2f} seconds.")
1405
1406	# Start another thread to wait for engine core process to start
1407	# and kill it - simulate fatal uncaught process exit.
1408
1409	def kill_first_child():
1410	while (child_pid := core_proc_pid[0]) is None:
1411	time.sleep(0.5)
1412	print(f"Killing child core process {child_pid}")
1413	assert isinstance(child_pid, int)
1414	os.kill(child_pid, signal.SIGKILL)
1415
1416	Thread(target=kill_first_child, daemon=True).start()
1417
1418	_core_client = EngineCoreClient.make_client(
1419	multiprocess_mode=True,
1420	asyncio_mode=True,
1421	vllm_config=vllm_config,
1422	executor_class=executor_class,
1423	log_stats=True,
1424	)
1425
1426	assert "Engine core initialization failed" in str(e_info.value)
1427
1428
1429	@create_new_process_for_each_test()
1430	def test_engine_core_proc_instantiation_cuda_empty(monkeypatch: pytest.MonkeyPatch):
1431	"""Test that EngineCoreProc can be instantiated when CUDA_VISIBLE_DEVICES
1432	is empty. This ensures the engine frontend does not need access to GPUs.
1433	"""
1434	from vllm.v1.engine.core import EngineCoreProc
1435	from vllm.v1.executor.abstract import Executor
1436
1437	# Create a simple mock executor instead of a complex custom class
1438	mock_executor_class = MagicMock(spec=Executor)
1439
1440	def create_mock_executor(vllm_config):
1441	mock_executor = MagicMock()
1442
1443	# Only implement the methods that are actually called during init
1444	from vllm.v1.kv_cache_interface import FullAttentionSpec
1445
1446	mock_spec = FullAttentionSpec(
1447	block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float16
1448	)
1449
1450	mock_executor.get_kv_cache_specs.return_value = [{"default": mock_spec}]
1451	mock_executor.get_supported_kv_cache_layouts.return_value = [["LBNHC"]]
1452	mock_executor.determine_available_memory.return_value = [1024 * 1024 * 1024]
1453	mock_executor.initialize_from_config.return_value = None
1454	mock_executor.supports_draft_weight_updates.return_value = False
1455
1456	return mock_executor
1457
1458	mock_executor_class.side_effect = create_mock_executor
1459
1460	with monkeypatch.context() as m:
1461	m.setenv("CUDA_VISIBLE_DEVICES", "") # No CUDA devices
1462
1463	from vllm.v1.engine.utils import EngineZmqAddresses
1464
1465	def mock_startup_handshake(
1466	self, handshake_socket, local_client, headless, parallel_config
1467	):
1468	return EngineZmqAddresses(
1469	inputs=["tcp://127.0.0.1:5555"],
1470	outputs=["tcp://127.0.0.1:5556"],
1471	coordinator_input=None,
1472	coordinator_output=None,
1473	)
1474
1475	# Background processes are not important here
1476	m.setattr(EngineCoreProc, "startup_handshake", mock_startup_handshake)
1477
1478	vllm_config = EngineArgs(
1479	model="deepseek-ai/DeepSeek-V2-Lite", trust_remote_code=True
1480	).create_engine_config()
1481	engine_core_proc = EngineCoreProc(
1482	vllm_config=vllm_config,
1483	local_client=True,
1484	handshake_address="tcp://127.0.0.1:12345",
1485	executor_class=mock_executor_class,
1486	log_stats=False,
1487	engine_index=0,
1488	)
1489
1490	engine_core_proc.shutdown()
1491
1492
1493	def _ready_response_with_weight_transfer(weight_transfer_config):
1494	"""Build a ready response against a mock executor."""
1495	from vllm.v1.engine.core import EngineCoreProc
1496
1497	executor = MagicMock()
1498	executor.supports_draft_weight_updates.return_value = True
1499
1500	proc = MagicMock()
1501	proc.model_executor = executor
1502	proc.vllm_config = MagicMock()
1503	proc.vllm_config.weight_transfer_config = weight_transfer_config
1504	proc.vllm_config.lora_config = None
1505	proc.vllm_config.model_config.enable_sleep_mode = False
1506	proc.scheduler = MagicMock()
1507
1508	response = EngineCoreProc._make_ready_response(proc)
1509	return response, executor
1510
1511
1512	def test_ready_response_skips_draft_rpc_without_weight_transfer():
1513	"""No weight-transfer config: the executor must not be consulted."""
1514	response, executor = _ready_response_with_weight_transfer(None)
1515
1516	executor.supports_draft_weight_updates.assert_not_called()
1517	assert response.supports_draft_weight_updates is False
1518
1519
1520	def test_ready_response_queries_executor_with_weight_transfer():
1521	"""Weight transfer configured: the executor is still consulted."""
1522	weight_transfer_config = MagicMock()
1523	weight_transfer_config.backend = "nixl"
1524
1525	response, executor = _ready_response_with_weight_transfer(weight_transfer_config)
1526
1527	executor.supports_draft_weight_updates.assert_called_once()
1528	assert response.supports_draft_weight_updates is True
1529
1530
1531	def test_apply_ready_response_retains_kv_event_sources():
1532	"""Each engine's resolved publisher config is kept by DP rank; engines
1533	without KV events contribute nothing."""
1534	import msgspec
1535
1536	from vllm.config.kv_events import KVEventsConfig
1537
1538	client = object.__new__(MPClient)
1539	client._effective_attention_block_sizes = set()
1540	client._kv_event_sources = {}
1541	client.vllm_config = SimpleNamespace(
1542	cache_config=SimpleNamespace(block_size=16, num_gpu_blocks=0),
1543	model_config=SimpleNamespace(max_model_len=8192),
1544	)
1545
1546	def ready(rank: int, config: KVEventsConfig \| None) -> bytes:
1547	return msgspec.msgpack.encode(
1548	EngineCoreReadyResponse(
1549	max_model_len=8192,
1550	num_gpu_blocks=100,
1551	block_size=16,
1552	dp_stats_address=None,
1553	dtype="bfloat16",
1554	vllm_version="test",
1555	world_size=2,
1556	data_parallel_size=2,
1557	tensor_parallel_size=1,
1558	pipeline_parallel_size=1,
1559	decode_context_parallel_size=1,
1560	data_parallel_rank=rank,
1561	max_num_seqs=256,
1562	max_num_batched_tokens=8192,
1563	instance_id="test-instance",
1564	supports_lora=False,
1565	max_loras=0,
1566	kv_events_config=config,
1567	)
1568	)
1569
1570	rank1 = KVEventsConfig(
1571	enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:41233"
1572	)
1573	client._apply_ready_response(ready(1, rank1))
1574	client._apply_ready_response(ready(0, None))
1575	assert client.get_kv_event_sources() == {1: rank1}
1576
1577	# A rank that comes back (elastic re-add) reports its new port.
1578	rank1_again = KVEventsConfig(
1579	enable_kv_cache_events=True, publisher="zmq", endpoint="tcp://10.0.0.2:47001"
1580	)
1581	client._apply_ready_response(ready(1, rank1_again))
1582	assert client.get_kv_event_sources() == {1: rank1_again}