fix: count amdgpu devices across every PCI domain

src/hyperloom/common/platform_probe.py

@@ -46,6 +46,7 @@
 
 _CPU_ROOT = "sys/devices/system/cpu"
 _NODE_ROOT = "sys/devices/system/node"
+_AMDGPU_DRIVER_ROOT = "sys/bus/pci/drivers/amdgpu"
 
 
 def read_kernel_file(path: Path | str, *, root: Path = DEFAULT_ROOT) -> str:
@@ -135,6 +136,25 @@ def kernel_release(*, root: Path = DEFAULT_ROOT) -> str:
     return read_kernel_file("proc/sys/kernel/osrelease", root=root) or "unknown"
 
 
+def amdgpu_device_count(*, root: Path = DEFAULT_ROOT) -> int | None:
+    """PCI devices bound to ``amdgpu``, or ``None`` when none are readable.
+
+    Counts every PCI domain rather than ``0000`` alone. A host with more
+    devices than one domain can address puts its GPUs under ``0002:``,
+    ``0003:`` and so on -- which is the normal layout on an MI300/MI350-class
+    node -- and a ``0000:``-only match counts zero of them there, i.e. it fails
+    on exactly the hardware this record is written to describe.
+
+    The ``*:*:*.*`` shape is what separates a device entry from the driver's
+    own ``bind``/``unbind``/``module`` siblings, which share the directory and
+    are not PCI addresses.
+    """
+    try:
+        return len(list((root / _AMDGPU_DRIVER_ROOT).glob("*:*:*.*"))) or None
+    except OSError:
+        return None
+
+
 @dataclass(frozen=True)
 class CpuPlatform:
     """Host CPU tuning state. Every field degrades independently."""
@@ -229,7 +249,7 @@ def platform_fingerprint(
                 # run could see. *_VISIBLE_DEVICES masking does not change this
                 # number, so it is named for the host to keep it from being read
                 # as the run's device count.
-                "host_count": len(list(Path("/sys/bus/pci/drivers/amdgpu").glob("0000:*"))) or None,
+                "host_count": amdgpu_device_count(),
                 # probe=False: gpu_type already answers this, and report
                 # generation runs in-process under unit tests that must not
                 # spawn rocminfo.
@@ -255,6 +275,7 @@ def platform_fingerprint(
 __all__ = [
     "CpuPlatform",
     "DEFAULT_ROOT",
+    "amdgpu_device_count",
     "boost_state",
     "cpu_model",
     "cpufreq_governor",

src/hyperloom/common/tests/test_platform_probe.py

@@ -13,6 +13,7 @@
 from pathlib import Path
 
 from hyperloom.common.platform_probe import (
+    amdgpu_device_count,
     boost_state,
     cpu_model,
     cpufreq_governor,
@@ -121,3 +122,42 @@ def test_read_kernel_file_never_raises(tmp_path):
     d = tmp_path / "adir"
     d.mkdir()
     assert read_kernel_file("adir", root=tmp_path) == ""
+
+
+def _amdgpu(root: Path, *addresses: str) -> Path:
+    """A driver directory holding the given devices next to its own entries."""
+    d = root / "sys/bus/pci/drivers/amdgpu"
+    d.mkdir(parents=True, exist_ok=True)
+    for address in addresses:
+        (d / address).mkdir()
+    for entry in ("bind", "unbind", "uevent", "new_id", "remove_id"):
+        (d / entry).write_text("")
+    (d / "module").mkdir()
+    return root
+
+
+def test_gpus_outside_pci_domain_zero_are_counted(tmp_path):
+    """A host large enough to need several PCI domains puts its GPUs outside
+    ``0000:``, and that is the hardware this record exists to describe. Matching
+    ``0000:`` alone reported no accelerators on exactly those machines."""
+    root = _amdgpu(
+        tmp_path,
+        "0002:00:01.0", "0002:00:02.0", "0002:00:03.0", "0002:00:04.0",
+        "0003:00:01.0", "0003:00:02.0", "0003:00:03.0", "0003:00:04.0",
+    )
+    assert amdgpu_device_count(root=root) == 8
+
+
+def test_the_ordinary_single_domain_host_still_counts(tmp_path):
+    root = _amdgpu(tmp_path, "0000:63:00.0", "0000:a3:00.0")
+    assert amdgpu_device_count(root=root) == 2
+
+
+def test_driver_control_entries_are_not_devices(tmp_path):
+    """``bind``, ``module`` and friends share the directory with the devices."""
+    assert amdgpu_device_count(root=_amdgpu(tmp_path)) is None
+
+
+def test_a_host_without_the_amdgpu_driver_reports_nothing(tmp_path):
+    """None, not zero: the driver dir is absent, which is not a count of zero."""
+    assert amdgpu_device_count(root=tmp_path) is None