| """CPU-only resume/provenance and two-worker scheduling checks.""" |
| import io |
| import json |
| from pathlib import Path |
| import sys |
| import tempfile |
| import threading |
| import unittest |
| from contextlib import redirect_stdout |
| from dataclasses import replace |
| from unittest.mock import patch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "scripts")) |
| import run_elf_remaining as q |
| from elf_common import SOURCE_COMMIT, default_sampling, sha256, write_json |
|
|
|
|
| class QueueTests(unittest.TestCase): |
| def setUp(self): |
| self.temp = tempfile.TemporaryDirectory() |
| self.addCleanup(self.temp.cleanup) |
| self.root = Path(self.temp.name) |
| self.assets = {t: {"checkpoint": "model", "sha256": {"model": "locked"}} |
| for t in ("owt", "wmt14", "xsum")} |
| write_json(self.root / "artifacts/elf_lock.json", {}) |
| write_json(self.root / "data/elf/assets.json", { |
| "assets": self.assets, "lock_sha256": sha256(self.root / "artifacts/elf_lock.json")}) |
| self.asset_hash = sha256(self.root / "data/elf/assets.json") |
|
|
| def input(self, task, split, count): |
| prefix = task == "owt-prefix" |
| path = self.root / "data/elf" / ("owt-prefix.jsonl" if prefix else f"{task}-{split}.jsonl") |
| path.write_text(''.join(json.dumps({"id": i}) + '\n' for i in range(count))) |
| meta = {"task": task, "split": split, "count": count, "sha256": sha256(path)} |
| if prefix: |
| source = self.root / "data/conditional/owt-c64/prompts.jsonl" |
| source.parent.mkdir(parents=True, exist_ok=True) |
| source.write_text("source\n") |
| meta.update(protocol="c64_text_t5_v1", source_sha256=sha256(source)) |
| write_json(path.with_suffix(".manifest.json"), meta) |
| return path |
|
|
| def request(self, cell): |
| return { |
| "task": cell.task, "steps": cell.steps, "seed": 42, "compile": False, "batch_size": 8, |
| "source_commit": SOURCE_COMMIT, "assets_manifest_sha256": self.asset_hash, |
| "checkpoint_task": "owt" if cell.task == "owt-prefix" else cell.task, |
| "checkpoint_selection": "ema_params1", "checkpoint_sha256": "locked", |
| "sampling": default_sampling(cell.task, cell.steps), |
| "model_canvas_t5_tokens": 1088 if cell.task == "xsum" else 128 if cell.task == "wmt14" else 1024, |
| "protocol": "c64_text_t5_v1" if cell.task == "owt-prefix" else "elf_native_v1", |
| "condition_tokenization_policy": q.EOS_POLICY if cell.task in ("wmt14", "xsum") else "none", |
| "input": {"sha256": cell.input_sha256} if cell.input_sha256 else None, |
| } |
|
|
| def result(self, path, cell, status="quality_valid"): |
| request = self.request(cell) |
| write_json(path / "request.json", request) |
| if status == "partial": |
| return |
| pairs = [(i, 0) for i in range(cell.count)] |
| if cell.task == "owt-prefix": |
| pairs += [(i, c) for c in range(1, 5) for i in range(min(256, cell.count))] |
| samples = path / "samples.jsonl" |
| samples.write_text(''.join(json.dumps({"id": j, "prompt_id": i, "completion_id": c}) + '\n' |
| for j, (i, c) in enumerate(pairs))) |
| gen = {**request, "sample_count": len(pairs), "prompt_count": cell.count, |
| "diversity": cell.task == "owt-prefix", "samples_sha256": sha256(samples)} |
| write_json(path / "generation.json", gen) |
| if status == "needs_evaluation": |
| return |
| metric = {"task": cell.task, "sample_count": cell.count, "samples_sha256": sha256(samples), |
| "generation_manifest_sha256": sha256(path / "generation.json"), |
| "valid": status == "quality_valid", "metrics": {"entropy_gpt2_nats": 5.1}} |
| if not metric["valid"]: |
| metric["reason"] = "At least one completion has fewer than two GPT-2 tokens; no samples were dropped" |
| write_json(path / "metrics.json", metric) |
|
|
| def reusable(self, cell): |
| return q.reusable_result(self.root, cell, self.asset_hash, self.assets) |
|
|
| def test_full_matrix_and_input_block_isolation(self): |
| cells, blocked = q.build_matrix(self.root, q.TASKS) |
| self.assertEqual(len(cells), 7) |
| self.assertEqual(len(blocked), 5) |
| for task, split, count in (("owt-prefix", "c64-text-t5", 1024), |
| ("wmt14", "validation", 3000), ("wmt14", "test", 3003), |
| ("xsum", "validation", 11332), ("xsum", "test", 11334)): |
| self.input(task, split, count) |
| cells, blocked = q.build_matrix(self.root, q.TASKS) |
| self.assertFalse(blocked) |
| self.assertEqual(len(cells), 32) |
| self.assertEqual(len({c.key for c in cells}), 32) |
| self.assertEqual([c.steps for c in cells if c.task == "owt"], [1, 2, 4, 8, 16, 32, 1024]) |
| source = self.root / "data/conditional/owt-c64/prompts.jsonl" |
| source.write_text("changed\n") |
| cells, blocked = q.build_matrix(self.root, q.TASKS) |
| self.assertEqual(len(cells), 25) |
| self.assertEqual(blocked[0]["task"], "owt-prefix") |
|
|
| def test_reuse_scores_or_generation_but_not_sanity_or_tampered_samples(self): |
| cell = q.Cell("owt", "unconditional", 8, 1024) |
| sanity = self.root / "results/elf-sanity/quality" / cell.key |
| self.result(sanity, replace(cell, count=1000)) |
| self.assertIsNone(self.reusable(cell)[0]) |
| path = self.root / "results/elf-old/quality" / cell.key |
| self.result(path, cell, "needs_evaluation") |
| self.assertEqual(self.reusable(cell)[0]["status"], "needs_evaluation") |
| self.result(path, cell) |
| self.assertEqual(self.reusable(cell)[0]["status"], "quality_valid") |
| with (path / "samples.jsonl").open("a") as f: |
| f.write('{}\n') |
| found, rejected = self.reusable(cell) |
| self.assertIsNone(found) |
| self.assertTrue(rejected) |
|
|
| def test_sampling_batch_and_no_eos_mismatch_rejected(self): |
| cell = q.Cell("wmt14", "validation", 8, 3000, "data", "hash") |
| for key, value in (("batch_size", 16), ("condition_tokenization_policy", "old"), |
| ("sampling", {}), ("checkpoint_sha256", "different")): |
| request = self.request(cell) |
| request[key] = value |
| self.assertFalse(q.matches_request(request, cell, self.asset_hash, self.assets)) |
|
|
| def test_invalid_quality_not_retried_and_conflicts_rejected(self): |
| cell = q.Cell("owt", "unconditional", 1, 2) |
| path = self.root / "results/elf-old/quality" / cell.key |
| self.result(path, cell, "quality_invalid") |
| self.assertEqual(self.reusable(cell)[0]["status"], "quality_invalid") |
| self.result(self.root / "results/elf-other/quality" / cell.key, cell) |
| with self.assertRaisesRegex(ValueError, "Conflicting"): |
| self.reusable(cell) |
|
|
| def test_partial_runs_preserved_and_prefix_requests_diversity(self): |
| cell = q.Cell("owt-prefix", "c64-text-t5", 32, 1024, "inputs", "digest") |
| root = self.root / "results/elf-gpu01" |
| path = root / "quality" / cell.key |
| self.result(path, cell, "partial") |
| retry = q.fresh_path(root, cell) |
| self.assertIn("elf-gpu01-retry1", str(retry)) |
| self.assertTrue((path / "request.json").exists()) |
| cmds = q.commands(self.root, sys.executable, cell, retry, True) |
| self.assertEqual([stage for stage, _ in cmds], ["generate", "evaluate"]) |
| self.assertIn("--diversity", cmds[0][1]) |
| self.assertNotIn("timing", str(cmds)) |
| self.assertEqual([s for s, _ in q.commands(self.root, sys.executable, cell, path, False)], ["evaluate"]) |
| self.result(path, cell) |
| self.assertEqual(self.reusable(cell)[0]["split"], "c64-text-t5") |
|
|
| def test_plan_does_not_query_gpu_or_launch_processes(self): |
| with patch.object(sys, "argv", ["queue", "--root", str(self.root), "--plan"]), \ |
| patch.object(q.subprocess, "check_output") as query, \ |
| patch.object(q.subprocess, "Popen") as popen, redirect_stdout(io.StringIO()): |
| self.assertEqual(q.main(), 2) |
| query.assert_not_called() |
| popen.assert_not_called() |
| self.assertFalse((self.root / "results").exists()) |
|
|
| def test_two_workers_isolated_and_invalid_quality_does_not_stop_queue(self): |
| cells = [q.Cell("owt", "unconditional", s, 2) for s in (1, 2, 4)] |
| barrier = threading.Barrier(2, timeout=10) |
| launches, lock = [], threading.Lock() |
|
|
| def fake_popen(command, **kwargs): |
| stage = "generate" if "generate" in command else "evaluate" |
| path = Path(command[command.index("--output" if stage == "generate" else "--run") + 1]) |
| cell = next(c for c in cells if str(path).endswith(c.key)) |
| with lock: |
| launches.append((cell.steps, stage, kwargs["env"]["CUDA_VISIBLE_DEVICES"])) |
|
|
| class Process: |
| def wait(inner): |
| if stage == "generate" and cell.steps in (1, 2): |
| barrier.wait() |
| status = "needs_evaluation" if stage == "generate" else "quality_invalid" if cell.steps == 1 else "quality_valid" |
| self.result(path, cell, status) |
| return int(status == "quality_invalid") |
| return Process() |
|
|
| with patch.object(sys, "argv", ["queue", "--root", str(self.root)]), \ |
| patch.object(sys, "platform", "linux"), patch.object(q, "build_matrix", return_value=(cells, [])), \ |
| patch.object(q.subprocess, "check_output", side_effect=lambda cmd, **kw: f"GPU-{cmd[2]}, H200 NVL"), \ |
| patch.object(q.subprocess, "Popen", side_effect=fake_popen), redirect_stdout(io.StringIO()): |
| self.assertEqual(q.main(), 0) |
| self.assertEqual(len(launches), 6) |
| self.assertEqual({gpu for _, _, gpu in launches}, {"GPU-0", "GPU-1"}) |
| for cell in cells: |
| self.assertEqual(len({gpu for step, _, gpu in launches if step == cell.steps}), 1) |
| report = json.loads(next((self.root / "results/elf-gpu01/logs").glob("*/queue.json")).read_text()) |
| self.assertEqual([r["status"] for r in report["runs"]], ["quality_invalid", "quality_valid", "quality_valid"]) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
