| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import pytest |
|
|
| from deepgenopix.config import ExperimentConfig |
| from deepgenopix.notebook_support import ( |
| NotebookPaths, |
| TrainNotebookContext, |
| ensure_train_artifacts, |
| prepare_train_context, |
| ) |
|
|
|
|
| def _make_context(tmp_path: Path, *, run_etl: bool = True): |
| repo_root = tmp_path / "repo" |
| run_root = repo_root / "data" / "output" / "baseline_v1" / "baseline_v1__px12" |
| processed_root = repo_root / "data" / "processed" / "te_visuals" / "baseline_v1__px12" |
| raw_root = repo_root / "data" / "raw" |
| raw_fasta = raw_root / "input.fa" |
|
|
| for directory in [run_root, processed_root, raw_root]: |
| directory.mkdir(parents=True, exist_ok=True) |
| raw_fasta.write_text(">seq0\nACGTACGTACGT\n", encoding="utf-8") |
|
|
| paths = NotebookPaths( |
| repo_root=repo_root, |
| run_root=run_root, |
| processed_root=processed_root, |
| checkpoint_dir=run_root / "checkpoints", |
| report_dir=run_root / "reports", |
| embedding_dir=run_root / "embeddings", |
| config_path=run_root / "config.json", |
| metrics_path=run_root / "metrics.csv", |
| manifest_path=run_root / "run_manifest.json", |
| notes_path=run_root / "run_notes.md", |
| report_path=run_root / "run_report.md", |
| raw_input_dir=raw_root, |
| raw_fasta=raw_fasta, |
| raw_parquet=None, |
| raw_split_parquet_root=None, |
| raw_split_summary=None, |
| ) |
|
|
| return TrainNotebookContext( |
| config=ExperimentConfig(), |
| paths=paths, |
| device="cpu", |
| in_colab=False, |
| run_mode="train", |
| resume_from_existing=False, |
| run_etl=run_etl, |
| copy_data_manifest=False, |
| export_embeddings=False, |
| ) |
|
|
|
|
| def test_ensure_train_artifacts_rebuilds_partial_dataset(monkeypatch, tmp_path): |
| context = _make_context(tmp_path) |
| processed_root = context.paths.processed_root |
| run_root = context.paths.run_root |
|
|
| (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8") |
| (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8") |
|
|
| calls: list[str] = [] |
|
|
| def fake_run_incremental_etl(fasta_path, output_dir, **kwargs): |
| calls.append(str(fasta_path)) |
| assert not (processed_root / "registry.csv").exists() |
| assert not (processed_root / "classes.json").exists() |
| assert not (processed_root / "tensors.lmdb").exists() |
| (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8") |
| (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8") |
| (processed_root / "tensors.lmdb").write_bytes(b"lmdb") |
|
|
| monkeypatch.setattr("deepgenopix.notebook_support.run_incremental_etl", fake_run_incremental_etl) |
|
|
| artifacts = ensure_train_artifacts(context) |
|
|
| assert calls == [str(context.paths.raw_fasta)] |
| assert artifacts.registry_path == processed_root / "registry.csv" |
| assert artifacts.classes_path == processed_root / "classes.json" |
| assert (processed_root / "tensors.lmdb").exists() |
| assert (run_root / "run_manifest.json").exists() |
|
|
|
|
| def test_ensure_train_artifacts_raises_when_lmdb_missing_and_reuse_only(monkeypatch, tmp_path): |
| context = _make_context(tmp_path, run_etl=False) |
| processed_root = context.paths.processed_root |
|
|
| (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8") |
| (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8") |
|
|
| called = False |
|
|
| def fake_run_incremental_etl(*args, **kwargs): |
| nonlocal called |
| called = True |
|
|
| monkeypatch.setattr("deepgenopix.notebook_support.run_incremental_etl", fake_run_incremental_etl) |
|
|
| try: |
| ensure_train_artifacts(context) |
| except FileNotFoundError as exc: |
| assert "tensors.lmdb" in str(exc) |
| else: |
| raise AssertionError("expected FileNotFoundError") |
|
|
| assert called is False |
|
|
|
|
| def test_ensure_train_artifacts_uses_pre_split_parquets(monkeypatch, tmp_path): |
| context = _make_context(tmp_path) |
| processed_root = context.paths.processed_root |
| split_root = context.paths.raw_input_dir |
| context.paths.raw_fasta = None |
| context.paths.raw_split_parquet_root = split_root |
|
|
| for split in ("train", "val", "test"): |
| split_dir = split_root / split |
| split_dir.mkdir(parents=True, exist_ok=True) |
| (split_dir / "te_seqdata.parquet").write_bytes(b"parquet") |
|
|
| calls: list[str] = [] |
|
|
| def fake_run_split_parquet_etl(split_input_root, output_dir, **kwargs): |
| calls.append(str(split_input_root)) |
| (processed_root / "registry.csv").write_text("original_index,lmdb_key,label,split\n0,0,A,train\n", encoding="utf-8") |
| (processed_root / "classes.json").write_text(json.dumps(["A"]), encoding="utf-8") |
| (processed_root / "tensors.lmdb").write_bytes(b"lmdb") |
|
|
| monkeypatch.setattr("deepgenopix.notebook_support.run_split_parquet_etl", fake_run_split_parquet_etl) |
|
|
| artifacts = ensure_train_artifacts(context) |
|
|
| assert calls == [str(split_root)] |
| assert artifacts.registry_path == processed_root / "registry.csv" |
|
|
|
|
| def test_prepare_train_context_refuses_to_overwrite_completed_run(tmp_path): |
| repo_root = tmp_path / "repo" |
| raw_root = repo_root / "data" / "raw" |
| run_root = repo_root / "data" / "output" / "baseline_v1" / ExperimentConfig().run_id |
| raw_root.mkdir(parents=True, exist_ok=True) |
| run_root.mkdir(parents=True, exist_ok=True) |
| (raw_root / "input.fa").write_text(">seq0\nACGTACGTACGT\n", encoding="utf-8") |
| (run_root / "metrics.csv").write_text("epoch,train_loss,val_accuracy\n1,1.0,0.5\n", encoding="utf-8") |
| (run_root / "run_manifest.json").write_text(json.dumps({"status": "complete"}), encoding="utf-8") |
|
|
| with pytest.raises(FileExistsError, match="Run outputs already exist"): |
| prepare_train_context( |
| ExperimentConfig(), |
| run_mode="train", |
| resume_from_existing=False, |
| run_etl=True, |
| copy_data_manifest=False, |
| export_embeddings=False, |
| repo_root=repo_root, |
| raw_input_dir=raw_root, |
| ) |
|
|
|
|
| def test_prepare_train_context_uses_pixel_signature_processed_root_and_split_summary(tmp_path): |
| repo_root = tmp_path / "repo" |
| raw_root = repo_root / "data" / "raw" |
|
|
| for split in ("train", "val", "test"): |
| split_dir = raw_root / split |
| split_dir.mkdir(parents=True, exist_ok=True) |
| (split_dir / "te_seqdata.parquet").write_bytes(b"parquet") |
| (raw_root / "split_summary.json").write_text("{}", encoding="utf-8") |
|
|
| context = prepare_train_context( |
| ExperimentConfig(pixel_stride_bp=9, val_frac=0.2, test_frac=0.1, split_seed=7, min_family_size=2), |
| run_mode="train", |
| resume_from_existing=False, |
| run_etl=True, |
| copy_data_manifest=False, |
| export_embeddings=False, |
| repo_root=repo_root, |
| raw_input_dir=raw_root, |
| ) |
|
|
| assert context.paths.processed_root == ( |
| repo_root / "data" / "processed" / "te_visuals" / "px12_ps9_fl500-500_vf0p2_tf0p1_seed7_mf2" |
| ) |
| assert context.paths.raw_split_parquet_root == raw_root |
| assert context.paths.raw_split_summary == raw_root / "split_summary.json" |
| manifest = json.loads(context.paths.manifest_path.read_text(encoding="utf-8")) |
| assert manifest["raw_split_summary"] == str(raw_root / "split_summary.json") |
|
|
|
|
| def test_training_loop_defers_amp_choice_to_trainer_env(): |
| source = Path("src/deepgenopix/notebook_support.py").read_text() |
| assert "use_amp=None" in source |
| assert 'use_amp=str(context.device).startswith("cuda")' not in source |
|
|
