| import json |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from dlb.gpu_matrix import active_gpus |
| from dlb.matrix import MatrixTask, write_matrix |
|
|
|
|
| ROOT = Path(__file__).parents[1] |
|
|
|
|
| def _task(index: int) -> MatrixTask: |
| steps = 2**index |
| return MatrixTask( |
| task_id=f"flm-lm1b-steps-{steps}", |
| category="many", |
| model="flm", |
| dataset="lm1b", |
| steps=steps, |
| sample_count=1024, |
| seed=42, |
| environment="dlb-flm", |
| adapter="flm", |
| source="flm", |
| provenance="official", |
| sample_dir=f"/tmp/dlb/samples/lm1b/flm/steps_{steps}", |
| metrics_path=f"/tmp/dlb/metrics/lm1b/flm/steps_{steps}/metrics.json", |
| timing_path=f"/tmp/dlb/timing/lm1b/flm/steps_{steps}/timing.json", |
| ) |
|
|
|
|
| def test_gpu_matrix_dry_run_binds_tasks_to_requested_gpus(tmp_path): |
| matrix = write_matrix(tmp_path / "generation.tsv", [_task(0), _task(1), _task(2)]) |
|
|
| completed = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "dlb.gpu_matrix", |
| "--root", |
| str(ROOT), |
| "--stage", |
| "generate", |
| "--matrix", |
| str(matrix), |
| "--gpus", |
| "0,1", |
| "--max-jobs", |
| "2", |
| "--dry-run", |
| ], |
| cwd=ROOT, |
| text=True, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| check=False, |
| ) |
|
|
| assert completed.returncode == 0, completed.stderr |
| records = [json.loads(line) for line in completed.stdout.splitlines()] |
| assert [record["gpu"] for record in records] == ["0", "1", "0"] |
| assert all(record["cuda_visible_devices"] == record["gpu"] for record in records) |
| assert all(record["stage"] == "generate" for record in records) |
| assert records[0]["command"][:2] == ["bash", str(ROOT / "scripts/run_one.sh")] |
| assert "--num-samples" in records[0]["command"] |
|
|
|
|
| def test_gpu_matrix_rejects_empty_gpu_list(tmp_path): |
| matrix = write_matrix(tmp_path / "generation.tsv", [_task(0)]) |
|
|
| completed = subprocess.run( |
| [ |
| sys.executable, |
| "-m", |
| "dlb.gpu_matrix", |
| "--root", |
| str(ROOT), |
| "--stage", |
| "benchmark", |
| "--matrix", |
| str(matrix), |
| "--gpus", |
| "", |
| "--dry-run", |
| ], |
| cwd=ROOT, |
| text=True, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| check=False, |
| ) |
|
|
| assert completed.returncode == 2 |
| assert "at least one GPU" in completed.stderr |
|
|
|
|
| def test_gpu_matrix_can_oversubscribe_gpus_when_max_jobs_exceeds_gpu_count(): |
| assert active_gpus(("0", "1"), 5) == ("0", "1", "0", "1", "0") |
|
|
|
|
| @pytest.mark.parametrize( |
| "script", |
| [ |
| "scripts/smoke_4gpu.sh", |
| "scripts/run_4gpu.sh", |
| "scripts/evaluate_4gpu.sh", |
| "scripts/benchmark_4gpu.sh", |
| ], |
| ) |
| def test_4gpu_wrappers_do_not_regenerate_explicit_custom_matrix( |
| tmp_path: Path, script: str |
| ): |
| custom = write_matrix(tmp_path / "custom.tsv", [_task(0)]) |
| before = custom.read_text(encoding="utf-8") |
|
|
| completed = subprocess.run( |
| [ |
| "bash", |
| str(ROOT / script), |
| "--root", |
| str(ROOT), |
| "--matrix", |
| str(custom), |
| "--gpus", |
| "0,1", |
| "--max-jobs", |
| "1", |
| "--dry-run", |
| ], |
| cwd=ROOT, |
| text=True, |
| stdout=subprocess.PIPE, |
| stderr=subprocess.PIPE, |
| check=False, |
| ) |
|
|
| assert completed.returncode == 0, completed.stderr |
| assert custom.read_text(encoding="utf-8") == before |
| records = [json.loads(line) for line in completed.stdout.splitlines()] |
| assert [record["task_id"] for record in records] == ["flm-lm1b-steps-1"] |
|
|
