feat(lora-worker): add optional train_id to POST /train (explicit publish-path namespace)

comfy-dev's explicit-over-implicit call: arbo now sends train_id, so the
worker no longer derives the loras/trained/{train_id}/ namespace from
output_dir.parent (which coupled it to arbo's handoff layout). train_id is
optional + path-safe-validated; when present it wins, else the path
derivation remains as the fallback. Wired through TrainRequest ->
validate_request -> published_relative_path -> _publish_lora. 18 tests green.
This commit is contained in:
2026-07-07 01:49:24 -07:00
parent 74dbfafdf1
commit b617a8b674
4 changed files with 34 additions and 6 deletions
@@ -107,6 +107,21 @@ def test_published_relative_path():
) == "trained/392cf898ac03/sindra_lora.safetensors"
def test_published_relative_path_explicit_train_id_wins():
# explicit train_id (arbo now sends it) takes precedence over the path-derived fallback
assert published_relative_path(
"/worktank/arbo/train/whatever/output", "sindra_lora", train_id="abc123"
) == "trained/abc123/sindra_lora.safetensors"
def test_train_id_optional_and_validated():
# absent -> OK (falls back to derivation); present+safe -> OK; present+unsafe -> 422
build_command(_req()) # no train_id key
build_command(_req(train_id="392cf898ac03"))
with pytest.raises(InvalidTrainRequest, match="train_id"):
build_command(_req(train_id="../../etc"))
def test_storetank_checkpoint_allowed():
# the canonical SDXL store (2026-06-13 move) — arbo dispatches base_model_path from here
argv, _, _ = build_command(
@@ -39,6 +39,7 @@ class TrainRequest(BaseModel):
tier: str
device_index: int = Field(ge=0, le=1)
seed: int = Field(default=42, ge=0)
train_id: str | None = None # Phase 2: arbo's id, the publish-path namespace (optional)
@app.post("/train")
@@ -84,6 +84,14 @@ def validate_request(req: dict) -> dict:
seed = req.get("seed", 42)
_require(isinstance(seed, int) and 0 <= seed <= 2**31 - 1, "seed must be a non-negative int32")
# ---- train_id (optional; arbo's id — the publish-path namespace, Phase 2) -------
# Explicit over the derive-from-output_dir.parent fallback (decouples the worker from
# arbo's handoff layout). Path-safe token since it lands in the loras/trained/{train_id}/ path.
train_id = req.get("train_id")
if train_id is not None:
_require(isinstance(train_id, str) and bool(_SAFE_OUTPUT_NAME.match(train_id)),
"train_id is not a safe token")
# ---- paths (containment-checked) ------------------------------------------------
dataset_dir = _validate_under(req.get("dataset_dir", ""), (config.HANDOFF_ROOT,), "dataset_dir")
output_dir = _validate_under(req.get("output_dir", ""), (config.HANDOFF_ROOT,), "output_dir")
@@ -103,17 +111,19 @@ def validate_request(req: dict) -> dict:
"dataset_dir": str(dataset_dir),
"output_dir": str(output_dir),
"base_model_path": str(base_model_path),
"train_id": train_id,
}
def published_relative_path(output_dir: str, output_name: str) -> str:
def published_relative_path(output_dir: str, output_name: str, train_id: str | None = None) -> 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`,
Uses the explicit `train_id` (arbo now sends it — decoupled from the path layout); falls
back to deriving it from the handoff layout `HANDOFF_ROOT/{train_id}/output` (the name of the
dir the LoRA was written into) when absent. 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"
tid = train_id or Path(output_dir).parent.name
return f"{config.PUBLISH_SUBDIR}/{tid}/{output_name}.safetensors"
def build_command(req: dict) -> tuple[list[str], dict[str, str], dict]:
+3 -1
View File
@@ -298,7 +298,9 @@ class JobManager:
if not job.lora_path:
return
try:
rel = published_relative_path(job.params["output_dir"], job.params["output_name"])
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: