Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import sys | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "scribe")) | |
| from gec.env import ( | |
| bundle_adapter, | |
| gpu_report, | |
| in_colab, | |
| latest_checkpoint, | |
| restore_artifacts, | |
| save_artifacts, | |
| ) | |
| from gec.notebook import StageContext | |
| class EnvDetectionTests(unittest.TestCase): | |
| def test_in_colab_is_false_in_this_environment(self) -> None: | |
| self.assertFalse(in_colab()) | |
| def test_gpu_report_returns_a_string_without_crashing(self) -> None: | |
| # Must be safe whether or not torch/CUDA is installed. | |
| self.assertIsInstance(gpu_report(), str) | |
| class CheckpointTests(unittest.TestCase): | |
| def test_no_checkpoint_returns_none(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| self.assertIsNone(latest_checkpoint(tmp)) | |
| def test_missing_dir_returns_none(self) -> None: | |
| self.assertIsNone(latest_checkpoint("/no/such/dir/xyz")) | |
| def test_latest_checkpoint_picks_highest_step(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| root = Path(tmp) | |
| for step in (50, 200, 100): | |
| (root / f"checkpoint-{step}").mkdir() | |
| (root / "checkpoint-notanumber").mkdir() # ignored | |
| self.assertEqual(Path(latest_checkpoint(root)).name, "checkpoint-200") | |
| class BackupHelperTests(unittest.TestCase): | |
| def test_local_restore_confirms_present_and_errors_when_missing(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| present = Path(tmp) / "have.jsonl" | |
| present.write_text("x", encoding="utf-8") | |
| # backup=None => local mode: present file is fine. | |
| restore_artifacts(None, [str(present)]) | |
| # missing file => clear SystemExit pointing back to Part 1. | |
| with self.assertRaises(SystemExit): | |
| restore_artifacts(None, [str(Path(tmp) / "missing.jsonl")]) | |
| def test_save_to_backup_copies_files(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| src = Path(tmp) / "metrics.json" | |
| src.write_text("{}", encoding="utf-8") | |
| backup = Path(tmp) / "drive" | |
| backup.mkdir() | |
| save_artifacts(backup, [str(src)]) | |
| self.assertTrue((backup / "metrics.json").exists()) | |
| def test_save_local_is_noop(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| src = Path(tmp) / "metrics.json" | |
| src.write_text("{}", encoding="utf-8") | |
| save_artifacts(None, [str(src)]) # should not raise | |
| class DurableOutputTests(unittest.TestCase): | |
| def _ctx(self, backup): | |
| # durable() only reads self.backup; the rest can be None for this test. | |
| return StageContext(profile=None, paths=None, backup=backup, in_colab=backup is not None) | |
| def test_local_durable_is_the_repo_path(self) -> None: | |
| rel = "artifacts/gec_pairs/pairs.jsonl" | |
| self.assertEqual(self._ctx(None).durable(rel), str(Path(rel))) | |
| def test_colab_durable_points_into_backup_by_basename(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| backup = Path(tmp) / "drive" | |
| backup.mkdir() | |
| out = Path(self._ctx(backup).durable("artifacts/gec_pairs/pairs.jsonl")) | |
| # same basename save()/restore() use, so downstream restore still finds it | |
| self.assertEqual(out, backup / "pairs.jsonl") | |
| self.assertTrue(out.parent.exists()) | |
| class BundleAdapterTests(unittest.TestCase): | |
| def test_bundle_keeps_adapter_files_and_drops_checkpoints(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| adapter = Path(tmp) / "qwen3" / "full" | |
| adapter.mkdir(parents=True) | |
| (adapter / "adapter_config.json").write_text("{}", encoding="utf-8") | |
| (adapter / "adapter_model.safetensors").write_text("weights", encoding="utf-8") | |
| (adapter / "darag_variant.json").write_text( | |
| '{"use_retrieval": true}', encoding="utf-8" | |
| ) | |
| ckpt = adapter / "checkpoint-50" | |
| ckpt.mkdir() | |
| (ckpt / "optimizer.pt").write_text("big", encoding="utf-8") | |
| dest = Path(tmp) / "share" | |
| bundle, zip_path = bundle_adapter(adapter, dest, name="carepath-gec") | |
| self.assertTrue((bundle / "adapter_config.json").exists()) | |
| self.assertTrue((bundle / "darag_variant.json").exists()) | |
| self.assertFalse((bundle / "checkpoint-50").exists()) # training-only | |
| self.assertIsNotNone(zip_path) | |
| self.assertTrue(zip_path.exists()) | |
| def test_bundle_errors_without_an_adapter(self) -> None: | |
| with tempfile.TemporaryDirectory() as tmp: | |
| empty = Path(tmp) / "empty" | |
| empty.mkdir() | |
| with self.assertRaises(SystemExit): | |
| bundle_adapter(empty, Path(tmp) / "share") | |
| if __name__ == "__main__": | |
| unittest.main() | |