vh
74dbfafdf1
On a train reaching succeeded, IN ADDITION to output/{name}.safetensors
(unchanged download source), COPY it into ComfyUI's loras search path at
/storetank/arbo/models/loras/trained/{train_id}/{name}.safetensors and
return published_lora_name (the ComfyUI-relative LoraLoader string) in the
terminal GET /train/{id} payload (arbo Phase 2 auto-registration, §4.1/§7).
- Copy not move; a publish failure NEVER fails the train (keeps succeeded,
omits published_lora_name, logs the reason to the tailable run log).
- INV-T7-safe: a copy to a fixed computed path, no new free-form args.
- train_id derived from the handoff layout (output_dir.parent.name).
- Provisions loras/trained/ (arbotrain 2775, group-write per the Phase-1
lesson; world-readable/traversable for ComfyUI) via the deploy playbook.
- ComfyUI verified to resolve nested loras subfolders (no flat fallback).
- Pure path helper unit-tested; 16 tests green.
148 lines
5.4 KiB
Python
148 lines
5.4 KiB
Python
"""Tests for the fixed-invocation builder — the INV-T7 enforcement surface.
|
|
|
|
These assert that (a) a valid request produces exactly the proven §4.6 command shape, and
|
|
(b) every out-of-bounds / unsafe / mis-fit request is REJECTED before any argv is produced.
|
|
Pure functions, no GPU, no subprocess — runnable anywhere.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from worker import config
|
|
from worker.invocation import InvalidTrainRequest, build_command, published_relative_path
|
|
|
|
|
|
def _req(**over):
|
|
base = {
|
|
"dataset_dir": "/worktank/arbo/train/t123/dataset",
|
|
"base_model_path": "/opt/fluxgym/models/sdxl/base.safetensors",
|
|
"output_dir": "/worktank/arbo/train/t123/out",
|
|
"output_name": "char_t123",
|
|
"trigger": "ohwx",
|
|
"subject_class": "woman",
|
|
"repeats": 10,
|
|
"tier": "balanced",
|
|
"device_index": 0,
|
|
"seed": 42,
|
|
}
|
|
base.update(over)
|
|
return base
|
|
|
|
|
|
def _flag_value(argv, flag):
|
|
return argv[argv.index(flag) + 1]
|
|
|
|
|
|
# ---- happy path ------------------------------------------------------------------------
|
|
|
|
def test_balanced_on_3090_builds_proven_shape():
|
|
argv, env, params = build_command(_req(tier="balanced", device_index=0))
|
|
assert str(config.ACCELERATE_BIN) == argv[0]
|
|
assert "launch" == argv[1]
|
|
assert str(config.SDXL_TRAIN_SCRIPT) in argv
|
|
# tier table (§4.6): balanced -> 1500 steps, dim 32, res 768
|
|
assert _flag_value(argv, "--max_train_steps") == "1500"
|
|
assert _flag_value(argv, "--network_dim") == "32"
|
|
assert _flag_value(argv, "--resolution") == "768,768"
|
|
# the load-bearing lean flags + caption gotcha
|
|
assert "--network_train_unet_only" in argv
|
|
assert "--gradient_checkpointing" in argv
|
|
assert _flag_value(argv, "--caption_extension") == ".txt"
|
|
assert _flag_value(argv, "--optimizer_type") == "adamw8bit"
|
|
# env: PCI_BUS_ID ordering + the assigned device + allocator
|
|
assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID"
|
|
assert env["CUDA_VISIBLE_DEVICES"] == "0"
|
|
assert env["PYTORCH_CUDA_ALLOC_CONF"] == "expandable_segments:True"
|
|
assert params["steps"] == 1500
|
|
|
|
|
|
def test_quality_on_a6000_ok_1024():
|
|
argv, env, _ = build_command(_req(tier="quality", device_index=1))
|
|
assert _flag_value(argv, "--resolution") == "1024,1024"
|
|
assert _flag_value(argv, "--max_train_steps") == "3000"
|
|
assert env["CUDA_VISIBLE_DEVICES"] == "1"
|
|
|
|
|
|
def test_paths_and_names_flow_through():
|
|
argv, _, _ = build_command(_req(output_name="my_char", trigger="ohwx"))
|
|
assert _flag_value(argv, "--output_name") == "my_char"
|
|
assert _flag_value(argv, "--train_data_dir") == "/worktank/arbo/train/t123/dataset"
|
|
|
|
|
|
# ---- rejections (INV-T7 / §4.6 tier-device fit) ----------------------------------------
|
|
|
|
def test_quality_on_3090_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="quality"):
|
|
build_command(_req(tier="quality", device_index=0))
|
|
|
|
|
|
def test_unknown_tier_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="tier"):
|
|
build_command(_req(tier="ultra"))
|
|
|
|
|
|
def test_bad_device_index_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="device_index"):
|
|
build_command(_req(device_index=3))
|
|
|
|
|
|
def test_dataset_dir_outside_handoff_root_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="dataset_dir"):
|
|
build_command(_req(dataset_dir="/etc/passwd"))
|
|
|
|
|
|
def test_path_traversal_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="dataset_dir"):
|
|
build_command(_req(dataset_dir="/worktank/arbo/train/../../etc/shadow"))
|
|
|
|
|
|
def test_base_model_outside_allowed_roots_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="base_model_path"):
|
|
build_command(_req(base_model_path="/home/someone/evil.safetensors"))
|
|
|
|
|
|
def test_published_relative_path():
|
|
# Phase 2 publish: derive {train_id} from the handoff output_dir -> ComfyUI-relative loras path
|
|
assert published_relative_path(
|
|
"/worktank/arbo/train/392cf898ac03/output", "sindra_lora"
|
|
) == "trained/392cf898ac03/sindra_lora.safetensors"
|
|
|
|
|
|
def test_storetank_checkpoint_allowed():
|
|
# the canonical SDXL store (2026-06-13 move) — arbo dispatches base_model_path from here
|
|
argv, _, _ = build_command(
|
|
_req(base_model_path="/storetank/arbo/models/checkpoints/albedobaseXL_v31Large.safetensors")
|
|
)
|
|
assert _flag_value(argv, "--pretrained_model_name_or_path") == \
|
|
"/storetank/arbo/models/checkpoints/albedobaseXL_v31Large.safetensors"
|
|
|
|
|
|
def test_unsafe_output_name_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="output_name"):
|
|
build_command(_req(output_name="a; rm -rf /"))
|
|
|
|
|
|
def test_unsafe_trigger_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="trigger"):
|
|
build_command(_req(trigger="$(whoami)"))
|
|
|
|
|
|
def test_repeats_out_of_range_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="repeats"):
|
|
build_command(_req(repeats=0))
|
|
with pytest.raises(InvalidTrainRequest, match="repeats"):
|
|
build_command(_req(repeats=1000))
|
|
|
|
|
|
def test_relative_dataset_dir_rejected():
|
|
with pytest.raises(InvalidTrainRequest, match="absolute"):
|
|
build_command(_req(dataset_dir="relative/path"))
|
|
|
|
|
|
def test_alpha_scales_with_dim():
|
|
_, _, params = build_command(_req(tier="fast"))
|
|
# fast -> dim 16; alpha = dim * ratio (0.5, matching Sindra) = 8
|
|
argv, _, _ = build_command(_req(tier="fast"))
|
|
assert _flag_value(argv, "--network_dim") == "16"
|
|
assert _flag_value(argv, "--network_alpha") == "8"
|
|
assert params["dim"] == 16
|