esh-pfi-infrastructure/services/lora-training-worker/tests/test_invocation.py

"""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_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(
        _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