Pin this process to the NUMA node of its GPU (CPU + memory).
(rank: int)
| 47 | |
| 48 | |
| 49 | def _numa_bind(rank: int) -> None: |
| 50 | """Pin this process to the NUMA node of its GPU (CPU + memory).""" |
| 51 | node = _gpu_numa_node(rank) |
| 52 | if node is None: |
| 53 | print(f"[rank {rank}] NUMA node unavailable, skipping binding", flush=True) |
| 54 | return |
| 55 | try: |
| 56 | # Read which CPUs belong to this NUMA node |
| 57 | cpus_path = f"/sys/devices/system/node/node{node}/cpulist" |
| 58 | cpu_list_str = open(cpus_path).read().strip() |
| 59 | # Parse "0-47,96-143" format into a set of CPU ids |
| 60 | numa_cpus = set() |
| 61 | for part in cpu_list_str.split(","): |
| 62 | if "-" in part: |
| 63 | lo, hi = part.split("-") |
| 64 | numa_cpus.update(range(int(lo), int(hi) + 1)) |
| 65 | else: |
| 66 | numa_cpus.add(int(part)) |
| 67 | # Intersect with current affinity (cgroup may restrict) |
| 68 | allowed = os.sched_getaffinity(0) |
| 69 | target = numa_cpus & allowed |
| 70 | if not target: |
| 71 | print(f"[rank {rank}] No CPUs in NUMA node {node} within cgroup, skipping", flush=True) |
| 72 | return |
| 73 | os.sched_setaffinity(0, target) |
| 74 | print(f"[rank {rank}] NUMA-bound to node {node}, CPUs {sorted(target)}", flush=True) |
| 75 | except Exception as e: |
| 76 | print(f"[rank {rank}] NUMA binding failed: {e}", flush=True) |
| 77 | |
| 78 | |
| 79 | def _print_numa_info(rank: str) -> None: |
no test coverage detected