chore(ci): reclaim scratch from cancelled development jobs

Signed-off-by: functionstackx <47992694+functionstackx@users.noreply.github.com>

Author: functionstackx

.github/maintenance/owned_scratch_driver.py

@@ -0,0 +1,221 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Draft transport for one protected Actions-only maintenance run."""
+
+import argparse
+import hashlib
+import json
+import os
+import re
+import subprocess
+import sys
+import tarfile
+import tempfile
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts/ci"))
+from slurm_client import PYTHON, ActionsContext, Connection
+
+BRANCH = "refs/heads/maintenance/owned-scratch-cleanup-1018-1019"
+FILES = {"controller.json", "srun.log", "result.json"}
+SCRATCH_ROOTS = {
+    "1018": "/tmp/dynamo-rocm-1018-LYcWorev",
+    "1019": "/tmp/dynamo-rocm-1019-hN5nn17x",
+}
+
+WORKER = r"""
+import json, os, re, socket, subprocess, sys
+job = os.environ.get("SLURM_JOB_ID", "")
+if os.getuid() != 20011 or socket.gethostname().split(".")[0] != "smci300x-ccs-aus-e07-40":
+    raise SystemExit("Wrong allocation user or host")
+if not re.fullmatch(r"[1-9][0-9]*", job) or job in ("1018", "1019"):
+    raise SystemExit("Expected a new maintenance job")
+print(json.dumps({"event":"allocation", "job_id":job, "node":socket.gethostname(), "uid":os.getuid()}), flush=True)
+payload = sys.stdin.read()
+result = subprocess.run(["/usr/bin/sudo", "-n", "/usr/bin/python3", "-I", "-", job], input=payload, text=True, timeout=480, cwd="/")
+raise SystemExit(result.returncode)
+"""
+
+REMOTE = r"""
+import fcntl, hashlib, json, os, re, stat, subprocess, sys, time
+from pathlib import Path
+run_key, sha = sys.argv[1:]
+if os.getuid() != 20011 or not re.fullmatch(r"gh-[1-9][0-9]*-[1-9][0-9]*", run_key) or not re.fullmatch(r"[a-f0-9]{40}", sha):
+    raise SystemExit("Invalid maintenance identity")
+os.umask(0o077)
+base = Path.home().resolve() / "dynamo-rocm-ci" / "maintenance"
+base.mkdir(mode=0o700, parents=True, exist_ok=True)
+if base.resolve() != base or base.stat().st_uid != 20011:
+    raise SystemExit("Untrusted maintenance directory")
+lock = os.fdopen(os.open(base / "jobs-1018-1019.lock", os.O_WRONLY | os.O_CREAT | os.O_NOFOLLOW | os.O_NONBLOCK, 0o600), "a")
+lock_info = os.fstat(lock.fileno())
+if not stat.S_ISREG(lock_info.st_mode) or lock_info.st_uid != 20011 or lock_info.st_nlink != 1:
+    raise SystemExit("Untrusted maintenance lock")
+fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
+run = base / run_key
+run.mkdir(mode=0o700)
+data = json.load(sys.stdin)
+controller = {"run_key":run_key,"source_sha":sha,"started_at":time.time(),"payload_sha256":hashlib.sha256(data["payload"].encode()).hexdigest()}
+(run / "controller.json").write_text(json.dumps(controller)+"\n")
+print(json.dumps({"run_dir":str(run),**controller}), flush=True)
+result = {"status":"failed","run_key":run_key,"source_sha":sha}
+try:
+    queue = subprocess.run(["squeue","--noheader","--user=20011","--format=%A"], check=True, capture_output=True, text=True, timeout=30)
+    rows = queue.stdout.split()
+    if any(not value.isdecimal() for value in rows) or {"1018","1019"}.intersection(rows):
+        raise RuntimeError("Old jobs are active or queue identity is ambiguous")
+    command = ["srun","--partition=compute-1","--nodelist=smci300x-ccs-aus-e07-40","--nodes=1","--ntasks=1","--cpus-per-task=1","--mem=1G","--gres=none","--time=00:10:00","--immediate=60","--kill-on-bad-exit=1","--chdir=/tmp","--job-name=owned-scratch-"+run_key,"/usr/bin/python3","-I","-c",data["worker"]]
+    with (run / "srun.log").open("w") as output:
+        completed = subprocess.run(command, input=data["payload"], text=True, stdout=output, stderr=subprocess.STDOUT, timeout=720)
+    result["returncode"] = completed.returncode
+    result["status"] = "passed" if completed.returncode == 0 else "failed"
+    completed.check_returncode()
+finally:
+    result["finished_at"] = time.time()
+    (run / "result.json").write_text(json.dumps(result)+"\n")
+print(json.dumps(result), flush=True)
+"""
+
+COLLECT = r"""
+import os, re, stat, sys, tarfile
+from pathlib import Path
+key = sys.argv[1]
+if os.getuid() != 20011 or not re.fullmatch(r"gh-[1-9][0-9]*-[1-9][0-9]*", key):
+    raise SystemExit("Invalid collection identity")
+root = Path.home().resolve() / "dynamo-rocm-ci" / "maintenance" / key
+if root.exists() and (root.resolve() != root or root.stat().st_uid != 20011):
+    raise RuntimeError("Untrusted maintenance directory")
+with tarfile.open(fileobj=sys.stdout.buffer, mode="w|") as archive:
+    for name in ("controller.json", "srun.log", "result.json"):
+        path = root / name
+        if not path.exists():
+            continue
+        with os.fdopen(os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK), "rb") as source:
+            info = os.fstat(source.fileno())
+            if not stat.S_ISREG(info.st_mode) or info.st_uid != 20011 or info.st_size > 10_000_000:
+                raise RuntimeError("Untrusted maintenance evidence")
+            entry = tarfile.TarInfo(name)
+            entry.size = info.st_size
+            entry.mode = 0o600
+            archive.addfile(entry, source)
+"""
+
+
+def verify_success(context, seen):
+    if seen != FILES:
+        raise RuntimeError("Successful maintenance is missing required evidence")
+    controller = json.loads((context.output / "controller.json").read_text())
+    result = json.loads((context.output / "result.json").read_text())
+    for value in (controller, result):
+        if (
+            value.get("run_key") != context.run_key
+            or value.get("source_sha") != context.sha
+        ):
+            raise RuntimeError("Maintenance evidence identity mismatch")
+    payload = Path(__file__).with_name("owned_scratch_payload.py").read_bytes()
+    if controller.get("payload_sha256") != hashlib.sha256(payload).hexdigest():
+        raise RuntimeError("Maintenance payload identity mismatch")
+    if result.get("status") != "passed" or result.get("returncode") != 0:
+        raise RuntimeError("Maintenance evidence does not report success")
+    events = []
+    for line in (context.output / "srun.log").read_text().splitlines():
+        if line.startswith("{"):
+            value = json.loads(line)
+            if isinstance(value, dict) and "event" in value:
+                events.append(value)
+
+    def single(name):
+        found = [value for value in events if value["event"] == name]
+        if len(found) != 1:
+            raise RuntimeError(f"Missing or repeated maintenance event: {name}")
+        return found[0]
+
+    allocation, before, after = (
+        single(name) for name in ("allocation", "before", "after")
+    )
+    job = allocation.get("job_id", "")
+    if not re.fullmatch(r"[1-9][0-9]*", job) or job in SCRATCH_ROOTS:
+        raise RuntimeError("Invalid maintenance allocation evidence")
+    if (
+        allocation.get("uid") != 20011
+        or allocation.get("node", "").split(".")[0] != "smci300x-ccs-aus-e07-40"
+    ):
+        raise RuntimeError("Unexpected maintenance allocation user or node")
+    if (
+        before.get("maintenance_job") != job
+        or before.get("node") != "smci300x-ccs-aus-e07-40"
+    ):
+        raise RuntimeError("Root helper allocation evidence mismatch")
+    if after.get("status") != "passed" or any(
+        value["event"] == "failure" for value in events
+    ):
+        raise RuntimeError("Root helper evidence does not report success")
+    for phase in (before, after):
+        targets = phase.get("targets", {})
+        if set(targets) != set(SCRATCH_ROOTS) or any(
+            targets[key].get("scratch") != path for key, path in SCRATCH_ROOTS.items()
+        ):
+            raise RuntimeError("Maintenance target evidence mismatch")
+    if any(after["targets"][key].get("exists") is not False for key in SCRATCH_ROOTS):
+        raise RuntimeError("Maintenance targets remain after reported success")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("action", choices=("run", "collect"))
+    args = parser.parse_args()
+    context = ActionsContext.from_environment()
+    if (
+        os.environ.get("GITHUB_REPOSITORY") != "SemiAnalysisAI/dynamo"
+        or os.environ.get("GITHUB_REF") != BRANCH
+        or context.expected_uid != 20011
+    ):
+        raise RuntimeError(
+            "One-off maintenance is restricted to its reviewed fork branch"
+        )
+    context.output.mkdir(parents=True, exist_ok=True)
+    connection = Connection(context)
+    if args.action == "run":
+        payload = Path(__file__).with_name("owned_scratch_payload.py").read_text()
+        (context.output / "payload-sha256.txt").write_text(
+            hashlib.sha256(payload.encode()).hexdigest() + "\n"
+        )
+        with (context.output / "transport.log").open("w") as output:
+            connection.call(
+                [*PYTHON, "-c", REMOTE, context.run_key, context.sha],
+                input=json.dumps({"payload": payload, "worker": WORKER}),
+                text=True,
+                stdout=output,
+                stderr=subprocess.STDOUT,
+                timeout=780,
+            )
+    else:
+        with tempfile.TemporaryFile() as stream:
+            connection.call(
+                [*PYTHON, "-c", COLLECT, context.run_key], stdout=stream, timeout=60
+            )
+            stream.seek(0)
+            with tarfile.open(fileobj=stream, mode="r:") as archive:
+                seen = set()
+                for member in archive:
+                    if (
+                        member.name not in FILES
+                        or member.name in seen
+                        or not member.isfile()
+                        or member.size > 10_000_000
+                    ):
+                        raise RuntimeError("Unexpected maintenance artifact")
+                    seen.add(member.name)
+                    (context.output / member.name).write_bytes(
+                        archive.extractfile(member).read()
+                    )
+        outcome = os.environ.get("MAINTENANCE_OUTCOME")
+        if outcome not in {"success", "failure", "cancelled", "skipped"}:
+            raise RuntimeError("Expected the Actions maintenance step outcome")
+        if outcome == "success":
+            verify_success(context, seen)
+
+
+if __name__ == "__main__":
+    main()