space_runner.py

#!/usr/bin/env python3
import json
import os
import signal
import subprocess
import sys
import threading
from collections import deque
from http.server import BaseHTTPRequestHandler, HTTPServer


def main():
    cmd = sys.argv[1:]
    if not cmd:
        print("No command provided to space_runner.py", flush=True)
        sys.exit(2)

    # Start workload with captured output so we can expose failure tracebacks.
    child = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
    )

    recent_logs = deque(maxlen=400)

    state = {
        "status": "running",
        "pid": child.pid,
        "returncode": None,
        "command": cmd,
    }

    def stream_output():
        if child.stdout is None:
            return
        for line in child.stdout:
            line = line.rstrip("\n")
            recent_logs.append(line)
            print(line, flush=True)

    class Handler(BaseHTTPRequestHandler):
        def do_GET(self):
            payload = json.dumps({**state, "recent_logs": list(recent_logs)}).encode("utf-8")
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", str(len(payload)))
            self.end_headers()
            self.wfile.write(payload)

        def log_message(self, fmt, *args):
            return

    server = HTTPServer(("0.0.0.0", int(os.environ.get("PORT", "7860"))), Handler)

    log_thread = threading.Thread(target=stream_output, daemon=True)
    log_thread.start()

    def wait_child():
        rc = child.wait()
        state["returncode"] = rc
        state["status"] = "completed" if rc == 0 else "failed"

    watcher = threading.Thread(target=wait_child, daemon=True)
    watcher.start()

    def forward_signal(signum, _frame):
        if child.poll() is None:
            child.send_signal(signum)

    signal.signal(signal.SIGTERM, forward_signal)
    signal.signal(signal.SIGINT, forward_signal)

    print("Health server listening on 0.0.0.0:{}".format(os.environ.get("PORT", "7860")), flush=True)
    print("Executing workload: {}".format(" ".join(cmd)), flush=True)

    server.serve_forever()


if __name__ == "__main__":
    main()