from __future__ import annotations
import logging
import os
import re
import torch
from . import envs
logger = logging.getLogger("modelexpress.ucx_utils")
def _read_int_file(path: str) -> int | None:
try:
with open(path, "r") as f:
return int(f.read().strip())
except (OSError, ValueError):
return None
def _read_str_file(path: str) -> str | None:
try:
with open(path, "r") as f:
return f.read().strip()
except OSError:
return None
def _parse_ib_rate_gbps(rate_str: str) -> float | None:
if not rate_str:
return None
parts = rate_str.strip().split()
if not parts:
return None
try:
return float(parts[0])
except ValueError:
return None
def _gpu_pci_bdf(device_id: int) -> str | None:
try:
props = torch.cuda.get_device_properties(device_id)
domain = int(getattr(props, "pci_domain_id", 0))
bus = int(props.pci_bus_id)
dev = int(props.pci_device_id)
except (AttributeError, RuntimeError, AssertionError, TypeError) as e:
logger.warning(f"NIC pin probe: unable to read PCI BDF for device {device_id}: {e}")
return None
return f"{domain:04x}:{bus:02x}:{dev:02x}.0"
def _gpu_numa_node(device_id: int) -> int | None:
bdf = _gpu_pci_bdf(device_id)
if bdf is None:
return None
return _read_int_file(f"/sys/bus/pci/devices/{bdf}/numa_node")
def _pci_path_components(bdf: str) -> list[str]:
try:
rp = os.path.realpath(f"/sys/bus/pci/devices/{bdf}")
except OSError:
return []
bdf_re = re.compile(r"^[0-9a-f]{4}:[0-9a-f]{2}:[0-9a-f]{2}\.[0-9a-f]$")
return [p for p in rp.split("/") if bdf_re.match(p)]
def _pci_common_depth(a: list[str], b: list[str]) -> int:
n = min(len(a), len(b))
for i in range(n):
if a[i] != b[i]:
return i
return n
def _nic_pci_bdf(nic_name: str) -> str | None:
try:
target = os.readlink(f"/sys/class/infiniband/{nic_name}/device")
except OSError:
return None
return os.path.basename(target.rstrip("/"))
def _nic_has_accessible_verbs_device(nic_name: str) -> bool:
verbs_dir = f"/sys/class/infiniband/{nic_name}/device/infiniband_verbs"
try:
verbs = os.listdir(verbs_dir)
except OSError:
return False
return any(
name.startswith("uverbs")
and os.path.exists(f"/dev/infiniband/{name}")
for name in verbs
)
def _list_compute_ib_nics(
min_rate_gbps: float | None = None,
) -> list[tuple[str, int, float, list[str]]]:
base = "/sys/class/infiniband"
if not os.path.isdir(base):
return []
try:
names = sorted(os.listdir(base))
except OSError:
return []
candidates: list[tuple[str, float, str]] = []
for name in names:
if "bond" in name:
continue
if not _nic_has_accessible_verbs_device(name):
continue
port_dir = f"{base}/{name}/ports/1"
if not os.path.isdir(port_dir):
continue
rate_str = _read_str_file(f"{port_dir}/rate")
rate = _parse_ib_rate_gbps(rate_str) if rate_str else None
if rate is None:
continue
candidates.append((name, rate, port_dir))
if not candidates:
return []
if min_rate_gbps is None:
threshold = max(r for _, r, _ in candidates)
else:
threshold = min_rate_gbps
out: list[tuple[str, int, float, list[str]]] = []
for name, rate, _port_dir in candidates:
if rate < threshold:
continue
numa = _read_int_file(f"{base}/{name}/device/numa_node")
if numa is None:
numa = -1
bdf = _nic_pci_bdf(name)
path = _pci_path_components(bdf) if bdf else []
out.append((name, numa, rate, path))
return out
def probe_nic_pin_for_device(
device_id: int, min_rate_gbps: float | None = None
) -> str | None:
nics = _list_compute_ib_nics(min_rate_gbps)
if not nics:
rate_desc = (
"max-rate auto-detect"
if min_rate_gbps is None
else f"rate >= {min_rate_gbps} Gb/s"
)
logger.warning(
f"MX_RDMA_NIC_PIN auto-probe: no compute IB-class NICs found "
f"under /sys/class/infiniband ({rate_desc}); skipping pin. "
f"This is expected on hosts without IB / RoCE / OPA / EFA, "
f"or on fabrics outside the kernel verbs API (e.g. Slingshot)."
)
return None
try:
num_gpus = torch.cuda.device_count()
except Exception:
num_gpus = 0
gpu_paths: dict[int, list[str]] = {}
gpu_numa: dict[int, int] = {}
for gi in range(num_gpus):
bdf = _gpu_pci_bdf(gi)
if bdf is None:
continue
gpu_paths[gi] = _pci_path_components(bdf)
numa = _read_int_file(f"/sys/bus/pci/devices/{bdf}/numa_node")
gpu_numa[gi] = numa if numa is not None else -1
if device_id not in gpu_paths:
logger.warning(
f"MX_RDMA_NIC_PIN auto-probe: GPU {device_id} not found among "
f"visible CUDA devices ({sorted(gpu_paths.keys())}); skipping pin"
)
return None
assigned_count: dict[str, int] = {n[0]: 0 for n in nics}
assignments: dict[int, tuple[str, int]] = {}
for gi in sorted(gpu_paths.keys()):
gpu_path = gpu_paths[gi]
ranked: list[tuple[int, int, str]] = []
for nic_name, _nic_numa, _nic_rate, nic_path in nics:
score = _pci_common_depth(gpu_path, nic_path)
ranked.append((-score, assigned_count[nic_name], nic_name))
ranked.sort()
chosen_name = ranked[0][2]
chosen_score = -ranked[0][0]
assignments[gi] = (chosen_name, chosen_score)
assigned_count[chosen_name] += 1
chosen_name, chosen_score = assignments[device_id]
nic_numa_map = {n[0]: n[1] for n in nics}
nic_rate_map = {n[0]: n[2] for n in nics}
same_numa_nics = [
n[0] for n in nics if n[1] == gpu_numa.get(device_id, -2) and n[1] >= 0
]
full_map = {gi: a[0] for gi, a in sorted(assignments.items())}
cross_socket = (
gpu_numa.get(device_id, -1) >= 0
and nic_numa_map.get(chosen_name, -1) >= 0
and gpu_numa[device_id] != nic_numa_map[chosen_name]
)
if cross_socket:
logger.warning(
f"MX_RDMA_NIC_PIN auto-probe: GPU {device_id} -> {chosen_name}:1 "
f"is CROSS-SOCKET (GPU NUMA {gpu_numa[device_id]}, NIC NUMA "
f"{nic_numa_map[chosen_name]}); single-flow bandwidth will be "
f"capped by UPI / Infinity Fabric. PCIe common-depth {chosen_score}, "
f"same-NUMA NICs available: {same_numa_nics}, full GPU->NIC map: "
f"{full_map}"
)
else:
logger.info(
f"MX_RDMA_NIC_PIN auto-probe: GPU {device_id} -> {chosen_name}:1 "
f"(PCIe common-depth {chosen_score}; GPU NUMA "
f"{gpu_numa.get(device_id)}, NIC NUMA {nic_numa_map.get(chosen_name)}, "
f"NIC rate {nic_rate_map.get(chosen_name)} Gb/s; "
f"same-NUMA NICs: {same_numa_nics}; full GPU->NIC map: {full_map})"
)
return f"{chosen_name}:1"
def _resolve_nic_pin(device_id: int) -> str | None:
raw = envs.MX_RDMA_NIC_PIN
if raw == "" or raw.lower() in ("off", "0", "false", "no"):
return None
if "," in raw:
nic_list = [n.strip() for n in raw.split(",") if n.strip()]
if 0 <= device_id < len(nic_list):
pinned = nic_list[device_id]
logger.info(
f"MX_RDMA_NIC_PIN explicit list: device {device_id} -> {pinned}"
)
return pinned
logger.warning(
f"MX_RDMA_NIC_PIN explicit list: device_id {device_id} out of "
f"range for list of length {len(nic_list)}; skipping pin"
)
return None
raw_min = envs.MX_RDMA_NIC_PIN_MIN_RATE_GBPS
if raw_min is None or raw_min.strip() == "":
min_rate: float | None = None
else:
try:
min_rate = float(raw_min)
except ValueError:
logger.warning(
f"MX_RDMA_NIC_PIN_MIN_RATE_GBPS={raw_min!r} not a float; "
f"falling back to max-rate auto-detect"
)
min_rate = None
return probe_nic_pin_for_device(device_id, min_rate_gbps=min_rate)
def apply_nic_pin_for_device(device_id: int) -> None:
pinned = _resolve_nic_pin(device_id)
if pinned:
prev = envs.UCX_NET_DEVICES
os.environ["UCX_NET_DEVICES"] = pinned
logger.info(
f"NIXL NIC pin: device {device_id} -> "
f"UCX_NET_DEVICES={pinned} (was: {prev})"
)