|
1 | 1 | """Headroom admission, not synthetic activity or a board power limit.""" |
2 | 2 |
|
| 3 | +import json |
| 4 | +import math |
| 5 | +import os |
| 6 | +from pathlib import Path |
3 | 7 |
|
4 | | -def local_budget() -> dict: |
| 8 | + |
| 9 | +def available_cpus( |
| 10 | + cgroup_root: Path = Path("/sys/fs/cgroup"), membership: Path = Path("/proc/self/cgroup") |
| 11 | +) -> int: |
| 12 | + """Count process affinity and Linux cgroup v2 quota; admission reserves CPU separately.""" |
| 13 | + count = ( |
| 14 | + len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count() or 1 |
| 15 | + ) |
| 16 | + if membership.exists(): |
| 17 | + for line in membership.read_text().splitlines(): |
| 18 | + if line.startswith("0::"): |
| 19 | + base = cgroup_root.resolve() |
| 20 | + folder = (base / line[3:].lstrip("/")).resolve() |
| 21 | + if not folder.is_relative_to(base): |
| 22 | + # A cgroup namespace can report a path outside the mounted subtree. |
| 23 | + # Its mounted root still exposes the namespace's effective root quota. |
| 24 | + folder = base |
| 25 | + while folder.is_relative_to(base): |
| 26 | + quota = folder / "cpu.max" |
| 27 | + if quota.exists(): |
| 28 | + limit, period = quota.read_text().split() |
| 29 | + if limit != "max": |
| 30 | + count = min(count, max(1, math.floor(int(limit) / int(period)))) |
| 31 | + if folder == base: |
| 32 | + break |
| 33 | + folder = folder.parent |
| 34 | + return max(1, count) |
| 35 | + |
| 36 | + |
| 37 | +def local_budget(root: Path | None = None) -> dict: |
5 | 38 | import psutil |
6 | 39 |
|
7 | 40 | cpu = psutil.cpu_percent(interval=0.05) |
8 | 41 | free = psutil.virtual_memory().available / 2**30 |
9 | | - slots = 0 if cpu >= 85 or free < 2 else 1 if cpu >= 65 or free < 4 else 2 |
| 42 | + maximum = 2 |
| 43 | + reserve_ram, ram_per_slot = 0, 2 |
| 44 | + path = root / "research/state/resource-policy.json" if root is not None else None |
| 45 | + if path is not None and path.exists(): |
| 46 | + policy = json.loads(path.read_text()) |
| 47 | + maximum = policy["max_cpu_slots"] |
| 48 | + if type(maximum) is not int or not 1 <= maximum <= 4: |
| 49 | + raise ValueError("Local CPU policy must allow 1-4 slots") |
| 50 | + reserve_ram, ram_per_slot = ( |
| 51 | + policy.get("reserve_ram_gib", 0), |
| 52 | + policy.get("ram_per_slot_gib", 2), |
| 53 | + ) |
| 54 | + if ( |
| 55 | + type(reserve_ram) is not int |
| 56 | + or not 0 <= reserve_ram <= 8 |
| 57 | + or type(ram_per_slot) is not int |
| 58 | + or not 2 <= ram_per_slot <= 4 |
| 59 | + ): |
| 60 | + raise ValueError("Invalid local RAM reservation") |
| 61 | + cpus = available_cpus() |
| 62 | + maximum = min(maximum, max(1, (cpus - 1) // 2)) |
| 63 | + ram_slots = max(0, int((free - reserve_ram) // ram_per_slot)) |
| 64 | + slots = 0 if cpu >= 85 else min(1 if cpu >= 65 else maximum, ram_slots) |
10 | 65 | return { |
11 | 66 | "cpu_slots": slots, |
| 67 | + "max_cpu_slots": maximum, |
| 68 | + "available_cpus": cpus, |
| 69 | + "reserve_ram_gib": reserve_ram, |
| 70 | + "ram_per_slot_gib": ram_per_slot, |
12 | 71 | "host_cpu_percent": cpu, |
13 | 72 | "free_ram_gib": round(free, 2), |
14 | 73 | "reason": "Host pressure; queued work retained" |
15 | 74 | if not slots |
16 | 75 | else "Measured CPU/RAM headroom", |
17 | | - "scope": "At most two local CPU/network drones. GPU serving/training use existing exclusive admission. Windows uses its own lower limits; RTX remains operator-disabled.", |
| 76 | + "scope": "Bounded local CPU/network admission; one model job. GPU serving/training keep exclusive gates. Windows keeps its own pressure, RAM and game guards. Slots are concurrency limits, not measured speedup.", |
18 | 77 | } |
| 78 | + |
| 79 | + |
| 80 | +def cpu_assignment(active: list[list[int]]) -> list[int]: |
| 81 | + """Allocate distinct affinity slots for concurrent owned Python experiments.""" |
| 82 | + cpus = sorted(os.sched_getaffinity(0))[: available_cpus()] |
| 83 | + if len(cpus) > 2: |
| 84 | + cpus = cpus[1:] |
| 85 | + used = {cpu for assignment in active for cpu in assignment} |
| 86 | + return [cpu for cpu in cpus if cpu not in used][:2] |
0 commit comments