social-data-attribution-w2/tests/test_slurm_wrapper_paths.py

from __future__ import annotations

from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]

WRAPPER_DIRS = [
    REPO_ROOT / "scripts" / "slurm" / "eda",
    REPO_ROOT / "scripts" / "slurm" / "enrichment",
    REPO_ROOT / "scripts" / "slurm" / "launcher",
    REPO_ROOT / "scripts" / "slurm" / "manifest",
    REPO_ROOT / "scripts" / "slurm" / "r2_sunset",
    REPO_ROOT / "scripts" / "slurm" / "sampling",
    REPO_ROOT / "scripts" / "slurm" / "upload",
]


def test_nested_slurm_wrappers_do_not_use_two_level_repo_root() -> None:
    stale_patterns = [
        '$(cd "$(dirname "$0")/../.." && pwd)',
        '$(cd "$SCRIPT_DIR/../.." && pwd)',
    ]
    for wrapper_dir in WRAPPER_DIRS:
        for path in sorted(wrapper_dir.iterdir()):
            if not path.is_file():
                continue
            text = path.read_text(encoding="utf-8")
            for pattern in stale_patterns:
                assert pattern not in text, (
                    f"{path} still contains stale root path {pattern}"
                )


def test_slurm_wrappers_do_not_reference_removed_cli_modules() -> None:
    scripts_dir = REPO_ROOT / "scripts" / "slurm"
    stale_modules = [
        "python3 -m dolma.pool_sample_cli",
        "python3 -m dolma.corpus_stats_cli",
    ]
    for path in sorted(scripts_dir.rglob("*")):
        if not path.is_file() or path.suffix not in {".sh", ".sbatch", ".py"}:
            continue
        text = path.read_text(encoding="utf-8")
        for pattern in stale_modules:
            assert pattern not in text, (
                f"{path} still references removed module {pattern}"
            )


def test_pool_sample_launcher_uses_current_paths() -> None:
    path = REPO_ROOT / "scripts" / "slurm" / "sampling" / "launch_pool_sample.sh"
    text = path.read_text(encoding="utf-8")

    assert 'SHARD_FILE="$REPO_DIR/scripts/slurm/data/pool_shards_job_${i}.txt"' in text
    assert (
        "sbatch --export=ALL,JOB_INDEX=$i scripts/slurm/sampling/sample_pool_150B.sh"
        in text
    )


def test_pool_sample_and_corpus_stats_wrappers_use_current_modules() -> None:
    sample_text = (
        REPO_ROOT / "scripts" / "slurm" / "sampling" / "sample_pool_150B.sh"
    ).read_text(encoding="utf-8")
    smoke_text = (
        REPO_ROOT / "scripts" / "slurm" / "sampling" / "smoke_pool_sample.sh"
    ).read_text(encoding="utf-8")
    corpus_stats_text = (
        REPO_ROOT / "scripts" / "slurm" / "manifest" / "corpus_stats.sh"
    ).read_text(encoding="utf-8")

    assert "python3 -m dolma.pool_sample.cli \\" in sample_text
    assert "python3 -m dolma.pool_sample.cli \\" in smoke_text
    assert "python3 -m dolma.corpus_stats.cli \\" in corpus_stats_text
    assert "scripts/slurm/manifest/corpus_stats.sh" in corpus_stats_text


def test_enrichment_manifest_generator_and_array_script_use_current_paths() -> None:
    manifest_text = (
        REPO_ROOT / "scripts" / "manifests" / "make_enrich_manifest.py"
    ).read_text(encoding="utf-8")
    enrich_text = (
        REPO_ROOT / "scripts" / "slurm" / "enrichment" / "enrich_pool_150B_gpu.sh"
    ).read_text(encoding="utf-8")

    assert "scripts/slurm/enrichment/enrich_manifest.txt" in manifest_text
    assert "scripts/slurm/enrichment/enrich_manifest.txt" in enrich_text
    assert "scripts/slurm/enrichment/enrich_pool_150B_gpu.sh" in manifest_text


def test_attribution_slurm_wrappers_use_current_bergson_cli_flags() -> None:
    attribution_dir = REPO_ROOT / "scripts" / "slurm" / "attribution"
    stale_patterns = [
        "--method mean",
        "--score mean",
    ]
    required_patterns = {
        "query_index.sh": "--aggregation mean",
        "query_index_instruct.sh": "--aggregation mean",
        "score.sh": 'SCORE="individual"',
        "data_index.sh": "--score individual",
        "trackstar_reduce.sbatch": "--aggregation mean",
        "trackstar_score.sbatch": 'SCORE="${SCORE:-individual}"',
    }

    for path in sorted(attribution_dir.iterdir()):
        if not path.is_file():
            continue
        text = path.read_text(encoding="utf-8")
        for pattern in stale_patterns:
            assert pattern not in text, (
                f"{path} still contains stale Bergson flag {pattern}"
            )

        expected = required_patterns.get(path.name)
        if expected is not None:
            assert expected in text, (
                f"{path} is missing expected Bergson flag {expected}"
            )


def test_attribution_score_wrappers_use_explicit_fp32_score_precision() -> None:
    expected_patterns = {
        "data_index.sh": [
            "--index_cfg.precision fp32",
            "--score_cfg.precision fp32",
        ],
        "score.sh": [
            "--index_cfg.precision fp32",
            "--score_cfg.precision fp32",
        ],
        "trackstar_score.sbatch": [
            '--index_cfg.precision "$PRECISION"',
            '--score_cfg.precision "$PRECISION"',
        ],
    }
    stale_patterns = {
        "data_index.sh": ["--precision fp32"],
        "score.sh": ["--precision fp32"],
        "trackstar_score.sbatch": ['--precision "$PRECISION"'],
    }
    attribution_dir = REPO_ROOT / "scripts" / "slurm" / "attribution"

    for path in sorted(attribution_dir.iterdir()):
        if not path.is_file():
            continue
        text = path.read_text(encoding="utf-8")
        for pattern in stale_patterns.get(path.name, []):
            assert pattern not in text, (
                f"{path} still contains stale score precision flag {pattern}"
            )
        for pattern in expected_patterns.get(path.name, []):
            assert pattern in text, (
                f"{path} is missing explicit score precision flag {pattern}"
            )