code/tests/test_alpha_grid.py

from __future__ import annotations

import json

from scripts.launch_alpha_grid import (
    ALPHAS,
    BLOCKS,
    METHODS,
    SEEDS,
    EXPECTED_CELLS_PER_JOB,
    _commands,
    _inspect,
    _can_resume,
    GPU_DEVICES,
)


def test_alpha_grid_has_28_sequential_jobs_and_560_cells(tmp_path):
    commands = _commands(tmp_path, tmp_path / "grid")

    assert len(commands) == len(ALPHAS) * len(BLOCKS)
    assert len(commands) * EXPECTED_CELLS_PER_JOB == 560
    assert all(method in command for command in commands.values() for method in METHODS)
    assert all(str(seed) in command for command in commands.values() for seed in SEEDS)
    assert all("--alpha-fpeft" in command and "--alpha-baseline" in command for command in commands.values())


def test_inspect_uses_all_methods_and_requested_seeds(tmp_path):
    path = tmp_path / "results.json"
    path.write_text(json.dumps({
        f"{method}__rank32__seed{seed}": {
            "status": "complete",
            "prediction_distribution": [8, 2],
        }
        for method in METHODS
        for seed in [42]
    }))

    state = _inspect(path, [42])

    assert state["complete"] == len(METHODS)
    assert state["errors"] == 0
    assert state["finished"] is True
    assert _can_resume(state)


def test_llama_rte_uses_low_memory_batching(tmp_path):
    command = _commands(tmp_path, tmp_path / "grid")["alpha0.25__llama_rte"]
    assert command[command.index("--batch-size") + 1] == "1"
    assert command[command.index("--grad-accum") + 1] == "16"
    assert command[command.index("--max-length") + 1] == "128"
    assert "--gradient-checkpointing" in command


def test_gpu2_device_is_configured():
    assert 2 in GPU_DEVICES
    assert GPU_DEVICES[2] != GPU_DEVICES[1]