aios/services/aios-spark-runner/runner.py

#!/usr/bin/env python3
"""AIOS Spark Runner — GitOps-managed isolated task executor.

Runs with least privilege (aios-spark-runner user).
Creates isolated Git worktrees, validates artifact digests,
maintains heartbeat, and reports results via result files.
"""
import os
import sys
import time
import json
import stat
import hashlib
import subprocess
import shutil
import tempfile
import signal
from pathlib import Path

STATE_DIR = Path(os.environ.get("AIOS_SPARK_STATE_DIR", "/var/lib/aios-spark-runner"))
RUNTIME_DIR = Path(os.environ.get("AIOS_SPARK_RUNTIME_DIR", "/run/aios-spark-runner"))
CONFIG_DIR = Path(os.environ.get("AIOS_SPARK_CONFIG_DIR", "/etc/aios-spark-runner"))
WORKTREE_BASE = Path(os.environ.get("AIOS_SPARK_WORKTREE_BASE", str(STATE_DIR / "worktrees")))
REPO_PATH = Path(os.environ.get("AIOS_REPO_PATH", "/data/aios-gitops/repo.git"))
POLL_INTERVAL = int(os.environ.get("AIOS_SPARK_POLL_INTERVAL", "10"))
HEARTBEAT_FRESHNESS = int(os.environ.get("AIOS_SPARK_HEARTBEAT_FRESHNESS", "60"))
TASK_TIMEOUT = int(os.environ.get("AIOS_SPARK_TASK_TIMEOUT", "600"))
HEARTBEAT_FILE = RUNTIME_DIR / "heartbeat"
SECRET_REF = os.environ.get("AIOS_SPARK_SECRET_REF", "")

def log(msg):
    ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    print(f"[{ts}] {msg}", flush=True)

def die(msg, code=1):
    log(f"FATAL: {msg}")
    sys.exit(code)

def heartbeat():
    RUNTIME_DIR.mkdir(parents=True, exist_ok=True)
    HEARTBEAT_FILE.write_text(str(time.time()))

def verify_artifact(digest: str, path: Path):
    """Verify artifact SHA-256 digest before use."""
    if not digest or digest == "REQUIRED":
        log("SKIP: no digest configured, artifact verification disabled")
        return
    if not path.exists():
        die(f"artifact missing: {path}")
    h = hashlib.sha256()
    h.update(path.read_bytes())
    actual = h.hexdigest()
    if actual != digest:
        die(f"artifact digest mismatch: {path} expected={digest} actual={actual}")
    log(f"artifact verified: {path}")

def provision_worktree(commit_or_branch: str) -> Path:
    """Create an isolated Git worktree from the canonical repository."""
    WORKTREE_BASE.mkdir(parents=True, exist_ok=True)
    wt_dir = WORKTREE_BASE / f"spark-{int(time.time())}"
    subprocess.run(
        ["git", "-C", str(REPO_PATH), "worktree", "add", str(wt_dir), commit_or_branch],
        check=True, capture_output=True, timeout=120)
    log(f"worktree created: {wt_dir} @ {commit_or_branch}")
    return wt_dir

def cleanup_worktree(wt_dir: Path):
    """Remove a Git worktree cleanly."""
    if not wt_dir.exists():
        return
    subprocess.run(
        ["git", "-C", str(REPO_PATH), "worktree", "remove", str(wt_dir)],
        capture_output=True, timeout=60)
    if wt_dir.exists():
        shutil.rmtree(wt_dir, ignore_errors=True)
    log(f"worktree removed: {wt_dir}")

def run_task(task_cmd: list, workdir: Path) -> dict:
    """Run a task with sandboxing and timeout, return result dict."""
    log(f"running task: {task_cmd} in {workdir}")
    try:
        r = subprocess.run(
            task_cmd, cwd=workdir, capture_output=True, text=True,
            timeout=TASK_TIMEOUT)
        return {"exit": r.returncode, "stdout": r.stdout, "stderr": r.stderr}
    except subprocess.TimeoutExpired:
        return {"exit": -1, "stdout": "", "stderr": "TIMEOUT"}
    except Exception as exc:
        return {"exit": -1, "stdout": "", "stderr": str(exc)}

def main():
    log("AIOS Spark Runner starting")
    if not REPO_PATH.exists():
        die(f"repository not found: {REPO_PATH}")

    while True:
        heartbeat()
        # Poll for task files in state directory
        tasks = sorted(STATE_DIR.glob("task-*.json"))
        for tf in tasks:
            try:
                task = json.loads(tf.read_text())
            except Exception as e:
                log(f"BAD TASK FILE: {tf} {e}")
                tf.unlink(missing_ok=True)
                continue
            commit = task.get("commit", "HEAD")
            cmd = task.get("command", [])
            wt = None
            try:
                wt = provision_worktree(commit)
                result = run_task(cmd, wt)
                result_file = STATE_DIR / f"result-{tf.stem}.json"
                result_file.write_text(json.dumps(result, indent=2))
                log(f"task complete: {tf} exit={result['exit']}")
            finally:
                if wt:
                    cleanup_worktree(wt)
                tf.unlink(missing_ok=True)
        time.sleep(POLL_INTERVAL)

if __name__ == "__main__":
    main()