koreallmdev's picture
Upload anonymized DGX harness engineering export V1.4 (part 2)
6fda4b7 verified
Raw History Blame Contribute Delete
14.6 kB
from __future__ import annotations
import hashlib
import json
import os
import secrets
import signal
import socket
import subprocess
import time
import urllib.request
from pathlib import Path
from typing import Any, Dict, List, Optional
from harness_core import AuditLog, write_json_atomic
from launch_plan import validate_launch_plan
def canonical_sha(obj: Any) -> str:
raw = json.dumps(
obj,
sort_keys=True,
ensure_ascii=False,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(raw).hexdigest()
def proc_start_ticks(pid: int) -> int:
# /proc/<pid>/stat field 22. comm can contain spaces inside parentheses,
# so split only after the final ') '.
text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
tail = text[text.rfind(") ") + 2:].split()
# Tail starts at field 3, so field 22 is tail index 19.
return int(tail[19])
def proc_cmdline(pid: int) -> List[str]:
raw = Path(f"/proc/{pid}/cmdline").read_bytes()
return [
x.decode("utf-8", errors="replace")
for x in raw.split(b"\0")
if x
]
def proc_state(pid: int):
try:
text = Path(f"/proc/{pid}/stat").read_text(encoding="utf-8")
except FileNotFoundError:
return None
tail = text[text.rfind(") ") + 2:].split()
if not tail:
return None
return tail[0]
def pid_exists(pid: int) -> bool:
"""
Return True only for a live, non-zombie process.
kill(pid, 0) alone is insufficient because a terminated child can remain
briefly as a zombie until its parent reaps it. For lifecycle ownership,
zombie means execution has ended and is therefore considered stopped.
"""
try:
os.kill(pid, 0)
except ProcessLookupError:
return False
state = proc_state(pid)
if state is None:
return False
if state == "Z":
return False
return True
def check_port_free(host: str, port: int) -> Dict[str, Any]:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 0)
s.bind((host, port))
return {"status": "PASS", "free": True}
except OSError as e:
return {
"status": "BLOCK",
"free": False,
"error": repr(e),
}
finally:
s.close()
def choose_free_port(host: str = "127.0.0.1") -> int:
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
try:
s.bind((host, 0))
return int(s.getsockname()[1])
finally:
s.close()
class RunOwnedRuntimeController:
def __init__(
self,
*,
run_dir: Path,
audit: AuditLog,
launch_plan: Dict[str, Any],
readiness_url: Optional[str] = None,
readiness_nonce: Optional[str] = None,
startup_timeout: float = 15.0,
stop_timeout: float = 10.0,
):
self.run_dir = Path(run_dir)
self.run_dir.mkdir(parents=True, exist_ok=True)
self.audit = audit
self.plan = dict(launch_plan)
self.readiness_url = readiness_url
self.readiness_nonce = readiness_nonce
self.startup_timeout = float(startup_timeout)
self.stop_timeout = float(stop_timeout)
self.state_path = self.run_dir / "RUNTIME_STATE.json"
self.stdout_path = self.run_dir / "runtime.stdout.log"
self.stderr_path = self.run_dir / "runtime.stderr.log"
def _read_state(self):
if not self.state_path.is_file():
return None
return json.loads(
self.state_path.read_text(encoding="utf-8")
)
def _write_state(self, state):
write_json_atomic(self.state_path, state)
def _clear_state(self):
if self.state_path.exists():
self.state_path.unlink()
def validate_plan(self):
result = validate_launch_plan(self.plan)
if result["status"] != "PASS":
node-7.example.invalid RuntimeError("LAUNCH_PLAN_BLOCK")
return result
def verify_ownership(self) -> Dict[str, Any]:
state = self._read_state()
if not state:
return {
"status": "BLOCK",
"reason": "NO_RUNTIME_STATE",
}
pid = int(state["pid"])
if not pid_exists(pid):
return {
"status": "BLOCK",
"reason": "PID_NOT_ALIVE",
"pid": pid,
}
actual_pgid = os.getpgid(pid)
if actual_pgid != int(state["pgid"]):
return {
"status": "BLOCK",
"reason": "PGID_MISMATCH",
"expected": state["pgid"],
"actual": actual_pgid,
}
actual_ticks = proc_start_ticks(pid)
if actual_ticks != int(state["proc_start_ticks"]):
return {
"status": "BLOCK",
"reason": "PROC_START_TICKS_MISMATCH",
"expected": state["proc_start_ticks"],
"actual": actual_ticks,
}
actual_cmdline = proc_cmdline(pid)
expected_cmdline = state["expected_cmdline"]
if actual_cmdline != expected_cmdline:
return {
"status": "BLOCK",
"reason": "CMDLINE_MISMATCH",
"expected": expected_cmdline,
"actual": actual_cmdline,
}
if canonical_sha(actual_cmdline) != state["cmdline_sha256"]:
return {
"status": "BLOCK",
"reason": "CMDLINE_SHA_MISMATCH",
}
return {
"status": "PASS",
"pid": pid,
"pgid": actual_pgid,
"proc_start_ticks": actual_ticks,
"cmdline_sha256": state["cmdline_sha256"],
"ownership_nonce": state["ownership_nonce"],
}
def probe(self) -> Dict[str, Any]:
own = self.verify_ownership()
if own["status"] != "PASS":
return {
"status": "BLOCK",
"ownership": own,
}
if not self.readiness_url:
return {
"status": "PASS",
"ownership": own,
"readiness": "NOT_CONFIGURED",
}
try:
with urllib.request.urlopen(
self.readiness_url,
timeout=2.0,
) as r:
body = json.loads(
r.read().decode("utf-8")
)
except Exception as e:
return {
"status": "BLOCK",
"ownership": own,
"readiness": "HTTP_FAIL",
"error": repr(e),
}
if self.readiness_nonce is not None:
if body.get("nonce") != self.readiness_nonce:
return {
"status": "BLOCK",
"ownership": own,
"readiness": "NONCE_MISMATCH",
"body": body,
}
return {
"status": "PASS",
"ownership": own,
"readiness": "PASS",
"body": body,
}
def start(self) -> Dict[str, Any]:
plan_check = self.validate_plan()
existing = self._read_state()
if existing:
own = self.verify_ownership()
if own["status"] == "PASS":
self.audit.append(
"RUNTIME_START_REUSE",
{"ownership": own},
)
return {
"status": "PASS",
"action": "ALREADY_RUNNING_RUN_OWNED",
"ownership": own,
}
# Fail closed on live-but-not-owned state. Clear only when PID is dead.
pid = int(existing.get("pid", -1))
if pid > 0 and pid_exists(pid):
node-7.example.invalid RuntimeError(
"EXISTING_RUNTIME_STATE_NOT_OWNED:"
+ json.dumps(own, ensure_ascii=False)
)
self.audit.append(
"STALE_RUNTIME_STATE_CLEARED_BEFORE_START",
{"old_state": existing, "ownership": own},
)
self._clear_state()
executable = str(self.plan["executable"])
argv = [str(x) for x in self.plan["argv"]]
cmd = [executable] + argv
env = dict(os.environ)
for k, v in self.plan["env"].items():
env[str(k)] = str(v)
self.stdout_path.parent.mkdir(
parents=True,
exist_ok=True,
)
stdout_f = self.stdout_path.open("ab", buffering=0)
stderr_f = self.stderr_path.open("ab", buffering=0)
try:
proc = subprocess.Popen(
cmd,
cwd=str(self.plan["cwd"]),
env=env,
stdin=subprocess.DEVNULL,
stdout=stdout_f,
stderr=stderr_f,
shell=False,
start_new_session=True,
close_fds=True,
)
finally:
stdout_f.close()
stderr_f.close()
pid = int(proc.pid)
pgid = int(os.getpgid(pid))
ticks = proc_start_ticks(pid)
cmdline = proc_cmdline(pid)
nonce = str(
self.plan.get(
"ownership_nonce",
secrets.token_hex(16),
)
)
state = {
"schema_version": 1,
"pid": pid,
"pgid": pgid,
"proc_start_ticks": ticks,
"expected_cmdline": cmdline,
"cmdline_sha256": canonical_sha(cmdline),
"launch_plan_sha256": canonical_sha(self.plan),
"ownership_nonce": nonce,
"ownership_role": self.plan["ownership_role"],
"started_unix": int(time.time()),
"readiness_url": self.readiness_url,
"stdout_path": str(self.stdout_path),
"stderr_path": str(self.stderr_path),
}
self._write_state(state)
self.audit.append(
"RUNTIME_STARTED",
{
"pid": pid,
"pgid": pgid,
"proc_start_ticks": ticks,
"cmdline_sha256": state["cmdline_sha256"],
"launch_plan": plan_check,
},
)
deadline = time.time() + self.startup_timeout
last = None
while time.time() < deadline:
last = self.probe()
if last["status"] == "PASS":
self.audit.append(
"RUNTIME_READY",
last,
)
return {
"status": "PASS",
"action": "STARTED",
"probe": last,
}
if not pid_exists(pid):
break
time.sleep(0.1)
self.audit.append(
"RUNTIME_START_NOT_READY",
{"last_probe": last},
)
# Do not broad-kill. Exact-stop only if ownership still verifies.
own = self.verify_ownership()
if own["status"] == "PASS":
self.stop()
node-7.example.invalid RuntimeError(
"RUNTIME_READINESS_TIMEOUT:"
+ json.dumps(last, ensure_ascii=False)
)
def stop(self) -> Dict[str, Any]:
state = self._read_state()
if not state:
return {
"status": "PASS",
"action": "NO_STATE",
}
own = self.verify_ownership()
if own["status"] != "PASS":
pid = int(state.get("pid", -1))
if pid > 0 and not pid_exists(pid):
self.audit.append(
"STALE_RUNTIME_STATE_CLEARED_ON_STOP",
{"state": state, "ownership": own},
)
self._clear_state()
return {
"status": "PASS",
"action": "STALE_STATE_CLEARED",
}
node-7.example.invalid RuntimeError(
"STOP_REFUSED_NOT_RUN_OWNED:"
+ json.dumps(own, ensure_ascii=False)
)
pgid = int(own["pgid"])
pid = int(own["pid"])
os.killpg(pgid, signal.SIGTERM)
self.audit.append(
"RUNTIME_SIGTERM_SENT",
{
"pid": pid,
"pgid": pgid,
"signal": "SIGTERM",
},
)
deadline = time.time() + self.stop_timeout
while time.time() < deadline:
if not pid_exists(pid):
self._clear_state()
self.audit.append(
"RUNTIME_STOPPED",
{
"pid": pid,
"pgid": pgid,
"method": "SIGTERM",
},
)
return {
"status": "PASS",
"action": "SIGTERM_EXACT_RUN_OWNED_PGID",
"pid": pid,
"pgid": pgid,
}
time.sleep(0.1)
# No SIGKILL escalation by design.
self.audit.append(
"RUNTIME_STOP_TIMEOUT",
{
"pid": pid,
"pgid": pgid,
"sigkill_escalation": False,
},
)
return {
"status": "BLOCK",
"action": "SIGTERM_TIMEOUT_NO_SIGKILL",
"pid": pid,
"pgid": pgid,
}
def recover(self) -> Dict[str, Any]:
state = self._read_state()
if not state:
return {
"status": "PASS",
"action": "NO_STATE",
}
own = self.verify_ownership()
if own["status"] == "PASS":
return {
"status": "PASS",
"action": "RUN_OWNED_PROCESS_STILL_ALIVE",
"ownership": own,
}
pid = int(state.get("pid", -1))
if pid > 0 and pid_exists(pid):
return {
"status": "BLOCK",
"action": "LIVE_PROCESS_OWNERSHIP_MISMATCH",
"ownership": own,
}
self.audit.append(
"STALE_RUNTIME_STATE_RECOVERED",
{
"state": state,
"ownership": own,
},
)
self._clear_state()
return {
"status": "PASS",
"action": "STALE_STATE_CLEARED",
}