feat(lora-worker): Phase 2 publish-step — copy succeeded LoRA into ComfyUI loras + published_lora_name
On a train reaching succeeded, IN ADDITION to output/{name}.safetensors
(unchanged download source), COPY it into ComfyUI's loras search path at
/storetank/arbo/models/loras/trained/{train_id}/{name}.safetensors and
return published_lora_name (the ComfyUI-relative LoraLoader string) in the
terminal GET /train/{id} payload (arbo Phase 2 auto-registration, §4.1/§7).
- Copy not move; a publish failure NEVER fails the train (keeps succeeded,
omits published_lora_name, logs the reason to the tailable run log).
- INV-T7-safe: a copy to a fixed computed path, no new free-form args.
- train_id derived from the handoff layout (output_dir.parent.name).
- Provisions loras/trained/ (arbotrain 2775, group-write per the Phase-1
lesson; world-readable/traversable for ComfyUI) via the deploy playbook.
- ComfyUI verified to resolve nested loras subfolders (no flat fallback).
- Pure path helper unit-tested; 16 tests green.
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user