feat(data_plane): mooncake cpu rdma backend

.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 \

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
 

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

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

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"])

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)

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"):

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.

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.