A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 23:45 UTC
[P/D] Move FakeNixlWrapper to test dir (#21328)
Signed-off-by: Rui Qiao <ruisearch42@gmail.com>
Rui Qiao committed
Jul 24, 2025 at 08:53 UTC
1e9ea8e69d6da725e3fb5f287360c986a7901d09
4 files changed
+140
-115
tests/v1/kv_connector/unit/test_nixl_connector.py
+136
-33
@@ -1,10 +1,15 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import contextlib
5
+import inspect
6
import os
7
import tempfile
8
import textwrap
9
import time
10
+import uuid
11
+from collections import defaultdict
12
+from typing import Optional
13
from unittest.mock import patch
14
15
import pytest
@@ -16,30 +21,118 @@ from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
21
KVConnectorRole, NixlAgentMetadata, NixlConnector, NixlConnectorMetadata,
22
NixlConnectorWorker)
23
from vllm.forward_context import ForwardContext
19
-from vllm.mocks.mock_nixl_connector import FakeNixlWrapper
24
from vllm.sampling_params import SamplingParams
25
26
from .utils import create_request, create_scheduler, create_vllm_config
27
28
25
-def _make_stub_pkg() -> str:
26
- """Return a directory that makes
27
- `from nixl._api import nixl_agent` resolve to our FakeNixlWrapper."""
28
- td = tempfile.mkdtemp()
29
- pkg_root = os.path.join(td, "nixl", "_api")
30
- os.makedirs(pkg_root, exist_ok=True)
29
+class FakeNixlWrapper:
30
+ """Mock implementation of NixlWrapper for testing.
31
32
- stub = textwrap.dedent("""\
33
- # Forward the real FakeNixlWrapper that the driver already defined.
34
- print("In fake package")
35
- from vllm.mocks.mock_nixl_connector import FakeNixlWrapper as nixl_agent
36
- """)
37
- with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
38
- f.write(stub)
32
+ We don't inherit from nixl._api.nixl_agent because nixl may not be
33
+ installed.
34
+
35
+ Note: The complete source of this class is also used in the
36
+ `_make_fake_nixl_pkg` function to create a fake nixl package
37
+ for Ray workers.
38
+ """
39
+
40
+ AGENT_METADATA = b"fake_agent_metadata"
41
+ REMOTE_AGENT_NAME = "remote_agent"
42
+
43
+ def __init__(self, agent_name: str, *args, **kwargs):
44
+ self._cycles_before_xfer_done = 0
45
+ self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(
46
+ lambda: 0)
47
+
48
+ def get_reg_descs(self, caches_data, memory_type: str) -> list:
49
+ return [str(uuid.uuid4()) for _ in caches_data]
50
+
51
+ def register_memory(self, descs) -> None:
52
+ pass
53
+
54
+ def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
55
+ return [str(uuid.uuid4()) for _ in blocks_data]
56
+
57
+ def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
58
+ return uuid.uuid4().int
59
+
60
+ def get_agent_metadata(self) -> bytes:
61
+ return self.AGENT_METADATA
62
+
63
+ def add_remote_agent(self, agent_metadata: bytes) -> str:
64
+ return self.REMOTE_AGENT_NAME
65
+
66
+ def get_new_notifs(self) -> dict[str, list[bytes]]:
67
+ # Used to collect done_sending, which we don't test yet.
68
+ return {}
69
+
70
+ def check_xfer_state(self, handle: int) -> str:
71
+ if self._check_xfer_state_cycles[
72
+ handle] >= self._cycles_before_xfer_done:
73
+ return "DONE"
74
+ self._check_xfer_state_cycles[handle] += 1
75
+ return "PROC"
76
+
77
+ def release_xfer_handle(self, handle: int) -> None:
78
+ pass
79
+
80
+ def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
81
+ pass
82
+
83
+ def make_prepped_xfer(self,
84
+ xfer_type: str,
85
+ local_xfer_side_handle: int,
86
+ local_block_descs_ids: list[int],
87
+ remote_xfer_side_handle: int,
88
+ remote_block_descs_ids: list[int],
89
+ notif_msg: Optional[bytes] = None) -> int:
90
+ return uuid.uuid4().int
91
40
- # touch parent package
41
- open(os.path.join(td, "nixl", "__init__.py"), "w").close()
42
- return td
92
+ def transfer(self, handle: int) -> str:
93
+ return "PROC"
94
+
95
+ ############################################################
96
+ # Follow are for changing the behavior during testing.
97
+ ############################################################
98
+
99
+ def set_cycles_before_xfer_done(self, cycles: int):
100
+ """Set the number of cycles before a transfer is considered done."""
101
+ self._cycles_before_xfer_done = cycles
102
+
103
+
104
+@contextlib.contextmanager
105
+def _make_fake_nixl_pkg():
106
+ """Context manager that creates a temporary package making
107
+ `from nixl._api import nixl_agent` resolve to our FakeNixlWrapper.
108
+
109
+ Automatically cleans up the temporary directory when done.
110
+ """
111
+ with tempfile.TemporaryDirectory() as td:
112
+ pkg_root = os.path.join(td, "nixl", "_api")
113
+ os.makedirs(pkg_root, exist_ok=True)
114
+
115
+ # Get the source code of FakeNixlWrapper class and dedent it
116
+ fake_nixl_source = inspect.getsource(FakeNixlWrapper)
117
+ fake_nixl_source = textwrap.dedent(fake_nixl_source)
118
+
119
+ stub = f"""\
120
+# Copy of FakeNixlWrapper implementation for Ray workers
121
+import uuid
122
+from collections import defaultdict
123
+from typing import Optional
124
+
125
+{fake_nixl_source}
126
+
127
+# Export as nixl_agent
128
+nixl_agent = FakeNixlWrapper
129
+"""
130
+ with open(os.path.join(pkg_root, "__init__.py"), "w") as f:
131
+ f.write(stub)
132
+
133
+ # touch parent package
134
+ open(os.path.join(td, "nixl", "__init__.py"), "w").close()
135
+ yield td
136
137
138
def test_basic_interface():
@@ -351,27 +444,37 @@ def test_abort_timeout_on_prefiller(monkeypatch, distributed_executor_backend):
444
kv_connector="NixlConnector",
445
kv_role="kv_both",
446
)
447
+ llm_kwargs = {
448
+ "model": model_name,
449
+ "enforce_eager": True,
450
+ "gpu_memory_utilization": 0.5,
451
+ "kv_transfer_config": kv_transfer_config,
452
+ "distributed_executor_backend": distributed_executor_backend,
453
+ }
454
+
455
timeout = 6
456
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
457
monkeypatch.setenv("VLLM_NIXL_ABORT_REQUEST_TIMEOUT", str(timeout))
458
358
- # Build runtime_env only if we’re using Ray
459
+ # Build runtime_env only if we're using Ray
460
if distributed_executor_backend == "ray":
360
- runtime_env = {
361
- "working_dir": _make_stub_pkg(), # ship stub package
362
- "env_vars": {
363
- "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),
364
- },
365
- }
366
- ray.init(runtime_env=runtime_env)
367
-
368
- llm = LLM(
369
- model=model_name,
370
- enforce_eager=True,
371
- gpu_memory_utilization=0.5,
372
- kv_transfer_config=kv_transfer_config,
373
- distributed_executor_backend=distributed_executor_backend,
374
- )
461
+ with _make_fake_nixl_pkg() as working_dir:
462
+ runtime_env = {
463
+ "working_dir": working_dir, # ship fake nixl package
464
+ "env_vars": {
465
+ "VLLM_NIXL_ABORT_REQUEST_TIMEOUT": str(timeout),
466
+ },
467
+ }
468
+ ray.init(runtime_env=runtime_env)
469
+
470
+ _run_abort_timeout_test(llm_kwargs, timeout)
471
+ else:
472
+ _run_abort_timeout_test(llm_kwargs, timeout)
473
+
474
+
475
+def _run_abort_timeout_test(llm_kwargs: dict, timeout: int):
476
+ """Helper function to run the abort timeout test logic."""
477
+ llm = LLM(**llm_kwargs)
478
remote_prefill_opts = {
479
"do_remote_decode": True,
480
"do_remote_prefill": False,
vllm/distributed/kv_transfer/kv_connector/utils.py
+4
-6
@@ -120,8 +120,8 @@ class KVOutputAggregator:
120
output corresponding to Rank 0 for scheduler."""
121
122
def __init__(self, world_size: int):
123
- # Complete transfer tracker. Used by to track finished requests
124
- # [req_id -> n_finished_workers]
123
+ # Complete transfer tracker. Used to track finished requests
124
+ # [req_id -> n_remaining_workers]
125
self._recv_remaining_count = defaultdict[str, int](lambda: world_size)
126
self._send_remaining_count = defaultdict[str, int](lambda: world_size)
127
@@ -134,12 +134,10 @@ class KVOutputAggregator:
134
remaining_count_dict: dict[str, int],
135
finished_set: set[str]) -> None:
136
for req_id in req_ids or ():
137
- new_count = remaining_count_dict[req_id] - 1
138
- if new_count == 0:
137
+ remaining_count_dict[req_id] -= 1
138
+ if remaining_count_dict[req_id] == 0:
139
finished_set.add(req_id)
140
del remaining_count_dict[req_id]
141
- else:
142
- remaining_count_dict[req_id] = new_count
141
142
finished_sending = set[str]()
143
finished_recving = set[str]()
vllm/mocks/__init__.py
vllm/mocks/mock_nixl_connector.py
deleted
-76
@@ -1,76 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import uuid
4
-from collections import defaultdict
5
-from typing import Optional
6
-
7
-
8
-class FakeNixlWrapper:
9
- """Mock implementation of NixlWrapper for testing.
10
-
11
- We don't inherit from nixl._api.nixl_agent because nixl may not be
12
- installed.
13
- """
14
-
15
- AGENT_METADATA = b"fake_agent_metadata"
16
- REMOTE_AGENT_NAME = "remote_agent"
17
-
18
- def __init__(self, agent_name: str, *args, **kwargs):
19
- self._cycles_before_xfer_done = 0
20
- self._check_xfer_state_cycles: defaultdict[int, int] = defaultdict(
21
- lambda: 0)
22
-
23
- def get_reg_descs(self, caches_data, memory_type: str) -> list:
24
- return [str(uuid.uuid4()) for _ in caches_data]
25
-
26
- def register_memory(self, descs) -> None:
27
- pass
28
-
29
- def get_xfer_descs(self, blocks_data, memory_type: str) -> list:
30
- return [str(uuid.uuid4()) for _ in blocks_data]
31
-
32
- def prep_xfer_dlist(self, agent_name: str, descs: list) -> int:
33
- return uuid.uuid4().int
34
-
35
- def get_agent_metadata(self) -> bytes:
36
- return self.AGENT_METADATA
37
-
38
- def add_remote_agent(self, agent_metadata: bytes) -> str:
39
- return self.REMOTE_AGENT_NAME
40
-
41
- def get_new_notifs(self) -> dict[str, list[bytes]]:
42
- # Used to collect done_sending, which we don't test yet.
43
- return {}
44
-
45
- def check_xfer_state(self, handle: int) -> str:
46
- if self._check_xfer_state_cycles[
47
- handle] >= self._cycles_before_xfer_done:
48
- return "DONE"
49
- self._check_xfer_state_cycles[handle] += 1
50
- return "PROC"
51
-
52
- def release_xfer_handle(self, handle: int) -> None:
53
- pass
54
-
55
- def send_notif(self, agent_name: str, notif_msg: bytes) -> None:
56
- pass
57
-
58
- def make_prepped_xfer(self,
59
- xfer_type: str,
60
- local_xfer_side_handle: int,
61
- local_block_descs_ids: list[int],
62
- remote_xfer_side_handle: int,
63
- remote_block_descs_ids: list[int],
64
- notif_msg: Optional[bytes] = None) -> int:
65
- return uuid.uuid4().int
66
-
67
- def transfer(self, handle: int) -> str:
68
- return "PROC"
69
-
70
- ############################################################
71
- # Follow are for changing the behavior during testing.
72
- ############################################################
73
-
74
- def set_cycles_before_xfer_done(self, cycles: int):
75
- """Set the number of cycles before a transfer is considered done."""
76
- self._cycles_before_xfer_done = cycles