giant/tests/test_workflow_spec.py

"""Workflow spec parsing, validation, and spec hashes (gitea #83)."""

import pytest

from giant.workflow.spec import (
    WorkflowSpecError,
    epoch_milestones,
    load_spec,
    parse_spec,
    spec_hash,
)

MINIMAL = {
    "workflow": {"name": "wf", "result_dir": "/tmp/wf"},
    "condor": {"accounting_group": "cms", "repo_dir": "/work/lbogner/giant"},
    "dataset": {"steps": "/data/train", "reference": "/data/holdout"},
    "train": [{"name": "a", "epochs": 3}],
    "rollout": [{"name": "a", "train": "a"}],
    "analysis": [{"name": "cmp", "rollouts": ["a"], "chunks": 4}],
}


def _spec(**patch):
    raw = {k: (v.copy() if isinstance(v, dict) else list(v)) for k, v in MINIMAL.items()}
    raw.update(patch)
    return parse_spec(raw)


def test_parses_minimal_spec():
    spec = _spec()
    assert spec.name == "wf"
    assert spec.log_dir == "/tmp/wf/logs"  # derived from result_dir
    assert spec.train("a").epochs == 3
    assert spec.rollout("a").train == "a"
    assert spec.analysis("cmp").rollouts == ("a",)
    # defaults come from the dataclasses, not the file
    assert spec.geometry.method == "slab"
    assert spec.condor.docker_image_gpu == "mschnepf/slc7-condocker"


def test_example_config_is_valid():
    spec = load_spec("configs/workflow_example.toml")
    assert {t.name for t in spec.trains} == {"baseline", "router-balanced"}
    assert spec.analysis("baseline-vs-router").rollouts == ("baseline", "router-balanced")


@pytest.mark.parametrize(
    "patch, message",
    [
        ({"train": [{"name": "a"}, {"name": "a"}]}, "unique"),
        ({"rollout": [{"name": "r", "train": "nope"}]}, "names no"),
        ({"analysis": [{"name": "c", "rollouts": ["nope"]}]}, "not defined"),
        ({"analysis": [{"name": "c", "rollouts": []}]}, "at least one"),
        ({"analysis": [{"name": "c", "rollouts": ["a"], "chunks": 0}]}, "chunks must be"),
        ({"train": [{"name": "a", "epochs": 0}]}, "epochs must be"),
        ({"train": [{"name": "a", "epchs": 3}]}, "unknown key"),
        ({"geometry": {"methd": "slab"}}, "unknown key"),
    ],
)
def test_validation_errors(patch, message):
    with pytest.raises(WorkflowSpecError, match=message):
        _spec(**patch)


def test_unknown_top_level_table_rejected():
    with pytest.raises(WorkflowSpecError, match="unknown top-level"):
        _spec(nonsense={})


def test_missing_required_table_rejected():
    raw = {k: v for k, v in MINIMAL.items() if k != "dataset"}
    with pytest.raises(WorkflowSpecError, match=r"missing required \[dataset\]"):
        parse_spec(raw)


def test_unknown_lookup_names_are_explicit():
    spec = _spec()
    with pytest.raises(WorkflowSpecError, match="no \\[\\[train\\]\\] named 'zzz'"):
        spec.train("zzz")


def test_hash_is_stable_and_order_independent():
    a = _spec()
    b = parse_spec(
        {
            "dataset": MINIMAL["dataset"],
            "condor": MINIMAL["condor"],
            "workflow": MINIMAL["workflow"],
            "train": MINIMAL["train"],
            "rollout": MINIMAL["rollout"],
            "analysis": MINIMAL["analysis"],
        }
    )
    assert a.train_hash("a") == b.train_hash("a")
    assert a.analysis_hash("cmp") == b.analysis_hash("cmp")
    assert len(a.train_hash("a")) == 8


def test_hash_changes_with_own_settings():
    base = _spec()
    changed = _spec(train=[{"name": "a", "epochs": 4}])
    assert base.train_hash("a") != changed.train_hash("a")


def test_hash_propagates_from_parents():
    """A dataset change must move every downstream task's directory."""
    base = _spec()
    changed = _spec(dataset={"steps": "/data/other", "reference": "/data/holdout"})
    assert base.train_hash("a") != changed.train_hash("a")
    assert base.rollout_hash("a") != changed.rollout_hash("a")
    assert base.analysis_hash("cmp") != changed.analysis_hash("cmp")

    # ... and so must a change to a training the analysis transitively uses.
    retrained = _spec(train=[{"name": "a", "epochs": 9}])
    assert retrained.analysis_hash("cmp") != base.analysis_hash("cmp")
    # while an unrelated knob on the analysis leaves the training alone
    rebinned = _spec(analysis=[{"name": "cmp", "rollouts": ["a"], "chunks": 4, "bins": 99}])
    assert rebinned.train_hash("a") == base.train_hash("a")
    assert rebinned.analysis_hash("cmp") != base.analysis_hash("cmp")


def test_warm_cache_hash_ignores_epochs():
    """Epoch count doesn't change the setup cache, so it must not re-warm it."""
    base = _spec()
    longer = _spec(train=[{"name": "a", "epochs": 50}])
    assert base.warm_cache_hash("a") == longer.warm_cache_hash("a")
    other_cfg = _spec(train=[{"name": "a", "epochs": 3, "config": "configs/router.toml"}])
    assert base.warm_cache_hash("a") != other_cfg.warm_cache_hash("a")


def test_spec_hash_expands_dataclasses():
    spec = _spec()
    assert spec_hash(spec.dataset) == spec_hash(spec.dataset)
    assert spec_hash(spec.dataset) != spec_hash(spec.geometry)


@pytest.mark.parametrize(
    "epochs, per_job, expected",
    [
        (3, 1, [1, 2, 3]),
        (10, 3, [3, 6, 9, 10]),
        (9, 3, [3, 6, 9]),
        (1, 5, [1]),
    ],
)
def test_epoch_milestones(epochs, per_job, expected):
    spec = _spec(train=[{"name": "a", "epochs": epochs, "epochs_per_job": per_job}])
    assert epoch_milestones(spec.train("a")) == expected