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

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

Browse Source 
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:
vh
2026-07-07 01:49:24 -07:00
parent 74dbfafdf1
commit b617a8b674
4 changed files with 34 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
services/lora-training-worker/worker/invocation.py
+15 -5
View File
@@ -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]: