giant/tests/test_condor_gpu.py

"""Tests for GPU-job HTCondor submission (`giant train-submit` / `giant rollout-submit`)."""

from __future__ import annotations

from pathlib import Path

from typer.testing import CliRunner

from giant import config as gconfig
from giant.cli import app
from giant.condor import CondorJobMeta, GpuSubmitConfig, parse_cluster_id, write_gpu_submit

runner = CliRunner()

_MINIMAL_TOML = """\
[train]
mode = "flow"
epochs = 1

[model]
hidden_dim = 8
n_blocks = 2
"""


# ---------------------------------------------------------------------------
# default_out_dir
# ---------------------------------------------------------------------------


def test_default_out_dir_encodes_hyperparams():
    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {})
    out_dir = gconfig.default_out_dir(cfg)
    name = out_dir.name
    assert f"_h{cfg['model']['hidden_dim']}_" in name
    assert f"_b{cfg['model']['n_blocks']}_" in name
    assert f"_c{cfg['model']['conditioning']}_" in name


def test_default_out_dir_avoids_collisions():
    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {})
    a = gconfig.default_out_dir(cfg)
    b = gconfig.default_out_dir(cfg)
    assert a != b
    # same hyperparam-derived prefix, differing only in the uuid tag
    assert a.name.rsplit("_", 1)[0] == b.name.rsplit("_", 1)[0]


# ---------------------------------------------------------------------------
# giant/condor.py: GpuSubmitConfig / write_gpu_submit
# ---------------------------------------------------------------------------


def test_write_gpu_submit_description(tmp_path: Path):
    cfg = GpuSubmitConfig(
        run_dir=tmp_path / "condor",
        accounting_group="cms",
        repo_dir=tmp_path,
        command="uv run giant train data.parquet --config c.toml --out out --device cuda",
    )
    sub = write_gpu_submit(cfg, "#!/bin/bash\necho hi\n")
    txt = sub.read_text()

    assert "universe = docker" in txt
    assert "docker_image = mschnepf/slc7-condocker" in txt
    assert "+RemoteJob = True" in txt
    assert "RequestGPUs = 1" in txt
    assert "accounting_group = cms" in txt
    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in txt
    assert "ProvidesETPResources" not in txt
    assert "queue 1" in txt

    wrapper = cfg.run_dir / "run.sh"
    assert wrapper.exists()
    assert wrapper.stat().st_mode & 0o111
    assert wrapper.read_text() == "#!/bin/bash\necho hi\n"


def test_gpu_requirements_include_type_and_memory_pins(tmp_path: Path):
    cfg = GpuSubmitConfig(
        run_dir=tmp_path / "condor",
        accounting_group="cms",
        repo_dir=tmp_path,
        command="uv run giant train ...",
        request_gpus=2,
        gpu_type="Tesla V100-PCIE-32GB",
        gpu_memory_mb=16000,
    )
    txt = write_gpu_submit(cfg, "#!/bin/bash\n").read_text()

    assert "RequestGPUs = 2" in txt
    assert 'TARGET.GPUs_DeviceName =?= "Tesla V100-PCIE-32GB"' in txt
    assert "TARGET.GPUs_GlobalMemoryMb >= 16000" in txt


# ---------------------------------------------------------------------------
# CondorJobMeta / parse_cluster_id
# ---------------------------------------------------------------------------


def test_condor_job_meta_round_trip(tmp_path: Path):
    meta = CondorJobMeta(
        accounting_group="cms",
        request_gpus=1,
        gpu_type=None,
        gpu_memory_mb=None,
        request_memory_mb=16384,
        request_walltime_s=172800,
        docker_image="mschnepf/slc7-condocker",
        repo_dir="/ceph/lbogner/giant",
        command="uv run giant train ...",
        submitted_at="2026-07-24T00:00:00+00:00",
    )
    path = tmp_path / "submission.json"
    meta.save(path)
    loaded = CondorJobMeta.load(path)
    assert loaded == meta
    assert loaded.cluster_id is None

    loaded.cluster_id = 123456
    loaded.save(path)
    assert CondorJobMeta.load(path).cluster_id == 123456


def test_parse_cluster_id():
    assert parse_cluster_id("1 job(s) submitted to cluster 123456.\n") == 123456
    assert parse_cluster_id("ERROR: something went wrong\n") is None


# ---------------------------------------------------------------------------
# CLI: giant train-submit / giant rollout-submit (--dry-run only, no cluster contact)
# ---------------------------------------------------------------------------


def test_train_submit_dry_run_writes_condor_dir(tmp_path: Path):
    config = tmp_path / "config.toml"
    config.write_text(_MINIMAL_TOML)
    data = tmp_path / "train.parquet"
    out_dir = tmp_path / "ckpt"

    result = runner.invoke(
        app,
        [
            "train-submit",
            str(data),
            "--config",
            str(config),
            "--accounting-group",
            "cms",
            "--out",
            str(out_dir),
            "--dry-run",
        ],
    )
    assert result.exit_code == 0, result.output

    run_dir = out_dir / "condor"
    sub_txt = (run_dir / "job.sub").read_text()
    assert "+RemoteJob = True" in sub_txt
    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt

    wrapper = (run_dir / "run.sh").read_text()
    assert "--device cuda" in wrapper
    assert "last.pt" in wrapper
    assert "RESUME" in wrapper

    meta = CondorJobMeta.load(run_dir / "submission.json")
    assert meta.accounting_group == "cms"
    assert meta.cluster_id is None


def test_rollout_submit_dry_run_writes_condor_dir(tmp_path: Path):
    data = tmp_path / "seed.parquet"
    checkpoint = tmp_path / "best.pt"
    geometry = tmp_path / "oracle.pkl"
    out = tmp_path / "rollout.parquet"

    result = runner.invoke(
        app,
        [
            "rollout-submit",
            str(data),
            "--checkpoint",
            str(checkpoint),
            "--geometry",
            str(geometry),
            "--accounting-group",
            "cms",
            "--out",
            str(out),
            "--dry-run",
        ],
    )
    assert result.exit_code == 0, result.output

    condor_dirs = list(tmp_path.glob("condor_*"))
    assert len(condor_dirs) == 1
    run_dir = condor_dirs[0]

    sub_txt = (run_dir / "job.sub").read_text()
    assert "+RemoteJob = True" in sub_txt
    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt

    wrapper = (run_dir / "run.sh").read_text()
    assert "--device cuda" in wrapper
    assert "--prediction-id" in wrapper
    assert "last.pt" not in wrapper
    assert "RESUME" not in wrapper

    meta = CondorJobMeta.load(run_dir / "submission.json")
    assert meta.accounting_group == "cms"