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, sudo_password: str | None = None):
    job = JOBS[job_id]
    job["state"] = "running"
    try:
        actual_args = list(args)
        if sudo_password is not None:
            for i, arg in enumerate(actual_args):
                if isinstance(arg, str):
                    actual_args[i] = arg.replace("sudo -n", "sudo -S").replace("sudo ", "sudo -S ")

        proc = subprocess.Popen(
            actual_args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
            stdin=subprocess.PIPE if sudo_password is not None else None,
            bufsize=0,
            env={**os.environ, **(env or {})},
        )
        _PROCS[job_id] = proc

        if sudo_password is not None and proc.stdin:
            proc.stdin.write((sudo_password + "\n").encode("utf-8"))
            proc.stdin.flush()
            proc.stdin.close()

        _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")
        
        # Check if failed due to sudo authorization failure
        if proc.returncode != 0 and job["log"]:
            log_str = "\n".join(job["log"])
            if "a password is required" in log_str or "password" in log_str.lower() or "sudo:" in log_str:
                job["sudo_failed"] = True
    except Exception as exc:
        _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:
                _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:
                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,
              sudo_password: str | None = None, group: str | None = None) -> str:
    job_id = uuid.uuid4().hex[:12]
    # Mask password in log if present in args
    log_args = list(args)
    JOBS[job_id] = {
        "id": job_id, "label": label, "state": "queued", "group": group,
        "log": ["$ " + " ".join(shlex.quote(a) for a in log_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, sudo_password), daemon=True).start()
    return job_id


def active_in_group(group: str) -> dict | None:
    """Erster laufender/wartender Job einer Gruppe (z.B. 'maintenance'), sonst None.
    Basis für den Wartungs-Riegel: nur EIN System-Update gleichzeitig."""
    for j in JOBS.values():
        if j.get("group") == group and j.get("state") in ("running", "queued"):
            return j
    return None


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:
            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()]