diff --git a/.github/actions/test-template/action.yml b/.github/actions/test-template/action.yml index f4b1062f7a5..299d9ad43d6 100644 --- a/.github/actions/test-template/action.yml +++ b/.github/actions/test-template/action.yml @@ -142,10 +142,41 @@ runs: echo -e "\033[1;34m│ script : ${{ inputs.script }}\033[0m" echo -e "\033[1;34m│ runner : ${{ inputs.runner }}\033[0m" echo -e "\033[1;34m│ container : ${{ inputs.registry }}/${{ inputs.image }}:${{ inputs.image-tag || github.run_id }}\033[0m" + # Expose the host's RDMA devices so mooncake_cpu tests exercise the real + # transport instead of finding no verbs device and skipping. A container + # sees the host's /sys/class/infiniband but does NOT inherit + # /dev/infiniband/uverbs*, which is what libibverbs actually opens; + # ibv_reg_mr additionally needs IPC_LOCK and an unlimited memlock rlimit + # to pin its registrations. + # + # Requires an mlx5 device on the host, either fabric — rdma_devices() + # prefers InfiniBand and falls back to RoCE. When the host has one, + # NEMO_RL_REQUIRE_MOONCAKE promotes the tests' skip into a failure, so + # losing this passthrough surfaces as a red job rather than silently + # dropping mooncake coverage. + # + # --device against a missing path makes docker run itself fail, which + # would break every test rather than only the RDMA ones — hence the + # host-side check (see detect_rdma.sh for what it gates on). + source "${{ github.workspace }}/${{ github.run_id }}/${{ steps.uuid.outputs.id }}/nemo-rl/tests/scripts/detect_rdma.sh" + RDMA_FLAGS=() + # MC_ENABLE_DEST_DEVICE_AFFINITY is normally decided by + # configure_engine_env in nemo_rl/data_plane/adapters/transfer_queue_env.py, + # which sets it only on a RoCE-only fabric. Passing it here overrides that + # gate (an already-set value is left alone) and also covers processes that + # never run the data-plane factory, e.g. a bare pytest invocation. + if [[ "${{ inputs.is_doc_test }}" != "true" ]] && rdma_device_available; then + RDMA_FLAGS=(--device=/dev/infiniband --cap-add=IPC_LOCK --ulimit memlock=-1 + --env NEMO_RL_REQUIRE_MOONCAKE=1 + --env MC_ENABLE_DEST_DEVICE_AFFINITY=1) + fi + echo -e "\033[1;34m│ rdma : ${RDMA_FLAGS[*]:-none}\033[0m" echo -e "\033[1;34m└──────────────────────────────────────────────────────────────────────────┘\033[0m" + echo "::group::Logs" docker run --rm -u root --runtime=nvidia --gpus all \ --shm-size=64g \ + "${RDMA_FLAGS[@]}" \ --env TRANSFORMERS_OFFLINE=0 \ --env GHA_RUNNER=${{ inputs.runner }} \ --env HYDRA_FULL_ERROR=1 \ diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index afaf8a38af3..dabb1c50dc8 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -534,11 +534,19 @@ data_plane: enabled: false impl: transfer_queue backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') - storage_capacity: 1000000 # max samples retained per partition - num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" - local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" + # Backend-specific blocks: only the one named by `backend` is read, and an + # absent block means that backend's defaults (see SimpleStorageConfig / + # MooncakeCpuConfig in nemo_rl/data_plane/interfaces.py). + simple: + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: ${mul:2, ${cluster.num_nodes}} # TQ wants >= 2 per node + mooncake_cpu: + # Per client process — see MooncakeCpuConfig before raising these. + global_segment_size: 68719476736 # 64 GiB/process + local_buffer_size: 4294967296 # 4 GiB/process + reuse_registered_buffers: true # reuse RDMA-registered buffers + staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool # observability: # NotRequired # enabled: false diff --git a/examples/nemo_gym/run_grpo_nemo_gym.py b/examples/nemo_gym/run_grpo_nemo_gym.py index 7210879a725..02cb621df2c 100644 --- a/examples/nemo_gym/run_grpo_nemo_gym.py +++ b/examples/nemo_gym/run_grpo_nemo_gym.py @@ -40,6 +40,7 @@ ) from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_response_data +from nemo_rl.data_plane.factory import maybe_configure_data_plane_env from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.environments.nemo_gym import ( setup_nemo_gym_config, @@ -224,6 +225,8 @@ def main() -> None: pprint.pprint(config) with rl_init_timer.time("ray_connect"): + # Must precede init_ray() — see maybe_configure_data_plane_env's docstring. + maybe_configure_data_plane_env(config.data_plane) init_ray() # `is_trajectory_collection` is a NeMo-RL-side control-flow knob; pop it diff --git a/examples/run_grpo.py b/examples/run_grpo.py index 51b1f08698e..c6f06943b75 100644 --- a/examples/run_grpo.py +++ b/examples/run_grpo.py @@ -27,6 +27,7 @@ ) from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_response_data +from nemo_rl.data_plane.factory import maybe_configure_data_plane_env from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.models.generation import configure_generation_config from nemo_rl.utils.config import ( @@ -108,6 +109,8 @@ def main() -> None: ) with rl_init_timer.time("ray_connect"): + # Must precede init_ray() — see maybe_configure_data_plane_env's docstring. + maybe_configure_data_plane_env(config.data_plane) init_ray() # setup tokenizer diff --git a/examples/run_grpo_single_controller.py b/examples/run_grpo_single_controller.py index fef3b2294c1..8a440b4832d 100644 --- a/examples/run_grpo_single_controller.py +++ b/examples/run_grpo_single_controller.py @@ -36,6 +36,7 @@ setup_single_controller, ) from nemo_rl.algorithms.utils import get_tokenizer +from nemo_rl.data_plane.factory import maybe_configure_data_plane_env from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.environments.nemo_gym import setup_nemo_gym_config from nemo_rl.models.generation import configure_generation_config @@ -115,6 +116,8 @@ def main() -> None: f"📊 Using checkpoint directory: {config.checkpointing['checkpoint_dir']}" ) + # Must precede init_ray() — see maybe_configure_data_plane_env's docstring. + maybe_configure_data_plane_env(config.data_plane) init_ray() tokenizer = get_tokenizer(config.policy["tokenizer"]) diff --git a/examples/run_grpo_sliding_puzzle.py b/examples/run_grpo_sliding_puzzle.py index 72ad784ae6f..9d802d5332f 100644 --- a/examples/run_grpo_sliding_puzzle.py +++ b/examples/run_grpo_sliding_puzzle.py @@ -27,6 +27,7 @@ from nemo_rl.algorithms.grpo import MasterConfig, grpo_train, setup from nemo_rl.algorithms.utils import get_tokenizer, set_seed from nemo_rl.data.interfaces import DatumSpec, LLMMessageLogType +from nemo_rl.data_plane.factory import maybe_configure_data_plane_env from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.environments.games.sliding_puzzle import ( SlidingPuzzleConfig, @@ -232,6 +233,8 @@ def main(): ) with rl_init_timer.time("ray_connect"): + # Must precede init_ray() — see maybe_configure_data_plane_env's docstring. + maybe_configure_data_plane_env(config.data_plane) init_ray() set_seed(config.grpo.seed) diff --git a/examples/run_vlm_grpo.py b/examples/run_vlm_grpo.py index de37ba988a7..da363ea4080 100644 --- a/examples/run_vlm_grpo.py +++ b/examples/run_vlm_grpo.py @@ -22,6 +22,7 @@ from nemo_rl.algorithms.grpo import MasterConfig, async_grpo_train, grpo_train, setup from nemo_rl.algorithms.utils import get_tokenizer from nemo_rl.data.utils import setup_response_data +from nemo_rl.data_plane.factory import maybe_configure_data_plane_env from nemo_rl.distributed.virtual_cluster import init_ray from nemo_rl.models.generation import configure_generation_config from nemo_rl.utils.config import ( @@ -83,6 +84,8 @@ def main() -> None: ) with rl_init_timer.time("ray_connect"): + # Must precede init_ray() — see maybe_configure_data_plane_env's docstring. + maybe_configure_data_plane_env(config.data_plane) init_ray() with rl_init_timer.time("tokenizer"): diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 046b4e059d0..44d20e97d8e 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -409,29 +409,44 @@ global_forward_pad_seqlen = round_up(1320, 64) = 1344 ## Configuration The data plane is configured via a `data_plane:` block in the master -YAML (`examples/configs/...`). **YAML is the single source of truth -for defaults** — the adapter has no hidden `cfg.get(key, default)` -fallbacks. The canonical exemplar is +YAML (`examples/configs/...`). The canonical exemplar is `examples/configs/grpo_math_1B.yaml`. -All eight keys below are **required** when `enabled=true`. Recipes -under `examples/configs/recipes/**/*.yaml` inherit them via -`defaults:` from the exemplar. +`enabled`, `impl`, `backend` and `claim_meta_poll_interval_s` are +**required** when `enabled=true`. Backend sizing lives in a block named +for the backend that reads it; only the block named by `backend` is +consulted. An absent `mooncake_cpu:` block means that backend's +defaults, declared on `MooncakeCpuConfig` in +`nemo_rl/data_plane/interfaces.py`. `simple:` is **not** optional — +`num_storage_units` has no static default, since no single value is +right across cluster sizes, so a `simple` run without the block fails +validation. Recipes under `examples/configs/recipes/**/*.yaml` inherit +all of it via `defaults:`. ```yaml data_plane: enabled: false # flip to true to engage grpo_train_sync impl: transfer_queue # only one impl today backend: "simple" # "simple" or "mooncake_cpu" - storage_capacity: 1000000 # max samples retained per partition - num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" - local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" + simple: + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: ${mul:2, ${cluster.num_nodes}} # TQ wants >= 2 per node + mooncake_cpu: + global_segment_size: 68719476736 # 64 GiB/process + local_buffer_size: 4294967296 # 4 GiB/process + reuse_registered_buffers: true # reuse RDMA-registered buffers + staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool # observability: # NotRequired # enabled: false ``` +These keys used to sit directly under `data_plane:`. That spelling is not +rejected — it is simply never read. A config still using it silently gets +this backend's defaults instead of its own values: an inherited config +supplies the nested block, so a surviving flat key always loses the merge, +with no warning either way. + Backend choice: - **`simple`** — ZMQ-backed; lowest setup overhead. Default for tests and small runs. diff --git a/nemo_rl/data_plane/adapters/__init__.py b/nemo_rl/data_plane/adapters/__init__.py index 341a77c5bc6..b1839a12a62 100644 --- a/nemo_rl/data_plane/adapters/__init__.py +++ b/nemo_rl/data_plane/adapters/__init__.py @@ -11,3 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +# +# Keep this file import-free. ``transfer_queue_env`` is importable only because +# reaching it does not load ``transfer_queue`` (and therefore mooncake, which +# snapshots MC_* on the way in); re-exporting an adapter here would break that. diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 995cfa24c37..b5e803f1f6d 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -22,23 +22,33 @@ from __future__ import annotations +import contextlib +import glob import ipaddress import os +import resource import socket -import subprocess +import threading import time import warnings +import weakref from importlib import resources +from queue import Empty, SimpleQueue from typing import Any, cast import torch + +# Loading this loads mooncake, which snapshots MC_* on the way in. Configure the +# engine before this import — see nemo_rl.data_plane.adapters.transfer_queue_env. import transfer_queue as tq from tensordict import TensorDict +from nemo_rl.data_plane.adapters.transfer_queue_env import rail_link_layers from nemo_rl.data_plane.interfaces import ( DataPlaneClient, DataPlaneConfig, KVBatchMeta, + backend_config, ) from nemo_rl.data_plane.schema import PROMOTE_1D_FIELDS @@ -72,136 +82,380 @@ def _get_local_node_ip() -> str: return "" +def rdma_devices() -> str: + """Return this host's RDMA devices as mooncake's comma-separated list. + + ``MC_MOONCAKE_DEVICE`` wins and is passed through verbatim (one device or + a list). Otherwise every rail is offered: the NICs are split across NUMA + domains, so naming only one device makes the other domain's ranks cross + the socket on every transfer. + + Offering every rail is only safe because + ``MC_ENABLE_DEST_DEVICE_AFFINITY`` pins each transfer's peer rail to the + local one by name, so a cross-rail pair is never formed — see + :mod:`nemo_rl.data_plane.adapters.transfer_queue_env`. Without it, on a + fabric where each rail is its own subnet, a cross-rail draw has no route and + dies with "transport retry counter exceeded". + + IB and RoCE are never mixed; InfiniBand is preferred when present. + + Also the skip predicate for the mooncake tests — ``mooncake_cpu`` is + RDMA-only, so they cannot run without a device. + """ + override = os.environ.get("MC_MOONCAKE_DEVICE", "") + if override: + return override + # sysfs lists devices the kernel knows about; libibverbs can only open the + # ones exposed as /dev/infiniband/uverbs*. Containers routinely have the + # former without the latter, where mooncake fails with "No available RNIC" + # well after setup has begun — so treat a missing verbs node as no device. + if not glob.glob("/dev/infiniband/uverbs*"): + return "" + layers = rail_link_layers() + ib = [n for n, layer in layers.items() if layer == "InfiniBand"] + roce = [n for n, layer in layers.items() if layer == "Ethernet"] + # No space after the comma: mooncake splits on "," only. + return ",".join(ib or roce) + + def _mooncake_transport_config() -> dict: - protocol = os.environ.get("MC_MOONCAKE_PROTOCOL", "tcp") - if protocol != "rdma": - return {"protocol": "tcp"} - device = os.environ.get("MC_MOONCAKE_DEVICE", "") - if not device: - try: - out = subprocess.run( - [ - "sh", - "-c", - "for d in /sys/class/infiniband/mlx5_*/ports/1/link_layer; do " - " test -f $d && grep -q Ethernet $d && basename $(dirname $(dirname $d)); " - "done | head -1", - ], - check=False, - capture_output=True, - text=True, - ).stdout.strip() - device = out or "" - except Exception: - device = "" - if device: - os.environ.setdefault("MC_GID_INDEX", os.environ.get("MC_GID_INDEX", "3")) - return {"protocol": "rdma", "device_name": device} + # mooncake_cpu exists for the zero-copy RDMA MooncakeStore path (TQ v0.1.8), + # so RDMA is the only transport it runs: there is no TCP fallback, and a + # host without an RDMA device fails here rather than quietly degrading. + # Runs on the driver only, so it assumes homogeneous nodes — the device it + # finds is broadcast to every client. + devices = rdma_devices() + if not devices: + raise RuntimeError( + "data_plane.backend='mooncake_cpu' requires RDMA, but no usable " + "mlx5 device was found. Check that /dev/infiniband/uverbs* exists " + "(a container does not inherit it from the host even though it " + "does see /sys/class/infiniband) — name a device with " + "MC_MOONCAKE_DEVICE=, or use data_plane.backend='simple'." + ) + return {"protocol": "rdma", "device_name": devices} -def _connect_existing() -> None: - """Worker-process path: connect this process's client to the Ray cluster. +# A slot is held for exactly one transfer, so waiting minutes for one means +# this process runs more concurrent transfers than the pool has slots — not +# that a transfer is slow. Fail with that diagnosis rather than block forever. +_STAGING_SLOT_TIMEOUT_S = 600.0 - Connects to the already-running named controller actor. Mirrors - rl-arena/arena/dataplane_client.py's `tq.init()` (no args) call. + +def _memlock_limit() -> str: + """Return this process's RLIMIT_MEMLOCK soft limit, for error messages.""" + soft, _ = resource.getrlimit(resource.RLIMIT_MEMLOCK) + return "unlimited" if soft == resource.RLIM_INFINITY else f"{soft} bytes" + + +def _register_checked(store: Any, ptr: int, nbytes: int) -> None: + """``store.register_buffer`` with its status actually checked. + + Mooncake returns a status int here, and TQ drops it at every call site + (``mooncake_client.py``'s ``_register_all_buffers``). A registration that + fails is then invisible: the transfer into that unmapped region comes + back as the generic ``TRANSFER_FAIL`` (-800), which carries no root + cause, and burns its three retries against the same unmapped memory. + Registration pins pages with ``ibv_reg_mr`` once per RDMA rail, so it is + exactly the call that a memlock rlimit or a missing ``IPC_LOCK`` breaks. + + ``None`` counts as success — the binding's return type has varied across + mooncake wheels, so only an explicit non-zero status is a failure. """ - tq.init() + status = store.register_buffer(ptr, nbytes) + if status is not None and status != 0: + raise RuntimeError( + f"mooncake register_buffer(0x{ptr:x}, {nbytes} bytes) failed with " + f"status {status}. Registration pins the pages with ibv_reg_mr " + f"once per rail (devices={rdma_devices() or 'none'}), so it needs " + f"IPC_LOCK and a high memlock rlimit — RLIMIT_MEMLOCK is " + f"{_memlock_limit()} here. Lower data_plane.mooncake_cpu.global_segment_size / local_buffer_size if the limit is the bound." + ) + + +class _StagingPool: + """RDMA-registered host buffers, owned by one mooncake client. + Not thread-local: the ``ThreadPoolExecutor`` is rebuilt inside each + get/put, so thread-local buffers would be discarded every call. Sized to + the executor width so no worker normally waits for a slot. + + A slot's buffer is registered for as long as the pool holds it. The + invariant that matters is the converse: **no buffer is ever freed while + still registered**, because mooncake would keep a mapping over an address + the allocator immediately hands to the next caller. + """ -_TQ_RUNTIME_ENV_PATCHED = False + def __init__(self, store: Any, n_slots: int, max_bytes: int) -> None: + self._store = store + self._free: SimpleQueue = SimpleQueue() + for _ in range(n_slots): + self._free.put(None) # allocated on first use + self._n_slots = n_slots + self._max_bytes = max_bytes + + @contextlib.contextmanager + def buffer(self, nbytes: int): + # Outliers bypass the pool: slots only ever grow, so admitting one + # long-sequence sample would pin that size in every slot for the + # rest of the run. Registering it transiently is the cheaper trade. + if nbytes > self._max_bytes: + tmp = torch.empty(nbytes, dtype=torch.uint8) + _register_checked(self._store, tmp.data_ptr(), tmp.nbytes) + try: + yield tmp + finally: + self._store.unregister_buffer(tmp.data_ptr()) + return + try: + buf = self._free.get(timeout=_STAGING_SLOT_TIMEOUT_S) + except Empty: + raise RuntimeError( + f"No mooncake staging slot free after {_STAGING_SLOT_TIMEOUT_S}s. " + f"The pool has {self._n_slots} slots, sized to one TQ worker " + "pool, so this means overlapping put/get calls in this process. " + "Set data_plane.mooncake_cpu.reuse_registered_buffers=false to " + "fall back to upstream's per-call registration." + ) from None + try: + if buf is None or buf.nbytes < nbytes: + if buf is not None: + status = self._store.unregister_buffer(buf.data_ptr()) + if status is not None and status != 0: + # Dropping it now would hand memory the NIC may still + # map back to the allocator — see _register_checked. + raise RuntimeError( + f"mooncake unregister_buffer(0x{buf.data_ptr():x}) " + f"failed with status {status}; refusing to free a " + "buffer that may still be registered." + ) + # Empty the slot before allocating: if the registration + # below fails, the slot must come back empty rather than + # holding a buffer the NIC no longer maps. + buf = None + grown = torch.empty(nbytes, dtype=torch.uint8) + _register_checked(self._store, grown.data_ptr(), grown.nbytes) + buf = grown + yield buf + finally: + self._free.put(buf) + + +class _StagingPoolRegistry: + """Owns each client's staging pool, keyed weakly so it dies with the client. + + Weak keys because the registry is reachable from the patched class for the + process lifetime; a strong table would pin every client's registered + buffers for that long. + """ + + def __init__(self, n_slots: int, max_bytes: int) -> None: + self._n_slots = n_slots + self._max_bytes = max_bytes + self._lock = threading.Lock() + self._pools: weakref.WeakKeyDictionary[Any, _StagingPool] = ( + weakref.WeakKeyDictionary() + ) + + def pool_for(self, client: Any) -> _StagingPool: + """Return ``client``'s pool, building it at most once across threads. + + Locked because ``put``/``get`` drive the thread workers from a + ``ThreadPoolExecutor``, so two of them reach a cold client at once + whenever a call splits into more than one ``BATCH_SIZE_LIMIT`` batch. + Unsynchronized, the loser's pool is dropped on the floor and its buffers + are freed while still registered — see :func:`_register_checked` for why + that surfaces as a bare ``TRANSFER_FAIL``. The lock is taken on every + lookup rather than double-checked: it is uncontended after the first + transfer, and nanoseconds against a millisecond RDMA transfer is not + worth reasoning about visibility. + """ + with self._lock: + pool = self._pools.get(client) + if pool is None: + pool = self._pools[client] = _StagingPool( + client._store, self._n_slots, self._max_bytes + ) + return pool -def _resolve_tq_pin() -> str: - """Return the ``TransferQueue`` requirement string from nemo-rl metadata. +def _tq_shape_drift_error( + missing: str, consequence: str, target: str, *, opt_out: bool = False +) -> RuntimeError: + """Build the error shared by the monkey-patch shape guards below. - Single source of truth is ``pyproject.toml`` — we read it back via - ``importlib.metadata.requires`` so the runtime_env injection cannot - drift from the dependency declaration. + All three guards fire for the same reason — the pinned ``transfer_queue`` + revision no longer has the internals a patch depends on — so they share + this message shape rather than each hand-rolling it. """ - from importlib.metadata import requires - - for req in requires("nemo-rl") or []: - spec = req.split(";")[0].strip() - if spec.lower().startswith("transferqueue"): - return spec - raise RuntimeError( - "Could not resolve TransferQueue dependency from nemo-rl metadata. " - "Check pyproject.toml under [project.dependencies]." + remedy = f"re-point the patch at the new {target}" + if opt_out: + remedy += ( + ", or set data_plane.mooncake_cpu.reuse_registered_buffers=false " + "to run on upstream's per-call registration deliberately" + ) + return RuntimeError( + f"transfer_queue's {missing}, so {consequence}. The TQ pin in " + f"pyproject.toml has moved: {remedy}." ) -def _patch_tq_actor_runtime_env() -> None: - """Inject a per-actor ``runtime_env`` pin into TQ's actor ``.options()``. - - TQ spawns ``SimpleStorageUnit`` and ``TransferQueueController`` via - ``Cls.options(...).remote(...)`` without a runtime_env, so they - inherit the job-level env. In a multi-node container deployment - where each node has its own ``/opt/nemo_rl_venv``, the driver's - ``uv sync`` only updates ray-head's venv and a worker-node actor - fails with ``ModuleNotFoundError``. This monkey-patch makes Ray - pip-install TQ into a per-actor runtime_env on first spawn (cached - per-node by Ray afterwards). Idempotent. Couples us to TQ's internal - class layout — if TQ restructures, this becomes a no-op with a - logged warning and we fall back to per-node ``uv sync``. - - The pin is sourced from nemo-rl's installed metadata via - :func:`_resolve_tq_pin` so it cannot drift from ``pyproject.toml``. - - TODO(zhiyul): remove this patch once the nightly container image - is published with ``TransferQueue`` baked in via ``pyproject.toml``. - When every node starts from that image, the base env already has TQ - and Ray actors inherit it — this injection then becomes pure - overhead (Ray builds a redundant per-actor pip env on top of the - container's existing TQ install). Drop the call from - ``TQDataPlaneClient.__init__`` and delete this function. +def _patch_mooncake_register_check() -> None: + """Make a failed RDMA registration fail at the registration. + + Upstream's ``_register_all_buffers`` ignores ``register_buffer``'s + status, so every worker that uses it — including the two bytes workers + :func:`_patch_mooncake_staging_buffers` leaves alone — transfers into + memory the NIC may never have mapped and reports only ``TRANSFER_FAIL`` + (-800). Applied for every ``mooncake_cpu`` client, independent of + ``reuse_registered_buffers``, so the check survives disabling the pool. + + Raises if ``_register_all_buffers`` is missing, rather than returning + early: this check has no ``reuse_registered_buffers``-style opt-out, so + silently skipping it would put a failed registration back to surfacing + only as a bare ``TRANSFER_FAIL``, which is exactly the diagnosability + this patch exists to add. """ - global _TQ_RUNTIME_ENV_PATCHED - if _TQ_RUNTIME_ENV_PATCHED: + try: + from transfer_queue.storage.clients import mooncake_client as _mc + except ImportError: return - runtime_env = {"pip": [_resolve_tq_pin()]} + cls = getattr(_mc, "MooncakeStoreClient", None) + if cls is None or getattr(cls, "_nrl_register_checked", False): + return + if not hasattr(cls, "_register_all_buffers"): + raise _tq_shape_drift_error( + "MooncakeStoreClient no longer has _register_all_buffers", + "a failed RDMA registration would go back to surfacing only as " + "a bare TRANSFER_FAIL (-800) with no root cause", + "call site", + ) - def _install(cls) -> bool: - if not hasattr(cls, "options"): - return False - original = cls.options + def _register_all_buffers(self, ptrs, sizes): # type: ignore[no-untyped-def] + for ptr, size in zip(ptrs, sizes, strict=True): + _register_checked(self._store, ptr, size) - def patched(*args, **kwargs): - kwargs.setdefault("runtime_env", runtime_env) - return original(*args, **kwargs) + cls._register_all_buffers = _register_all_buffers + cls._nrl_register_checked = True - cls.options = patched # type: ignore[method-assign] - return True - unpatched_classes: list[str] = [] - try: - from transfer_queue.storage.simple_storage import SimpleStorageUnit +def _patch_mooncake_staging_buffers(max_bytes: int) -> None: + """Reuse RDMA-registered host buffers for mooncake tensor GETs and PUTs. - if not _install(SimpleStorageUnit): - unpatched_classes.append("SimpleStorageUnit") - except ImportError: - unpatched_classes.append("SimpleStorageUnit") - try: - from transfer_queue.controller import TransferQueueController + Upstream's thread workers allocate a fresh destination per call and + register/unregister it on the critical path. Pinning pages for DMA costs + several times the wire time for the same bytes, and because the buffers + are freed each call the pointers are always new, so nothing can be + cached. This keeps a small pool of registered buffers alive instead. - if not _install(TransferQueueController): - unpatched_classes.append("TransferQueueController") + Monkey-patched because TransferQueue is pinned by git SHA in + ``pyproject.toml``. Raises if the internals it drives are not shaped as + expected, rather than returning early: a silent return would leave + ``reuse_registered_buffers: true`` reading as on while the pool is + never built, with no symptom besides lost throughput. + """ + try: + from transfer_queue.storage.clients import mooncake_client as _mc + from transfer_queue.utils.mooncake_utils import _aligned_offsets, split_by_bytes + from transfer_queue.utils.tensor_utils import get_nbytes except ImportError: - unpatched_classes.append("TransferQueueController") - - if unpatched_classes: - # Soft-fail: TQ may have moved its actor classes. The driver will - # still work; multi-node TQ may need the per-node `uv sync` workaround. - warnings.warn( - "Could not patch every TQ actor class for runtime_env injection: " - f"unpatched={unpatched_classes}. " - "Multi-node TQ may fail with ModuleNotFoundError: 'transfer_queue' " - "on worker nodes. Workaround: run `uv sync` inside each node's " - "container before the driver runs.", - RuntimeWarning, - stacklevel=2, + return + + cls = getattr(_mc, "MooncakeStoreClient", None) + if cls is None or getattr(cls, "_nrl_staging_patched", False): + return + if not all( + hasattr(cls, a) + for a in ( + "_get_tensors_thread_worker", + "_batch_get_into_with_retry", + "_put_tensors_thread_worker", + "_batch_upsert_with_retry", + ) + ): + raise _tq_shape_drift_error( + "MooncakeStoreClient no longer has the methods the staging pool patches", + "reuse_registered_buffers cannot be honoured and every transfer " + "would silently re-register its buffers", + "call sites", + opt_out=True, + ) + + _n_slots_raw = getattr(_mc, "MAX_BATCH_WORKER_THREADS", None) + if not isinstance(_n_slots_raw, int): + raise _tq_shape_drift_error( + "mooncake_client module no longer exposes MAX_BATCH_WORKER_THREADS " + "as an int", + "the staging pool cannot be sized and reuse_registered_buffers " + "cannot be honoured", + "constant", + opt_out=True, ) - _TQ_RUNTIME_ENV_PATCHED = True + n_slots: int = _n_slots_raw + registry = _StagingPoolRegistry(n_slots, max_bytes) + + def _get_tensors_thread_worker( + self, batch_keys, batch_shapes, batch_dtypes, indexes + ): # type: ignore[no-untyped-def] + pool = registry.pool_for(self) + batch_nbytes = get_nbytes(batch_dtypes, batch_shapes) + tensors: list[Any] = [None] * len(batch_keys) + # Split the payload to fit a bounded buffer rather than sizing the + # buffer to the payload — this is what keeps the pool footprint fixed. + for idxs in split_by_bytes(batch_nbytes, max_bytes): + g_keys = [batch_keys[i] for i in idxs] + g_nbytes = [batch_nbytes[i] for i in idxs] + offsets, total = _aligned_offsets(g_nbytes) + with pool.buffer(total) as buf: + base = buf.data_ptr() + self._batch_get_into_with_retry( + g_keys, [base + off for off in offsets], g_nbytes + ) + # Clone: the buffer is reused by the next group and next call. + for pos, off, nb in zip(idxs, offsets, g_nbytes, strict=True): + tensors[pos] = ( + buf[off : off + nb] + .view(batch_dtypes[pos]) + .reshape(tuple(batch_shapes[pos])) + .clone() + ) + return tensors, indexes + + def _put_tensors_thread_worker(self, batch_keys, batch_tensors): # type: ignore[no-untyped-def] + """PUT direction of the GET patch: stage into the pooled buffer, then transfer.""" + pool = registry.pool_for(self) + contiguous = [t.contiguous() for t in batch_tensors] + nbytes = [t.nbytes for t in contiguous] + for idxs in split_by_bytes(nbytes, max_bytes): + g_nbytes = [nbytes[i] for i in idxs] + offsets, total = _aligned_offsets(g_nbytes) + with pool.buffer(total) as buf: + base = buf.data_ptr() + for i, off, nb in zip(idxs, offsets, g_nbytes, strict=True): + # reshape(-1) before view(uint8): view() resizes the last + # dim and raises on the 0-d scalars real payloads carry. + buf[off : off + nb].copy_( + contiguous[i].reshape(-1).view(torch.uint8) + ) + self._batch_upsert_with_retry( + [batch_keys[i] for i in idxs], + [base + off for off in offsets], + g_nbytes, + ) + + cls._get_tensors_thread_worker = _get_tensors_thread_worker + cls._put_tensors_thread_worker = _put_tensors_thread_worker + cls._nrl_staging_patched = True + + +def _connect_existing() -> None: + """Worker-process path: connect this process's client to the Ray cluster. + + Connects to the already-running named controller actor. Mirrors + rl-arena/arena/dataplane_client.py's `tq.init()` (no args) call. + """ + tq.init() def _init_tq(cfg: DataPlaneConfig) -> None: @@ -211,8 +465,6 @@ def _init_tq(cfg: DataPlaneConfig) -> None: base = OmegaConf.load(str(resources.files("transfer_queue") / "config.yaml")) backend = cfg["backend"] - storage_capacity = cfg["storage_capacity"] - num_storage_units = cfg["num_storage_units"] # polling_mode=True: controller returns empty BatchMeta instead of raising # TimeoutError when no samples are ready yet. The client-side blocking @@ -220,13 +472,18 @@ def _init_tq(cfg: DataPlaneConfig) -> None: controller_overlay = {"controller": {"polling_mode": True}} if backend == "simple": + # Resolved here, not above: MooncakeStore has no unit count and no + # sample cap — it sizes from global_segment_size/local_buffer_size — + # so both keys are SimpleStorage-only and reading them at the top + # implied otherwise. + simple_cfg = backend_config(cfg) overlay = { **controller_overlay, "backend": { "storage_backend": "SimpleStorage", "SimpleStorage": { - "total_storage_size": storage_capacity, - "num_data_storage_units": num_storage_units, + "total_storage_size": simple_cfg.storage_capacity, + "num_data_storage_units": simple_cfg.num_storage_units, }, }, } @@ -266,21 +523,16 @@ def _init_tq(cfg: DataPlaneConfig) -> None: "Mooncake backend requires a local node IP; " "_get_local_node_ip() returned empty." ) - # Mooncake virtual segment / local buffer sizing. Defaults sized - # for production-scale rollouts (multi-iter DAPO, large - # message_log object payloads); under-sized values cause - # ``batch_get_tensor returned None`` once mooncake exhausts its - # internal allocator headroom. Lazy-mmap'd, so RSS is bounded - # by actual traffic. Override per-recipe via - # ``data_plane.global_segment_size`` / - # ``data_plane.local_buffer_size`` (bytes). + # Sizes are per client process and RDMA-pinned — see MooncakeCpuConfig + # in nemo_rl/data_plane/interfaces.py for the per-node arithmetic. + mooncake_cfg = backend_config(cfg) overlay = { **controller_overlay, "backend": { "storage_backend": "MooncakeStore", "MooncakeStore": { - "global_segment_size": int(cfg["global_segment_size"]), - "local_buffer_size": int(cfg["local_buffer_size"]), + "global_segment_size": int(mooncake_cfg.global_segment_size), + "local_buffer_size": int(mooncake_cfg.local_buffer_size), # _init_tq runs on the driver only — driver IS the # head, so local_ip here is also the head's IP that # mooncake_master + the metadata server bind to. @@ -295,11 +547,6 @@ def _init_tq(cfg: DataPlaneConfig) -> None: conf = OmegaConf.merge(base, overlay) - # Inject runtime_env into TQ's actor spawn so SimpleStorageUnit / - # TransferQueueController land on workers with transfer_queue available - # — see _patch_tq_actor_runtime_env() docstring for the why. - _patch_tq_actor_runtime_env() - # pyrefly: ignore # bad-argument-type tq.init(conf=conf) @@ -439,21 +686,18 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: """ # mooncake_cpu setup must run BEFORE _init_tq / _connect_existing # — once tq.init/connect runs, Mooncake's engine.so reads the - # env vars and they can't be changed. Three per-process knobs + # env vars and they can't be changed. Two per-process knobs are # needed in EVERY process that builds a TQ client (driver, # SyncRolloutActor, every MegatronPolicyWorker rank): # 1. MC_TCP_BIND_ADDRESS — Mooncake engine.so writes this into # desc.ip_or_host_name, the address peers receive from the # metadata service. Without it, getifaddrs()[0] picks usb0 # (169.254.x APIPA) and peers fail to connect. - # 2. MC_STORE_MEMCPY=0 — Mooncake LOCAL_MEMCPY fast-path - # reinterpret_casts cross-process pointers, segfaulting - # MemcpyWorkerPool. PR #1995 (merged 2026-04-30) fixes the - # root cause but isn't in any published wheel yet - # (mooncake-transfer-engine 0.3.10.post2 was bumped before - # that merge). Drop this once the wheel includes the fix. - # 3. KV-path 1D promotion — works around TQ's + # 2. KV-path 1D promotion — works around TQ's # extract_field_schema schema/data mismatch for 1D fields. + # The cluster-wide MC_* knobs are NOT among them; they are set + # once on the driver, before this module is importable — see + # nemo_rl.data_plane.adapters.transfer_queue_env. if cfg["backend"] == "mooncake_cpu": local_ip = _get_local_node_ip() if local_ip: @@ -462,7 +706,19 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: # be a no-op and the actor would announce the driver's # IP — peers fail with "connection refused". os.environ["MC_TCP_BIND_ADDRESS"] = local_ip - os.environ.setdefault("MC_STORE_MEMCPY", "0") + # Do not add MC_* setup here — mooncake snapshotted its config when + # this module imported, so a write now is silently ignored. + # Both must run before the first get, in every process with a TQ + # client. The registration check is unconditional: it also covers + # the two bytes workers the staging patch leaves untouched, which + # is where an unchecked registration surfaces as TRANSFER_FAIL. + _patch_mooncake_register_check() + # Opt-out flag, defaulted on MooncakeCpuConfig rather than here: + # an absent mooncake_cpu block means "this backend's defaults", + # so the pool is on unless a config deliberately turns it off. + mooncake_cfg = backend_config(cfg) + if mooncake_cfg.reuse_registered_buffers: + _patch_mooncake_staging_buffers(mooncake_cfg.staging_buffer_size) # Workaround for TQ KVStorageManager's 1D-field schema/data # mismatch (only `mooncake_cpu` goes through that path; `simple` @@ -477,6 +733,10 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: _connect_existing() self._poll_interval_s = cfg["claim_meta_poll_interval_s"] self._closed = False + # Fields whose schema this process has already warmed, per partition. + # See register_partition: the controller's field map is append-only, + # so a field only ever needs warming once. + self._warmed_fields: dict[str, set[str]] = {} # ── (A) task-mediated ─────────────────────────────────────────────── @@ -503,12 +763,17 @@ def register_partition( # Registering everything from a single driver thread before any # client request races with a put removes the trigger entirely. # + # Only new field names need warming: the controller's + # ``field_name_mapping`` is append-only (never deleted, and our + # ``clear_samples`` zeroes rows without popping the partition). + already = self._warmed_fields.setdefault(partition_id, set()) + fields = [f for f in fields if f not in already] + if not fields: + return # Use a unique KV key instead of ``client.put``'s default row id # (``0@field`` at the Mooncake storage layer). Mooncake does not # support upsert, so repeated schema warmups can collide with # stale metadata from a previous registration. - if not fields: - return schema_key = ( f"__schema__:{partition_id}:{os.getpid()}:{id(self)}:{time.time_ns()}" ) @@ -523,6 +788,10 @@ def register_partition( tags=[{}], ) tq.kv_clear(keys=[schema_key], partition_id=partition_id) + # Only mark warmed once the write actually landed — otherwise a + # failed put (mooncake's own retries already exhausted) poisons the + # cache and a future retry of this call would wrongly skip warmup. + already.update(fields) def claim_meta( self, @@ -688,8 +957,6 @@ def clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None sample_ids = list(listing.get(partition_id, {}).keys()) if not sample_ids: if cleared_via_none: - import warnings - warnings.warn( f"clear_samples(sample_ids=None, partition_id={partition_id!r}) " "found nothing to clear — TQ's kv_list returned no keys for " diff --git a/nemo_rl/data_plane/adapters/transfer_queue_env.py b/nemo_rl/data_plane/adapters/transfer_queue_env.py new file mode 100644 index 00000000000..7b73adadce4 --- /dev/null +++ b/nemo_rl/data_plane/adapters/transfer_queue_env.py @@ -0,0 +1,129 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Mooncake engine environment, configured *before* the engine is imported. + +Mooncake snapshots its whole ``MC_*`` configuration as its extension loads, so +these variables only take effect if they are in ``os.environ`` beforehand. That +makes import order load-bearing, and the failure is silent: a late write lands +in ``os.environ`` — where it still reads back correctly — while the engine keeps +the value it captured. On a rail-isolated RoCE fabric that silence costs a run, +with every transfer dying as "transport retry counter exceeded". + +This module therefore deliberately imports **nothing** from ``transfer_queue`` +or ``mooncake``, so importing it can never be the thing that loads the engine. +Keep it that way — a convenience import here, or one added to this package's +``__init__``, would defeat the whole point. :func:`configure_engine_env` turns +the ordering violation into an error rather than a silent misconfiguration. +""" + +from __future__ import annotations + +import glob +import os +import sys +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from nemo_rl.data_plane.interfaces import DataPlaneConfig + +# Importing either loads mooncake's extension (transfer_queue's storage.clients +# package eagerly imports mooncake_client, which does `from mooncake.store +# import ...` at module scope), which is the point after which MC_* can no +# longer be configured. Submodules need no entry of their own: importing +# mooncake.store registers the mooncake package first. +_ENGINE_MODULES = ("transfer_queue", "mooncake") + + +def _engine_already_imported() -> str | None: + """Return the first engine module already in ``sys.modules``, else None.""" + return next((m for m in _ENGINE_MODULES if m in sys.modules), None) + + +def rail_link_layers() -> dict[str, str]: + """Map each mlx5 rail to its port-1 link layer, read from sysfs.""" + layers: dict[str, str] = {} + for path in sorted(glob.glob("/sys/class/infiniband/mlx5_*/ports/1/link_layer")): + try: + layers[Path(path).parents[2].name] = Path(path).read_text().strip() + except OSError: + continue + return layers + + +def fabric_is_roce_only() -> bool: + """True when the host has RoCE rails and no InfiniBand. + + Deliberately requires *seeing* InfiniBand to answer False, so an empty or + unreadable sysfs cannot silently opt a RoCE host out of the pairing hint. + """ + layers = set(rail_link_layers().values()) + return "Ethernet" in layers and "InfiniBand" not in layers + + +def _wanted_engine_env() -> dict[str, str]: + """The ``MC_*`` values this backend needs, for this host's fabric.""" + # LOCAL_MEMCPY reinterpret_casts cross-process pointers and segfaults + # MemcpyWorkerPool; upstream PR #1995 fixes it but is not in a published + # wheel. Drop once the pinned wheel includes it. + wanted = {"MC_STORE_MEMCPY": "0"} + if fabric_is_roce_only(): + # Pin each transfer's peer rail to the local one by name. Mooncake + # otherwise picks the peer independently (Topology::selectDevice), and + # on RoCE each rail is its own subnet, so a cross-rail pair has no + # route. Measured on the gb200 CI runners: every cross-rail pair failed, + # no same-rail pair ever did. InfiniBand routes cross-rail, so it is + # left alone. Mooncake reads this presence-only (config.cpp:318), so + # `=0` enables it too — unsetting it is the only way to disable it. + wanted["MC_ENABLE_DEST_DEVICE_AFFINITY"] = "1" + return wanted + + +def configure_engine_env(cfg: DataPlaneConfig) -> None: + """Set the mooncake knobs that must be identical in every process. + + No-op unless the backend is ``mooncake_cpu``; ``simple`` has no engine. + Values already present in the environment are left alone, so a launcher can + override any of them — except ``MC_ENABLE_DEST_DEVICE_AFFINITY``, which + mooncake reads presence-only, so a launcher trying to override it to + ``"0"`` enables it instead; unsetting it is the only way to disable it. + + Call this before anything imports ``transfer_queue`` or ``mooncake``. + :func:`nemo_rl.data_plane.factory.maybe_configure_data_plane_env` does, on + the driver before ``init_ray``, and Ray hands the result to every worker. + + Raises: + RuntimeError: if a variable still needs setting but the engine is + already imported, i.e. the value can no longer reach it. Fatal on + purpose — the alternative is a run that looks configured and is not. + """ + if cfg["backend"] != "mooncake_cpu": + return + + missing = {k: v for k, v in _wanted_engine_env().items() if k not in os.environ} + if not missing: + # Already set by the launcher, or by an earlier call in this process. + return + + imported = _engine_already_imported() + if imported is not None: + raise RuntimeError( + f"mooncake's engine was already imported (via {imported!r}) before " + f"{sorted(missing)} could be set, so the engine would never see " + "them: it reads its MC_* configuration once, as the extension " + "loads. Configure before anything imports transfer_queue or " + "mooncake, or export the variables in the launcher environment." + ) + + os.environ.update(missing) diff --git a/nemo_rl/data_plane/factory.py b/nemo_rl/data_plane/factory.py index 86b5a944813..08f454442b1 100644 --- a/nemo_rl/data_plane/factory.py +++ b/nemo_rl/data_plane/factory.py @@ -18,6 +18,44 @@ from nemo_rl.data_plane.interfaces import DataPlaneClient, DataPlaneConfig +def maybe_configure_data_plane_env(cfg: DataPlaneConfig | None) -> None: + """Set backend env vars that must be identical in every process. + + Call this on the driver **before** ``init_ray()``: ``init_ray`` snapshots the + driver's environment into ``runtime_env["env_vars"]`` and hands it to every + Ray worker, which are fresh processes, so the value is in place before they + run any engine code. + + The binding constraint is not ``init_ray`` but that this must run before + anything imports the backend's engine, which snapshots its configuration as + it loads. :func:`~nemo_rl.data_plane.adapters.transfer_queue_env.configure_engine_env` + raises rather than silently no-op'ing if that is violated. + + Subprocesses the driver did not spawn through Ray inherit whatever the + driver's environment held at fork, so they are covered as long as this ran + first. + + No-op when the data plane is disabled or the backend has no such knobs. + + Args: + cfg: Data-plane config, or ``None`` when the data plane is off. + """ + if cfg is None or not cfg["enabled"]: + return + + impl = cfg["impl"] + if impl == "transfer_queue": + # transfer_queue_env, not the adapter — importing the adapter loads + # mooncake, which is what this call has to precede. + from nemo_rl.data_plane.adapters.transfer_queue_env import ( + configure_engine_env, + ) + + configure_engine_env(cfg) + else: + raise ValueError(f"unknown data_plane impl: {impl!r}") + + def build_data_plane_client( cfg: DataPlaneConfig | None, *, bootstrap: bool = True ) -> DataPlaneClient: diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 41a98f0c0ed..8ee07141c27 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -39,9 +39,57 @@ from dataclasses import dataclass, field from typing import Any, Callable, Literal, NotRequired, Sequence, TypedDict +from pydantic import BaseModel from tensordict import TensorDict +class SimpleStorageConfig(BaseModel, extra="allow"): + """Sizing for ``backend="simple"``. Ignored by every other backend. + + ``num_storage_units`` scales with the cluster: TQ round-robins storage + units over Ray nodes and recommends ``>= 2 x`` the node count. No static + default is correct across cluster sizes, so this is required rather than + defaulted — a class field cannot see ``cluster.num_nodes``, only the + exemplar YAML can, via ``${mul:2, ${cluster.num_nodes}}``. Every recipe + inherits that from the exemplar; set a plain int to pin it. + """ + + storage_capacity: int = 1000000 # max samples retained per partition + num_storage_units: int + + +class MooncakeCpuConfig(BaseModel, extra="allow"): + """Sizing and RDMA knobs for ``backend="mooncake_cpu"``. Ignored otherwise. + + ``global_segment_size`` / ``local_buffer_size`` are per client *process* + (one per GPU), so a node pays ``gpus_per_node x (segment + buffer)``. Under + RDMA that memory is pinned and resident from setup, so keep the per-node + product in mind when raising them. Under-sizing surfaces as + ``batch_get_tensor returned None``. + + ``reuse_registered_buffers`` keeps a pool of RDMA-registered buffers alive + instead of registering a fresh one per transfer; set false to fall back to + upstream's per-call registration. + + ``staging_buffer_size`` is that pool's per-slot ceiling. It is a pooling + threshold, not a size limit: a bigger payload still transfers, just with a + transient registration. Slots ratchet — they grow to the largest payload + admitted and never shrink — so raise it only when a per-key payload (one + sample of one field) genuinely exceeds it, not for headroom. + + Every RDMA rail on the host is offered to mooncake (see ``rdma_devices``). + That is only safe with ``MC_ENABLE_DEST_DEVICE_AFFINITY=1``, which pins each + transfer's peer rail to the local one by name; on a rail-isolated RoCE + fabric a cross-rail pair has no route. It is set on RoCE-only hosts by + ``nemo_rl.data_plane.adapters.transfer_queue_env.configure_engine_env``. + """ + + global_segment_size: int = 68719476736 # 64 GiB per client process + local_buffer_size: int = 4294967296 # 4 GiB per client process + reuse_registered_buffers: bool = True + staging_buffer_size: int = 268435456 # 256 MiB per pool slot + + class DataPlaneConfig(TypedDict): """Feature-gated config; defaults to disabled. @@ -49,30 +97,60 @@ class DataPlaneConfig(TypedDict): the TQ adapter, not by NeMo-RL. ``impl`` selects which adapter we go through. - Required keys (always set in exemplar YAML — never defaulted in code): - ``enabled``, ``impl``, ``backend``, ``storage_capacity``, - ``num_storage_units``, ``claim_meta_poll_interval_s``, - ``global_segment_size``, ``local_buffer_size``. - - ``global_segment_size`` / ``local_buffer_size`` are only *read* when - ``backend == "mooncake_cpu"``; the simple backend ignores them. - They are required (not NotRequired) so the YAML carries the full - schema and there are no hidden Python defaults. + Backend-specific knobs live under a block named for the backend that reads + them — ``simple:`` and ``mooncake_cpu:`` — mirroring TransferQueue's own + ``config.yaml`` and the per-backend overlay :func:`_init_tq` builds. Only + the block named by ``backend`` is consulted, so a config selecting + ``simple`` never has to mention mooncake's RDMA sizing at all. An absent + ``mooncake_cpu:`` block means "use :class:`MooncakeCpuConfig`'s + defaults" — but ``simple:`` is **not** optional: ``num_storage_units`` + has no static default, since no single value is right across cluster + sizes, so a ``simple`` run without the block fails validation. + + Required keys (always set in the exemplar YAML): ``enabled``, ``impl``, + ``backend``, ``claim_meta_poll_interval_s``. + + ``storage_capacity`` / ``num_storage_units`` / ``global_segment_size`` / + ``local_buffer_size`` used to sit at this level. A config still using that + spelling is not rejected — the flat key is simply never read, and + :func:`backend_config` resolves the nested block (or its defaults) as if it + were absent. See there. """ enabled: bool impl: Literal["transfer_queue"] backend: Literal["simple", "mooncake_cpu"] - storage_capacity: int - num_storage_units: int claim_meta_poll_interval_s: float - global_segment_size: int - local_buffer_size: int + simple: NotRequired[SimpleStorageConfig] + mooncake_cpu: NotRequired[MooncakeCpuConfig] controller_address: NotRequired[str] ack_timeout_ms: NotRequired[int] observability: NotRequired["ObservabilityConfig"] +_BACKEND_MODELS: dict[str, type[BaseModel]] = { + "simple": SimpleStorageConfig, + "mooncake_cpu": MooncakeCpuConfig, +} + + +def backend_config(cfg: DataPlaneConfig) -> Any: + """Return the validated sizing block for ``cfg["backend"]``. + + Reads the nested block and lets the model supply anything it omits, so no + caller ever writes a fallback. Works whether ``cfg`` came through pydantic + (block already coerced to a model) or as a plain dict from a test. + + Sizing is read only from the nested block. A config still using the + pre-nesting flat spelling gets this backend's defaults, not its own values. + """ + backend = cfg["backend"] + nested = cfg.get(backend) or {} + if isinstance(nested, BaseModel): + nested = nested.model_dump(exclude_unset=True) + return _BACKEND_MODELS[backend].model_validate(nested) + + class ObservabilityConfig(TypedDict): """Optional middleware that records per-op metrics on the client. diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index 0604c17359b..a51f11fc72b 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -286,6 +286,14 @@ def init_ray(log_dir: Optional[str] = None) -> None: If that cluster uses the same CUDA_VISIBLE_DEVICES or Slurm managed tag we will reuse it. Otherwise, we will detach and start a fresh local cluster. + Any process env var that must reach every worker (e.g. a backend engine + knob such as data_plane's) has to be set before this call, in the caller + — this function snapshots ``dict(os.environ)`` below into + ``runtime_env["env_vars"]``, which is the only point such a setting + becomes cluster-wide. See + :func:`~nemo_rl.data_plane.factory.maybe_configure_data_plane_env`, which + a data-plane-enabled launcher calls immediately before this one. + Args: log_dir: Optional directory to store Ray logs and temp files. """ diff --git a/pyrefly.toml b/pyrefly.toml index aee1da0635f..a3a7dae9305 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -118,6 +118,7 @@ project-includes = [ "nemo_rl/data_plane/adapters/__init__.py", "nemo_rl/data_plane/adapters/noop.py", "nemo_rl/data_plane/adapters/transfer_queue.py", + "nemo_rl/data_plane/adapters/transfer_queue_env.py", "nemo_rl/data_plane/codec.py", "nemo_rl/data_plane/column_io.py", "nemo_rl/data_plane/factory.py", diff --git a/tests/functional/grpo_async_gym_single_controller.sh b/tests/functional/grpo_async_gym_single_controller.sh index 178eb202029..79de252d87c 100755 --- a/tests/functional/grpo_async_gym_single_controller.sh +++ b/tests/functional/grpo_async_gym_single_controller.sh @@ -94,11 +94,9 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE ++data_plane.enabled=true \ ++data_plane.impl=transfer_queue \ ++data_plane.backend=simple \ - ++data_plane.storage_capacity=1000000 \ - ++data_plane.num_storage_units=2 \ + ++data_plane.simple.storage_capacity=1000000 \ + ++data_plane.simple.num_storage_units=2 \ ++data_plane.claim_meta_poll_interval_s=0.5 \ - ++data_plane.global_segment_size=549755813888 \ - ++data_plane.local_buffer_size=68719476736 \ ++async_rl.sampler.name=in_order \ ++async_rl.sampler.max_lookahead_versions=0 \ ++async_rl.min_groups_for_streaming_train=4 \ diff --git a/tests/functional/grpo_dp_mooncake.sh b/tests/functional/grpo_dp_mooncake.sh index a5f6e3d0bd9..9992f0d4593 100755 --- a/tests/functional/grpo_dp_mooncake.sh +++ b/tests/functional/grpo_dp_mooncake.sh @@ -10,6 +10,15 @@ git config --global --add safe.directory $PROJECT_ROOT set -eou pipefail +# mooncake_cpu is RDMA-only, so this needs an mlx5 device libibverbs can open +# (either fabric). Skip rather than fail on hosts that have none. +source "$SCRIPT_DIR/../scripts/detect_rdma.sh" +if [[ -z "${MC_MOONCAKE_DEVICE:-}" ]] && ! rdma_device_available; then + echo "[SKIP] no usable mlx5 RDMA device; mooncake_cpu requires RDMA." \ + "Set MC_MOONCAKE_DEVICE= to override." + exit 0 +fi + EXP_NAME=$(basename $0 .sh) EXP_DIR=$SCRIPT_DIR/$EXP_NAME LOG_DIR=$EXP_DIR/logs @@ -38,8 +47,8 @@ uv run coverage run -a --data-file=$PROJECT_ROOT/tests/.coverage --source=$PROJE data_plane.enabled=true \ data_plane.impl=transfer_queue \ data_plane.backend=mooncake_cpu \ - data_plane.global_segment_size=4294967296 \ - data_plane.local_buffer_size=1073741824 \ + data_plane.mooncake_cpu.global_segment_size=4294967296 \ + data_plane.mooncake_cpu.local_buffer_size=1073741824 \ $@ \ 2>&1 | tee $RUN_LOG diff --git a/tests/functional/grpo_sc_gym_router_failover.sh b/tests/functional/grpo_sc_gym_router_failover.sh index f3be22c560c..e5e7d838fc9 100755 --- a/tests/functional/grpo_sc_gym_router_failover.sh +++ b/tests/functional/grpo_sc_gym_router_failover.sh @@ -167,11 +167,9 @@ PYTHONUNBUFFERED=1 uv run "$PROJECT_ROOT/examples/run_grpo_single_controller.py" ++data_plane.enabled=true \ ++data_plane.impl=transfer_queue \ ++data_plane.backend=simple \ - ++data_plane.storage_capacity=1000000 \ - ++data_plane.num_storage_units=2 \ + ++data_plane.simple.storage_capacity=1000000 \ + ++data_plane.simple.num_storage_units=2 \ ++data_plane.claim_meta_poll_interval_s=0.5 \ - ++data_plane.global_segment_size=549755813888 \ - ++data_plane.local_buffer_size=68719476736 \ ++async_rl.sampler.name=in_order \ ++async_rl.sampler.max_lookahead_versions=0 \ ++async_rl.min_groups_for_streaming_train=4 \ diff --git a/tests/run_functional_in_docker.sh b/tests/run_functional_in_docker.sh index 793ca9ed8f6..40b27469d6d 100755 --- a/tests/run_functional_in_docker.sh +++ b/tests/run_functional_in_docker.sh @@ -45,8 +45,17 @@ fi # The workaround is we launch the job but set umask 000 so all files created as root are rwxrwxrwx. # We have found that 111 does not always work and can leave the filesystem permissions in a bad state. +# Expose RDMA devices when the host has them, matching CI (see +# .github/actions/test-template/action.yml). Without these the mooncake_cpu +# tests skip locally while running in CI, so "passes locally" means less. +source "$SCRIPT_DIR/scripts/detect_rdma.sh" +RDMA_FLAGS=() +if rdma_device_available; then + RDMA_FLAGS=(--device=/dev/infiniband --cap-add=IPC_LOCK) +fi + # Run the script inside the Docker container with GPU support -docker run -u root $INTERACTIVE_FLAG --ulimit memlock=-1 --ulimit stack=67108864 --rm --gpus '"device=0,1"' \ +docker run -u root $INTERACTIVE_FLAG --ulimit memlock=-1 --ulimit stack=67108864 "${RDMA_FLAGS[@]}" --rm --gpus '"device=0,1"' \ -v "$PROJECT_ROOT:$PROJECT_ROOT" \ -v $HF_HOME:/hf_home \ -v $HF_DATASETS_CACHE:/hf_datasets_cache \ diff --git a/tests/run_unit_in_docker.sh b/tests/run_unit_in_docker.sh index 2ed46abdf62..c30bdbb2c3a 100755 --- a/tests/run_unit_in_docker.sh +++ b/tests/run_unit_in_docker.sh @@ -35,5 +35,14 @@ fi # The workaround is we launch the job but set umask 000 so all files created as root are rwxrwxrwx. # We have found that 111 does not always work and can leave the filesystem permissions in a bad state. +# Expose RDMA devices when the host has them, matching CI (see +# .github/actions/test-template/action.yml). Without these the mooncake_cpu +# fixtures skip locally while running in CI, so "passes locally" means less. +source "$SCRIPT_DIR/scripts/detect_rdma.sh" +RDMA_FLAGS=() +if rdma_device_available; then + RDMA_FLAGS=(--device=/dev/infiniband --cap-add=IPC_LOCK) +fi + # Run the script inside the Docker container with GPU support -docker run -u root $INTERACTIVE_FLAG --ulimit memlock=-1 --ulimit stack=67108864 --cap-add=SYS_PTRACE --rm --gpus '"device=0,1"' -v "$(realpath $SCRIPT_DIR/..):/workspace" -v $HF_HOME:/hf_home -e HF_TOKEN -e HF_HOME=/hf_home -e HOME=/tmp/ -w /workspace/tests "$CONTAINER" -- bash -x -c "umask 000 && uv run --group test bash -x ./run_unit.sh $@" +docker run -u root $INTERACTIVE_FLAG --ulimit memlock=-1 --ulimit stack=67108864 --cap-add=SYS_PTRACE "${RDMA_FLAGS[@]}" --rm --gpus '"device=0,1"' -v "$(realpath $SCRIPT_DIR/..):/workspace" -v $HF_HOME:/hf_home -e HF_TOKEN -e HF_HOME=/hf_home -e HOME=/tmp/ -w /workspace/tests "$CONTAINER" -- bash -x -c "umask 000 && uv run --group test bash -x ./run_unit.sh $@" diff --git a/tests/scripts/detect_rdma.sh b/tests/scripts/detect_rdma.sh new file mode 100644 index 00000000000..1c5144ff8cd --- /dev/null +++ b/tests/scripts/detect_rdma.sh @@ -0,0 +1,10 @@ +#!/bin/bash +# True if the host exposes an mlx5 RDMA device libibverbs can open (RoCE or +# InfiniBand) — what mooncake_cpu's rdma_devices() (nemo_rl/data_plane/adapters/ +# transfer_queue.py) requires. Gate on uverbs* specifically, not just the +# /dev/infiniband directory: a host can have the directory without a verbs +# node, which is what libibverbs actually opens. +rdma_device_available() { + compgen -G "/dev/infiniband/uverbs*" >/dev/null && + compgen -G "/sys/class/infiniband/mlx5_*/ports/1/link_layer" >/dev/null +} diff --git a/tests/unit/data_plane/_rollout_shapes.py b/tests/unit/data_plane/_rollout_shapes.py index 3bb2e614522..147f0a70611 100644 --- a/tests/unit/data_plane/_rollout_shapes.py +++ b/tests/unit/data_plane/_rollout_shapes.py @@ -242,3 +242,22 @@ def mooncake_available() -> bool: raise return False return True + + +def rdma_available() -> bool: + """Return True if a usable mlx5 RDMA device is present. + + Set ``NEMO_RL_REQUIRE_MOONCAKE=1`` to promote a missing device into a + loud ``RuntimeError`` instead of returning False — same promotion rule + as :func:`mooncake_available`, for the "RDMA device absent" precondition. + """ + from nemo_rl.data_plane.adapters.transfer_queue import rdma_devices + + if rdma_devices(): + return True + if os.environ.get("NEMO_RL_REQUIRE_MOONCAKE") == "1": + raise RuntimeError( + "no usable mlx5 RDMA device — mooncake_cpu requires RDMA " + "(set MC_MOONCAKE_DEVICE= to override)" + ) + return False diff --git a/tests/unit/data_plane/conftest.py b/tests/unit/data_plane/conftest.py index 18158484079..0f7e2da3019 100644 --- a/tests/unit/data_plane/conftest.py +++ b/tests/unit/data_plane/conftest.py @@ -33,7 +33,7 @@ from nemo_rl.data_plane import build_data_plane_client -from ._rollout_shapes import mooncake_available +from ._rollout_shapes import mooncake_available, rdma_available def _make_tq_cfg(backend: str) -> dict: @@ -41,11 +41,15 @@ def _make_tq_cfg(backend: str) -> dict: "enabled": True, "impl": "transfer_queue", "backend": backend, - "storage_capacity": 1024, - "num_storage_units": 1, "claim_meta_poll_interval_s": 0.5, - "global_segment_size": 8589934592, # 8 GiB — sized for CI host RAM - "local_buffer_size": 1073741824, # 1 GiB + "simple": {"storage_capacity": 1024, "num_storage_units": 1}, + "mooncake_cpu": { + "global_segment_size": 8589934592, # 8 GiB — sized for CI host RAM + "local_buffer_size": 1073741824, # 1 GiB + # reuse_registered_buffers omitted on purpose: absent must mean on, + # so the fixture exercises the default the same way a user config + # that never mentions the flag does. + }, } @@ -67,6 +71,15 @@ def _session_tq_client_mooncake_cpu(): "mooncake not installed — skipping mooncake_cpu " "(set NEMO_RL_REQUIRE_MOONCAKE=1 to fail loud)" ) + # mooncake_cpu is RDMA-only, so it cannot run without an RDMA device. CI + # sets NEMO_RL_REQUIRE_MOONCAKE=1 on runners that have one, which turns + # this skip into a failure — otherwise losing the device passthrough would + # silently drop mooncake coverage and still go green. + if not rdma_available(): + pytest.skip( + "no usable mlx5 RDMA device — mooncake_cpu requires RDMA " + "(set MC_MOONCAKE_DEVICE= to override)" + ) client = build_data_plane_client(_make_tq_cfg("mooncake_cpu")) yield client client.close() diff --git a/tests/unit/data_plane/test_backend_config.py b/tests/unit/data_plane/test_backend_config.py new file mode 100644 index 00000000000..2054a2b2da9 --- /dev/null +++ b/tests/unit/data_plane/test_backend_config.py @@ -0,0 +1,120 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Resolution of the per-backend sizing block. + +``data_plane`` carries one block per backend (``simple:`` / ``mooncake_cpu:``) +and only the selected one is read, falling back to that backend's defaults when +absent. Getting this wrong would silently run a job at the wrong RDMA segment +size or with the staging pool off, neither of which fails loudly. +""" + +from __future__ import annotations + +import pydantic +import pytest +from pydantic import TypeAdapter + +from nemo_rl.data_plane.interfaces import ( + DataPlaneConfig, + MooncakeCpuConfig, + SimpleStorageConfig, + backend_config, +) + +_BASE = { + "enabled": True, + "impl": "transfer_queue", + "claim_meta_poll_interval_s": 0.5, +} + + +def _cfg(backend: str, **extra) -> dict: + return {**_BASE, "backend": backend, **extra} + + +def test_nested_block_is_used() -> None: + cfg = _cfg( + "mooncake_cpu", + mooncake_cpu={"global_segment_size": 111, "reuse_registered_buffers": False}, + ) + resolved = backend_config(cfg) + assert isinstance(resolved, MooncakeCpuConfig) + assert resolved.global_segment_size == 111 + assert resolved.reuse_registered_buffers is False + + +def test_absent_block_falls_back_to_model_defaults() -> None: + """The point of the nesting: a config need not mention a backend it isn't using. + + Pins the literals rather than comparing against MooncakeCpuConfig()'s own + attributes — that would hold for any value the class default was changed + to and couldn't catch a regression of the sizing itself. + """ + resolved = backend_config(_cfg("mooncake_cpu")) + assert resolved.global_segment_size == 68719476736 # 64 GiB per client process + assert resolved.local_buffer_size == 4294967296 # 4 GiB per client process + # The opt-out flag defaults on, so omitting it must not disable the pool. + assert resolved.reuse_registered_buffers is True + + +def test_accepts_an_already_coerced_model() -> None: + """Configs arriving via pydantic have the block coerced to a model already.""" + cfg = _cfg("mooncake_cpu", mooncake_cpu=MooncakeCpuConfig(global_segment_size=555)) + assert backend_config(cfg).global_segment_size == 555 + + +def test_partial_nested_block_keeps_other_defaults() -> None: + cfg = _cfg("mooncake_cpu", mooncake_cpu={"local_buffer_size": 7}) + resolved = backend_config(cfg) + assert resolved.local_buffer_size == 7 + assert resolved.global_segment_size == MooncakeCpuConfig().global_segment_size + + +def test_simple_backend_nested_block_is_used() -> None: + cfg = _cfg("simple", simple={"storage_capacity": 7, "num_storage_units": 3}) + resolved = backend_config(cfg) + assert isinstance(resolved, SimpleStorageConfig) + assert resolved.storage_capacity == 7 + assert resolved.num_storage_units == 3 + + +def test_simple_backend_num_storage_units_has_no_default() -> None: + """No static default is correct across cluster sizes (see the field's + docstring), so an absent/incomplete simple: block must raise rather than + silently run at a node count the config never chose.""" + with pytest.raises(pydantic.ValidationError, match="num_storage_units"): + backend_config(_cfg("simple")) + with pytest.raises(pydantic.ValidationError, match="num_storage_units"): + backend_config(_cfg("simple", simple={"storage_capacity": 7})) + + +def test_only_the_selected_backend_is_read() -> None: + """A mooncake block must not leak into a simple run, or vice versa.""" + cfg = _cfg( + "simple", + simple={"storage_capacity": 5, "num_storage_units": 1}, + mooncake_cpu={"global_segment_size": 999}, + ) + resolved = backend_config(cfg) + assert isinstance(resolved, SimpleStorageConfig) + assert not hasattr(resolved, "global_segment_size") + + +def test_schema_validates_without_any_backend_block() -> None: + """Regression guard: a required backend key is what broke SingleController CI. + + ``data_plane`` built from scratch — not inherited from the exemplar — must + validate, otherwise MasterConfig fails before training starts. + """ + TypeAdapter(DataPlaneConfig).validate_python(_cfg("simple")) diff --git a/tests/unit/data_plane/test_mooncake_staging_pool.py b/tests/unit/data_plane/test_mooncake_staging_pool.py new file mode 100644 index 00000000000..5d2cd96ff68 --- /dev/null +++ b/tests/unit/data_plane/test_mooncake_staging_pool.py @@ -0,0 +1,226 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""RDMA buffer-registration bookkeeping for the mooncake staging pool. + +Every failure guarded here reaches production as the same symptom — mooncake's +generic ``TRANSFER_FAIL`` (-800) out of ``batch_upsert_from``, carrying no root +cause and often on keys the pool never touched, because the damage is to +address-level registration state shared with the unpatched byte paths. None of +it needs an RDMA device to reproduce: the pool's contract is with +``register_buffer`` / ``unregister_buffer``, so a recording double is enough. +""" + +from __future__ import annotations + +import threading +import time + +import pytest + +from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter + + +class _FakeStore: + """Recording stand-in for ``MooncakeDistributedStore``. + + Tracks which addresses are currently registered so a test can assert the + invariant that actually matters: the pool never drops a buffer whose + memory mooncake still maps. + """ + + def __init__(self, fail_after: int | None = None) -> None: + # None: never fail. n: fail every registration after the first n. + self.fail_after = fail_after + self.registered: dict[int, int] = {} + self.register_calls = 0 + self.unregistered: list[int] = [] + + def register_buffer(self, ptr: int, nbytes: int) -> int: + self.register_calls += 1 + if self.fail_after is not None and self.register_calls > self.fail_after: + return -1 + self.registered[ptr] = nbytes + return 0 + + def unregister_buffer(self, ptr: int) -> int: + self.unregistered.append(ptr) + self.registered.pop(ptr, None) + return 0 + + +class _FakeClient: + """Just the attribute surface ``_StagingPoolRegistry.pool_for`` touches.""" + + def __init__(self, store: _FakeStore) -> None: + self._store = store + + +# ── register_buffer status checking ────────────────────────────────────────── + + +@pytest.mark.parametrize("status", [0, None], ids=["zero", "none"]) +def test_register_checked_accepts_success_statuses(status) -> None: + """``None`` must pass: the binding's return type varies across wheels.""" + store = _FakeStore() + store.register_buffer = lambda ptr, nbytes: status # type: ignore[method-assign] + tq_adapter._register_checked(store, 0x1000, 4096) + + +# Comfortably above every payload these tests stage, so the ceiling only +# matters in the test that sets it explicitly. +_MAX = 1 << 24 + + +def test_register_checked_raises_on_failed_registration() -> None: + """The whole point: fail at the registration, not three retries later.""" + store = _FakeStore(fail_after=0) + with pytest.raises(RuntimeError, match="register_buffer.*failed with status -1"): + tq_adapter._register_checked(store, 0x1000, 4096) + + +def test_register_all_buffers_patch_checks_upstream_call_site(monkeypatch) -> None: + """The patched path is the one in the -800 traceback. + + ``_put_bytes_thread_worker`` — non-tensor fields, untouched by the staging + pool — reaches ``register_buffer`` only through ``_register_all_buffers``, + which upstream calls for its side effect and never inspects. + """ + mc = pytest.importorskip("transfer_queue.storage.clients.mooncake_client") + cls = mc.MooncakeStoreClient + + # Restore both the method and the idempotence flag so the shared module is + # left exactly as found for other tests in this session. + monkeypatch.setattr(cls, "_register_all_buffers", cls._register_all_buffers) + monkeypatch.setattr(cls, "_nrl_register_checked", False, raising=False) + + tq_adapter._patch_mooncake_register_check() + + client = cls.__new__(cls) # __init__ would build a real mooncake store + client._store = _FakeStore(fail_after=1) + with pytest.raises(RuntimeError, match="register_buffer"): + client._register_all_buffers([0x1000, 0x2000], [4096, 4096]) + + +# ── pool slot bookkeeping ──────────────────────────────────────────────────── + + +def test_growing_a_slot_unregisters_before_dropping_the_old_buffer() -> None: + store = _FakeStore() + pool = tq_adapter._StagingPool(store, n_slots=1, max_bytes=_MAX) + + with pool.buffer(1024) as small: + small_ptr = small.data_ptr() + assert store.registered == {small_ptr: 1024} + + with pool.buffer(8 * 1024 * 1024) as big: + big_ptr = big.data_ptr() + + assert small_ptr in store.unregistered + assert store.registered == {big_ptr: 8 * 1024 * 1024} + + +def test_failed_growth_leaves_the_slot_empty_not_poisoned() -> None: + """A slot must never come back holding an unregistered buffer. + + Reusing one is the silent variant of this bug: every later transfer + through that slot writes into memory the NIC never mapped and returns + -800, which retrying cannot fix. + """ + store = _FakeStore(fail_after=0) + pool = tq_adapter._StagingPool(store, n_slots=1, max_bytes=_MAX) + + with pytest.raises(RuntimeError, match="register_buffer"): + with pool.buffer(1024): + pass + assert store.registered == {} + + store.fail_after = None + with pool.buffer(1024) as buf: + assert store.registered == {buf.data_ptr(): 1024} + + +def test_oversized_transfer_bypasses_the_pool_and_unregisters() -> None: + """Outliers get a transient registration; it must not outlive the call.""" + store = _FakeStore() + pool = tq_adapter._StagingPool(store, n_slots=1, max_bytes=4096) + + with pool.buffer(8192) as buf: + assert store.registered == {buf.data_ptr(): 8192} + oversized_ptr = buf.data_ptr() + + assert store.unregistered == [oversized_ptr] + assert store.registered == {} + + +def test_slot_exhaustion_fails_loudly_instead_of_hanging(monkeypatch) -> None: + """More concurrent transfers than slots must raise, not block forever.""" + monkeypatch.setattr(tq_adapter, "_STAGING_SLOT_TIMEOUT_S", 0.05) + pool = tq_adapter._StagingPool(_FakeStore(), n_slots=1, max_bytes=_MAX) + + with pool.buffer(1024): + with pytest.raises(RuntimeError, match="No mooncake staging slot free"): + with pool.buffer(1024): + pass + + +# ── lazy construction under concurrency ────────────────────────────────────── + + +def test_pool_is_constructed_once_under_concurrent_first_use(monkeypatch) -> None: + """The regression: a second pool's buffers are freed while still mapped. + + TQ submits one thread worker per ``BATCH_SIZE_LIMIT`` (400 keys) batch to + a shared ``ThreadPoolExecutor``, so a rollout step whose tensor fields + exceed that reaches a cold client from several threads at once. The loser + of an unsynchronized check-then-set has its pool overwritten and garbage + collected, freeing registered memory that the allocator immediately hands + to the next caller — including ``_put_bytes_thread_worker``'s receive + region, which is where the -800 surfaced. + + Slowing the constructor makes the interleaving deterministic rather than + relying on winning a race a fixed number of times. + """ + constructed: list[object] = [] + original_init = tq_adapter._StagingPool.__init__ + + def slow_init(self, store, n_slots, max_bytes): # type: ignore[no-untyped-def] + time.sleep(0.05) # widen the check-then-set window + original_init(self, store, n_slots, max_bytes) + constructed.append(self) + + monkeypatch.setattr(tq_adapter._StagingPool, "__init__", slow_init) + + # The production registry is a local of _patch_mooncake_staging_buffers, so + # build one here rather than reaching into the patch closure. + registry = tq_adapter._StagingPoolRegistry(4, _MAX) + client = _FakeClient(_FakeStore()) + n_threads = 8 + barrier = threading.Barrier(n_threads) + seen: list[object] = [] + + def worker() -> None: + barrier.wait() + seen.append(registry.pool_for(client)) + + threads = [threading.Thread(target=worker) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert len(constructed) == 1 + assert len(seen) == n_threads + # Identity, not storage location: every caller must get the one pool that + # was actually constructed. + assert seen == [constructed[0]] * n_threads diff --git a/tests/unit/data_plane/test_rdma_device_selection.py b/tests/unit/data_plane/test_rdma_device_selection.py new file mode 100644 index 00000000000..58eb9f69436 --- /dev/null +++ b/tests/unit/data_plane/test_rdma_device_selection.py @@ -0,0 +1,237 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Device selection for the mooncake transport: all IB rails, never RoCE +alongside them. A regression here still trains, just slower, so nothing +else would catch it. +""" + +import os +import sys + +import pytest + +from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter +from nemo_rl.data_plane.adapters import transfer_queue_env as tq_env + + +@pytest.fixture +def fake_fabric(monkeypatch): + """Install a synthetic device inventory. + + The sysfs scan lives in ``tq_env.rail_link_layers`` and the uverbs gate in + the adapter, but both reach it through the same stdlib ``glob`` module + object, so one patch covers both: the uverbs glob gates on device + availability and the link_layer glob enumerates. + """ + + def _install(layers: dict[str, str], *, uverbs: bool = True): + def fake_glob(pattern: str): + if pattern.startswith("/dev/infiniband/uverbs"): + return ["/dev/infiniband/uverbs0"] if uverbs else [] + return [f"/sys/class/infiniband/{d}/ports/1/link_layer" for d in layers] + + def fake_read_text(path, *args, **kwargs): + return layers[path.parents[2].name] + + monkeypatch.setattr(tq_env.glob, "glob", fake_glob) + monkeypatch.setattr(tq_env.Path, "read_text", fake_read_text) + monkeypatch.setattr(os, "environ", dict(os.environ)) + monkeypatch.delenv("MC_MOONCAKE_DEVICE", raising=False) + + return _install + + +# The real pool0 layout: eight 400 Gb/s IB rails plus one 100 Gb/s RoCE port. +_MIXED = { + "mlx5_0": "InfiniBand", + "mlx5_1": "InfiniBand", + "mlx5_2": "InfiniBand", + "mlx5_3": "Ethernet", + "mlx5_4": "InfiniBand", + "mlx5_5": "InfiniBand", + "mlx5_6": "InfiniBand", + "mlx5_7": "InfiniBand", + "mlx5_8": "InfiniBand", +} + + +def test_prefers_infiniband_and_excludes_roce(fake_fabric): + """The regression this guards: mlx5_3 was chosen over eight IB rails. + + Exact equality also pins the three things the format depends on: all + eight rails (not one), no space after the comma (mooncake splits on "," + only), and no RoCE device mixed in. + """ + fake_fabric(_MIXED) + assert ( + tq_adapter.rdma_devices() + == "mlx5_0,mlx5_1,mlx5_2,mlx5_4,mlx5_5,mlx5_6,mlx5_7,mlx5_8" + ) + + +def test_falls_back_to_roce_only_when_no_ib(fake_fabric): + fake_fabric({"mlx5_0": "Ethernet", "mlx5_1": "Ethernet"}) + assert tq_adapter.rdma_devices() == "mlx5_0,mlx5_1" + + +def test_every_roce_rail_is_offered(fake_fabric): + """The gb200 CI layout — 4 RoCE rails, two per NUMA domain — yields all 4. + + Correctness does not come from thinning this list; it comes from + MC_ENABLE_DEST_DEVICE_AFFINITY pinning each transfer's peer rail to the + local one. Keeping one rail per domain does NOT make this safe: the pair + that survives such a filter (mlx5_0, mlx5_2) is itself cross-rail and was + measured failing on that fleet. + """ + fake_fabric( + { + "mlx5_0": "Ethernet", + "mlx5_1": "Ethernet", + "mlx5_2": "Ethernet", + "mlx5_3": "Ethernet", + } + ) + assert tq_adapter.rdma_devices() == "mlx5_0,mlx5_1,mlx5_2,mlx5_3" + + +def test_empty_without_verbs_node(fake_fabric): + """Containers see /sys without /dev/infiniband; mooncake fails late there.""" + fake_fabric(_MIXED, uverbs=False) + assert tq_adapter.rdma_devices() == "" + + +def test_env_override_wins_verbatim(fake_fabric, monkeypatch): + fake_fabric(_MIXED) + monkeypatch.setenv("MC_MOONCAKE_DEVICE", "mlx5_9,mlx5_10") + assert tq_adapter.rdma_devices() == "mlx5_9,mlx5_10" + + +def test_transport_config_is_rdma_and_carries_all_rails(fake_fabric): + """The device list must reach mooncake, and the transport stays RDMA.""" + fake_fabric(_MIXED) + cfg = tq_adapter._mooncake_transport_config() + assert cfg["protocol"] == "rdma" + assert cfg["device_name"] == tq_adapter.rdma_devices() + + +def test_raises_when_no_device_since_mooncake_is_rdma_only(fake_fabric): + fake_fabric(_MIXED, uverbs=False) + with pytest.raises(RuntimeError, match="requires RDMA"): + tq_adapter._mooncake_transport_config() + + +# ── Peer-rail pairing ──────────────────────────────────────────────────────── +# +# Mooncake picks the peer rail at random unless told otherwise. Where each rail +# is its own subnet (the RoCE-only gb200 CI runners) a cross-rail pair has no +# route, which was 100% of the failures observed there. + + +def _mooncake_cfg() -> dict: + return { + "enabled": True, + "impl": "transfer_queue", + "backend": "mooncake_cpu", + "claim_meta_poll_interval_s": 0.5, + } + + +@pytest.fixture +def clean_env(monkeypatch): + """Isolate os.environ and pretend the engine has not been imported yet. + + The real ``sys.modules`` always has ``transfer_queue`` in it here — this + test module imports the adapter — so the "not yet imported" case has to be + injected rather than arranged. + """ + monkeypatch.setattr(os, "environ", dict(os.environ)) + monkeypatch.delenv("MC_ENABLE_DEST_DEVICE_AFFINITY", raising=False) + monkeypatch.delenv("MC_STORE_MEMCPY", raising=False) + monkeypatch.setattr(tq_env, "_engine_already_imported", lambda: None) + + +@pytest.fixture +def engine_imported(clean_env, monkeypatch): + """Flip ``clean_env``'s verdict: pretend the engine is already loaded.""" + monkeypatch.setattr(tq_env, "_engine_already_imported", lambda: "transfer_queue") + + +def test_affinity_pinned_on_a_roce_only_fabric(clean_env, fake_fabric): + """Same-rail pairing is what makes offering every rail safe.""" + fake_fabric({"mlx5_0": "Ethernet", "mlx5_1": "Ethernet"}) + tq_env.configure_engine_env(_mooncake_cfg()) + assert os.environ["MC_ENABLE_DEST_DEVICE_AFFINITY"] == "1" + + +def test_affinity_left_alone_on_infiniband(clean_env, fake_fabric): + """IB routes cross-rail, so the hint is not our call to make there. + + Scoped deliberately: the cross-rail failure was only ever measured on RoCE, + and this cluster cannot test IB. + """ + fake_fabric({"mlx5_0": "InfiniBand", "mlx5_1": "InfiniBand"}) + tq_env.configure_engine_env(_mooncake_cfg()) + assert "MC_ENABLE_DEST_DEVICE_AFFINITY" not in os.environ + + +def test_roce_gate_does_not_fail_open_when_sysfs_is_empty(clean_env, fake_fabric): + """No rails at all must not read as "InfiniBand, skip the hint".""" + fake_fabric({}) + assert tq_env.fabric_is_roce_only() is False + + +def test_existing_value_is_not_clobbered(clean_env, fake_fabric): + """A launcher-supplied value wins; we only fill the gap.""" + fake_fabric({"mlx5_0": "Ethernet"}) + os.environ["MC_ENABLE_DEST_DEVICE_AFFINITY"] = "0" + tq_env.configure_engine_env(_mooncake_cfg()) + assert os.environ["MC_ENABLE_DEST_DEVICE_AFFINITY"] == "0" + + +def test_not_applied_to_simple_backend(clean_env, fake_fabric): + """The knob is mooncake-only; `simple` never touches RDMA.""" + fake_fabric({"mlx5_0": "Ethernet"}) + tq_env.configure_engine_env({**_mooncake_cfg(), "backend": "simple"}) + assert "MC_ENABLE_DEST_DEVICE_AFFINITY" not in os.environ + + +def test_raises_when_the_engine_was_already_imported(engine_imported, fake_fabric): + """The whole point of the split: too-late must be loud, not silent. + + Mooncake snapshots MC_* as its extension loads, so a value set after that + reads back fine from os.environ while the engine ignores it — which is how + this cost several CI runs before being caught. + """ + fake_fabric({"mlx5_0": "Ethernet"}) + with pytest.raises(RuntimeError, match="already imported"): + tq_env.configure_engine_env(_mooncake_cfg()) + + +def test_no_raise_when_already_set_even_if_engine_imported( + engine_imported, fake_fabric +): + """The normal worker path: Ray handed down the driver's environment.""" + fake_fabric({"mlx5_0": "Ethernet"}) + os.environ["MC_ENABLE_DEST_DEVICE_AFFINITY"] = "1" + os.environ["MC_STORE_MEMCPY"] = "0" + tq_env.configure_engine_env(_mooncake_cfg()) # must not raise + + +def test_engine_already_imported_detects_a_loaded_module(monkeypatch): + """Pin the detector itself, since every other test stubs it out.""" + monkeypatch.setitem(sys.modules, "transfer_queue", object()) + assert tq_env._engine_already_imported() == "transfer_queue" + for name in tq_env._ENGINE_MODULES: + monkeypatch.delitem(sys.modules, name, raising=False) + assert tq_env._engine_already_imported() is None diff --git a/tests/unit/data_plane/test_tq_lifecycle.py b/tests/unit/data_plane/test_tq_lifecycle.py index 3e79a50ff84..4ab3b51abab 100644 --- a/tests/unit/data_plane/test_tq_lifecycle.py +++ b/tests/unit/data_plane/test_tq_lifecycle.py @@ -49,20 +49,42 @@ def fake_clear(**kwargs): monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", fake_put) monkeypatch.setattr(tq_adapter.tq, "kv_clear", fake_clear) + # bootstrap=False only connects to an existing controller; stubbing that + # lets the real __init__ run, so this test cannot drift from it. + monkeypatch.setattr(tq_adapter, "_connect_existing", lambda: None) - client = object.__new__(tq_adapter.TQDataPlaneClient) + client = tq_adapter.TQDataPlaneClient( + { + "enabled": True, + "impl": "transfer_queue", + "backend": "simple", + "claim_meta_poll_interval_s": 0.5, + }, + bootstrap=False, + ) client.register_partition( partition_id="obj-backend", fields=["msg_log"], num_samples=8, consumer_tasks=["read"], ) + # Same fields again: already warmed, so no second put/clear. client.register_partition( partition_id="obj-backend", fields=["msg_log"], num_samples=8, consumer_tasks=["read"], ) + assert len(put_calls) == 1 + + # A genuinely new field warms only that field, under a fresh key -- + # mooncake has no upsert, so a reused key would hit stale metadata. + client.register_partition( + partition_id="obj-backend", + fields=["msg_log", "rewards"], + num_samples=8, + consumer_tasks=["read"], + ) assert len(put_calls) == 2 schema_keys = [call["keys"][0] for call in put_calls] @@ -74,7 +96,7 @@ def fake_clear(**kwargs): ] assert [list(call["fields"].keys()) for call in put_calls] == [ ["msg_log"], - ["msg_log"], + ["rewards"], ] assert clear_calls == [ {"keys": [schema_keys[0]], "partition_id": "obj-backend"}, diff --git a/tests/unit/distributed/test_virtual_cluster.py b/tests/unit/distributed/test_virtual_cluster.py index 72bb2adc3ee..6df9b2de683 100644 --- a/tests/unit/distributed/test_virtual_cluster.py +++ b/tests/unit/distributed/test_virtual_cluster.py @@ -211,6 +211,72 @@ def test_ray_uses_same_cluster_for_permuted_cuda_devices(): assert mock_ray_shutdown.call_count == 0 +def test_maybe_configure_data_plane_env_then_init_ray_threads_env_vars(): + """This is the pattern every data-plane-enabled launcher uses -- + maybe_configure_data_plane_env(config.data_plane) immediately before + init_ray() -- because init_ray's env_vars snapshot is the only point a + backend engine knob becomes cluster-wide (see both functions' + docstrings). init_ray itself has no data-plane awareness; the ordering + at the call site is what makes this work, so exercise the two calls + together rather than mocking either one out. + + The fabric probe and the already-imported guard ARE mocked: both read + process/host state (sysfs, sys.modules) that would otherwise make this + test's outcome depend on the machine it runs on and on whether another + test already imported the transfer_queue adapter. + """ + from nemo_rl.data_plane.factory import maybe_configure_data_plane_env + from nemo_rl.distributed.virtual_cluster import init_ray + + env_mod = "nemo_rl.data_plane.adapters.transfer_queue_env" + with ( + patch("ray.init") as mock_ray_init, + patch("ray.cluster_resources") as mock_cluster_resources, + patch(f"{env_mod}.fabric_is_roce_only", return_value=True), + patch(f"{env_mod}._engine_already_imported", return_value=None), + ): + # Matching tag -> init_ray takes the "reuse existing cluster" path + # and returns after exactly one ray.init call, the one whose + # runtime_env we need to inspect. + mock_cluster_resources.return_value = {"nrl_tag_0": 1} + env = {"CUDA_VISIBLE_DEVICES": "0"} + with patch.dict(os.environ, env, clear=True): + maybe_configure_data_plane_env( + { + "enabled": True, + "impl": "transfer_queue", + "backend": "mooncake_cpu", + "claim_meta_poll_interval_s": 0.5, + } + ) + init_ray() + + assert mock_ray_init.call_count == 1 + env_vars = mock_ray_init.call_args_list[0][1]["runtime_env"]["env_vars"] + assert env_vars["MC_STORE_MEMCPY"] == "0" + assert env_vars["MC_ENABLE_DEST_DEVICE_AFFINITY"] == "1" + + +def test_init_ray_alone_has_no_data_plane_awareness(): + """Every non-data-plane launcher's call (bare init_ray(), no preceding + maybe_configure_data_plane_env) must not touch mooncake env vars -- + init_ray does not know the data plane exists.""" + with ( + patch("ray.init") as mock_ray_init, + patch("ray.cluster_resources") as mock_cluster_resources, + ): + mock_cluster_resources.return_value = {"nrl_tag_0": 1} + env = {"CUDA_VISIBLE_DEVICES": "0"} + with patch.dict(os.environ, env, clear=True): + from nemo_rl.distributed.virtual_cluster import init_ray + + init_ray() + + env_vars = mock_ray_init.call_args_list[0][1]["runtime_env"]["env_vars"] + assert "MC_STORE_MEMCPY" not in env_vars + assert "MC_ENABLE_DEST_DEVICE_AFFINITY" not in env_vars + + def test_mcore_py_executable(): # The temporary directory is created within the project. # For some reason, creating a virtual environment outside of the project diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 5fd3709096a..f557a085e0c 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -515,11 +515,15 @@ data_plane: enabled: false impl: transfer_queue backend: "simple" # TQ storage backend ('simple' or 'mooncake_cpu') - storage_capacity: 1000000 # max samples retained per partition - num_storage_units: 2 # storage shards claim_meta_poll_interval_s: 0.5 # blocking-claim poll cadence - global_segment_size: 549755813888 # 512 GiB — used when backend == "mooncake_cpu" - local_buffer_size: 68719476736 # 64 GiB — used when backend == "mooncake_cpu" + simple: + storage_capacity: 1000000 # max samples retained per partition + num_storage_units: ${mul:2, ${cluster.num_nodes}} # TQ wants >= 2 per node + mooncake_cpu: + global_segment_size: 68719476736 # 64 GiB/process + local_buffer_size: 4294967296 # 4 GiB/process + reuse_registered_buffers: true # reuse RDMA-registered buffers + staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool # observability: # NotRequired # enabled: false diff --git a/tests/unit/single_controller/test_run_grpo_single_controller.py b/tests/unit/single_controller/test_run_grpo_single_controller.py index c0b5371a57e..5b5d0330276 100644 --- a/tests/unit/single_controller/test_run_grpo_single_controller.py +++ b/tests/unit/single_controller/test_run_grpo_single_controller.py @@ -35,7 +35,7 @@ def main_context(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: "megatron_cfg": {"mtp_num_layers": 2}, }, env={}, - data_plane={"enabled": True}, + data_plane={"enabled": True, "impl": "transfer_queue", "backend": "simple"}, logger={"log_dir": "/tmp/logs"}, checkpointing={"enabled": False}, async_rl=SimpleNamespace( diff --git a/tests/unit/single_controller/test_train_pump.py b/tests/unit/single_controller/test_train_pump.py index 1df434d9860..e68a0546643 100644 --- a/tests/unit/single_controller/test_train_pump.py +++ b/tests/unit/single_controller/test_train_pump.py @@ -69,11 +69,12 @@ def _simple_tq_cfg() -> dict: "enabled": True, "impl": "transfer_queue", "backend": "simple", - "storage_capacity": 1024, - "num_storage_units": 1, "claim_meta_poll_interval_s": 0.5, - "global_segment_size": 8589934592, # 8 GiB - "local_buffer_size": 1073741824, # 1 GiB + "simple": {"storage_capacity": 1024, "num_storage_units": 1}, + "mooncake_cpu": { + "global_segment_size": 8589934592, # 8 GiB + "local_buffer_size": 1073741824, # 1 GiB + }, }