giant/tests/test_workflow_tasks.py

"""Workflow task graph: dependencies, output paths, condor settings (gitea #83)."""

import pytest

from giant.analysis.catalog import catalog_ids, get_spec as get_plot_spec
from giant.workflow import tasks
from giant.workflow.spec import parse_spec

CONDOR = {
    "accounting_group": "cms",
    "repo_dir": "/work/lbogner/giant",
    "env_script": "/work/lbogner/giant/condor_env.sh",
}

RAW = {
    "workflow": {"name": "wf", "result_dir": "/results/wf"},
    "condor": CONDOR,
    "dataset": {"steps": "/data/train", "reference": "/data/holdout"},
    "train": [
        {"name": "base", "epochs": 3, "gpu_memory_mb": 20000},
        {"name": "router", "epochs": 2},
    ],
    "rollout": [
        {"name": "base", "train": "base"},
        {"name": "router", "train": "router"},
    ],
    "analysis": [{"name": "cmp", "rollouts": ["base", "router"], "chunks": 4}],
}


@pytest.fixture
def spec():
    s = parse_spec(RAW)
    tasks.set_spec(s)
    return s


def _requires(task):
    return list(task.requires() or [])


def test_epoch_chain_is_linear_and_rooted_at_warm_cache(spec):
    h = spec.train_hash("base")
    third = tasks.TrainEpochTask(name="base", spec_hash=h, milestone=3)
    second = _requires(third)
    assert [type(t) for t in second] == [tasks.TrainEpochTask]
    assert second[0].milestone == 2
    first = _requires(second[0])[0]
    assert first.milestone == 1
    root = _requires(first)
    assert [type(t) for t in root] == [tasks.WarmCacheTask]
    # the warm cache is keyed by its own hash, not the training's
    assert root[0].spec_hash == spec.warm_cache_hash("base")


def test_epoch_task_outputs_last_pt_per_milestone(spec):
    h = spec.train_hash("base")
    path = tasks.TrainEpochTask(name="base", spec_hash=h, milestone=2).output().path
    assert path == f"/results/wf/train_epoch/name=base/spec_hash={h}/epochs=2/last.pt"


def test_train_task_requires_final_epoch_and_publishes_canonical_outputs(spec):
    h = spec.train_hash("base")
    train = tasks.TrainTask(name="base", spec_hash=h)
    (dep,) = _requires(train)
    assert isinstance(dep, tasks.TrainEpochTask) and dep.milestone == 3
    out = train.output()
    assert set(out) == {"best.pt", "last.pt", "metrics.csv"}
    assert out["best.pt"].path == f"/results/wf/train/name=base/spec_hash={h}/best.pt"
    # local: it only copies files around, no reason to queue a job for it
    assert train.batch_system == "local"


def test_rollout_requires_training_geometry_and_reference(spec):
    ro = tasks.RolloutTask(name="base", spec_hash=spec.rollout_hash("base"))
    deps = _requires(ro)
    assert [type(d) for d in deps] == [tasks.TrainTask, tasks.GeometryOracleTask, tasks.DatasetTask]
    assert deps[0].name == "base"
    assert deps[2].path == "/data/holdout"
    out = ro.output()
    assert out["rollout.yaml"].path.endswith("rollout.yaml")
    # the sidecar sits next to the parquet — the deterministic path
    # `giant rollout --out` now produces
    assert out["rollout.yaml"].path[: -len(".yaml")] == out["rollout.parquet"].path[: -len(".parquet")]


def test_analysis_prep_requires_every_named_rollout(spec):
    prep = tasks.AnalysisPrepTask(name="cmp", spec_hash=spec.analysis_hash("cmp"))
    deps = _requires(prep)
    assert [d.name for d in deps] == ["base", "router"]
    assert all(isinstance(d, tasks.RolloutTask) for d in deps)
    assert prep.batch_system == "local"


def test_compute_job_enumeration_collapses_non_chunkable_specs(spec):
    jobs = tasks.analysis_jobs(spec, "cmp")
    non_chunkable = [i for i in catalog_ids() if not get_plot_spec(i).chunkable]
    expected = (len(catalog_ids()) - len(non_chunkable)) * 4 + len(non_chunkable)
    assert len(jobs) == expected
    assert non_chunkable, "expected some chunkable=False specs in the catalog"
    for spec_id in non_chunkable:
        assert [c for i, c in jobs if i == spec_id] == [0]


def test_compute_output_matches_the_on_disk_contract(spec):
    h = spec.analysis_hash("cmp")
    task = tasks.AnalysisComputeTask(name="cmp", spec_hash=h, plot_id="event_mean_length", chunk=2)
    assert task.output().path == (
        f"/results/wf/analysis/name=cmp/spec_hash={h}/reduced_partial/event_mean_length__2.json"
    )
    (dep,) = _requires(task)
    assert isinstance(dep, tasks.AnalysisPrepTask)


def test_render_requires_every_compute_job_and_runs_locally(spec):
    render = tasks.AnalysisRenderTask(name="cmp", spec_hash=spec.analysis_hash("cmp"))
    deps = _requires(render)
    assert len(deps) == len(tasks.analysis_jobs(spec, "cmp"))
    assert render.batch_system == "local"  # the only step importing plotstyle/LaTeX
    assert render.output().path.endswith("/plots/metadata.yaml")


def test_workflow_task_wraps_every_analysis(spec):
    deps = _requires(tasks.WorkflowTask(workflow_name="wf"))
    assert [(type(d), d.name) for d in deps] == [(tasks.AnalysisRenderTask, "cmp")]


def test_workflow_without_analysis_falls_back_to_rollouts():
    raw = {k: v for k, v in RAW.items() if k != "analysis"}
    tasks.set_spec(parse_spec(raw))
    deps = _requires(tasks.WorkflowTask(workflow_name="wf"))
    assert [type(d) for d in deps] == [tasks.RolloutTask, tasks.RolloutTask]


def test_gpu_settings_carry_remote_ceph_and_pins(spec):
    settings = tasks.TrainEpochTask(name="base", spec_hash=spec.train_hash("base"), milestone=1).htcondor_settings
    assert settings["+RemoteJob"] == "True"
    assert settings["RequestGPUs"] == 1
    assert "TARGET.ProvidesEtpCeph =?= True" in settings["requirements"]
    assert "TARGET.GPUs_GlobalMemoryMb >= 20000" in settings["requirements"]
    assert settings["accounting_group"] == "cms"
    assert settings["docker_image"] == "mschnepf/slc7-condocker"


def test_cpu_settings_used_for_analysis_compute(spec):
    task = tasks.AnalysisComputeTask(
        name="cmp", spec_hash=spec.analysis_hash("cmp"), plot_id="event_mean_length", chunk=0
    )
    settings = task.htcondor_settings
    assert settings["docker_image"] == "cverstege/alma9-gridjob"
    assert "RequestGPUs" not in settings
    # no run_meta.json yet (prep hasn't run), so no walltime is claimed
    assert "+RequestWalltime" not in settings


def test_cpu_settings_local_files_use_provides_etp_resources():
    raw = {**RAW, "condor": {**CONDOR, "remote": False}}
    spec = parse_spec(raw)
    tasks.set_spec(spec)
    settings = tasks.GeometryOracleTask(spec_hash=spec.geometry_hash()).htcondor_settings
    assert settings["requirements"] == "TARGET.ProvidesETPResources"
    assert "+RemoteJob" not in settings


def test_missing_dataset_fails_immediately(spec):
    with pytest.raises(FileNotFoundError, match="/ceph"):
        tasks.DatasetTask(path="/data/train").complete()


def test_dataset_that_exists_is_complete(tmp_path, spec):
    (tmp_path / "steps.parquet").write_text("")
    assert tasks.DatasetTask(path=str(tmp_path / "steps.parquet")).complete()