Look up the NUMA node for a GPU from sysfs.
(gpu_index: int)
| 19 | |
| 20 | |
| 21 | def _gpu_numa_node(gpu_index: int) -> int | None: |
| 22 | """Look up the NUMA node for a GPU from sysfs.""" |
| 23 | try: |
| 24 | pci_bus = ( |
| 25 | subprocess.check_output( |
| 26 | [ |
| 27 | "nvidia-smi", |
| 28 | "-i", |
| 29 | str(gpu_index), |
| 30 | "--query-gpu=pci.bus_id", |
| 31 | "--format=csv,noheader", |
| 32 | ], |
| 33 | text=True, |
| 34 | ) |
| 35 | .strip() |
| 36 | .lower() |
| 37 | ) |
| 38 | # nvidia-smi may return "00000000:XX:YY.Z", sysfs uses "0000:XX:YY.Z" |
| 39 | if pci_bus.startswith("00000000:"): |
| 40 | pci_bus = pci_bus[4:] |
| 41 | numa_path = f"/sys/bus/pci/devices/{pci_bus}/numa_node" |
| 42 | node = int(open(numa_path).read().strip()) |
| 43 | return node if node >= 0 else None |
| 44 | except Exception as e: |
| 45 | warnings.warn(f"[gpu {gpu_index}] Failed to look up NUMA node: {e}", stacklevel=2) |
| 46 | return None |
| 47 | |
| 48 | |
| 49 | def _numa_bind(rank: int) -> None: |