nemo_rl.distributed.numa_utils

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.

Module Contents

Functions

bind_to_gpu_numa	Pin the current process to the NUMA-local CPUs and memory of the given GPU.
resolve_visible_gpu_id	Map a process-local CUDA device index to its node-global physical GPU id.
_load_libnuma	Load libnuma, returning None if unavailable.
_get_numa_node	Return the NUMA node for the given CPU set, or -1 on failure.
_set_numa_membind	Hard-bind memory allocations to the NUMA node of the given CPUs.
_parse_cpulist	Parse a Linux cpulist string like ‘0-71’ into a set of ints.

Data

logger
GPU_CPU_AFFINITY_PATH

API

nemo_rl.distributed.numa_utils.logger

‘getLogger(…)’

nemo_rl.distributed.numa_utils.GPU_CPU_AFFINITY_PATH

‘get(…)’

nemo_rl.distributed.numa_utils.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.

Parameters:

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.

nemo_rl.distributed.numa_utils.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).

nemo_rl.distributed.numa_utils._load_libnuma() → ctypes.CDLL | None

Load libnuma, returning None if unavailable.

nemo_rl.distributed.numa_utils._get_numa_node(libnuma: ctypes.CDLL, cpus: set[int]) → int

Return the NUMA node for the given CPU set, or -1 on failure.

nemo_rl.distributed.numa_utils._set_numa_membind(cpus: set[int]) → bool

Hard-bind memory allocations to the NUMA node of the given CPUs.

nemo_rl.distributed.numa_utils._parse_cpulist(cpulist: str) → set[int]

Parse a Linux cpulist string like ‘0-71’ into a set of ints.