"""Pre-flight and serialization guard for runs that allocate a Metal expert arena.
Every `strata_window_cli` prompt run mmaps `STRATA_METAL_ARENA_GB` of expert slots and, with
`STRATA_METAL_WIRE_ARENA=1`, `mlock`s them. Wired memory cannot be evicted, so two of those on a
machine that is already serving from a resident engine does not degrade gracefully: the kernel swaps
everything else out and Metal stalls, which reads as "the GPU is dying". This module exists because
that happened: three arena runs were started before any was awaited, on top of a live 9.27 GB
`--serve` engine, and it took the machine down.
The fix is structural, not a promise to be careful:
import arena_preflight as ap
job = await ap.arena_run(*argv, env=env) # one at a time, env-gated, reported
`arena_run` holds a process-wide asyncio lock, so a concurrent call blocks instead of oversubscribing,
and refuses to start at all when headroom is short unless `force=True`.
"""
import asyncio
import os
import re
import sys
PAGE = 16384 # Apple Silicon
_DEFAULT_ARENA_GB = 6.5
# an arena run also wants its dense weights, the PLE/MTP files, KV, and the framework itself
_OVERHEAD_GB = 4.0
_lock = asyncio.Lock()
def _sh(*argv):
import subprocess
return subprocess.run(argv, capture_output=True, text=True).stdout
def processes():
"""Live arena-allocating or arena-holding processes, newest state from ps."""
out = _sh("ps", "-ax", "-o", "pid,rss,etime,command")
found = []
for line in out.splitlines()[1:]:
parts = line.split(None, 3)
if len(parts) < 4:
continue
pid, rss, etime, cmd = parts
if "ps -ax" in cmd or "ps -p" in cmd:
continue
# arena holders: the engine itself and the serve supervisor that spawns it
if "strata_window_cli" in cmd or "server.py" in cmd:
found.append({"pid": int(pid), "rss_gb": int(rss) / 1e6, "etime": etime,
"serve": "--serve" in cmd or "server.py" in cmd, "cmd": cmd[:110]})
return found
def memory():
out = _sh("vm_stat")
pages = {}
for line in out.splitlines():
m = re.match(r"(.+?):\s+(\d+)\.?$", line.strip())
if m:
pages[m.group(1)] = int(m.group(2))
free_gb = (pages.get("Pages free", 0) + pages.get("Pages speculative", 0)) * PAGE / 1e9
wired_gb = pages.get("Pages wired down", 0) * PAGE / 1e9
swapped = _sh("sysctl", "-n", "vm.swapusage").strip()
used = re.search(r"used = ([\d.]+)M", swapped)
total = re.search(r"total = ([\d.]+)M", swapped)
return {
"free_gb": free_gb,
"wired_gb": wired_gb,
"swap_used_gb": float(used.group(1)) / 1024 if used else -1.0,
"swap_total_gb": float(total.group(1)) / 1024 if total else -1.0,
"swap_encrypted_unusable": "encrypted" in swapped,
}
def snapshot(arena_gb=None):
if arena_gb is None:
arena_gb = float(os.environ.get("STRATA_METAL_ARENA_GB", _DEFAULT_ARENA_GB))
procs, mem = processes(), memory()
held = sum(p["rss_gb"] for p in procs)
need = arena_gb + _OVERHEAD_GB
return {"arena_gb": arena_gb, "need_gb": need, "held_gb": held,
"headroom_gb": mem["free_gb"], "free_after_gb": mem["free_gb"] - need,
"ok": mem["free_gb"] >= need and mem["swap_used_gb"] < 1.0,
"processes": procs, "memory": mem}
def report(snap=None):
s = snap or snapshot()
print(f"arena preflight: need {s['need_gb']:.1f} GB "
f"(arena {s['arena_gb']:.1f} + overhead {_OVERHEAD_GB:.1f}), "
f"free {s['headroom_gb']:.1f} GB, after {s['free_after_gb']:+.1f} GB, "
f"swap {s['memory']['swap_used_gb']:.2f}/{s['memory']['swap_total_gb']:.2f} GB, "
f"wired {s['memory']['wired_gb']:.1f} GB")
if s["processes"]:
for p in s["processes"]:
tag = "SERVE" if p["serve"] else "cli"
print(f" [{tag}] pid {p['pid']} {p['rss_gb']:.2f} GB up {p['etime']}")
else:
print(" no live arena process")
print(f" verdict: {'OK to start one arena run' if s['ok'] else 'REFUSE - not enough headroom'}")
return s["ok"]
async def arena_run(*argv, env=None, arena_gb=None, force=False, quiet=False, tail=40):
"""Run one arena-allocating program to completion, serialized and pre-flighted."""
from albedo import run # the session's process launcher
async with _lock:
snap = snapshot(arena_gb if arena_gb is not None else
(float(env["STRATA_METAL_ARENA_GB"]) if env and "STRATA_METAL_ARENA_GB" in env else None))
if not quiet:
report(snap)
if not snap["ok"] and not force:
raise RuntimeError(
"arena_preflight: refusing to start; headroom "
f"{snap['headroom_gb']:.1f} GB < need {snap['need_gb']:.1f} GB or swap "
f"{snap['memory']['swap_used_gb']:.2f} GB in use. Free memory, shrink "
"STRATA_METAL_ARENA_GB, or pass force=True deliberately.")
job = run(*argv, env=env)
await job
if not quiet:
out = job.read() if job.exit_code == 0 else job.read()
for line in out.splitlines():
if "STRATA_RESULT" in line or "expert stream" in line or "error" in line.lower():
print(" " + line[:200])
return job