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:
@@ -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]:
|
||||
|
||||
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user