Merge 1692ef3f589c9cb83b7c469b3a67c0105fa41a3a into 553741557a57df44a46fb3ab72b6aafeaf945337

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

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)