diff --git a/playbooks/deploy-lora-training-worker.yaml b/playbooks/deploy-lora-training-worker.yaml index 392f8a5..5830745 100644 --- a/playbooks/deploy-lora-training-worker.yaml +++ b/playbooks/deploy-lora-training-worker.yaml @@ -15,6 +15,7 @@ vars: stage_dir: /tmp/lora-training-worker-stage install_dir: /opt/lora-training-worker handoff_dir: /worktank/arbo/train + loras_publish_dir: /storetank/arbo/models/loras/trained worker_user: llmuser arbo_user: lkraven # arbo container runs as uid 1000 = host lkraven group: arbotrain @@ -46,6 +47,14 @@ steps: sudo: true changed_when: "false" + - name: Create the Phase-2 LoRA publish dir (ComfyUI loras/trained, group-writable) + # 2775 (not 2770): world-readable + traversable so ComfyUI (uid 1025 comfytoo) can list + + # load; group arbotrain + group-WRITE so the worker (llmuser) can publish into it. setgid + # propagates the group to per-train subdirs (the Phase-1 group-write lesson). + shell: mkdir -p {{ loras_publish_dir }} && chgrp {{ group }} {{ loras_publish_dir }} && chmod 2775 {{ loras_publish_dir }} + sudo: true + changed_when: "false" + - name: Create the install dir owned by the worker user shell: mkdir -p {{ install_dir }} && chown {{ worker_user }}:{{ worker_user }} {{ install_dir }} sudo: true diff --git a/services/lora-training-worker/tests/test_invocation.py b/services/lora-training-worker/tests/test_invocation.py index bd5e866..29f9d7a 100644 --- a/services/lora-training-worker/tests/test_invocation.py +++ b/services/lora-training-worker/tests/test_invocation.py @@ -8,7 +8,7 @@ Pure functions, no GPU, no subprocess — runnable anywhere. import pytest from worker import config -from worker.invocation import InvalidTrainRequest, build_command +from worker.invocation import InvalidTrainRequest, build_command, published_relative_path def _req(**over): @@ -100,6 +100,13 @@ def test_base_model_outside_allowed_roots_rejected(): build_command(_req(base_model_path="/home/someone/evil.safetensors")) +def test_published_relative_path(): + # Phase 2 publish: derive {train_id} from the handoff output_dir -> ComfyUI-relative loras path + assert published_relative_path( + "/worktank/arbo/train/392cf898ac03/output", "sindra_lora" + ) == "trained/392cf898ac03/sindra_lora.safetensors" + + def test_storetank_checkpoint_allowed(): # the canonical SDXL store (2026-06-13 move) — arbo dispatches base_model_path from here argv, _, _ = build_command( diff --git a/services/lora-training-worker/worker/config.py b/services/lora-training-worker/worker/config.py index 0127bb8..c29cf6c 100644 --- a/services/lora-training-worker/worker/config.py +++ b/services/lora-training-worker/worker/config.py @@ -51,6 +51,17 @@ ALLOWED_MODEL_ROOTS = tuple( if p ) +# ---- LoRA publish step (Phase 2 — auto-registration into ComfyUI Generate) ------------- +# On a train reaching `succeeded`, IN ADDITION to output/{name}.safetensors (the unchanged +# download source), the worker COPIES the LoRA into ComfyUI's loras search path under +# PUBLISH_SUBDIR/{train_id}/ and reports `published_lora_name` (the path RELATIVE to the +# loras root — the exact string a ComfyUI LoraLoader.lora_name widget takes). ComfyUI +# (verified 2026-07-07) resolves loras SUBfolders, so the nested scheme works. A publish +# failure NEVER fails the train — it just omits published_lora_name (download still works). +# Stays inside INV-T7 (a copy to a fixed computed path; no new free-form args). +LORAS_PUBLISH_ROOT = Path(os.environ.get("LORA_WORKER_LORAS_ROOT", "/storetank/arbo/models/loras")) +PUBLISH_SUBDIR = "trained" # loras/trained/{train_id}/{name}.safetensors + # ---- Worker state + logs (survives a worker restart for boot reconciliation, §4.3) ----- STATE_DIR = Path(os.environ.get("LORA_WORKER_STATE_DIR", "/opt/lora-training-worker/state")) LOG_DIR = Path(os.environ.get("LORA_WORKER_LOG_DIR", "/opt/lora-training-worker/logs")) diff --git a/services/lora-training-worker/worker/invocation.py b/services/lora-training-worker/worker/invocation.py index 00a9e07..b818c50 100644 --- a/services/lora-training-worker/worker/invocation.py +++ b/services/lora-training-worker/worker/invocation.py @@ -106,6 +106,16 @@ def validate_request(req: dict) -> dict: } +def published_relative_path(output_dir: str, output_name: str) -> str: + """The ComfyUI-relative loras path for a succeeded LoRA (Phase 2 publish step). + + Derives {train_id} from the handoff layout `HANDOFF_ROOT/{train_id}/output` — i.e. the + name of the dir the LoRA was written into. Returns e.g. `trained/392cf898ac03/name.safetensors`, + the exact string a ComfyUI `LoraLoader.lora_name` widget takes. Pure — no filesystem touch.""" + train_id = Path(output_dir).parent.name + return f"{config.PUBLISH_SUBDIR}/{train_id}/{output_name}.safetensors" + + def build_command(req: dict) -> tuple[list[str], dict[str, str], dict]: """Validate `req` and return `(argv, env_overlay, params)` for the fixed sd-scripts invocation. diff --git a/services/lora-training-worker/worker/jobs.py b/services/lora-training-worker/worker/jobs.py index 0387178..8f2e790 100644 --- a/services/lora-training-worker/worker/jobs.py +++ b/services/lora-training-worker/worker/jobs.py @@ -15,6 +15,7 @@ from __future__ import annotations import json import os import re +import shutil import signal import subprocess import threading @@ -24,7 +25,7 @@ from pathlib import Path from typing import Optional from . import config -from .invocation import build_command # raises InvalidTrainRequest (→422), propagated by caller +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+)") @@ -44,6 +45,7 @@ class Job: 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 @@ -58,6 +60,7 @@ class Job: "loss": self.loss, "eta_s": self.eta_s, "lora_path": self.lora_path, + "published_lora_name": self.published_lora_name, "error": self.error, } @@ -80,6 +83,7 @@ class Job: 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") @@ -244,6 +248,7 @@ class JobManager: 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 @@ -251,6 +256,7 @@ class JobManager: 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: @@ -258,6 +264,8 @@ class JobManager: 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). @@ -266,6 +274,7 @@ class JobManager: 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) @@ -273,8 +282,41 @@ class JobManager: 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"]) + 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)