diff --git a/deep_cluster/topology/bind.py b/deep_cluster/topology/bind.py index 7d180fd..9178570 100644 --- a/deep_cluster/topology/bind.py +++ b/deep_cluster/topology/bind.py @@ -39,6 +39,38 @@ def _libnuma() -> ctypes.CDLL: return library +def _nvidia_gpu_uuid(entry: Path) -> str | None: + try: + text = (entry / "information").read_text() + except OSError: + return None + for line in text.splitlines(): + key, sep, value = line.partition(":") + if sep and key.strip() == "GPU UUID": + return value.strip() or None + return None + + +def _resolve_nvidia_identifier(identifier: str, *, _root: Path) -> int: + """Map a visible NVIDIA device identifier to a physical index. + + CUDA_VISIBLE_DEVICES accepts device ordinals, PCI bus ids, and GPU UUIDs + (including unambiguous UUID prefixes); NUMA lookup must handle all of them. + """ + value = identifier.strip() + if value.isdigit(): + return int(value) + gpu_dir = _root / "proc/driver/nvidia/gpus" + entries = tuple(sorted((path for path in gpu_dir.iterdir() if path.is_dir()), key=lambda path: path.name)) if gpu_dir.is_dir() else () + matches = [index for index, entry in enumerate(entries) if entry.name.lower() == value.lower()] + if not matches and value.upper().startswith("GPU-"): + wanted = value.upper() + matches = [index for index, entry in enumerate(entries) if (_nvidia_gpu_uuid(entry) or "").upper().startswith(wanted)] + if len(matches) == 1: + return matches[0] + raise IndexError(f"cannot resolve visible NVIDIA device {identifier!r} to a physical index") + + def which_numa( i_gpu: int | None = None, i_ib: int | None = None, @@ -57,7 +89,7 @@ def which_numa( match = re.search(r"(\d+)$", identifier) physical_index = int(match.group(1)) if match is not None else int(identifier) else: - physical_index = int(identifier) + physical_index = _resolve_nvidia_identifier(identifier, _root=_root) return device_numa_node(physical_index, _root=_root) except IndexError as error: raise AssertionError("no such GPU found") from error diff --git a/tests/deep_cluster/test_bind_numa.py b/tests/deep_cluster/test_bind_numa.py new file mode 100644 index 0000000..6e80c65 --- /dev/null +++ b/tests/deep_cluster/test_bind_numa.py @@ -0,0 +1,74 @@ +"""CPU-only tests: NUMA lookup must accept every identifier CUDA_VISIBLE_DEVICES allows. + +Operators pin devices by PCI bus id or GPU UUID just as often as by ordinal; the NUMA +binding path runs unconditionally at engine startup, so any of those forms must resolve +(or fail with the documented "no such GPU found") instead of dying in int(). +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from deep_cluster.topology.bind import which_numa + +UUID_A = "GPU-aaaa1111-2222-3333-4444-555555555555" +UUID_B = "GPU-bbbb1111-2222-3333-4444-555555555555" +BUS_A = "0000:ab:00.0" +BUS_B = "0000:cd:00.0" + + +@pytest.fixture +def fake_root(tmp_path: Path) -> Path: + """A procfs/sysfs-shaped tree advertising two NVIDIA GPUs on distinct NUMA nodes.""" + for bus_id, uuid, numa in ((BUS_A, UUID_A, 0), (BUS_B, UUID_B, 1)): + gpu_dir = tmp_path / "proc/driver/nvidia/gpus" / bus_id + gpu_dir.mkdir(parents=True) + (gpu_dir / "information").write_text(f"Device Minor: 0\nGPU UUID: {uuid}\n", encoding="utf-8") + pci_dir = tmp_path / "sys/bus/pci/devices" / bus_id + pci_dir.mkdir(parents=True) + (pci_dir / "numa_node").write_text(f"{numa}\n", encoding="utf-8") + (tmp_path / "dev").mkdir() + (tmp_path / "dev/nvidia0").touch() + (tmp_path / "dev/nvidia1").touch() + return tmp_path + + +def _visible(value: str | None, monkeypatch: pytest.MonkeyPatch) -> None: + if value is None: + monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False) + else: + monkeypatch.setenv("CUDA_VISIBLE_DEVICES", value) + + +def test_ordinal_baseline(fake_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _visible(None, monkeypatch) + assert which_numa(0, _root=fake_root) == 0 + assert which_numa(1, _root=fake_root) == 1 + _visible("1", monkeypatch) + assert which_numa(0, _root=fake_root) == 1 + + +@pytest.mark.parametrize("value", [BUS_B, BUS_B.upper(), UUID_B, "GPU-bbb"]) +def test_non_ordinal_identifiers_resolve(fake_root: Path, monkeypatch: pytest.MonkeyPatch, value: str) -> None: + _visible(value, monkeypatch) + assert which_numa(0, _root=fake_root) == 1 + + +def test_unknown_identifier_reports_no_such_gpu(fake_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _visible("0000:ff:00.0", monkeypatch) + with pytest.raises(AssertionError, match="no such GPU found"): + which_numa(0, _root=fake_root) + + +def test_ambiguous_uuid_prefix_reports_no_such_gpu(fake_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _visible("GPU-", monkeypatch) + with pytest.raises(AssertionError, match="no such GPU found"): + which_numa(0, _root=fake_root) + + +def test_negative_identifier_reports_no_such_gpu(fake_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: + _visible("-1", monkeypatch) + with pytest.raises(AssertionError, match="no such GPU found"): + which_numa(0, _root=fake_root)