owui-tools/vmagent/app/runner.py

#!/usr/bin/env python3
"""Claims VM jobs from a shared volume and runs them.

WHY FILES AND NOT HTTP

Exactly the reasoning in octave/run.sh, and the same protocol. The api
container lives inside owuitools-netguard, behind a default-deny filter
whose entire outbound allowlist is squid and valkey. Reaching a new
service from there means editing netguard/firewall.sh - a file whose
whole value is that it is short, verified at start, and rarely touched.

A shared volume needs none of that:

  /vmjobs/in/<id>/request.json   what to do
  /vmjobs/in/<id>/meta           the nonce, written by the api
  /vmjobs/in/<id>.ready          created LAST, so its appearance means
                                 the directory is complete. Without it
                                 this races the api and reads half a job.
  /vmjobs/run/<id>/CANCEL        the api's stop signal, mid-flight
  /vmjobs/out/<id>.out           the result, moved into place atomically

THE NONCE

Same purpose as the Octave runner's: the api stamps each job and checks
the stamp on the way back, so a result cannot be forged by anything else
that can write the volume. Unlike the Octave case the executed code never
sees this volume at all - it runs inside a VM with no access to it - so
this is defence in depth rather than the load-bearing control.

SERIAL, ON PURPOSE

One job at a time, because there is one VM at a time. A queue would only
let a second caller wait in a different place.
"""

from __future__ import annotations

import json
import logging
import os
import shutil
import signal
import sys
import threading
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import loop as loop_mod
from client import Vmd, VmdError, VmdUnavailable

JOBS = os.environ.get("OWUITOOLS_VM_JOBS_DIR", "/vmjobs")
POLL_INTERVAL = 0.25
REAP_MINUTES = int(os.environ.get("OWUITOOLS_VM_REAP_MINUTES", "180"))

# stdout and ./logs/vmagent.log both - see app/owuitools/logging_setup.py
# for why. Inlined rather than shared because this container deliberately
# copies almost nothing from the api image.
_handlers: list[logging.Handler] = [logging.StreamHandler(sys.stdout)]
_log_dir = os.environ.get("OWUITOOLS_LOG_DIR", "/logs")
_log_problem = ""
if _log_dir:
    try:
        os.makedirs(_log_dir, exist_ok=True)
        _handlers.append(logging.FileHandler(os.path.join(_log_dir, "vmagent.log"),
                                             encoding="utf-8"))
    except OSError as exc:
        _log_problem = f"{_log_dir} is not writable ({exc}); logging to stdout only"

logging.basicConfig(
    level=os.environ.get("OWUITOOLS_LOG_LEVEL", "INFO").upper(),
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    handlers=_handlers,
    force=True,
)
log = logging.getLogger("vmagent.runner")
if _log_problem:
    log.warning(_log_problem)

_stop = threading.Event()


def _watch_for_cancel(run_dir: str, vmd_path: str) -> threading.Thread:
    """Turn the api's CANCEL file into a real stop.

    The archived stack's cancel verb recorded an intention and killed
    nothing, so a cancelled run went to completion anyway. Here the file
    reaches vmd, which runs `systemctl stop owuivm.service` - and that
    kills qemu whatever any Python loop happens to be doing.
    """
    marker = os.path.join(run_dir, "CANCEL")

    def watch() -> None:
        while not _stop.is_set():
            if os.path.exists(marker):
                log.warning("cancellation requested; stopping the machine")
                try:
                    with Vmd(vmd_path) as vmd:
                        vmd.call("session_cancel")
                except (VmdError, VmdUnavailable) as exc:
                    log.warning("could not deliver the cancellation: %s", exc)
                return
            time.sleep(0.5)

    thread = threading.Thread(target=watch, daemon=True)
    thread.start()
    return thread


def run_job(job_id: str) -> None:
    run_dir = os.path.join(JOBS, "run", job_id)
    out_tmp = os.path.join(JOBS, "out", f"{job_id}.tmp")
    out_path = os.path.join(JOBS, "out", f"{job_id}.out")

    nonce = ""
    meta_path = os.path.join(run_dir, "meta")
    if os.path.exists(meta_path):
        for line in open(meta_path, encoding="utf-8", errors="replace"):
            if line.startswith("NONCE="):
                nonce = line.strip().split("=", 1)[1]
    # Deleted before anything else runs, exactly as the Octave runner
    # does: the nonce should live only in this process's memory.
    try:
        os.unlink(meta_path)
    except OSError:
        pass

    try:
        with open(os.path.join(run_dir, "request.json"), encoding="utf-8") as handle:
            request = json.load(handle)
    except (OSError, ValueError) as exc:
        _write_result(out_tmp, out_path, nonce, "failed",
                      f"FAILED: the job could not be read ({exc}).")
        return

    job = loop_mod.Job(
        task=str(request.get("task", ""))[:20000],
        os_name=str(request.get("os", "ubuntu")),
        interface=str(request.get("interface", "cli")),
        project=request.get("project") or None,
        chat_id=str(request.get("chat_id", "")),
        message_id=str(request.get("message_id", "")),
        owner=str(request.get("owner", "")),
        answer=str(request.get("answer", ""))[:4000],
    )

    _stop.clear()
    watcher = _watch_for_cancel(run_dir, os.environ.get("OWUITOOLS_VMD_SOCKET",
                                                        "/run/owuivm/agent/agent.sock"))
    started = time.monotonic()
    try:
        outcome = loop_mod.run(job)
        body = outcome.report
        state = outcome.state
        steps = outcome.steps
    except Exception as exc:  # noqa: BLE001
        log.exception("job %s failed outright", job_id)
        body = (f"FAILED: the VM agent stopped unexpectedly ({type(exc).__name__}). "
                "Nothing further is known about what it had done. Do not retry "
                "automatically; tell the user.")
        state, steps = "failed", 0
    finally:
        _stop.set()
        watcher.join(timeout=2.0)

    _write_result(out_tmp, out_path, nonce, state, body,
                  steps=steps, seconds=round(time.monotonic() - started, 1))
    shutil.rmtree(run_dir, ignore_errors=True)


def _write_result(tmp: str, final: str, nonce: str, state: str, body: str,
                  steps: int = 0, seconds: float = 0.0) -> None:
    payload = {"nonce": nonce, "state": state, "steps": steps,
               "seconds": seconds, "report": body}
    with open(tmp, "w", encoding="utf-8") as handle:
        json.dump(payload, handle)
    os.replace(tmp, final)          # atomic: the api never sees a partial file


def reap() -> None:
    """Nothing else reclaims these.

    A result the api gave up waiting for, or a job whose .ready never
    arrived, would sit on the volume until the container restarted.
    """
    cutoff = time.time() - REAP_MINUTES * 60
    for sub in ("out", "in", "run", "progress"):
        directory = os.path.join(JOBS, sub)
        if not os.path.isdir(directory):
            continue
        for name in os.listdir(directory):
            path = os.path.join(directory, name)
            try:
                if os.path.getmtime(path) > cutoff:
                    continue
                shutil.rmtree(path, ignore_errors=True) if os.path.isdir(path) else os.unlink(path)
            except OSError:
                pass


def main() -> int:
    for sub in ("in", "run", "out"):
        os.makedirs(os.path.join(JOBS, sub), mode=0o700, exist_ok=True)

    # Anything left over from a previous container start would otherwise
    # wedge its caller until timeout.
    shutil.rmtree(os.path.join(JOBS, "run"), ignore_errors=True)
    os.makedirs(os.path.join(JOBS, "run"), mode=0o700, exist_ok=True)
    try:
        with Vmd() as vmd:
            reaped = vmd.call("reap").get("reaped") or []
            if reaped:
                log.info("vmd reaped %d leftover(s) at start", len(reaped))
    except (VmdError, VmdUnavailable) as exc:
        log.warning("could not reach vmd at start: %s", exc)

    running = True

    def shutdown(_sig, _frame):
        nonlocal running
        running = False
        _stop.set()

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    log.info("vmagent runner ready, watching %s", JOBS)
    ticks = 0
    while running:
        claimed = False
        in_dir = os.path.join(JOBS, "in")
        for name in sorted(os.listdir(in_dir)) if os.path.isdir(in_dir) else []:
            if not name.endswith(".ready"):
                continue
            job_id = name[: -len(".ready")]
            if not job_id.isalnum():
                _unlink(os.path.join(in_dir, name))
                continue
            source = os.path.join(in_dir, job_id)
            target = os.path.join(JOBS, "run", job_id)
            try:
                os.rename(source, target)
            except OSError:
                _unlink(os.path.join(in_dir, name))
                continue
            _unlink(os.path.join(in_dir, name))
            claimed = True
            log.info("claimed job %s", job_id)
            run_job(job_id)
            log.info("finished job %s", job_id)

        ticks += 1
        if ticks >= 2400:               # roughly every ten minutes
            ticks = 0
            reap()
        _touch(os.path.join(JOBS, ".heartbeat"))
        if not claimed:
            time.sleep(POLL_INTERVAL)
    return 0


def _unlink(path: str) -> None:
    try:
        os.unlink(path)
    except OSError:
        pass


def _touch(path: str) -> None:
    try:
        with open(path, "a"):
            os.utime(path, None)
    except OSError:
        pass


if __name__ == "__main__":
    raise SystemExit(main())