"""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()