vh
b617a8b674
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.
191 lines
9.0 KiB
Python
191 lines
9.0 KiB
Python
"""Fixed-invocation command builder — the enforcement point for INV-T7.
|
|
|
|
arbo hands the worker a set of BOUNDED parameters; this module turns them into the ONE
|
|
`accelerate launch sdxl_train_network.py …` command the worker is allowed to run. Every
|
|
parameter is validated (type, range, allowlist, path-containment) before a single argv
|
|
element is produced. There is NO path from an arbo request to a free-form argument — the
|
|
flag set is a constant, only the *values* of vetted parameters vary.
|
|
|
|
`build_command` is a pure function (no I/O, no subprocess) so it is trivially unit-testable:
|
|
feed it a request dict, assert the argv + env. `InvalidTrainRequest` maps to HTTP 422.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from pathlib import Path
|
|
|
|
from . import config
|
|
|
|
# Safe token for names that land in a filename / kohya folder path. No shell metachars,
|
|
# no path separators, no leading dot — even though we never use a shell (argv list), this
|
|
# also protects the on-disk paths built from `output_name` / `trigger` / `subject_class`.
|
|
_SAFE_TOKEN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9 _.\-]{0,63}$")
|
|
_SAFE_OUTPUT_NAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_\-]{0,63}$") # filename stem, no spaces
|
|
|
|
|
|
class InvalidTrainRequest(ValueError):
|
|
"""A train request failed validation → HTTP 422 (never a silent downgrade)."""
|
|
|
|
|
|
def _require(cond: bool, msg: str) -> None:
|
|
if not cond:
|
|
raise InvalidTrainRequest(msg)
|
|
|
|
|
|
def _validate_under(path_str: str, roots: tuple[Path, ...], field: str) -> Path:
|
|
"""Resolve `path_str` and require it to sit under one of `roots` (no traversal escape)."""
|
|
_require(isinstance(path_str, str) and path_str.startswith("/"), f"{field} must be an absolute path")
|
|
p = Path(path_str)
|
|
# Reject traversal explicitly even before resolve() (belt): no '..' components.
|
|
_require(".." not in p.parts, f"{field} must not contain '..'")
|
|
resolved = p.resolve()
|
|
for root in roots:
|
|
try:
|
|
resolved.relative_to(root.resolve())
|
|
return resolved
|
|
except ValueError:
|
|
continue
|
|
allowed = ", ".join(str(r) for r in roots)
|
|
raise InvalidTrainRequest(f"{field} {resolved} is not under an allowed root ({allowed})")
|
|
|
|
|
|
def validate_request(req: dict) -> dict:
|
|
"""Validate + normalize an arbo /train request. Returns a clean param dict or raises."""
|
|
# ---- tier -> (steps, dim, res) --------------------------------------------------
|
|
tier_raw = req.get("tier")
|
|
_require(tier_raw in config.TIERS, f"tier must be one of {sorted(config.TIERS)}; got {tier_raw!r}")
|
|
tier = str(tier_raw)
|
|
spec = config.TIERS[tier]
|
|
|
|
# ---- device_index (int, 0=3090 / 1=A6000 under PCI_BUS_ID) ----------------------
|
|
device_index = req.get("device_index")
|
|
_require(isinstance(device_index, int) and device_index in (config.DEVICE_3090, config.DEVICE_A6000),
|
|
f"device_index must be {config.DEVICE_3090} (3090) or {config.DEVICE_A6000} (A6000)")
|
|
|
|
# ---- tier/device fit: quality (1024) is A6000-only (§4.6) -----------------------
|
|
if tier == "quality":
|
|
_require(device_index in config.QUALITY_ONLY_DEVICES,
|
|
"tier 'quality' (1024) requires the A6000 (device_index=1); the 3090 cannot fit it")
|
|
|
|
# ---- names (safe tokens; used in argv AND on-disk paths) ------------------------
|
|
trigger = req.get("trigger")
|
|
subject_class = req.get("subject_class")
|
|
output_name = req.get("output_name")
|
|
_require(isinstance(trigger, str) and bool(_SAFE_TOKEN.match(trigger)), "trigger is not a safe token")
|
|
_require(isinstance(subject_class, str) and bool(_SAFE_TOKEN.match(subject_class)),
|
|
"subject_class is not a safe token")
|
|
_require(isinstance(output_name, str) and bool(_SAFE_OUTPUT_NAME.match(output_name)),
|
|
"output_name is not a safe filename stem")
|
|
|
|
# ---- repeats + seed (bounded ints) ----------------------------------------------
|
|
repeats = req.get("repeats")
|
|
_require(isinstance(repeats, int) and 1 <= repeats <= 100, "repeats must be an int in [1, 100]")
|
|
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")
|
|
base_model_path = _validate_under(req.get("base_model_path", ""), config.ALLOWED_MODEL_ROOTS, "base_model_path")
|
|
|
|
return {
|
|
"tier": tier,
|
|
"steps": spec["steps"],
|
|
"dim": spec["dim"],
|
|
"resolution": spec["resolution"],
|
|
"device_index": device_index,
|
|
"trigger": trigger,
|
|
"subject_class": subject_class,
|
|
"output_name": output_name,
|
|
"repeats": repeats,
|
|
"seed": seed,
|
|
"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, train_id: str | None = None) -> str:
|
|
"""The ComfyUI-relative loras path for a succeeded LoRA (Phase 2 publish step).
|
|
|
|
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."""
|
|
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]:
|
|
"""Validate `req` and return `(argv, env_overlay, params)` for the fixed sd-scripts invocation.
|
|
|
|
argv is a LIST (never a shell string) — no metachar interpretation is possible. env_overlay
|
|
is merged onto os.environ by the caller (the CUDA ordering + allocator knobs from §4.6).
|
|
`params` is the vetted, normalized parameter dict (so the caller need not re-validate).
|
|
"""
|
|
p = validate_request(req)
|
|
h = config.SDXL_HPARAMS
|
|
batch = h["train_batch_size_full"] if p["device_index"] == config.DEVICE_A6000 else h["train_batch_size_lean"]
|
|
alpha = int(round(p["dim"] * float(h["network_alpha_ratio"])))
|
|
res = f"{p['resolution']},{p['resolution']}"
|
|
|
|
argv = [
|
|
str(config.ACCELERATE_BIN), "launch",
|
|
"--num_processes", "1",
|
|
"--num_cpu_threads_per_process", "1",
|
|
"--mixed_precision", h["mixed_precision"],
|
|
"--dynamo_backend", "no",
|
|
str(config.SDXL_TRAIN_SCRIPT),
|
|
"--pretrained_model_name_or_path", p["base_model_path"],
|
|
"--train_data_dir", p["dataset_dir"],
|
|
"--output_dir", p["output_dir"],
|
|
"--output_name", p["output_name"],
|
|
"--save_model_as", "safetensors",
|
|
"--save_precision", h["save_precision"],
|
|
"--mixed_precision", h["mixed_precision"],
|
|
# LoRA network (§4.6 tier dim; unet-only per the lean recipe)
|
|
"--network_module", "networks.lora",
|
|
"--network_dim", str(p["dim"]),
|
|
"--network_alpha", str(alpha),
|
|
"--network_train_unet_only",
|
|
# steps + resolution (§4.6 tier table)
|
|
"--max_train_steps", str(p["steps"]),
|
|
"--resolution", res,
|
|
"--enable_bucket",
|
|
"--train_batch_size", batch,
|
|
# memory-lean base (§4.6 — the proven Sindra flag set)
|
|
"--gradient_checkpointing",
|
|
"--cache_latents", "--cache_latents_to_disk",
|
|
"--cache_text_encoder_outputs", "--cache_text_encoder_outputs_to_disk",
|
|
"--optimizer_type", "adamw8bit",
|
|
"--no_half_vae",
|
|
"--sdpa",
|
|
"--caption_extension", ".txt", # CRITICAL: sd-scripts defaults to .caption → would skip our captions
|
|
# optimizer schedule (agent-discretion defaults; confirm vs Sindra)
|
|
"--learning_rate", h["learning_rate"],
|
|
"--lr_scheduler", h["lr_scheduler"],
|
|
"--lr_warmup_steps", h["lr_warmup_steps"],
|
|
"--min_snr_gamma", h["min_snr_gamma"],
|
|
"--noise_offset", h["noise_offset"],
|
|
"--max_data_loader_n_workers", h["max_data_loader_n_workers"],
|
|
"--persistent_data_loader_workers",
|
|
"--seed", str(p["seed"]),
|
|
]
|
|
|
|
env_overlay = {
|
|
"CUDA_DEVICE_ORDER": "PCI_BUS_ID", # index 0=3090, 1=A6000 (recipe-pinned ordering)
|
|
"CUDA_VISIBLE_DEVICES": str(p["device_index"]),
|
|
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
|
|
}
|
|
return argv, env_overlay, p
|