diff --git a/services/lora-training-worker/tests/test_invocation.py b/services/lora-training-worker/tests/test_invocation.py index 29f9d7a..73a2b98 100644 --- a/services/lora-training-worker/tests/test_invocation.py +++ b/services/lora-training-worker/tests/test_invocation.py @@ -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( diff --git a/services/lora-training-worker/worker/app.py b/services/lora-training-worker/worker/app.py index fc91438..7d0a1a1 100644 --- a/services/lora-training-worker/worker/app.py +++ b/services/lora-training-worker/worker/app.py @@ -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") diff --git a/services/lora-training-worker/worker/invocation.py b/services/lora-training-worker/worker/invocation.py index b818c50..9e0ee59 100644 --- a/services/lora-training-worker/worker/invocation.py +++ b/services/lora-training-worker/worker/invocation.py @@ -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]: diff --git a/services/lora-training-worker/worker/jobs.py b/services/lora-training-worker/worker/jobs.py index 8f2e790..bbd8fd1 100644 --- a/services/lora-training-worker/worker/jobs.py +++ b/services/lora-training-worker/worker/jobs.py @@ -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: