esh-pfi-infrastructure/services/lora-training-worker/worker/jobs.py

"""Training-job lifecycle: one job at a time, durable, with boot reconciliation.

The worker runs at most ONE sd-scripts process (§4.1 — 1-at-a-time; arbo's lease is the
real serializer, the worker's 409 is a backstop). A monitor thread owns the subprocess:
it spawns `accelerate launch …`, tails the log for step/loss, and lands a terminal status.

Durability (§4.3 boot reconciliation): job records persist to JSON. On startup any job
left non-terminal whose OS process is gone is marked `failed` ("worker restarted") — so
arbo's `GET /train/{id}` never sees a phantom `training` after a worker blip; arbo then
releases its lease per its boot rule.
"""

from __future__ import annotations

import json
import os
import re
import shutil
import signal
import subprocess
import threading
import time
import uuid
from pathlib import Path
from typing import Optional

from . import config
from .invocation import build_command, published_relative_path  # build_command raises InvalidTrainRequest (→422)

# sd-scripts / tqdm progress: `steps:  12%|█▏ | 50/400 [00:30<03:30, 1.66it/s, avr_loss=0.123]`
_STEP_RE = re.compile(r"(\d+)\s*/\s*(\d+)")
_LOSS_RE = re.compile(r"avr_loss=([0-9]*\.?[0-9]+)")

TERMINAL = {"succeeded", "failed", "cancelled"}


class Job:
    def __init__(self, worker_job_id: str, params: dict, log_path: Path):
        self.id = worker_job_id
        self.params = params
        self.log_path = log_path
        self.status = "queued"  # queued|preparing|training|succeeded|failed|cancelled
        self.step = 0
        self.total_steps = params.get("steps", 0)
        self.loss: Optional[float] = None
        self.eta_s: Optional[int] = None
        self.lora_path: Optional[str] = None
        self.published_lora_name: Optional[str] = None  # Phase 2: ComfyUI-relative loras path
        self.error: Optional[str] = None
        self.pid: Optional[int] = None
        self.started_at: Optional[float] = None
        self.finished_at: Optional[float] = None

    def public(self) -> dict:
        return {
            "worker_job_id": self.id,
            "status": self.status,
            "step": self.step,
            "total_steps": self.total_steps,
            "loss": self.loss,
            "eta_s": self.eta_s,
            "lora_path": self.lora_path,
            "published_lora_name": self.published_lora_name,
            "error": self.error,
        }

    def to_record(self) -> dict:
        return {
            **self.public(),
            "params": self.params,
            "log_path": str(self.log_path),
            "pid": self.pid,
            "started_at": self.started_at,
            "finished_at": self.finished_at,
        }

    @classmethod
    def from_record(cls, rec: dict) -> "Job":
        job = cls(rec["worker_job_id"], rec.get("params", {}), Path(rec.get("log_path", "/dev/null")))
        job.status = rec.get("status", "failed")
        job.step = rec.get("step", 0)
        job.total_steps = rec.get("total_steps", 0)
        job.loss = rec.get("loss")
        job.eta_s = rec.get("eta_s")
        job.lora_path = rec.get("lora_path")
        job.published_lora_name = rec.get("published_lora_name")
        job.error = rec.get("error")
        job.pid = rec.get("pid")
        job.started_at = rec.get("started_at")
        job.finished_at = rec.get("finished_at")
        return job


def _pid_alive(pid: Optional[int]) -> bool:
    if not pid:
        return False
    try:
        os.kill(pid, 0)
        return True
    except (OSError, ProcessLookupError):
        return False


class JobManager:
    """Singleton owner of the one active job + retained history. Thread-safe."""

    def __init__(self):
        self._lock = threading.Lock()
        self._jobs: dict[str, Job] = {}  # id -> Job (bounded to MAX_RETAINED_JOBS)
        self._active_id: Optional[str] = None
        self._proc: Optional[subprocess.Popen] = None
        config.STATE_DIR.mkdir(parents=True, exist_ok=True)
        config.LOG_DIR.mkdir(parents=True, exist_ok=True)
        self._load_and_reconcile()

    # ---- persistence ----------------------------------------------------------------
    def _persist(self) -> None:
        recs = [self._jobs[i].to_record() for i in list(self._jobs)]
        tmp = config.JOBS_FILE.with_suffix(".tmp")
        tmp.write_text(json.dumps(recs, indent=2))
        tmp.replace(config.JOBS_FILE)  # atomic

    def _load_and_reconcile(self) -> None:
        if not config.JOBS_FILE.exists():
            return
        try:
            recs = json.loads(config.JOBS_FILE.read_text())
        except (json.JSONDecodeError, OSError):
            return
        for rec in recs:
            job = Job.from_record(rec)
            # A non-terminal job from a prior worker life whose process is gone → failed.
            if job.status not in TERMINAL and not _pid_alive(job.pid):
                job.status = "failed"
                job.error = "worker restarted; training process not re-attachable"
                job.finished_at = job.finished_at or time.time()
            self._jobs[job.id] = job
        # If a job somehow survived as alive, keep it active (best-effort re-attach of state).
        for job in self._jobs.values():
            if job.status not in TERMINAL and _pid_alive(job.pid):
                self._active_id = job.id
                self._proc = None  # can't re-wrap the Popen; monitor via pid liveness
                threading.Thread(target=self._reattach_monitor, args=(job.id,), daemon=True).start()
        self._persist()

    # ---- public API -----------------------------------------------------------------
    def submit(self, req: dict) -> dict:
        """Start a train. Raises InvalidTrainRequest (→422) or Busy (→409)."""
        argv, env_overlay, params = build_command(req)  # validates; raises InvalidTrainRequest (→422)
        with self._lock:
            if self._active_id and self._jobs[self._active_id].status not in TERMINAL:
                raise Busy(f"a training job is already running: {self._active_id}")
            job_id = "wjob_" + uuid.uuid4().hex[:12]
            log_path = config.LOG_DIR / f"{job_id}.log"
            job = Job(job_id, params, log_path)
            self._jobs[job_id] = job
            self._active_id = job_id
            self._trim()
            self._persist()
        threading.Thread(target=self._run, args=(job_id, argv, env_overlay), daemon=True).start()
        return {"worker_job_id": job_id}

    def get(self, job_id: str) -> Optional[dict]:
        with self._lock:
            job = self._jobs.get(job_id)
            return job.public() if job else None

    def log_tail(self, job_id: str, max_bytes: int = 16384) -> Optional[str]:
        job = self._jobs.get(job_id)
        if not job or not job.log_path.exists():
            return None
        with open(job.log_path, "rb") as fh:
            fh.seek(0, os.SEEK_END)
            size = fh.tell()
            fh.seek(max(0, size - max_bytes))
            return fh.read().decode("utf-8", "replace")

    def cancel(self, job_id: str) -> Optional[dict]:
        with self._lock:
            job = self._jobs.get(job_id)
            if not job:
                return None
            if job.status in TERMINAL:
                return job.public()
            pid = job.pid
        # best-effort kill of the whole process group (accelerate spawns children)
        if pid:
            try:
                os.killpg(os.getpgid(pid), signal.SIGTERM)
            except (OSError, ProcessLookupError):
                pass
        with self._lock:
            job.status = "cancelled"
            job.error = "cancelled by request"
            job.finished_at = time.time()
            self._persist()
            return job.public()

    def active(self) -> Optional[str]:
        with self._lock:
            if self._active_id and self._jobs[self._active_id].status not in TERMINAL:
                return self._active_id
            return None

    # ---- internals ------------------------------------------------------------------
    def _trim(self) -> None:
        if len(self._jobs) <= config.MAX_RETAINED_JOBS:
            return
        # drop oldest terminal jobs first
        terminal = [i for i, j in self._jobs.items() if j.status in TERMINAL]
        for i in terminal[: len(self._jobs) - config.MAX_RETAINED_JOBS]:
            self._jobs.pop(i, None)

    def _set(self, job_id: str, **kw) -> None:
        with self._lock:
            job = self._jobs.get(job_id)
            if not job:
                return
            for k, v in kw.items():
                setattr(job, k, v)
            self._persist()

    def _run(self, job_id: str, argv: list[str], env_overlay: dict[str, str]) -> None:
        job = self._jobs[job_id]
        self._set(job_id, status="preparing", started_at=time.time())
        env = {**os.environ, **env_overlay}
        try:
            with open(job.log_path, "wb") as logf:
                proc = subprocess.Popen(
                    argv, cwd=str(config.SD_SCRIPTS_DIR), env=env,
                    stdout=logf, stderr=subprocess.STDOUT,
                    start_new_session=True,  # own process group → cancel can killpg
                )
                with self._lock:
                    self._proc = proc
                    job.pid = proc.pid
                    job.status = "training"
                    self._persist()
                self._poll_until_exit(job_id, proc)
        except (OSError, ValueError) as exc:
            self._set(job_id, status="failed", error=f"spawn failed: {exc}", finished_at=time.time())
            return

    def _poll_until_exit(self, job_id: str, proc: subprocess.Popen) -> None:
        job = self._jobs[job_id]
        while proc.poll() is None:
            self._update_progress(job_id)
            time.sleep(3)
        rc = proc.returncode
        self._update_progress(job_id)
        succeeded = False
        with self._lock:
            if job.status == "cancelled":
                return  # cancel already set the terminal state
            if rc == 0:
                lora = self._find_lora(job)
                if lora:
                    job.status, job.lora_path = "succeeded", str(lora)
                    succeeded = True
                else:
                    job.status, job.error = "failed", "process exited 0 but no .safetensors found"
            else:
                job.status = "failed"
                job.error = job.error or f"training process exited {rc} (see log)"
            job.finished_at = time.time()
            self._persist()
        if succeeded:
            self._publish_lora(job)  # Phase 2: copy to ComfyUI loras (best-effort, outside the lock)

    def _reattach_monitor(self, job_id: str) -> None:
        # A job whose pid survived a worker restart: watch pid liveness (no Popen handle).
        job = self._jobs[job_id]
        while _pid_alive(job.pid):
            self._update_progress(job_id)
            time.sleep(3)
        self._update_progress(job_id)
        succeeded = False
        with self._lock:
            if job.status not in TERMINAL:
                lora = self._find_lora(job)
                job.status = "succeeded" if lora else "failed"
                job.lora_path = str(lora) if lora else None
                if not lora:
                    job.error = "reattached process ended without a .safetensors"
                else:
                    succeeded = True
                job.finished_at = time.time()
                self._persist()
        if succeeded:
            self._publish_lora(job)  # Phase 2: copy to ComfyUI loras (best-effort, outside the lock)

    def _publish_lora(self, job: Job) -> None:
        """Phase 2 publish step: COPY the succeeded LoRA into ComfyUI's loras search path and
        set `published_lora_name` (the ComfyUI-relative path). Best-effort — a failure NEVER
        changes the `succeeded` status (download via output/ still works); it just leaves
        published_lora_name unset so arbo won't auto-register, and logs why. INV-T7-safe: a copy
        to a fixed computed path, no new free-form args."""
        if not job.lora_path:
            return
        try:
            rel = published_relative_path(
                job.params["output_dir"], job.params["output_name"], job.params.get("train_id")
            )
            dest = config.LORAS_PUBLISH_ROOT / rel
            dest.parent.mkdir(parents=True, exist_ok=True)
            try:
                os.chmod(dest.parent, 0o2775)  # group(arbotrain)-writable + ComfyUI-traversable
            except OSError:
                pass
            shutil.copy2(job.lora_path, dest)
            os.chmod(dest, 0o644)  # world-readable so ComfyUI (comfytoo) can load it
            with self._lock:
                job.published_lora_name = rel
                self._persist()
        except (OSError, KeyError, ValueError) as exc:
            # keep succeeded; omit published_lora_name; record the reason in the tailable log
            try:
                with open(job.log_path, "a") as lf:
                    lf.write(f"\n[worker] LoRA publish skipped (train still succeeded): {exc}\n")
            except OSError:
                pass

    def _update_progress(self, job_id: str) -> None:
        job = self._jobs.get(job_id)
        if not job or not job.log_path.exists():
            return
        try:
            with open(job.log_path, "rb") as fh:
                fh.seek(0, os.SEEK_END)
                fh.seek(max(0, fh.tell() - 8192))
                tail = fh.read().decode("utf-8", "replace")
        except OSError:
            return
        # tqdm uses \r; split on both so we see the latest bar frame
        frames = re.split(r"[\r\n]", tail)
        step, loss = job.step, job.loss
        for frame in frames:
            m = _STEP_RE.search(frame)
            if m and int(m.group(2)) == job.total_steps:  # match the step bar, not a random N/M
                step = int(m.group(1))
            lm = _LOSS_RE.search(frame)
            if lm:
                loss = float(lm.group(1))
        eta = None
        if step > 0 and job.started_at and job.total_steps:
            elapsed = time.time() - job.started_at
            eta = int((job.total_steps - step) * (elapsed / step))
        self._set(job_id, step=step, loss=loss, eta_s=eta)

    def _find_lora(self, job: Job) -> Optional[Path]:
        out_dir = Path(job.params.get("output_dir", ""))
        name = job.params.get("output_name", "")
        candidate = out_dir / f"{name}.safetensors"
        if candidate.exists():
            return candidate
        # fall back to newest .safetensors in the output dir
        try:
            sfts = sorted(out_dir.glob("*.safetensors"), key=lambda p: p.stat().st_mtime)
            return sfts[-1] if sfts else None
        except OSError:
            return None


class Busy(RuntimeError):
    """A second train was dispatched while one is active → HTTP 409."""


# module-level singleton
manager = JobManager()