mission-control-v2/backend/services/jobengine.py

"""
Mini-Job-System: Hintergrund-Prozesse mit Live-Log + Download-Fortschritt.
Portiert aus Mission Control v1 (jobengine.py). In-Memory, ein Daemon-Thread je Job.
"""

import glob
import os
import shlex
import subprocess
import threading
import time
import uuid

JOBS: dict[str, dict] = {}
_PROCS: dict[str, subprocess.Popen] = {}
_LOG_CAP = 400


def _append_log(job: dict, line: str) -> None:
    job["log"].append(line)
    if len(job["log"]) > _LOG_CAP:
        del job["log"][0]


def _pump_output(job: dict, stream) -> None:
    """Liest byteweise; `\\r` (tqdm/hf-Fortschritt) überschreibt die letzte Zeile."""
    buf = b""
    overwrite = False
    pending_cr = False

    def commit():
        line = buf.decode("utf-8", "replace")
        if overwrite and job["log"]:
            job["log"][-1] = line
        else:
            _append_log(job, line)

    while True:
        ch = stream.read(1)
        if not ch:
            break
        if pending_cr:
            pending_cr = False
            if ch == b"\n":
                commit(); overwrite = False; buf = b""
                continue
            commit(); overwrite = True; buf = b""
        if ch == b"\r":
            pending_cr = True
        elif ch == b"\n":
            commit(); overwrite = False; buf = b""
        else:
            buf += ch
    if pending_cr:
        commit(); overwrite = True; buf = b""
    if buf:
        commit()


def _run_job(job_id: str, args: list[str], env: dict | None = None):
    job = JOBS[job_id]
    job["state"] = "running"
    try:
        proc = subprocess.Popen(
            args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0,
            env={**os.environ, **(env or {})},
        )
        _PROCS[job_id] = proc
        _pump_output(job, proc.stdout)
        proc.wait()
        job["returncode"] = proc.returncode
        job["state"] = "canceled" if job.get("canceled") else ("done" if proc.returncode == 0 else "failed")
    except Exception as exc:  # noqa: BLE001
        _append_log(job, f"[mc] Fehler: {exc}")
        job["state"] = "failed"
        job["returncode"] = -1
    finally:
        _PROCS.pop(job_id, None)
        job["finished_at"] = time.time()
        cb = job.pop("_on_done", None)
        if cb and job["state"] == "done":
            try:
                cb()
            except Exception as exc:  # noqa: BLE001
                _append_log(job, f"[mc] Nachbearbeitung-Fehler: {exc}")


def attach_download_progress(job_id: str, local_dir: str, total_bytes: int) -> None:
    """Fortschritt in % aus wachsenden *.incomplete-Dateien (hf schreibt sie)."""
    if not total_bytes or total_bytes <= 0:
        return
    job = JOBS.get(job_id)
    if job is not None:
        job["progress"] = 0
        job["total_bytes"] = total_bytes

    def _watch():
        pat = os.path.join(local_dir, ".cache", "huggingface", "download", "**", "*.incomplete")
        prev_t = prev_b = None
        rate = 0.0
        while True:
            j = JOBS.get(job_id)
            if not j or j["state"] in ("done", "failed", "canceled"):
                break
            try:
                inc = glob.glob(pat, recursive=True)
                cur = sum(os.path.getsize(f) for f in inc) if inc else 0
                if cur:
                    j["progress"] = min(99, int(cur * 100 / total_bytes))
                    j["done_bytes"] = cur
                    now = time.time()
                    if prev_t is not None and now > prev_t and cur >= prev_b:
                        inst = (cur - prev_b) / (now - prev_t)
                        rate = inst if rate == 0 else 0.3 * inst + 0.7 * rate
                        if rate > 0:
                            j["rate_bps"] = rate
                            j["eta_s"] = int((total_bytes - cur) / rate)
                    prev_t, prev_b = now, cur
            except Exception:  # noqa: BLE001
                pass
            time.sleep(1.0)
        j = JOBS.get(job_id)
        if j and j["state"] == "done":
            j["progress"] = 100
            j.pop("eta_s", None)

    threading.Thread(target=_watch, daemon=True).start()


def start_job(args: list[str], label: str, env: dict | None = None, on_done=None) -> str:
    job_id = uuid.uuid4().hex[:12]
    JOBS[job_id] = {
        "id": job_id, "label": label, "state": "queued",
        "log": ["$ " + " ".join(shlex.quote(a) for a in args)],
        "returncode": None, "started_at": time.time(), "finished_at": None,
    }
    if on_done:
        JOBS[job_id]["_on_done"] = on_done
    threading.Thread(target=_run_job, args=(job_id, args, env), daemon=True).start()
    return job_id


def cancel_job(job_id: str) -> bool:
    job = JOBS.get(job_id)
    if not job or job["state"] in ("done", "failed", "canceled"):
        return False
    job["canceled"] = True
    _append_log(job, "[mc] Abbruch angefordert…")
    proc = _PROCS.get(job_id)
    if proc is not None:
        try:
            proc.terminate()
        except Exception:  # noqa: BLE001
            pass
    else:
        job["state"] = "canceled"
        job["finished_at"] = time.time()
    return True


def public_jobs() -> list[dict]:
    """Jobs ohne interne Felder (_on_done) für die API."""
    return [{k: v for k, v in j.items() if not k.startswith("_")} for j in JOBS.values()]