diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index d62c2434c..5d60d53f2 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -675,7 +675,7 @@ Manages a NIXL agent and RDMA transfers for a single GPU worker: |--------|---------| | `__init__(agent_name, device_id, listen_port, accelerator_backend)` | Create NIXL agent with UCX backend; `listen_port` enables P2P listen thread; `accelerator_backend` owns torch device operations and accelerator capability gates | | `register_tensors(tensors)` | Register GPU tensors for RDMA, return serialized metadata. With `MX_POOL_REG=1` on a backend that supports pool registration, registers each unique cudaMalloc allocation backing the tensors instead of registering each tensor individually | -| `register_arena(arena, tensors)` | Register the used VMM arena range once through dmabuf when the active accelerator backend supports the VMM arena fast path, then publish every tensor descriptor against that single MR | +| `register_arena(arena, tensors)` | Register the used VMM arena range once through dmabuf when the active accelerator backend supports the VMM arena fast path, then publish every tensor descriptor against that single MR. Falls back to per-tensor registration when a tensor lies outside the arena range, or when the arena spans several `cuMemCreate` handles (a single MR cannot be addressed by cuda_ipc then; override with `MX_ARENA_SINGLE_MR=1`) | | `fetch_remote_and_wait(agent_name, ip, port)` | P2P: fetch remote NIXL metadata via listen thread (polls until loaded) | | `receive_from_source(source_metadata, source_tensors, ..., remote_agent_name)` | Execute RDMA read transfer; `remote_agent_name` skips `add_remote_agent` (P2P) | | `shutdown()` | Clean up NIXL agent and resources | @@ -848,7 +848,7 @@ After the load strategy succeeds and the engine has finished post-processing, `L The arena path does not require `MX_POOL_REG=1`. Pool-reg remains the optimization for normal cudaMalloc deployments. Arena deployments use the direct `register_arena` seam because they already know the contiguous VA range and do not need `cuMemGetAddressRange` to rediscover allocation boundaries. -Empirical validation on B200 + ConnectX on 2026-05-14 showed that a dmabuf MR over a multi-handle VMM range remains valid when holes exist inside the registered VA range. That is the property that lets `process_weights_after_loading` allocate replacement tensors, free discarded tensors, and still finish with one MR for the surviving used arena range. +Empirical validation on B200 + ConnectX showed that a dmabuf MR over a multi-handle VMM range remains valid when holes exist inside the registered VA range. That is the property that lets `process_weights_after_loading` allocate replacement tensors, free discarded tensors, and still finish with one MR for the surviving used arena range. That property is specific to the dmabuf/IB path: `cuda_ipc` cannot address a single MR spanning several `cuMemCreate` handles, so multi-allocation arenas fall back to per-tensor registration there. See [Multi-handle arenas](DEPLOYMENT.md#multi-handle-arenas). ## NIXL Integration @@ -953,7 +953,7 @@ See [`metadata.md`](metadata.md) for the full storage schema and debugging guide | `MX_METADATA_BACKEND` | (required on server; `""` on client) | Server: `redis` or `kubernetes`. Client: `""` / `server` / `redis` / `kubernetes` (central server) or `k8s-service` (decentralized via K8s Service routing) | | `MX_POOL_REG` | `0` | Discover cudaMalloc allocations via `cuMemGetAddressRange` and register each as a single NIXL block instead of registering tensors individually. Reduces NIXL registration count by 80-99% on typical vLLM models, cutting `ibv_reg_mr` time and metadata blob size; transfer semantics unchanged. Not required for `MX_VMM_ARENA=1`, which registers the arena directly | | `MX_VMM_ARENA` | `0` | Install a `CUDAPluggableAllocator` that routes weight-loading allocations into a CUDA VMM arena, then registers the used arena range once through dmabuf at end-of-load. Reserves 16.0 TiB of VA by default and commits physical memory only for mapped allocations. See [VMM Arena](#vmm-arena-cudapluggableallocator-hook) | -| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` until the upstream UCX `cuda_copy_md` length-truncation fix ships. | +| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` on any UCX predating the `cuda_copy_md` length-truncation fix (openucx/ucx#11461). Scoped to the `cuda_copy` transport; it does not affect `cuda_ipc`. | | `MX_P2P_METADATA` | `1` | Enable P2P metadata exchange on source workers. Set to `0` to publish full metadata through a central-coordinator backend; ignored on decentralized backends that require P2P metadata | | `MX_METADATA_PORT` | `5555` | Base NIXL listen port; effective port is `MX_METADATA_PORT + device_id` | | `MX_WORKER_GRPC_PORT` | `6555` | Base worker gRPC port for P2P tensor and artifact manifest serving; effective port is `MX_WORKER_GRPC_PORT + device_id` | diff --git a/docs/DEPLOYMENT.md b/docs/DEPLOYMENT.md index b89034ceb..9adf3d58a 100644 --- a/docs/DEPLOYMENT.md +++ b/docs/DEPLOYMENT.md @@ -581,7 +581,8 @@ See [`K8S_SERVICE_BACKEND.md`](K8S_SERVICE_BACKEND.md) for the design rationale, | `MX_METRICS_SCHEME` | `""` | Optional run/scheme label added to every metric, so multiple runs compare on one dashboard. | | `MX_POOL_REG` | `0` | Allocation-level NIXL registration via `cuMemGetAddressRange`. Registers each unique cudaMalloc block instead of each tensor, typically 80-99% fewer registrations, without changing transfer semantics. `MX_VMM_ARENA=1` uses direct arena registration and does not require pool-reg. | | `MX_VMM_ARENA` | `0` | Route weight allocations into a CUDA VMM arena via PyTorch's `CUDAPluggableAllocator`, then register the used arena range as one NIXL MR with dmabuf at end-of-load. Reserves 16.0 TiB of VA by default, with no physical commit until allocations are mapped. Requires the `modelexpress.vmm._alloc_ext` C extension to have built at install time; if it did not, this flag is a no-op with a warning and the loader falls back to the pool-reg path. See [VMM Arena](#vmm-arena-single-mr-registration). | -| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` until the upstream UCX `cuda_copy_md` length-truncation fix ships. | +| `MX_ARENA_SINGLE_MR` | `0` | Keep single-MR arena registration even when the arena spans several `cuMemCreate` handles. Only safe on transports that can register across handles (dmabuf/IB); cuda_ipc cannot, so the default falls back to per-tensor registration. See [VMM Arena](#vmm-arena-single-mr-registration). | +| `UCX_CUDA_COPY_REG_WHOLE_ALLOC` | (UCX default) | Set to `off` with `MX_VMM_ARENA=1` on any UCX predating the `cuda_copy_md` length-truncation fix (openucx/ucx#11461). Scoped to the `cuda_copy` transport; it does not affect `cuda_ipc`. | | `MX_NIXL_BACKEND` | `UCX` | NIXL backend for GPU-to-GPU RDMA. `UCX` (default) for InfiniBand / RoCE. `LIBFABRIC` for AWS EFA — see [NIXL Backend Selection](#nixl-backend-selection). | | `MX_RDMA_NIC_PIN` | (unset) | Per-rank IB NIC pinning. `auto` runs a topology probe; comma-separated NIC list is an explicit override. Workaround for openucx/ucx#11259. | | `MX_RDMA_NIC_PIN_MIN_RATE_GBPS` | (auto, max-rate filter) | Override the auto-detect rate filter with an explicit lower bound (Gb/s). | @@ -662,7 +663,9 @@ arena address. Frees unmap and release that handle, so replacement tensors created during post-processing can return physical memory before the final registration step. At end-of-load, ModelExpress registers the used arena range once through dmabuf and publishes all tensor descriptors against -that single MR. +that single MR — but only when the arena is backed by a single physical +allocation. An arena that spans several `cuMemCreate` handles falls back to +per-tensor registration; see [Multi-handle arenas](#multi-handle-arenas) below. Recommended source-worker setting: @@ -677,13 +680,48 @@ arena registration bypasses the pool-reg path and calls `register_arena` directly. The arena produces one MR for the used range regardless of the pool-reg setting. -Set `UCX_CUDA_COPY_REG_WHOLE_ALLOC=off` until the upstream UCX -`cuda_copy_md` length-truncation fix ships. Without it, UCX can truncate -a multi-handle VMM registration to the first physical handle, and RDMA -operations that cross into later handles fail. See the reproducer and -fix notes in this gist: +Set `UCX_CUDA_COPY_REG_WHOLE_ALLOC=off` on any UCX predating the +`cuda_copy_md` length-truncation fix (openucx/ucx#11461). Without it, UCX +can truncate a multi-handle VMM registration to the first physical +handle, and RDMA operations that cross into later handles fail. See the +reproducer and fix notes in this gist: . +That knob covers the `cuda_copy` transport only. It has no effect on +`cuda_ipc`, which has its own multi-handle limitation described in +[Multi-handle arenas](#multi-handle-arenas) below. + +#### Multi-handle arenas + +A CUDA fabric/IPC handle names exactly one `cuMemCreate` allocation. UCX +cuda_ipc resolves a registered region with `cuMemRetainAllocationHandle` and +`cuMemGetAddressRange`, both of which report the allocation holding the base +pointer rather than the whole reserve. Registering a multi-allocation arena as +one MR therefore publishes an rkey covering only its first chunk, and the peer +reads past what it mapped — measured on GB200 MNNVL as a segfault in +`cuMemcpyDtoDAsync_v2` with an arena spanning 1019 chunks. + +ModelExpress detects this (`live_allocation_count > 1`) and falls back to +per-tensor registration, which is correct because each arena allocation is one +handle, so every tensor lies wholly inside one. The log line names the count: + +```text +register_arena: arena spans 1019 physical allocations; a single MR would publish +an rkey covering only the first, which cuda_ipc cannot address. Falling back to +per-tensor registration ... +``` + +`UCX_CUDA_COPY_REG_WHOLE_ALLOC=off` does not cover this case: it applies to +cuda_copy, while the truncation above is in cuda_ipc, and UCX 1.21 has no +cuda_ipc equivalent. Upstream fixes are in flight for both sides — +[openucx/ucx#11283](https://github.com/openucx/ucx/pull/11283) for cuda_ipc and +[openucx/ucx#11461](https://github.com/openucx/ucx/pull/11461) for the +cuda_copy/dmabuf length truncation. + +Set `MX_ARENA_SINGLE_MR=1` to keep the single-MR path on deployments where it +was validated, i.e. dmabuf/IB, where `ibv_reg_dmabuf_mr` does span several +handles. + ### P2P Metadata Exchange P2P metadata exchange is enabled by default. Source workers expose their own per-worker gRPC `WorkerService` (the `WorkerGrpcServer` on `MX_WORKER_GRPC_PORT`) and their NIXL agent metadata directly on the worker's NIXL listen thread (`MX_METADATA_PORT`). Targets fetch tensor manifests or artifact manifests directly from the source worker rather than pulling them through the central store. For file-backed cache artifacts, targets also call `PrepareArtifactChunk` and `ReleaseArtifactChunk` on this worker service while bytes move through NIXL into target-local staging, then install the staged artifact into the runtime cache directory. The division of responsibility depends on which metadata backend is in use: diff --git a/modelexpress_client/python/modelexpress/envs.py b/modelexpress_client/python/modelexpress/envs.py index 56d4fc589..63016243a 100644 --- a/modelexpress_client/python/modelexpress/envs.py +++ b/modelexpress_client/python/modelexpress/envs.py @@ -94,6 +94,7 @@ MX_TRANSFER_LOG_DIR: str # VMM arena MX_VMM_ARENA: bool + MX_ARENA_SINGLE_MR: bool # Framework artifact (JIT cache) transfer MX_ARTIFACT_TRANSFER: bool MX_ARTIFACT_BUNDLE_ROOT: Optional[str] @@ -301,6 +302,7 @@ def _env_positive_float(name: str, default: float) -> float: "MX_TRANSFER_LOG_DIR": lambda: os.environ.get("MX_TRANSFER_LOG_DIR", "/tmp/mx_logs"), # ── VMM arena ────────────────────────────────────────────────────────── "MX_VMM_ARENA": lambda: os.environ.get("MX_VMM_ARENA") == "1", + "MX_ARENA_SINGLE_MR": lambda: os.environ.get("MX_ARENA_SINGLE_MR") == "1", # ── Framework artifact (JIT cache) transfer ──────────────────────────── "MX_ARTIFACT_TRANSFER": lambda: os.environ.get("MX_ARTIFACT_TRANSFER", "").strip().lower() in _TRUTHY, diff --git a/modelexpress_client/python/modelexpress/nixl_transfer.py b/modelexpress_client/python/modelexpress/nixl_transfer.py index 83c50b46c..bde3e280b 100644 --- a/modelexpress_client/python/modelexpress/nixl_transfer.py +++ b/modelexpress_client/python/modelexpress/nixl_transfer.py @@ -71,6 +71,16 @@ def _resolve_nixl_backend() -> str: return raw +def _arena_single_mr_forced() -> bool: + """Whether to keep the single-MR arena path on a multi-allocation arena. + + MX_ARENA_SINGLE_MR=1 forces it. Only safe on transports that can span + several cuMemCreate handles in one registration (dmabuf/IB); cuda_ipc + cannot. Read at call time so tests can toggle the env var. + """ + return envs.MX_ARENA_SINGLE_MR + + def _pool_reg_enabled() -> bool: """Whether allocation-level pool registration is enabled. @@ -393,12 +403,19 @@ def register_arena( consumes a dmabuf via `ibv_reg_dmabuf_mr` and produces ONE lkey/rkey covering all live tensors. - Empirically validated on Blackwell + ConnectX over InfiniBand - against a CUDA VMM range with multiple cuMemCreate handles and - mid-range holes (chunks unmapped + released after the export): - registration succeeds, the dmabuf attach pins the currently- - mapped physical pages, and the HCA translation table survives - subsequent CUDA-side unmaps. + The multi-handle case is validated on the dmabuf/IB path only. + On Blackwell + ConnectX over InfiniBand, against a CUDA VMM range + with multiple cuMemCreate handles and mid-range holes (chunks + unmapped + released after the export): registration succeeds, the + dmabuf attach pins the currently-mapped physical pages, and the + HCA translation table survives subsequent CUDA-side unmaps. + + It does NOT hold on UCX cuda_ipc, where a fabric handle names one + cuMemCreate allocation and a single MR would publish an rkey + covering only the first chunk. That is why this method falls back + to per-tensor registration when the arena spans several + allocations, unless MX_ARENA_SINGLE_MR overrides it. Upstream fix: + openucx/ucx#11283. Per-tensor descriptors are still built (tensor name -> addr, size, dtype) because the receiver matches by name and computes @@ -431,6 +448,35 @@ def register_arena( ) return self.register_tensors(tensors) + # A CUDA fabric/IPC handle names exactly one cuMemCreate allocation: + # UCX cuda_ipc resolves a region with cuMemRetainAllocationHandle and + # cuMemGetAddressRange, which report the FIRST allocation under the + # range rather than the whole reserve. Registering a multi-allocation + # arena as one MR therefore publishes an rkey covering only its first + # chunk, and the peer's cuMemcpyDtoDAsync_v2 reads past what it mapped. + # Measured on GB200 MNNVL: Kimi-K3 arena, 1019 chunks, segfault in + # uct_cuda_ipc_ep_get_zcopy. Per-tensor registration is correct because + # the arena does one cuMemCreate per allocation, so every tensor lies + # wholly inside one handle. + # + # dmabuf/IB registration does span several handles, so deployments that + # validated the single-MR path there can keep it with + # MX_ARENA_SINGLE_MR=1. + live_allocs = arena.live_allocation_count + if live_allocs > 1 and not _arena_single_mr_forced(): + logger.warning( + "register_arena: arena spans %d physical allocations; a single " + "MR would publish an rkey covering only the first, which " + "cuda_ipc cannot address. Falling back to per-tensor " + "registration for %d tensors over [0x%x, 0x%x). Set " + "MX_ARENA_SINGLE_MR=1 to force single-MR (dmabuf/IB only).", + live_allocs, + len(tensor_descriptors), + base, + base + used, + ) + return self.register_tensors(tensors, force_per_tensor=True) + # NIXL resolves descriptors by containment, so one tensor outside # [base, base+used) fails prep_xfer_dlist for the whole transfer. uncovered = [ diff --git a/modelexpress_client/python/modelexpress/vmm/README.md b/modelexpress_client/python/modelexpress/vmm/README.md index b95305fa5..22402e269 100644 --- a/modelexpress_client/python/modelexpress/vmm/README.md +++ b/modelexpress_client/python/modelexpress/vmm/README.md @@ -23,18 +23,37 @@ MX_VMM_ARENA=1 That is the only knob. The arena enables itself on supported devices and falls back with a warning when the underlying C extension is missing. -### Required deployment flag +### Deployment flag for the cuda_copy path UCX's `cuda_copy_md` path probes allocations with `cuMemGetAddressRange`, which returns per-handle bounds rather than the full reserve when called against a multi-handle VMM range. Without the override, UCX truncates -transfers to a single physical chunk. Until the upstream fix lands -(openucx/ucx#11461), set: +transfers to a single physical chunk. On any UCX predating the upstream +fix (openucx/ucx#11461), set: ```bash UCX_CUDA_COPY_REG_WHOLE_ALLOC=off ``` +This knob is scoped to the `cuda_copy` transport and does not affect +`cuda_ipc`. Multi-allocation arenas have a separate limitation on +`cuda_ipc`; see Transport support below before enabling the arena on an +NVLink or MNNVL deployment. + +### Transport support + +Single-MR arena registration is validated on the dmabuf/IB path, where +`ibv_reg_dmabuf_mr` genuinely spans several `cuMemCreate` handles. + +It does not hold on `cuda_ipc`. A CUDA fabric/IPC handle names exactly +one `cuMemCreate` allocation, so a single MR over a multi-allocation +arena publishes an rkey covering only the first chunk and the peer reads +past what it mapped. ModelExpress detects this and falls back to +per-tensor registration. `MX_ARENA_SINGLE_MR=1` overrides that fallback +and is only safe on dmabuf/IB. The upstream fix is openucx/ucx#11283. +The mechanism and the measured failure are in the Multi-handle arenas +section of `docs/DEPLOYMENT.md`. + ## Why this exists Without an arena, NIXL registers each tensor as its own MR. A @@ -148,12 +167,18 @@ Key invariants: Validated configurations: -- Blackwell B200 + ConnectX over InfiniBand, single-pod and P2P. +- Blackwell B200 + ConnectX over InfiniBand, single-pod and P2P. This is + the only transport on which single-MR registration is validated. - vLLM 0.x with `CUDAPluggableAllocator` (PyTorch 2.4+). - TP=8 ranks running concurrently against independent arenas. Known limitations: +- Single-MR registration does not work on the UCX `cuda_ipc` transport + when the arena spans more than one `cuMemCreate` handle. The published + rkey covers only the first chunk. ModelExpress falls back to per-tensor + registration in that case; see Transport support above. Upstream fix: + openucx/ucx#11283. - VMM arena is x86_64-Linux only (CUDA driver constraint). - Bump pointer is monotonic. Long-lived processes that load and unload many models within one arena lifetime will fragment the VA but never diff --git a/modelexpress_client/python/modelexpress/vmm/runtime.py b/modelexpress_client/python/modelexpress/vmm/runtime.py index 1a4467fdc..d52db7cec 100644 --- a/modelexpress_client/python/modelexpress/vmm/runtime.py +++ b/modelexpress_client/python/modelexpress/vmm/runtime.py @@ -244,10 +244,12 @@ def log_arena_post_load(ctx: "LoadContext") -> None: ``ibv_reg_dmabuf_mr`` over ``[base, base+used_bytes)`` already ran inside ``LoadStrategyChain`` via ``NixlTransferManager.register_arena``; this hook is purely - diagnostic. Empirically validated on Blackwell + ConnectX over - InfiniBand: the registration succeeds over a VA range with + diagnostic. Validated on the dmabuf/IB path, Blackwell + ConnectX + over InfiniBand: the registration succeeds over a VA range with mid-range holes from prior ``cuMemUnmap`` calls, and the dmabuf pin - keeps live tensor pages addressable to the HCA. + keeps live tensor pages addressable to the HCA. On ``cuda_ipc`` a + multi-allocation arena cannot be covered by one MR at all, and + ``register_arena`` falls back to per-tensor registration. """ arena = _vmm_arenas.get(ctx.device_id) if arena is None: diff --git a/modelexpress_client/python/tests/test_pool_registration.py b/modelexpress_client/python/tests/test_pool_registration.py index 46aa17cda..54e0e7407 100644 --- a/modelexpress_client/python/tests/test_pool_registration.py +++ b/modelexpress_client/python/tests/test_pool_registration.py @@ -14,6 +14,7 @@ from modelexpress.accelerators import NIXL_ACCELERATOR_MEM_TYPE from modelexpress.nixl_transfer import ( + _arena_single_mr_forced, NixlTransferManager, _pool_reg_enabled, ) @@ -222,6 +223,8 @@ def test_arena_registration_uses_vram_segment(self): used = tensor.numel() * tensor.element_size() class FakeArena: + live_allocation_count = 1 + def registered_range(self): return base, used @@ -237,6 +240,8 @@ def registered_range(self): def test_arena_registration_falls_back_when_tensor_uncovered(self): # A tensor outside [base, base+used) must not be served by the single MR. class FakeArena: + live_allocation_count = 1 + def registered_range(self): return 0x1000, 0x2000 @@ -254,6 +259,8 @@ def test_arena_fallback_bypasses_pool_registration(self, monkeypatch): monkeypatch.setenv("MX_POOL_REG", "1") class FakeArena: + live_allocation_count = 1 + def registered_range(self): return 0x1000, 0x2000 @@ -520,3 +527,119 @@ def test_raises_labeled_error(self): manager = self._manager(["ERROR"]) with pytest.raises(RuntimeError, match="test transfer failed"): manager._wait_for_xfer(object(), None, "test transfer") + + +class TestArenaSingleMrForced: + """MX_ARENA_SINGLE_MR keeps the single-MR path on a multi-allocation arena.""" + + def test_default_is_off(self, monkeypatch): + monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False) + assert _arena_single_mr_forced() is False + + def test_one_is_on(self, monkeypatch): + monkeypatch.setenv("MX_ARENA_SINGLE_MR", "1") + assert _arena_single_mr_forced() is True + + def test_read_at_call_time(self, monkeypatch): + monkeypatch.setenv("MX_ARENA_SINGLE_MR", "1") + assert _arena_single_mr_forced() is True + monkeypatch.setenv("MX_ARENA_SINGLE_MR", "0") + assert _arena_single_mr_forced() is False + + +class TestMultiAllocationArena: + """A single MR cannot describe an arena spanning several cuMemCreate handles. + + UCX cuda_ipc resolves a region with cuMemGetAddressRange, which reports only + the allocation holding the base pointer, so the published rkey would cover + just the first chunk and the peer would read past what it mapped. + """ + + @staticmethod + def _make_manager() -> NixlTransferManager: + mgr = NixlTransferManager(agent_name="test", device_id=0) + mgr._agent = MagicMock() + mgr._agent.get_agent_metadata.return_value = b"metadata" + return mgr + + @staticmethod + def _covering_arena(tensor, live_allocs): + base = tensor.data_ptr() + used = tensor.numel() * tensor.element_size() + + class FakeArena: + live_allocation_count = live_allocs + + def registered_range(self): + return base, used + + return FakeArena(), base, used + + def test_multi_allocation_arena_registers_per_tensor(self, monkeypatch): + monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False) + tensor = torch.zeros(4, dtype=torch.float32) + arena, _, _ = self._covering_arena(tensor, live_allocs=1019) + + mgr = self._make_manager() + assert mgr.register_arena(arena, {"w": tensor}) == b"metadata" + + # Per-tensor, not one MR over the arena range. + args, kwargs = mgr._agent.register_memory.call_args + assert args[0][0] is tensor + assert kwargs == {"backends": ["UCX"]} + + def test_single_allocation_arena_keeps_one_mr(self, monkeypatch): + monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False) + tensor = torch.zeros(4, dtype=torch.float32) + arena, base, used = self._covering_arena(tensor, live_allocs=1) + + mgr = self._make_manager() + assert mgr.register_arena(arena, {"w": tensor}) == b"metadata" + + mgr._agent.register_memory.assert_called_once_with( + [(base, used, 0, "")], + mem_type=NIXL_ACCELERATOR_MEM_TYPE, + backends=["UCX"], + ) + + def test_env_override_keeps_one_mr_on_multi_allocation(self, monkeypatch): + # dmabuf/IB can span several handles in one registration, so deployments + # that validated the single-MR path there can keep it. + monkeypatch.setenv("MX_ARENA_SINGLE_MR", "1") + tensor = torch.zeros(4, dtype=torch.float32) + arena, base, used = self._covering_arena(tensor, live_allocs=1019) + + mgr = self._make_manager() + assert mgr.register_arena(arena, {"w": tensor}) == b"metadata" + + mgr._agent.register_memory.assert_called_once_with( + [(base, used, 0, "")], + mem_type=NIXL_ACCELERATOR_MEM_TYPE, + backends=["UCX"], + ) + + def test_multi_allocation_fallback_bypasses_pool_registration(self, monkeypatch): + # Pool reg resolves the same per-handle bounds that were insufficient. + monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False) + monkeypatch.setenv("MX_POOL_REG", "1") + tensor = torch.zeros(4, dtype=torch.float32) + arena, _, _ = self._covering_arena(tensor, live_allocs=1019) + + mgr = self._make_manager() + assert mgr.register_arena(arena, {"w": tensor}) == b"metadata" + + args, kwargs = mgr._agent.register_memory.call_args + assert args[0][0] is tensor + assert kwargs == {"backends": ["UCX"]} + + def test_warning_names_the_allocation_count(self, monkeypatch, caplog): + monkeypatch.delenv("MX_ARENA_SINGLE_MR", raising=False) + tensor = torch.zeros(4, dtype=torch.float32) + arena, _, _ = self._covering_arena(tensor, live_allocs=1019) + + mgr = self._make_manager() + with caplog.at_level(logging.WARNING): + mgr.register_arena(arena, {"w": tensor}) + + assert "1019 physical allocations" in caplog.text + assert "MX_ARENA_SINGLE_MR=1" in caplog.text