[ROCm][Bugfix] Initialize Ray NIXL agents for sharded RDT (#56687)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Codex <noreply@openai.com>
Andreas Karatzas committed
Sep 14, 2026 at 19:34 UTC
a7576447b86454f5c5729958d921271392ced47f
5 files changed
+144
-7
tests/distributed/test_sharded_rdt_producer.py
+114
@@ -20,17 +20,23 @@ covered by the GPU tests in `test_sharded_rdt_trainer.py`.
20
21
import contextlib
22
import gc
23
+import os
24
+import sys
25
import threading
26
import time
27
from concurrent.futures import ThreadPoolExecutor
28
from concurrent.futures import TimeoutError as FuturesTimeout
29
+from types import ModuleType, SimpleNamespace
30
+from unittest.mock import Mock
31
32
import pytest
33
import torch
34
35
import vllm.distributed.weight_transfer.sharded_rdt_trainer as trainer_mod
36
+from vllm.distributed.nixl_utils import alias_nixl_for_ray
37
from vllm.distributed.weight_transfer.sharded_rdt_common import (
38
buffer_alloc_bytes,
39
+ register_nixl_memory,
40
)
41
from vllm.distributed.weight_transfer.sharded_rdt_trainer import (
42
DEFAULT_GATHER_LOOKAHEAD,
@@ -45,6 +51,114 @@ GROUP_A = ("model.layers.0.w", "model.layers.0.b")
51
GROUP_B = ("model.layers.1.w",)
52
53
54
+@pytest.fixture
55
+def ray_nixl_transport(monkeypatch):
56
+ import ray
57
+ from ray.experimental.rdt import util
58
+ from ray.experimental.rdt.nixl_tensor_transport import NixlTensorTransport
59
+
60
+ from vllm.distributed import nixl_utils
61
+ from vllm.platforms import current_platform
62
+
63
+ transport = NixlTensorTransport()
64
+ package = ModuleType("nixl_rocm")
65
+ api = ModuleType("nixl_rocm._api")
66
+ api.__dict__.update(nixl_agent=Mock(), nixl_agent_config=Mock())
67
+ package.__dict__["_api"] = api
68
+ modules = {package.__name__: package, api.__name__: api}
69
+ imported_rcache_limits = []
70
+
71
+ def import_nixl(name):
72
+ imported_rcache_limits.append(os.environ.get("UCX_RCACHE_MAX_UNRELEASED"))
73
+ monkeypatch.setitem(sys.modules, name, modules[name])
74
+ return modules[name]
75
+
76
+ nixl_names = ("nixl", "nixl._api", *modules)
77
+ for name in nixl_names:
78
+ monkeypatch.delitem(sys.modules, name, raising=False)
79
+ monkeypatch.delenv("UCX_RCACHE_MAX_UNRELEASED", raising=False)
80
+ importer = Mock(side_effect=import_nixl)
81
+ monkeypatch.setattr(
82
+ nixl_utils, "importlib", SimpleNamespace(import_module=importer)
83
+ )
84
+ monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: True)
85
+ monkeypatch.setattr(util, "transport_managers", {"NIXL": transport})
86
+ monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
87
+ monkeypatch.setattr(
88
+ ray,
89
+ "get_runtime_context",
90
+ lambda: SimpleNamespace(get_actor_id=lambda: "rdt-actor"),
91
+ )
92
+ try:
93
+ yield transport, api, importer, imported_rcache_limits
94
+ finally:
95
+ # The helper inserts aliases directly; remove them before monkeypatch
96
+ # restores any modules and environment values present before this test.
97
+ for name in nixl_names:
98
+ sys.modules.pop(name, None)
99
+ os.environ.pop("UCX_RCACHE_MAX_UNRELEASED", None)
100
+
101
+
102
+def test_ray_registers_memory_through_rocm_alias(ray_nixl_transport):
103
+ """Ray's registration must reach the ROCm agent through the alias."""
104
+ _, api, importer, _ = ray_nixl_transport
105
+ register_nixl_memory(torch.empty(4, dtype=torch.uint8))
106
+ api.nixl_agent.return_value.register_memory.assert_called()
107
+ assert sys.modules["nixl"] is sys.modules["nixl_rocm"]
108
+ assert sys.modules["nixl._api"] is api
109
+
110
+ importer.reset_mock()
111
+ alias_nixl_for_ray()
112
+ importer.assert_not_called()
113
+
114
+
115
+@pytest.mark.parametrize("rcache_limit", [None, "4096"])
116
+def test_ray_nixl_sets_ucx_limit_before_import(
117
+ monkeypatch, ray_nixl_transport, rcache_limit
118
+):
119
+ """UCX must see the default or the user's override during the first import."""
120
+ _, _, _, imported_rcache_limits = ray_nixl_transport
121
+ if rcache_limit is not None:
122
+ monkeypatch.setenv("UCX_RCACHE_MAX_UNRELEASED", rcache_limit)
123
+ alias_nixl_for_ray()
124
+ assert imported_rcache_limits == [rcache_limit or "1024"] * 2
125
+
126
+
127
+@pytest.mark.parametrize("existing_api", [False, True])
128
+def test_ray_nixl_preserves_non_rocm_or_existing_imports(
129
+ monkeypatch, ray_nixl_transport, existing_api
130
+):
131
+ from vllm.platforms import current_platform
132
+
133
+ _, _, importer, _ = ray_nixl_transport
134
+ if existing_api:
135
+ api = ModuleType("nixl._api")
136
+ monkeypatch.setitem(sys.modules, "nixl._api", api)
137
+ else:
138
+ monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
139
+ alias_nixl_for_ray()
140
+ importer.assert_not_called()
141
+ assert "UCX_RCACHE_MAX_UNRELEASED" not in os.environ
142
+ if existing_api:
143
+ assert sys.modules["nixl._api"] is api
144
+ else:
145
+ assert "nixl._api" not in sys.modules
146
+
147
+
148
+def test_ray_nixl_alias_reports_missing_rocm_package(monkeypatch, ray_nixl_transport):
149
+ from vllm.distributed import nixl_utils
150
+
151
+ transport, api, importer, _ = ray_nixl_transport
152
+ monkeypatch.setattr(nixl_utils, "is_nixl_available", lambda: False)
153
+ with pytest.raises(ImportError, match="nixl_rocm"):
154
+ register_nixl_memory(torch.empty(1))
155
+ assert transport._nixl_agent is None
156
+ api.nixl_agent.assert_not_called()
157
+ importer.assert_not_called()
158
+ assert "nixl" not in sys.modules
159
+ assert "nixl._api" not in sys.modules
160
+
161
+
162
@pytest.fixture
163
def server_factory(monkeypatch):
164
"""Build real `_RDTProducerServer` instances without a GPU.
vllm/distributed/nixl_utils.py
+17
@@ -91,9 +91,26 @@ def is_nixl_available() -> bool:
91
return pkg in sys.modules or importlib.util.find_spec(pkg) is not None
92
93
94
+def alias_nixl_for_ray() -> None:
95
+ """Let Ray's ``nixl._api`` import resolve to the ROCm implementation."""
96
+ if not current_platform.is_rocm() or "nixl._api" in sys.modules:
97
+ return
98
+
99
+ package_name = _get_nixl_package_name()
100
+ if not is_nixl_available():
101
+ raise ImportError(f"Ray's NIXL transport on ROCm requires {package_name}")
102
+
103
+ _maybe_set_ucx_rcache_limit()
104
+ package = importlib.import_module(package_name)
105
+ api = importlib.import_module(f"{package_name}._api")
106
+ sys.modules.setdefault("nixl", package)
107
+ sys.modules.setdefault("nixl._api", api)
108
+
109
+
110
__all__ = [
111
"NixlWrapper",
112
"nixl_agent_config",
113
"nixlXferTelemetry",
114
"is_nixl_available",
115
+ "alias_nixl_for_ray",
116
]
vllm/distributed/weight_transfer/sharded_rdt_common.py
+10
@@ -77,6 +77,16 @@ def check_ray_rdt_version() -> None:
77
)
78
79
80
+def register_nixl_memory(tensor: torch.Tensor) -> None:
81
+ """Register a tensor with Ray using the platform's NIXL package."""
82
+ from vllm.distributed.nixl_utils import alias_nixl_for_ray
83
+
84
+ alias_nixl_for_ray()
85
+ from ray.experimental import register_nixl_memory as ray_register_nixl_memory
86
+
87
+ ray_register_nixl_memory(tensor)
88
+
89
+
90
def assign_producer_indices(
91
num_producers: int, num_consumers: int, consumer_idx: int
92
) -> list[int]:
vllm/distributed/weight_transfer/sharded_rdt_engine.py
+2
-3
@@ -66,6 +66,7 @@ from vllm.distributed.weight_transfer.sharded_rdt_common import (
66
RdtRouter,
67
buffer_alloc_bytes,
68
check_ray_rdt_version,
69
+ register_nixl_memory,
70
)
71
from vllm.distributed.weight_transfer.sharded_rdt_fake import (
72
BakeSink,
@@ -588,8 +589,6 @@ class ShardedRDTWeightTransferEngine(
589
plan = self._cached_plan
590
if plan is None or not plan.chunks:
591
return
591
- from ray.experimental import register_nixl_memory
592
-
592
# (a) consumer receive buffers — one per ring slot, at the largest chunk.
593
max_pack = max(c.pack_bytes for c in plan.chunks)
594
alloc = buffer_alloc_bytes(max_pack, self._buffer_presize)
@@ -740,7 +739,7 @@ class ShardedRDTWeightTransferEngine(
739
any time after the metadata push. See the doc's "slot generation
740
handshake".
741
"""
743
- from ray.experimental import register_nixl_memory, set_target_for_ref
742
+ from ray.experimental import set_target_for_ref
743
744
if self._slot_read_done:
745
with self._slot_cv:
vllm/distributed/weight_transfer/sharded_rdt_trainer.py
+1
-4
@@ -47,6 +47,7 @@ from vllm.distributed.weight_transfer.sharded_rdt_common import (
47
ALLOWED_OPS,
48
buffer_alloc_bytes,
49
check_ray_rdt_version,
50
+ register_nixl_memory,
51
)
52
from vllm.logger import init_logger
53
@@ -334,8 +335,6 @@ class _RDTProducerServer:
335
doc). The warmup buffer stays registered so the agent's CUDA-HMEM path
336
stays initialized.
337
"""
337
- from ray.experimental import register_nixl_memory
338
-
338
self._nixl_warmup_buf = torch.zeros(1 << 20, dtype=torch.uint8, device="cuda")
339
with self._reg_lock:
340
register_nixl_memory(self._nixl_warmup_buf)
@@ -510,8 +509,6 @@ class _RDTProducerServer:
509
"""Allocate + NIXL-register one serve slot: the single allocation seam,
510
so registration cannot be skipped on either of the two paths that make
511
buffers (the init-time reservation and the serve-path backstop)."""
513
- from ray.experimental import register_nixl_memory
514
-
512
t = torch.empty(nbytes, dtype=torch.uint8, device=self._serve_device)
513
with self._reg_lock:
514
register_nixl_memory(t)