vh/esh-pfi-infrastructure
1
0
Fork 0
Code Issues 7 Pull Requests Actions Packages Projects Releases Wiki Activity

feat(lora-worker): Phase 2 publish-step — copy succeeded LoRA into ComfyUI loras + published_lora_name

Browse Source 
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:
Vuong Hoang
2026-07-07 00:22:20 -07:00
parent f5c628c56d
commit 74dbfafdf1
5 changed files with 81 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
playbooks/deploy-lora-training-worker.yaml
+9
View File
@@ -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
services/lora-training-worker/tests/test_invocation.py
+8 -1
View File
@@ -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(
services/lora-training-worker/worker/config.py
+11
View File
@@ -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"))
services/lora-training-worker/worker/invocation.py
+10
View File
@@ -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.
services/lora-training-worker/worker/jobs.py
+43 -1
View File
@@ -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)