From 681b80da4aa42bfc71cf60321bb0e9b685ba973a Mon Sep 17 00:00:00 2001 From: Youngeun Kwon Date: Wed, 8 Jul 2026 22:39:15 -0700 Subject: [PATCH] feat: Numa aware binding (#2613) Signed-off-by: Youngeun Kwon Signed-off-by: Ananth Subramaniam Signed-off-by: Terry Kong Co-authored-by: Ananth Subramaniam Co-authored-by: Terry Kong (cherry picked from commit 9e01af64b3891e5bcc01885e10e9ca185b3e3690) --- nemo_rl/distributed/numa_utils.py | 181 ++++++ .../models/generation/vllm/vllm_backend.py | 19 + nemo_rl/models/generation/vllm/vllm_worker.py | 15 + .../generation/vllm/vllm_worker_async.py | 2 + .../policy/workers/dtensor_policy_worker.py | 8 + .../workers/dtensor_policy_worker_v2.py | 8 + .../policy/workers/megatron_policy_worker.py | 8 + .../value/workers/dtensor_value_worker_v2.py | 8 + .../value/workers/megatron_value_worker.py | 9 + pyrefly.toml | 1 + ray.sub | 28 + tests/unit/distributed/test_numa_utils.py | 519 ++++++++++++++++++ 12 files changed, 806 insertions(+) create mode 100644 nemo_rl/distributed/numa_utils.py create mode 100644 tests/unit/distributed/test_numa_utils.py diff --git a/nemo_rl/distributed/numa_utils.py b/nemo_rl/distributed/numa_utils.py new file mode 100644 index 00000000000..550bad710d7 --- /dev/null +++ b/nemo_rl/distributed/numa_utils.py @@ -0,0 +1,181 @@ +# 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. + +"""NUMA-aware CPU affinity and memory binding for GPU workers. + +Uses a GPU→cpulist mapping file written by topology_probe.sh (in ray.sub) +at node startup. The file path is communicated via the NRL_GPU_CPU_AFFINITY_FILE +environment variable. See ray.sub for the writer side. + +Disable all binding with NRL_DISABLE_NUMA_BINDING=1. +Disable only memory policy with NRL_DISABLE_NUMA_MEMBIND=1. +""" + +import ctypes +import ctypes.util +import logging +import os + +logger = logging.getLogger(__name__) + +# IMPORTANT: This default path must stay in sync with topology_probe.sh in ray.sub. +# The canonical path is set via the NRL_GPU_CPU_AFFINITY_FILE env var exported by ray.sub. +GPU_CPU_AFFINITY_PATH = os.environ.get( + "NRL_GPU_CPU_AFFINITY_FILE", "/tmp/nrl_gpu_cpu_affinity" +) + + +def bind_to_gpu_numa(gpu_id: int) -> bool: + """Pin the current process to the NUMA-local CPUs and memory of the given GPU. + + Reads the GPU→cpulist mapping written by topology_probe.sh at node + startup, then calls os.sched_setaffinity() for CPU pinning and + numa_set_membind() for memory policy. Best-effort: failures are + logged, never raised. + + Args: + gpu_id: Node-global physical GPU index (``nvidia-smi`` numbering), which + is how the affinity file is keyed. Passed explicitly because + ``CUDA_VISIBLE_DEVICES`` lists all devices on the node under + ``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1`` and so does not + identify a single worker's GPU. In a Ray actor this is + ``int(ray.get_gpu_ids()[0])``. + + Returns True if CPU binding succeeded, False if skipped or failed. + Memory binding is attempted independently and logged separately. + """ + if os.environ.get("NRL_DISABLE_NUMA_BINDING") == "1": + return False + + gpu = str(gpu_id) + try: + with open(GPU_CPU_AFFINITY_PATH) as f: + for line in f: + line = line.strip() + if not line: + continue + idx, cpulist = line.split(":", 1) + if idx == gpu: + cpus = _parse_cpulist(cpulist) + os.sched_setaffinity(0, cpus) + logger.info("NUMA CPU binding: GPU %s → CPUs %s", gpu, cpulist) + _set_numa_membind(cpus) + return True + logger.debug("NUMA binding: GPU %s not found in %s", gpu, GPU_CPU_AFFINITY_PATH) + except FileNotFoundError: + logger.debug("NUMA binding skipped: %s not found", GPU_CPU_AFFINITY_PATH) + except Exception as exc: + logger.debug("NUMA binding skipped: %s", exc) + return False + + +def resolve_visible_gpu_id(local_index: int) -> int | None: + """Map a process-local CUDA device index to its node-global physical GPU id. + + ``CUDA_VISIBLE_DEVICES`` lists the physical GPU ids visible to this process + in device-index order, and ``local_index`` (e.g. + ``torch.cuda.current_device()``) indexes into that list. The affinity file is + keyed by the physical id, so return ``CUDA_VISIBLE_DEVICES[local_index]``. + + ``CUDA_VISIBLE_DEVICES`` contents depend on the worker: + - vLLM TP>1 (``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1``): the + per-instance device subset, e.g. ``"4,5"``. + - vLLM TP=1: a single isolated device, so ``local_index`` is 0. + + Returns the physical GPU id, or None if it cannot be resolved (unset CVD, + index out of range, or non-integer entries such as MIG UUIDs). + """ + cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not cvd: + return None + devices = cvd.split(",") + if local_index < 0 or local_index >= len(devices): + return None + try: + return int(devices[local_index]) + except ValueError: + return None + + +def _load_libnuma() -> ctypes.CDLL | None: + """Load libnuma, returning None if unavailable.""" + try: + return ctypes.CDLL("libnuma.so.1") + except OSError: + return None + + +def _get_numa_node(libnuma: ctypes.CDLL, cpus: set[int]) -> int: + """Return the NUMA node for the given CPU set, or -1 on failure.""" + libnuma.numa_node_of_cpu.restype = ctypes.c_int + return libnuma.numa_node_of_cpu(min(cpus)) + + +def _set_numa_membind(cpus: set[int]) -> bool: + """Hard-bind memory allocations to the NUMA node of the given CPUs.""" + if os.environ.get("NRL_DISABLE_NUMA_MEMBIND") == "1": + return False + + libnuma = _load_libnuma() + if libnuma is None: + logger.debug("NUMA membind skipped: libnuma.so.1 not available") + return False + + try: + numa_node = _get_numa_node(libnuma, cpus) + if numa_node < 0: + logger.debug( + "NUMA membind skipped: numa_node_of_cpu(%d) returned %d", + min(cpus), + numa_node, + ) + return False + + libnuma.numa_allocate_nodemask.restype = ctypes.c_void_p + libnuma.numa_bitmask_setbit.argtypes = [ctypes.c_void_p, ctypes.c_uint] + libnuma.numa_bitmask_setbit.restype = ctypes.c_void_p + libnuma.numa_set_membind.argtypes = [ctypes.c_void_p] + libnuma.numa_bitmask_free.argtypes = [ctypes.c_void_p] + + nodemask = libnuma.numa_allocate_nodemask() + if not nodemask: + logger.debug("NUMA membind skipped: numa_allocate_nodemask returned NULL") + return False + + try: + libnuma.numa_bitmask_setbit(nodemask, numa_node) + libnuma.numa_set_membind(nodemask) + finally: + libnuma.numa_bitmask_free(nodemask) + + logger.info( + "NUMA membind: hard-bound to node %d (from CPU %d)", numa_node, min(cpus) + ) + return True + except Exception as exc: + logger.debug("NUMA membind skipped: %s", exc) + return False + + +def _parse_cpulist(cpulist: str) -> set[int]: + """Parse a Linux cpulist string like '0-71' into a set of ints.""" + cpus: set[int] = set() + for part in cpulist.split(","): + part = part.strip() + if "-" in part: + lo, hi = part.split("-", 1) + cpus.update(range(int(lo), int(hi) + 1)) + else: + cpus.add(int(part)) + return cpus diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index e2c7452da37..37acf8230a5 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -92,6 +92,25 @@ def _read_mtp_layer_weights_from_checkpoint( class VllmInternalWorkerExtension: + def bind_numa(self) -> bool: + """Pin this TP worker to its GPU's NUMA-local CPUs/memory. + + Invoked via ``collective_rpc`` on each vLLM TP worker once the engine + (and CUDA) is up, so the worker's physical GPU id is resolved from its + local device index (see ``resolve_visible_gpu_id``). + """ + import torch + + from nemo_rl.distributed.numa_utils import ( + bind_to_gpu_numa, + resolve_visible_gpu_id, + ) + + gpu_id = resolve_visible_gpu_id(torch.cuda.current_device()) + if gpu_id is None: + return False + return bind_to_gpu_numa(gpu_id) + def init_collective( self, rank_prefix: int, diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index 0d6a68b9f30..cd5df87ce93 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -182,6 +182,19 @@ def __init__( _load_model() later to perform the heavy model loading. This enables overlapping vLLM model loading with NeMo Gym init. """ + from nemo_rl.distributed.numa_utils import bind_to_gpu_numa + + # Only bind single-GPU workers to their GPU's NUMA node. + # For TP>1 workers, the parent process spans multiple NUMA nodes; + # binding it would incorrectly constrain the EngineCore subprocess + # (which inherits sched_setaffinity + numa_set_membind via fork). + # Individual TP workers get their own NUMA binding via collective_rpc + # in post_init / post_init_async. + # ray.get_gpu_ids()[0] is this worker's physical GPU index, which keys + # the affinity file. + if bundle_indices is not None and len(bundle_indices) == 1: + bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + self._init_config( config, bundle_indices, fraction_of_gpus, seed, extra_env_vars ) @@ -589,6 +602,8 @@ def _create_engine(self, llm_kwargs: dict[str, Any]) -> None: self.llm = vllm.LLM(**llm_kwargs) def post_init(self): + if self.llm is not None: + self.llm.collective_rpc("bind_numa", args=tuple()) self.vllm_device_ids = self.report_device_id() if self._mtp_load_from_disk: self.llm.collective_rpc( diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index fab2e1330e8..e7d2b1eaa6b 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -422,6 +422,8 @@ def clear_vllm_logger_metrics(self) -> None: self.generation_tokens = [] async def post_init_async(self): + if self.llm is not None: + await self.llm.collective_rpc("bind_numa", args=tuple()) self.vllm_device_ids = await self.report_device_id_async() if self._mtp_load_from_disk: await self.llm.collective_rpc( diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker.py b/nemo_rl/models/policy/workers/dtensor_policy_worker.py index 6f3b821f276..7e487cbc3ea 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker.py @@ -200,6 +200,14 @@ def __init__( **kwargs: Any, ): """Initialize the DTensorPolicyWorker.""" + from nemo_rl.distributed.numa_utils import bind_to_gpu_numa + + # Pin to this worker's GPU-local CPUs/memory before CUDA init or model + # load; FSDP's D2H paths (weight refit, optimizer/checkpoint offload) + # benefit. ray.get_gpu_ids()[0] is the physical GPU index that keys the + # affinity file, and reading it does not initialize CUDA. + bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + self.tokenizer = tokenizer self.processor = processor self.is_vlm = processor is not None diff --git a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py index 5f0530fc602..944fedc5234 100644 --- a/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py +++ b/nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py @@ -231,6 +231,14 @@ def __init__( # Apply TE patch until TE is upgraded to 2.10.0 apply_transformer_engine_patch() + from nemo_rl.distributed.numa_utils import bind_to_gpu_numa + + # Pin to this worker's GPU-local CPUs/memory before model load; FSDP's + # D2H paths (weight refit, optimizer/checkpoint offload) benefit. + # ray.get_gpu_ids()[0] is the physical GPU index that keys the affinity + # file, and reading it does not initialize CUDA. + bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + # Store configuration self.cfg = config diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index 7e3ef83a1c2..9ead07baf36 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -268,6 +268,14 @@ def __init__( # Apply patch from https://github.com/NVIDIA/TransformerEngine/pull/2286/files apply_transformer_engine_patch() + from nemo_rl.distributed.numa_utils import bind_to_gpu_numa + + # local_rank (== ray.get_gpu_ids()[0]) is the physical GPU index that + # keys the affinity file. Pass it explicitly: CUDA_VISIBLE_DEVICES lists + # all node devices here (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1, + # set by configure_worker), so it can't identify this worker's GPU. + bind_to_gpu_numa(local_rank) + self.cfg = config self._router_replay_enabled = router_replay_enabled(config) diff --git a/nemo_rl/models/value/workers/dtensor_value_worker_v2.py b/nemo_rl/models/value/workers/dtensor_value_worker_v2.py index b5870140310..0ad1b4fedf7 100644 --- a/nemo_rl/models/value/workers/dtensor_value_worker_v2.py +++ b/nemo_rl/models/value/workers/dtensor_value_worker_v2.py @@ -148,6 +148,14 @@ def __init__( # Apply patches apply_transformer_engine_patch() + from nemo_rl.distributed.numa_utils import bind_to_gpu_numa + + # Pin to this worker's GPU-local CPUs/memory before model load; FSDP's + # D2H paths (weight refit, optimizer/checkpoint offload) benefit. + # ray.get_gpu_ids()[0] is the physical GPU index that keys the affinity + # file, and reading it does not initialize CUDA. + bind_to_gpu_numa(int(ray.get_gpu_ids()[0])) + # Store configuration and tokenizer self.cfg = config self.tokenizer = tokenizer diff --git a/nemo_rl/models/value/workers/megatron_value_worker.py b/nemo_rl/models/value/workers/megatron_value_worker.py index 22d0f489f77..b4fd1c36737 100644 --- a/nemo_rl/models/value/workers/megatron_value_worker.py +++ b/nemo_rl/models/value/workers/megatron_value_worker.py @@ -395,6 +395,15 @@ def __init__( apply_transformer_engine_patch() + from nemo_rl.distributed.numa_utils import bind_to_gpu_numa + + # Pin to this worker's GPU-local CPUs/memory before model load, matching + # the policy workers. local_rank (== ray.get_gpu_ids()[0]) is the physical + # GPU index that keys the affinity file. Pass it explicitly: + # CUDA_VISIBLE_DEVICES lists all node devices here + # (RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1). + bind_to_gpu_numa(local_rank) + self.cfg = config self.rank = get_rank_safe() diff --git a/pyrefly.toml b/pyrefly.toml index 2db7ccb5c19..de05ad5a3e5 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -115,6 +115,7 @@ project-includes = [ "nemo_rl/distributed/__init__.py", "nemo_rl/distributed/collectives.py", "nemo_rl/distributed/named_sharding.py", + "nemo_rl/distributed/numa_utils.py", "nemo_rl/distributed/ray_actor_environment_registry.py", "nemo_rl/distributed/virtual_cluster.py", "nemo_rl/distributed/worker_group_utils.py", diff --git a/ray.sub b/ray.sub index fb407abc5e6..7237cbe1ada 100644 --- a/ray.sub +++ b/ray.sub @@ -387,6 +387,34 @@ else fi fi +# Write GPU→cpulist mapping for NUMA binding. +# IMPORTANT: The env var name NRL_GPU_CPU_AFFINITY_FILE and default path must stay +# in sync with GPU_CPU_AFFINITY_PATH in nemo_rl/distributed/numa_utils.py. +export NRL_GPU_CPU_AFFINITY_FILE="/tmp/nrl_gpu_cpu_affinity" +# nvidia-smi topo -m's CPU Affinity column is unreliable on GB200 (empty for +# GPUs not directly attached to the socket). Use NUMA Affinity (always +# populated, at field NF-1 since GPU NUMA ID is last) and look up the +# node-local CPU list from sysfs. On GB200 the NUMA Affinity column can be a +# list like "0,2-17" (the GPU-local CPU NUMA node plus the GPU's HBM NUMA +# nodes); take the first entry, which is the local CPU NUMA node. +nvidia-smi topo -m 2>/dev/null | awk '/^GPU[0-9]/ { + gpu = \$1; sub(/GPU/, "", gpu) + numa = \$(NF-1) + sub(/[,-].*/, "", numa) + if (numa ~ /^[0-9]+\$/) print gpu, numa +}' | while read -r _gpu _numa; do + _cpulist=\$(cat "/sys/devices/system/node/node\${_numa}/cpulist" 2>/dev/null) + if [[ -n "\$_cpulist" ]]; then + echo "\${_gpu}:\${_cpulist}" + fi +done > "\$NRL_GPU_CPU_AFFINITY_FILE" || true +if [[ -s "\$NRL_GPU_CPU_AFFINITY_FILE" ]]; then + echo "NUMA affinity map written to \$NRL_GPU_CPU_AFFINITY_FILE:" + cat "\$NRL_GPU_CPU_AFFINITY_FILE" +else + echo "WARNING: Could not generate NUMA affinity map (nvidia-smi topo unavailable)" +fi + # Use \\\" so that when --resources="\$RAY_RESOURCES" expands, we pass valid JSON to ray # IMPORTANT: The key names "nvlink_domain_" and "topo_rank" below must stay in sync # with the constants NVLINK_DOMAIN_PREFIX and TOPO_RANK_KEY defined in diff --git a/tests/unit/distributed/test_numa_utils.py b/tests/unit/distributed/test_numa_utils.py new file mode 100644 index 00000000000..2df58d706a5 --- /dev/null +++ b/tests/unit/distributed/test_numa_utils.py @@ -0,0 +1,519 @@ +# 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. +"""Tests for NUMA-aware CPU affinity and memory binding. + +Unit tests run anywhere (no GPU needed). The benchmark test +(TestNUMABindingBenchmark) requires a GPU and libnuma and is +skipped otherwise — run it inside the RL container on a DGX/GB200 +to validate that binding actually improves D2H performance. +""" + +import os +import subprocess +import tempfile + +import pytest + +from nemo_rl.distributed.numa_utils import ( + _get_numa_node, + _load_libnuma, + _parse_cpulist, + _set_numa_membind, + bind_to_gpu_numa, + resolve_visible_gpu_id, +) + +# --------------------------------------------------------------------------- +# Test helpers +# --------------------------------------------------------------------------- + + +def _write_affinity_file_from_topo() -> str | None: + """Parse nvidia-smi topo and write a GPU→cpulist affinity file. + + Mirrors the production probe in ray.sub: the "CPU Affinity" column is + unreliable on GB200 (empty for GPUs not directly attached to a socket), so + ray.sub reads the "NUMA Affinity" column and looks up the node-local CPU + list from sysfs. We do the same here so the benchmark can't drift from + production: + + nvidia-smi topo -m | awk '/^GPU[0-9]/ { numa=$(NF-1); ... }' + → cat /sys/devices/system/node/node/cpulist + + nvidia-smi topo -m output format (GB200 NVL72 example): + + GPU0 ... NIC5 CPU Affinity NUMA Affinity GPU NUMA ID + GPU0 X ... SYS 0-71 0 N/A + GPU2 ... NODE 72-143 1 N/A + NIC0 ... X + ... + + The header row is indented (starts with whitespace) and NIC rows start with + "NIC", so startswith("GPU") skips both. On a GPU row the trailing + whitespace-split columns are: + parts[-3] = CPU Affinity (unreliable on GB200 — NOT used) + parts[-2] = NUMA Affinity (e.g. "0"/"1") ← used; matches ray.sub's $(NF-1) + parts[-1] = GPU NUMA ID (e.g. "N/A") + + The resulting affinity file has one line per GPU, e.g. "2:72-143". + + Returns the path to the temp file, or None if parsing fails. + """ + try: + output = subprocess.run( + ["nvidia-smi", "topo", "-m"], + capture_output=True, + text=True, + timeout=10, + ).stdout + except (FileNotFoundError, subprocess.TimeoutExpired): + return None + + entries: list[str] = [] + for line in output.splitlines(): + if not line.startswith("GPU"): + continue + parts = line.split() + gpu_idx = parts[0].replace("GPU", "") + # NUMA Affinity column (ray.sub parses $(NF-1)). On GB200 this is a list + # like "0,2-17" (the GPU-local CPU NUMA node plus the GPU's HBM NUMA + # nodes); take the first entry, which is the local CPU NUMA node. + numa = parts[-2].split(",")[0].split("-")[0] + if not numa.isdigit(): + continue + try: + with open(f"/sys/devices/system/node/node{numa}/cpulist") as nf: + cpulist = nf.read().strip() + except OSError: + continue + if cpulist: + entries.append(f"{gpu_idx}:{cpulist}") + + if not entries: + return None + + f = tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) + f.write("\n".join(entries) + "\n") + f.flush() + f.close() + return f.name + + +def _reset_membind_to_all() -> None: + """Reset NUMA memory policy to allow allocations from all nodes. + + Used to establish the "unbound" baseline in benchmarks. In production + code we only bind (never unbind), so this helper is test-only. + """ + libnuma = _load_libnuma() + if libnuma is None: + return + + import ctypes + + libnuma.numa_set_membind.argtypes = [ctypes.c_void_p] + libnuma.numa_allocate_nodemask.restype = ctypes.c_void_p + libnuma.numa_bitmask_setbit.argtypes = [ctypes.c_void_p, ctypes.c_uint] + libnuma.numa_bitmask_setbit.restype = ctypes.c_void_p + libnuma.numa_bitmask_free.argtypes = [ctypes.c_void_p] + libnuma.numa_max_node.restype = ctypes.c_int + + max_node = libnuma.numa_max_node() + nodemask = libnuma.numa_allocate_nodemask() + for n in range(max_node + 1): + libnuma.numa_bitmask_setbit(nodemask, n) + libnuma.numa_set_membind(nodemask) + libnuma.numa_bitmask_free(nodemask) + + +def _reset_all_bindings() -> None: + """Reset both CPU affinity and memory policy to unbound defaults.""" + os.sched_setaffinity(0, set(range(os.cpu_count() or 256))) + _reset_membind_to_all() + + +def _patch_affinity_path(path: str): + """Temporarily override GPU_CPU_AFFINITY_PATH in the numa_utils module.""" + import nemo_rl.distributed.numa_utils as mod + + old = mod.GPU_CPU_AFFINITY_PATH + mod.GPU_CPU_AFFINITY_PATH = path + return old + + +def _split_available_cpus() -> tuple[list[int], list[int]] | tuple[None, None]: + """Split the CPUs available to this process into two non-empty groups. + + Lets the CPU-binding tests exercise a real ``os.sched_setaffinity`` on any + host instead of hard-coding a 144-CPU (GB200) layout that fails on smaller + machines. Returns ``(group0, group1)`` as sorted lists, or ``(None, None)`` + when there are fewer than 2 available CPUs to split. + """ + avail = sorted(os.sched_getaffinity(0)) + if len(avail) < 2: + return None, None + mid = len(avail) // 2 + return avail[:mid], avail[mid:] + + +# --------------------------------------------------------------------------- +# Pure unit tests (no GPU or libnuma required) +# --------------------------------------------------------------------------- + + +class TestParseCpulist: + def test_single_range(self): + assert _parse_cpulist("0-71") == set(range(72)) + + def test_multiple_ranges(self): + assert _parse_cpulist("0-3,8-11") == {0, 1, 2, 3, 8, 9, 10, 11} + + def test_single_values(self): + assert _parse_cpulist("0,4,8") == {0, 4, 8} + + def test_mixed(self): + assert _parse_cpulist("0-2,5,10-12") == {0, 1, 2, 5, 10, 11, 12} + + def test_whitespace(self): + assert _parse_cpulist(" 0-3 , 8 ") == {0, 1, 2, 3, 8} + + def test_single_cpu(self): + assert _parse_cpulist("42") == {42} + + +class TestResolveVisibleGpuId: + """Resolve a process-local CUDA device index to its physical GPU id. + + This is the logic the vLLM TP-worker bind path relies on. It is + hardware-independent (pure CUDA_VISIBLE_DEVICES string parsing), so it + exercises the multi-GPU / non-zero-based-instance cases that local 2-GPU + hosts cannot reproduce. + """ + + def test_noset_subset_maps_local_to_physical(self, monkeypatch): + # vLLM TP=2 EngineCore for an instance on GPUs [4,5]: CVD is the + # per-instance subset, so local index i maps to that subset's i-th GPU. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") + assert resolve_visible_gpu_id(0) == 4 + assert resolve_visible_gpu_id(1) == 5 + + def test_full_node_list(self, monkeypatch): + # NOSET full-node list (e.g. Megatron): local index == physical index. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") + assert resolve_visible_gpu_id(3) == 3 + + def test_isolated_single_device(self, monkeypatch): + # vLLM TP=1 / DTensor: CVD isolated to one GPU, local index is 0. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "5") + assert resolve_visible_gpu_id(0) == 5 + + def test_no_cvd_returns_none(self, monkeypatch): + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + assert resolve_visible_gpu_id(0) is None + + def test_empty_cvd_returns_none(self, monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "") + assert resolve_visible_gpu_id(0) is None + + def test_out_of_range_returns_none(self, monkeypatch): + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "4,5") + assert resolve_visible_gpu_id(2) is None + + def test_non_integer_entries_return_none(self, monkeypatch): + # e.g. MIG UUIDs — cannot map to a physical index, so skip binding. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "GPU-abc,GPU-def") + assert resolve_visible_gpu_id(0) is None + + +class TestBindToGpuNuma: + """Test bind_to_gpu_numa logic with mock affinity files.""" + + def test_disabled_via_env(self, monkeypatch): + monkeypatch.setenv("NRL_DISABLE_NUMA_BINDING", "1") + assert bind_to_gpu_numa(0) is False + + def test_missing_affinity_file(self, monkeypatch): + monkeypatch.delenv("NRL_DISABLE_NUMA_BINDING", raising=False) + + old_path = _patch_affinity_path("/nonexistent/path") + try: + assert bind_to_gpu_numa(0) is False + finally: + _patch_affinity_path(old_path) + + def test_gpu_not_in_file(self, monkeypatch): + monkeypatch.delenv("NRL_DISABLE_NUMA_BINDING", raising=False) + monkeypatch.setenv("NRL_DISABLE_NUMA_MEMBIND", "1") + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write("0:0-71\n1:0-71\n2:72-143\n3:72-143\n") + f.flush() + + old_path = _patch_affinity_path(f.name) + try: + assert bind_to_gpu_numa(7) is False + finally: + _patch_affinity_path(old_path) + os.unlink(f.name) + + def test_successful_cpu_binding(self, monkeypatch): + """Verify sched_setaffinity is called with the correct CPU set.""" + monkeypatch.delenv("NRL_DISABLE_NUMA_BINDING", raising=False) + monkeypatch.setenv("NRL_DISABLE_NUMA_MEMBIND", "1") + + # Derive the cpulist from the CPUs actually available to this process so + # the test is host-portable. GPUs 0/1 map to the first CPU group, GPUs + # 2/3 to the second; binding GPU 2 should land us on the second group. + group0, group1 = _split_available_cpus() + if group1 is None: + pytest.skip("need >= 2 available CPUs to exercise NUMA CPU binding") + cpus0 = ",".join(map(str, group0)) + cpus1 = ",".join(map(str, group1)) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write(f"0:{cpus0}\n1:{cpus0}\n2:{cpus1}\n3:{cpus1}\n") + f.flush() + + old_path = _patch_affinity_path(f.name) + try: + result = bind_to_gpu_numa(2) + assert result is True + bound_cpus = os.sched_getaffinity(0) + assert bound_cpus == set(group1) + finally: + _patch_affinity_path(old_path) + os.unlink(f.name) + _reset_all_bindings() + + def test_explicit_gpu_id_ignores_cuda_visible_devices(self, monkeypatch): + """Binding follows the explicit gpu_id, not CUDA_VISIBLE_DEVICES. + + Under ``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=1`` a worker sees the + full node device list in ``CUDA_VISIBLE_DEVICES``, which cannot identify + its own GPU. ``gpu_id=2`` must bind to GPU 2's CPUs even though CVD lists + all 8 devices. + """ + # Full-node CVD as seen under NOSET mode. + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1,2,3,4,5,6,7") + monkeypatch.delenv("NRL_DISABLE_NUMA_BINDING", raising=False) + monkeypatch.setenv("NRL_DISABLE_NUMA_MEMBIND", "1") + + group0, group1 = _split_available_cpus() + if group1 is None: + pytest.skip("need >= 2 available CPUs to exercise NUMA CPU binding") + cpus0 = ",".join(map(str, group0)) + cpus1 = ",".join(map(str, group1)) + + with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f: + f.write(f"0:{cpus0}\n1:{cpus0}\n2:{cpus1}\n3:{cpus1}\n") + f.flush() + + old_path = _patch_affinity_path(f.name) + try: + result = bind_to_gpu_numa(2) + assert result is True + # gpu_id=2 maps to the second CPU group in the affinity file. + assert os.sched_getaffinity(0) == set(group1) + finally: + _patch_affinity_path(old_path) + os.unlink(f.name) + _reset_all_bindings() + + +class TestSetNumaMembind: + """Test membind with libnuma (skipped if libnuma unavailable).""" + + @pytest.fixture(autouse=True) + def _check_libnuma(self): + if _load_libnuma() is None: + pytest.skip("libnuma.so.1 not available") + + def test_membind_disabled(self, monkeypatch): + monkeypatch.setenv("NRL_DISABLE_NUMA_MEMBIND", "1") + assert _set_numa_membind({0, 1, 2}) is False + + def test_membind_succeeds(self, monkeypatch): + monkeypatch.delenv("NRL_DISABLE_NUMA_MEMBIND", raising=False) + cpus = os.sched_getaffinity(0) + assert _set_numa_membind(cpus) is True + + def test_get_numa_node_valid(self): + libnuma = _load_libnuma() + node = _get_numa_node(libnuma, {0}) + assert node >= 0 + + +# --------------------------------------------------------------------------- +# GPU benchmark test — validates binding has measurable impact on D2H +# --------------------------------------------------------------------------- + +HAS_CUDA = False +try: + import torch + + HAS_CUDA = torch.cuda.is_available() +except ImportError: + pass + +HAS_LIBNUMA = _load_libnuma() is not None + + +@pytest.mark.skipif(not HAS_CUDA, reason="No CUDA GPU available") +@pytest.mark.skipif(not HAS_LIBNUMA, reason="libnuma.so.1 not available") +class TestNUMABindingBenchmark: + """D2H bandwidth benchmark to validate NUMA binding impact. + + Run inside the RL container on a multi-socket node (DGX or GB200): + + CUDA_VISIBLE_DEVICES=0 pytest tests/unit/distributed/test_numa_utils.py::TestNUMABindingBenchmark -v -s + + The test measures D2H copy time with and without NUMA binding and + reports the speedup. It also verifies the process is correctly bound + to the expected NUMA node for the assigned GPU. + """ + + TENSOR_ELEMENTS = 64 * 1024 * 1024 # 256 MB at float32 + WARMUP_ITERS = 10 + BENCH_ITERS = 50 + + def _d2h_bandwidth_ms(self) -> float: + """Measure average D2H copy time in milliseconds.""" + t_gpu = torch.randn(self.TENSOR_ELEMENTS, device="cuda", dtype=torch.float32) + + for _ in range(self.WARMUP_ITERS): + _ = t_gpu.to("cpu", non_blocking=True) + torch.cuda.synchronize() + + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + start.record() + for _ in range(self.BENCH_ITERS): + _ = t_gpu.to("cpu", non_blocking=True) + end.record() + torch.cuda.synchronize() + + return start.elapsed_time(end) / self.BENCH_ITERS + + def test_d2h_with_numa_binding(self): + """Measure D2H bandwidth before and after NUMA binding. + + This test does NOT assert a specific speedup (it varies by + platform) — it reports the numbers so you can compare. + The key assertion is that binding succeeds and the process + ends up on the expected NUMA node. + """ + cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "") + if not cvd: + pytest.skip( + "CUDA_VISIBLE_DEVICES not set — run with CUDA_VISIBLE_DEVICES=N" + ) + + # Baseline: reset to fully unbound (all CPUs, all NUMA nodes) + _reset_all_bindings() + unbound_ms = self._d2h_bandwidth_ms() + + # Generate the affinity file from live nvidia-smi topo + affinity_file = _write_affinity_file_from_topo() + if affinity_file is None: + pytest.skip("Could not parse nvidia-smi topo output") + + # Apply NUMA binding (CPU affinity + membind) + old_path = _patch_affinity_path(affinity_file) + gpu_str = cvd.split(",")[0] + try: + result = bind_to_gpu_numa(int(gpu_str)) + assert result is True, f"bind_to_gpu_numa() failed for GPU {gpu_str}" + finally: + _patch_affinity_path(old_path) + + bound_cpus = os.sched_getaffinity(0) + libnuma = _load_libnuma() + numa_node = _get_numa_node(libnuma, bound_cpus) + assert numa_node >= 0, ( + f"Could not determine NUMA node for CPU {min(bound_cpus)}" + ) + + bound_ms = self._d2h_bandwidth_ms() + speedup = unbound_ms / bound_ms if bound_ms > 0 else float("inf") + + print(f"\n{'=' * 60}") + print(f"NUMA Binding D2H Benchmark (GPU {gpu_str})") + print(f"{'=' * 60}") + print( + f" Tensor size: {self.TENSOR_ELEMENTS * 4 / 1024 / 1024:.0f} MB (float32)" + ) + print(f" Iterations: {self.BENCH_ITERS}") + print(f" Unbound D2H: {unbound_ms:.3f} ms/iter") + print(f" Bound D2H: {bound_ms:.3f} ms/iter") + print(f" Speedup: {speedup:.3f}x") + print(f" Bound CPUs: {min(bound_cpus)}-{max(bound_cpus)}") + print(f" NUMA node: {numa_node}") + print(f"{'=' * 60}") + + os.unlink(affinity_file) + _reset_all_bindings() + + def test_actor_mapping_correctness(self): + """Verify that bind_to_gpu_numa places the process on the correct NUMA node. + + For each GPU on the node, sets CUDA_VISIBLE_DEVICES, calls + bind_to_gpu_numa, and checks that the resulting CPU affinity + and NUMA node match nvidia-smi topo. + """ + affinity_file = _write_affinity_file_from_topo() + if affinity_file is None: + pytest.skip("Could not parse nvidia-smi topo output") + + # Read back the affinity file to get the expected mapping + gpu_map: dict[str, str] = {} + with open(affinity_file) as f: + for line in f: + line = line.strip() + if line: + idx, cpulist = line.split(":", 1) + gpu_map[idx] = cpulist + + libnuma = _load_libnuma() + old_path = _patch_affinity_path(affinity_file) + + results = [] + try: + for gpu_idx, expected_cpulist in gpu_map.items(): + _reset_all_bindings() + + result = bind_to_gpu_numa(int(gpu_idx)) + assert result is True, f"bind_to_gpu_numa() failed for GPU {gpu_idx}" + + bound_cpus = os.sched_getaffinity(0) + expected_cpus = _parse_cpulist(expected_cpulist) + assert bound_cpus == expected_cpus, ( + f"GPU {gpu_idx}: expected CPUs {expected_cpulist}, " + f"got {min(bound_cpus)}-{max(bound_cpus)}" + ) + + numa_node = _get_numa_node(libnuma, bound_cpus) + results.append((gpu_idx, expected_cpulist, numa_node)) + finally: + _patch_affinity_path(old_path) + _reset_all_bindings() + os.unlink(affinity_file) + + print(f"\n{'=' * 60}") + print("Actor → GPU → NUMA Mapping Verification") + print(f"{'=' * 60}") + for gpu_idx, cpulist, numa_node in results: + print(f" GPU {gpu_idx} → CPUs {cpulist} → NUMA node {numa_node}") + print(f"{'=' * 60}")