From 7d80a0d2a654ceea4e6c6b4b2a39ed18352b664b Mon Sep 17 00:00:00 2001
From: Florin Senoner <23100806+FlorinSenoner@users.noreply.github.com>
Date: Sat, 1 Aug 2026 18:47:02 +0200
Subject: [PATCH 1/5] fix(notebook): recompute projections on every generate
---
.../notebooks/ProtSpace_Preparation.ipynb | 3 +-
.../tests/test_issue_338_reproduction.py | 101 ++++++++++++++++++
.../.openspec.yaml | 2 +
.../fix-notebook-projection-cache/README.md | 3 +
.../fix-notebook-projection-cache/design.md | 49 +++++++++
.../fix-notebook-projection-cache/proposal.md | 25 +++++
.../notebook-projection-cache-safety/spec.md | 23 ++++
.../fix-notebook-projection-cache/tasks.md | 20 ++++
8 files changed, 225 insertions(+), 1 deletion(-)
create mode 100644 apps/protspace/tests/test_issue_338_reproduction.py
create mode 100644 openspec/changes/fix-notebook-projection-cache/.openspec.yaml
create mode 100644 openspec/changes/fix-notebook-projection-cache/README.md
create mode 100644 openspec/changes/fix-notebook-projection-cache/design.md
create mode 100644 openspec/changes/fix-notebook-projection-cache/proposal.md
create mode 100644 openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
create mode 100644 openspec/changes/fix-notebook-projection-cache/tasks.md
diff --git a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
index 16a37952..e1dc849b 100644
--- a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
+++ b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
@@ -654,7 +654,7 @@
"\n",
" n_proteins = len(embedding_sets[0].headers)\n",
"\n",
- " # Build pipeline with caching enabled\n",
+ " # Build pipeline with non-projection caching enabled\n",
" reducer_params = ReducerParams(\n",
" n_neighbors=pw[\"n_neighbors\"].value,\n",
" min_dist=pw[\"min_dist\"].value,\n",
@@ -669,6 +669,7 @@
" bundled=True,\n",
" keep_tmp=True,\n",
" intermediate_dir=cache_dir,\n",
+ " refetch_stages=frozenset({\"projections\"}),\n",
" annotations=ann,\n",
" reducer_params=reducer_params,\n",
" stats=compute_stats_cb.value,\n",
diff --git a/apps/protspace/tests/test_issue_338_reproduction.py b/apps/protspace/tests/test_issue_338_reproduction.py
new file mode 100644
index 00000000..1e970819
--- /dev/null
+++ b/apps/protspace/tests/test_issue_338_reproduction.py
@@ -0,0 +1,101 @@
+import ast
+import json
+from dataclasses import asdict
+from pathlib import Path
+
+import numpy as np
+
+from protspace.data.loaders import EmbeddingSet
+from protspace.data.processors.pipeline import (
+ PipelineConfig,
+ ReductionPipeline,
+ parse_methods_arg,
+)
+
+
+class InputRecordingBase:
+ def __init__(self, config):
+ self.config = config
+ self.reducers = {"umap": object()}
+ self.inputs = []
+
+ def process_reduction(self, data, method, dims):
+ self.inputs.append(data.copy())
+ return {
+ "name": f"{method}{dims}",
+ "dimensions": dims,
+ "info": {},
+ "data": data[:, :dims].copy(),
+ }
+
+
+def _preparation_notebook_projection_refetch_stages() -> frozenset[str]:
+ notebook_path = (
+ Path(__file__).parents[1] / "notebooks" / "ProtSpace_Preparation.ipynb"
+ )
+ notebook = json.loads(notebook_path.read_text())
+ code_sources = (
+ "".join(cell["source"])
+ for cell in notebook["cells"]
+ if cell["cell_type"] == "code"
+ )
+ generate_source = next(source for source in code_sources if "def _on_gen" in source)
+ tree = ast.parse(generate_source)
+ config_call = next(
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Name)
+ and node.func.id == "PipelineConfig"
+ )
+ refetch_keyword = next(
+ (
+ keyword
+ for keyword in config_call.keywords
+ if keyword.arg == "refetch_stages"
+ ),
+ None,
+ )
+ if refetch_keyword is None:
+ return frozenset()
+ expression = ast.Expression(refetch_keyword.value)
+ return eval(
+ compile(expression, filename=str(notebook_path), mode="eval"),
+ {"__builtins__": {}, "frozenset": frozenset},
+ )
+
+
+def test_notebook_cache_invalidates_when_input_embeddings_change(tmp_path):
+ cache_dir = tmp_path / "output" / "tmp"
+ cache_dir.mkdir(parents=True)
+ config = PipelineConfig(
+ methods=parse_methods_arg(["umap2"]),
+ output_path=tmp_path / "output" / "data.parquetbundle",
+ keep_tmp=True,
+ intermediate_dir=cache_dir,
+ annotations=None,
+ refetch_stages=_preparation_notebook_projection_refetch_stages(),
+ )
+ pipeline = object.__new__(ReductionPipeline)
+ pipeline.config = config
+ pipeline.base = InputRecordingBase(asdict(config.reducer_params))
+
+ headers = ["P1", "P2", "P3"]
+ first_input = EmbeddingSet(
+ name="prot_t5",
+ data=np.zeros((3, 3), dtype=np.float32),
+ headers=headers,
+ )
+ changed_input = EmbeddingSet(
+ name="prot_t5",
+ data=np.full((3, 3), 7.0, dtype=np.float32),
+ headers=headers,
+ )
+
+ pipeline._run_reductions([first_input])
+ changed = pipeline._run_reductions([changed_input])[0]
+
+ assert len(pipeline.base.inputs) == 2
+ np.testing.assert_array_equal(
+ changed["data"], np.full((3, 2), 7.0, dtype=np.float32)
+ )
diff --git a/openspec/changes/fix-notebook-projection-cache/.openspec.yaml b/openspec/changes/fix-notebook-projection-cache/.openspec.yaml
new file mode 100644
index 00000000..5849c2db
--- /dev/null
+++ b/openspec/changes/fix-notebook-projection-cache/.openspec.yaml
@@ -0,0 +1,2 @@
+schema: spec-driven
+created: 2026-08-01
diff --git a/openspec/changes/fix-notebook-projection-cache/README.md b/openspec/changes/fix-notebook-projection-cache/README.md
new file mode 100644
index 00000000..b1f655dd
--- /dev/null
+++ b/openspec/changes/fix-notebook-projection-cache/README.md
@@ -0,0 +1,3 @@
+# fix-notebook-projection-cache
+
+Invalidate cached notebook projections when input data changes.
diff --git a/openspec/changes/fix-notebook-projection-cache/design.md b/openspec/changes/fix-notebook-projection-cache/design.md
new file mode 100644
index 00000000..01718763
--- /dev/null
+++ b/openspec/changes/fix-notebook-projection-cache/design.md
@@ -0,0 +1,49 @@
+## Context
+
+`ProtSpace_Preparation.ipynb` keeps `output/tmp` so expensive FASTA downloads, embeddings, and annotations can survive repeated Generate actions. `ReductionPipeline` also stores projections there. Its projection cache key includes the logical embedding name, method, dimensions, and reducer parameters, but not the embedding matrix or headers. The notebook reuses generic embedding names such as `prot_t5`, so a changed input can collide with a prior projection. The issue's desired notebook behavior is simpler than the CLI's reusable-cache behavior: Generate must recompute projections.
+
+## Goals / Non-Goals
+
+**Goals:**
+
+- Guarantee that every Preparation-notebook Generate action reduces the current embedding data.
+- Preserve caching for the notebook's more expensive input, embedding, and annotation stages.
+- Cover changed input with an observable reducer-execution regression.
+
+**Non-Goals:**
+
+- Redesign projection cache identity for CLI users.
+- Disable every notebook cache or alter query/embedding/annotation refresh semantics.
+- Change reducer parameters, projection naming, bundle layout, or output paths.
+
+## Decisions
+
+### Request the existing projections refetch stage from the notebook
+
+The notebook will construct `PipelineConfig` with `refetch_stages=frozenset({"projections"})`. `ReductionPipeline._load_cached_projection` already treats that stage as an instruction to bypass cached coordinates, while the other retained intermediates remain available.
+
+This uses the pipeline's public configuration contract and keeps cache lifecycle in one place.
+
+**Alternative: delete `proj_*.npz` files before each run.** Rejected because it duplicates cache naming/lifecycle knowledge in the notebook and introduces an unnecessary destructive filesystem operation.
+
+**Alternative: hash all embedding bytes and headers in the core cache key.** Rejected for this issue because it broadens CLI cache semantics and adds hashing cost to all callers. The notebook explicitly wants fresh projections on Generate, so selecting the existing refetch stage is both clearer and narrower.
+
+### Exercise actual cache behavior in the regression
+
+The regression will use the real `ReductionPipeline._run_reductions` cache path with a deterministic fake reducer. It will run two same-name embedding sets with different data through a configuration that requests projection refresh, then assert the reducer sees both inputs and the second result reflects the second input.
+
+The notebook artifact will also be validated as a parseable notebook with parseable code cells, following existing notebook verification practice.
+
+## Risks / Trade-offs
+
+- **Projection reruns take longer even when nothing changed.** → This is the explicit notebook correctness contract; expensive embedding and annotation intermediates remain cached.
+- **The regression could test pipeline behavior without proving notebook wiring.** → Verification will additionally inspect the executed notebook configuration path and validate all notebook code cells.
+- **A future pipeline refetch API rename could break the notebook.** → The focused pipeline regression and notebook configuration verification make that failure visible.
+
+## Migration Plan
+
+No data migration is required. Existing projection cache files may remain in `output/tmp`; the notebook will stop reading them during Generate. Rollback is a one-line notebook configuration revert.
+
+## Open Questions
+
+None.
diff --git a/openspec/changes/fix-notebook-projection-cache/proposal.md b/openspec/changes/fix-notebook-projection-cache/proposal.md
new file mode 100644
index 00000000..9e2a677e
--- /dev/null
+++ b/openspec/changes/fix-notebook-projection-cache/proposal.md
@@ -0,0 +1,25 @@
+## Why
+
+The Preparation notebook keeps one intermediate directory across Generate runs, but projection cache identity does not include the input embeddings. A later run can therefore rebundle stale coordinates when its input changes while the embedding name, method, and reducer parameters remain the same.
+
+## What Changes
+
+- Make every Generate action in `ProtSpace_Preparation.ipynb` explicitly recompute dimensionality-reduction projections.
+- Continue retaining the notebook's expensive query, embedding, and annotation intermediates; only projection reuse changes.
+- Add regression coverage proving an explicitly refreshed projection does not reuse coordinates from changed input data.
+
+## Capabilities
+
+### New Capabilities
+
+- `notebook-projection-cache-safety`: Defines how the Preparation notebook treats cached projections across Generate actions.
+
+### Modified Capabilities
+
+None.
+
+## Impact
+
+- Affected notebook: `apps/protspace/notebooks/ProtSpace_Preparation.ipynb`.
+- Affected tests: Python pipeline regression coverage for notebook-equivalent projection refresh behavior.
+- No CLI defaults, bundle format, public Python API, or dependencies change.
diff --git a/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
new file mode 100644
index 00000000..510ec89f
--- /dev/null
+++ b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
@@ -0,0 +1,23 @@
+## ADDED Requirements
+
+### Requirement: Preparation notebook Generate actions use current projection inputs
+
+The Preparation notebook SHALL recompute dimensionality-reduction projections on every Generate action and SHALL NOT read cached projection coordinates from an earlier action. This projection refresh SHALL NOT disable caching for other intermediate stages.
+
+#### Scenario: Reducer parameters change between Generate actions
+
+- **WHEN** a user changes a dimensionality-reduction parameter and activates Generate again
+- **THEN** the selected reducer runs with the current parameter value
+- **AND** the downloaded bundle contains coordinates produced by that run
+
+#### Scenario: Input data changes without changing its logical embedding name
+
+- **WHEN** a user changes the input embeddings while the embedding name, method, and reducer parameters match an earlier Generate action
+- **THEN** the reducer runs against the current embedding matrix
+- **AND** cached coordinates from the earlier input are not used
+
+#### Scenario: Non-projection intermediates remain reusable
+
+- **WHEN** the notebook requests fresh projections
+- **THEN** only the projection stage is explicitly refreshed
+- **AND** retained query, embedding, and annotation intermediates remain eligible for their existing cache behavior
diff --git a/openspec/changes/fix-notebook-projection-cache/tasks.md b/openspec/changes/fix-notebook-projection-cache/tasks.md
new file mode 100644
index 00000000..0e860fd6
--- /dev/null
+++ b/openspec/changes/fix-notebook-projection-cache/tasks.md
@@ -0,0 +1,20 @@
+## 1. Regression coverage
+
+- [x] 1.1 Add the smallest pipeline regression that changes same-name input embeddings across retained-cache runs and asserts projection refresh processes the second input.
+- [x] 1.2 Run the regression before implementation and record the expected stale-cache failure.
+
+## 2. Notebook implementation
+
+- [x] 2.1 Configure `ProtSpace_Preparation.ipynb` to explicitly refresh only the projection stage on every Generate action.
+- [x] 2.2 Keep query, embedding, and annotation cache wiring unchanged.
+
+## 3. Focused verification
+
+- [x] 3.1 Run the regression after implementation and observe it pass.
+- [x] 3.2 Validate the notebook with `nbformat` and compile every code cell after removing Colab magics.
+- [x] 3.3 Verify the original two-run reproduction returns coordinates from the changed input and invokes the reducer twice.
+
+## 4. Repository gates
+
+- [x] 4.1 Run affected Python tests and Ruff checks.
+- [x] 4.2 Run `pnpm precommit` before commit and push.
From 3c4b9f4cdc325c651fc6917f9e666653f7c75f5d Mon Sep 17 00:00:00 2001
From: Florin Senoner <23100806+FlorinSenoner@users.noreply.github.com>
Date: Wed, 5 Aug 2026 14:08:47 +0200
Subject: [PATCH 2/5] fix(notebook): isolate retained caches by input
---
apps/protspace/CLAUDE.md | 2 +-
.../notebooks/ProtSpace_Preparation.ipynb | 23 +-
.../src/protspace/data/processors/pipeline.py | 35 +++
.../tests/test_issue_338_reproduction.py | 101 --------
apps/protspace/tests/test_pipeline_utils.py | 217 ++++++++++++++++++
.../fix-notebook-projection-cache/README.md | 2 +-
.../fix-notebook-projection-cache/design.md | 31 ++-
.../fix-notebook-projection-cache/proposal.md | 13 +-
.../notebook-projection-cache-safety/spec.md | 38 ++-
.../fix-notebook-projection-cache/tasks.md | 13 +-
10 files changed, 347 insertions(+), 128 deletions(-)
delete mode 100644 apps/protspace/tests/test_issue_338_reproduction.py
diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md
index 06c63fe2..53b95df1 100644
--- a/apps/protspace/CLAUDE.md
+++ b/apps/protspace/CLAUDE.md
@@ -269,7 +269,7 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
| `test_interpro_annotation_retriever.py` | InterPro API mocking, parsing |
| `test_settings_converter.py` | Settings table ↔ visualization state conversion |
| `test_uniprot_annotation_retriever.py` | UniProt API mocking, inactive entry resolution |
-| `test_pipeline_utils.py` | ReductionPipeline, EmbeddingSet, method parsing, multi-input merging, inline param overrides |
+| `test_pipeline_utils.py` | ReductionPipeline, notebook input/annotation/projection cache identity, EmbeddingSet, method parsing, multi-input merging, inline param overrides |
| `test_stats.py` | Projection statistics: elbow, annotation-based validity (silhouette/DBI/CH per annotation), auto-cluster ARI/NMI agreement, faithfulness (dual continuity + global metrics), cluster-selection (elbow/silhouette/both), subsample determinism/order-invariance, silhouette consistency, `_align` no-id guard, silhouette→elbow fallback |
| `test_stats_cli.py` | `protspace stats` CLI + `prepare` stats wiring, `--stats-annotation` (auto/list) wiring, `--settings-out` guard, `--cluster-selection` validation |
| `test_stats_carriage.py` | Routing rows to bundle parts (metadata quality, annotation columns, cluster legend) |
diff --git a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
index e1dc849b..e82a5789 100644
--- a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
+++ b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
@@ -57,7 +57,9 @@
" PipelineConfig,\n",
" ReducerParams,\n",
" ReductionPipeline,\n",
+ " _input_cache_dir,\n",
" parse_methods_arg,\n",
+ " _query_fasta_cache_path,\n",
")"
]
},
@@ -581,8 +583,8 @@
"\n",
" out_dir = Path(\"output\")\n",
" out_dir.mkdir(exist_ok=True)\n",
- " cache_dir = out_dir / \"tmp\"\n",
- " cache_dir.mkdir(exist_ok=True)\n",
+ " cache_root = out_dir / \"tmp\"\n",
+ " cache_root.mkdir(exist_ok=True)\n",
" output_path = out_dir / \"data.parquetbundle\"\n",
"\n",
" step_html = HTML(value=\"Step 1/4: Loading embeddings...\")\n",
@@ -598,7 +600,7 @@
" print(\"Select at least one embedder.\")\n",
" return\n",
" step_html.value = \"Step 1/6: Downloading FASTA...\"\n",
- " fasta_cache = cache_dir / \"sequences.fasta\"\n",
+ " fasta_cache = _query_fasta_cache_path(cache_root, inp[\"query\"])\n",
" if fasta_cache.exists() and fasta_cache.stat().st_size > 0:\n",
" from protspace.data.loaders.query import (\n",
" extract_identifiers_from_fasta,\n",
@@ -611,6 +613,8 @@
" if not headers:\n",
" print(f\"No sequences found for query: {inp['query']}\")\n",
" return\n",
+ " cache_dir = _input_cache_dir(cache_root, fasta_path)\n",
+ " cache_dir.mkdir(parents=True, exist_ok=True)\n",
" backend, _emb_cfg = _resolve_backend_and_config()\n",
" embs = _drop_incompatible(embs, backend)\n",
" if not embs:\n",
@@ -631,6 +635,9 @@
" if not embs:\n",
" print(\"Select at least one embedder.\")\n",
" return\n",
+ " fasta_path = Path(inp[\"path\"])\n",
+ " cache_dir = _input_cache_dir(cache_root, fasta_path)\n",
+ " cache_dir.mkdir(parents=True, exist_ok=True)\n",
" backend, _emb_cfg = _resolve_backend_and_config()\n",
" embs = _drop_incompatible(embs, backend)\n",
" if not embs:\n",
@@ -639,15 +646,17 @@
" for emb_name in embs:\n",
" step_html.value = f\"Step 1/5: Computing {emb_name} embeddings ({backend})...\"\n",
" emb_set = embed_fasta(\n",
- " Path(inp[\"path\"]), emb_name,\n",
+ " fasta_path, emb_name,\n",
" backend=backend,\n",
" embed_config=_emb_cfg,\n",
" embedding_cache=cache_dir / f\"{emb_name}.h5\",\n",
" )\n",
- " emb_set.fasta_path = Path(inp[\"path\"])\n",
+ " emb_set.fasta_path = fasta_path\n",
" embedding_sets.append(emb_set)\n",
" else:\n",
" h5_path = Path(inp[\"path\"])\n",
+ " cache_dir = _input_cache_dir(cache_root, h5_path)\n",
+ " cache_dir.mkdir(parents=True, exist_ok=True)\n",
" name_override = inp.get(\"name\")\n",
" emb_set = load_h5([h5_path], name_override=name_override)\n",
" embedding_sets.append(emb_set)\n",
@@ -678,7 +687,9 @@
"\n",
" # Step 2: Annotations (cached after first run)\n",
" step_html.value = \"Step 2/4: Fetching annotations...\"\n",
- " metadata = pipeline._fetch_annotations(embedding_sets[0].headers)\n",
+ " metadata = pipeline._fetch_annotations(\n",
+ " embedding_sets[0].headers, embedding_sets\n",
+ " )\n",
"\n",
" # Step 3: Dimensionality reduction\n",
" step_html.value = \"Step 3/4: Reducing dimensions...\"\n",
diff --git a/apps/protspace/src/protspace/data/processors/pipeline.py b/apps/protspace/src/protspace/data/processors/pipeline.py
index cf1821b6..6c963289 100644
--- a/apps/protspace/src/protspace/data/processors/pipeline.py
+++ b/apps/protspace/src/protspace/data/processors/pipeline.py
@@ -82,6 +82,21 @@ class PipelineConfig:
reducer_params: ReducerParams = field(default_factory=ReducerParams)
+def _query_fasta_cache_path(cache_root: Path, query: str) -> Path:
+ """Return the retained FASTA path owned by one exact UniProt query."""
+ digest = hashlib.sha256(query.encode()).hexdigest()[:12]
+ return cache_root / "queries" / f"{digest}.fasta"
+
+
+def _input_cache_dir(cache_root: Path, input_path: Path) -> Path:
+ """Return the retained intermediate directory owned by one input file."""
+ digest = hashlib.sha256()
+ with input_path.open("rb") as source:
+ for chunk in iter(lambda: source.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return cache_root / "inputs" / digest.hexdigest()[:12]
+
+
# Valid override parameter names (from ReducerParams fields)
_VALID_OVERRIDE_KEYS = {f.name for f in fields(ReducerParams)}
# Field types for coercion
@@ -396,6 +411,26 @@ def _fetch_annotations(
if cache_path.exists():
cached_df = pd.read_parquet(cache_path)
+ cached_identifiers = (
+ Counter(cached_df["identifier"].astype(str))
+ if "identifier" in cached_df.columns
+ else Counter()
+ )
+ requested_identifiers = Counter(map(str, headers))
+
+ if cached_identifiers != requested_identifiers:
+ logger.info(
+ "Annotation cache input changed; fetching annotations "
+ "for the current identifiers"
+ )
+ api_df = ProteinAnnotationManager(
+ headers=headers,
+ annotations=annotations_list,
+ output_path=cache_path,
+ sequences=sequences,
+ ).to_pd()
+ return self._merge_csv(api_df, csv_df)
+
cached_annotations = set(cached_df.columns) - {"identifier"}
if annotations_list is None:
diff --git a/apps/protspace/tests/test_issue_338_reproduction.py b/apps/protspace/tests/test_issue_338_reproduction.py
deleted file mode 100644
index 1e970819..00000000
--- a/apps/protspace/tests/test_issue_338_reproduction.py
+++ /dev/null
@@ -1,101 +0,0 @@
-import ast
-import json
-from dataclasses import asdict
-from pathlib import Path
-
-import numpy as np
-
-from protspace.data.loaders import EmbeddingSet
-from protspace.data.processors.pipeline import (
- PipelineConfig,
- ReductionPipeline,
- parse_methods_arg,
-)
-
-
-class InputRecordingBase:
- def __init__(self, config):
- self.config = config
- self.reducers = {"umap": object()}
- self.inputs = []
-
- def process_reduction(self, data, method, dims):
- self.inputs.append(data.copy())
- return {
- "name": f"{method}{dims}",
- "dimensions": dims,
- "info": {},
- "data": data[:, :dims].copy(),
- }
-
-
-def _preparation_notebook_projection_refetch_stages() -> frozenset[str]:
- notebook_path = (
- Path(__file__).parents[1] / "notebooks" / "ProtSpace_Preparation.ipynb"
- )
- notebook = json.loads(notebook_path.read_text())
- code_sources = (
- "".join(cell["source"])
- for cell in notebook["cells"]
- if cell["cell_type"] == "code"
- )
- generate_source = next(source for source in code_sources if "def _on_gen" in source)
- tree = ast.parse(generate_source)
- config_call = next(
- node
- for node in ast.walk(tree)
- if isinstance(node, ast.Call)
- and isinstance(node.func, ast.Name)
- and node.func.id == "PipelineConfig"
- )
- refetch_keyword = next(
- (
- keyword
- for keyword in config_call.keywords
- if keyword.arg == "refetch_stages"
- ),
- None,
- )
- if refetch_keyword is None:
- return frozenset()
- expression = ast.Expression(refetch_keyword.value)
- return eval(
- compile(expression, filename=str(notebook_path), mode="eval"),
- {"__builtins__": {}, "frozenset": frozenset},
- )
-
-
-def test_notebook_cache_invalidates_when_input_embeddings_change(tmp_path):
- cache_dir = tmp_path / "output" / "tmp"
- cache_dir.mkdir(parents=True)
- config = PipelineConfig(
- methods=parse_methods_arg(["umap2"]),
- output_path=tmp_path / "output" / "data.parquetbundle",
- keep_tmp=True,
- intermediate_dir=cache_dir,
- annotations=None,
- refetch_stages=_preparation_notebook_projection_refetch_stages(),
- )
- pipeline = object.__new__(ReductionPipeline)
- pipeline.config = config
- pipeline.base = InputRecordingBase(asdict(config.reducer_params))
-
- headers = ["P1", "P2", "P3"]
- first_input = EmbeddingSet(
- name="prot_t5",
- data=np.zeros((3, 3), dtype=np.float32),
- headers=headers,
- )
- changed_input = EmbeddingSet(
- name="prot_t5",
- data=np.full((3, 3), 7.0, dtype=np.float32),
- headers=headers,
- )
-
- pipeline._run_reductions([first_input])
- changed = pipeline._run_reductions([changed_input])[0]
-
- assert len(pipeline.base.inputs) == 2
- np.testing.assert_array_equal(
- changed["data"], np.full((3, 2), 7.0, dtype=np.float32)
- )
diff --git a/apps/protspace/tests/test_pipeline_utils.py b/apps/protspace/tests/test_pipeline_utils.py
index 5ec6fc29..c2053830 100644
--- a/apps/protspace/tests/test_pipeline_utils.py
+++ b/apps/protspace/tests/test_pipeline_utils.py
@@ -1,8 +1,12 @@
"""Tests for pipeline utility functions."""
+import ast
+import json
from collections import Counter
+from pathlib import Path
import numpy as np
+import pandas as pd
import pytest
from protspace.data.loaders.embedding_set import (
@@ -11,6 +15,7 @@
format_projection_name,
merge_same_name_sets,
)
+from protspace.data.processors import pipeline as pipeline_module
from protspace.data.processors.pipeline import (
MethodSpec,
PipelineConfig,
@@ -21,6 +26,36 @@
parse_methods_arg,
)
+
+def _preparation_notebook_projection_refetch_stages() -> frozenset[str]:
+ notebook_path = (
+ Path(__file__).parents[1] / "notebooks" / "ProtSpace_Preparation.ipynb"
+ )
+ notebook = json.loads(notebook_path.read_text())
+ code_sources = (
+ "".join(cell["source"])
+ for cell in notebook["cells"]
+ if cell["cell_type"] == "code"
+ )
+ generate_source = next(source for source in code_sources if "def _on_gen" in source)
+ tree = ast.parse(generate_source)
+ config_call = next(
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Name)
+ and node.func.id == "PipelineConfig"
+ )
+ refetch_keyword = next(
+ keyword for keyword in config_call.keywords if keyword.arg == "refetch_stages"
+ )
+ expression = ast.Expression(refetch_keyword.value)
+ return eval(
+ compile(expression, filename=str(notebook_path), mode="eval"),
+ {"__builtins__": {}, "frozenset": frozenset},
+ )
+
+
# ---------------------------------------------------------------------------
# parse_method_spec
# ---------------------------------------------------------------------------
@@ -642,3 +677,185 @@ def boom(data, method, dims):
assert id(pipeline.base.config) == original_config_id, (
"base.config reference should be the original dict, not a replacement"
)
+
+
+# ---------------------------------------------------------------------------
+# Preparation notebook cache identity
+# ---------------------------------------------------------------------------
+
+
+class TestPreparationNotebookCacheIdentity:
+ def test_query_cache_path_changes_with_query(self, tmp_path):
+ globin = pipeline_module._query_fasta_cache_path(
+ tmp_path, "(family:globin) AND (reviewed:true)"
+ )
+ phosphatase = pipeline_module._query_fasta_cache_path(
+ tmp_path, "(family:phosphatase) AND (reviewed:true)"
+ )
+
+ assert globin != phosphatase
+ assert globin.parent == phosphatase.parent == tmp_path / "queries"
+
+ def test_input_cache_dir_changes_for_disjoint_fasta_inputs(self, tmp_path):
+ first = tmp_path / "first.fasta"
+ second = tmp_path / "second.fasta"
+ first.write_text(">P1\nAAAA\n")
+ second.write_text(">P2\nCCCC\n")
+
+ first_cache = pipeline_module._input_cache_dir(tmp_path, first)
+ second_cache = pipeline_module._input_cache_dir(tmp_path, second)
+
+ assert first_cache != second_cache
+
+ def test_input_cache_dir_changes_for_same_id_changed_sequence(self, tmp_path):
+ fasta = tmp_path / "input.fasta"
+ fasta.write_text(">P1\nAAAA\n")
+ original_cache = pipeline_module._input_cache_dir(tmp_path, fasta)
+
+ fasta.write_text(">P1\nCCCC\n")
+ changed_cache = pipeline_module._input_cache_dir(tmp_path, fasta)
+
+ assert changed_cache != original_cache
+
+ def test_input_cache_dir_is_reused_for_identical_content(self, tmp_path):
+ first = tmp_path / "first.fasta"
+ renamed = tmp_path / "renamed.fasta"
+ first.write_text(">P1\nAAAA\n")
+ renamed.write_text(">P1\nAAAA\n")
+
+ assert pipeline_module._input_cache_dir(
+ tmp_path, first
+ ) == pipeline_module._input_cache_dir(tmp_path, renamed)
+
+
+# ---------------------------------------------------------------------------
+# Annotation cache identity
+# ---------------------------------------------------------------------------
+
+
+def test_annotation_cache_is_rebuilt_for_different_identifiers(tmp_path, monkeypatch):
+ from protspace.data.annotations.manager import ProteinAnnotationManager
+
+ cache_dir = tmp_path / "cache"
+ cache_dir.mkdir()
+ pd.DataFrame(
+ {
+ "identifier": ["OLD1", "OLD2"],
+ "protein_name": ["old", "old"],
+ "gene_name": ["old", "old"],
+ "uniprot_kb_id": ["old", "old"],
+ }
+ ).to_parquet(cache_dir / "all_annotations.parquet")
+
+ pipeline = ReductionPipeline(
+ PipelineConfig(
+ methods=[],
+ output_path=tmp_path / "output.parquetbundle",
+ keep_tmp=True,
+ intermediate_dir=cache_dir,
+ annotations=["protein_name"],
+ )
+ )
+ captured = {}
+
+ def fresh_annotations(manager):
+ captured["headers"] = manager.headers
+ captured["cached_data"] = manager.cached_data
+ return pd.DataFrame(
+ {
+ "identifier": ["NEW1", "NEW2"],
+ "protein_name": ["new", "new"],
+ "gene_name": ["new", "new"],
+ "uniprot_kb_id": ["new", "new"],
+ }
+ )
+
+ monkeypatch.setattr(ProteinAnnotationManager, "to_pd", fresh_annotations)
+
+ result = pipeline._fetch_annotations(["NEW1", "NEW2"])
+
+ assert captured["headers"] == ["NEW1", "NEW2"]
+ assert captured["cached_data"] is None
+ assert result["identifier"].tolist() == ["NEW1", "NEW2"]
+
+
+def test_annotation_cache_is_reused_for_matching_identifiers(tmp_path, monkeypatch):
+ from protspace.data.annotations.manager import ProteinAnnotationManager
+
+ cache_dir = tmp_path / "cache"
+ cache_dir.mkdir()
+ pd.DataFrame(
+ {
+ "identifier": ["P1", "P2"],
+ "protein_name": ["one", "two"],
+ "gene_name": ["gene-one", "gene-two"],
+ "uniprot_kb_id": ["id-one", "id-two"],
+ }
+ ).to_parquet(cache_dir / "all_annotations.parquet")
+ pipeline = ReductionPipeline(
+ PipelineConfig(
+ methods=[],
+ output_path=tmp_path / "output.parquetbundle",
+ keep_tmp=True,
+ intermediate_dir=cache_dir,
+ annotations=["protein_name"],
+ )
+ )
+
+ def unexpected_fetch(_manager):
+ pytest.fail("matching annotation identifiers should reuse the cache")
+
+ monkeypatch.setattr(ProteinAnnotationManager, "to_pd", unexpected_fetch)
+
+ result = pipeline._fetch_annotations(["P2", "P1"])
+
+ assert result["identifier"].tolist() == ["P1", "P2"]
+
+
+# ---------------------------------------------------------------------------
+# Preparation notebook projection refresh
+# ---------------------------------------------------------------------------
+
+
+def test_notebook_refreshes_same_name_changed_input_through_pipeline(tmp_path):
+ config = PipelineConfig(
+ methods=parse_methods_arg(["umap2"]),
+ output_path=tmp_path / "output" / "data.parquetbundle",
+ keep_tmp=True,
+ intermediate_dir=tmp_path / "output" / "tmp",
+ annotations=None,
+ refetch_stages=_preparation_notebook_projection_refetch_stages(),
+ )
+ config.intermediate_dir.mkdir(parents=True)
+ pipeline = ReductionPipeline(config)
+ inputs = []
+
+ def record_input(data, method, dims):
+ inputs.append(data.copy())
+ return {
+ "name": f"{method}{dims}",
+ "dimensions": dims,
+ "info": {},
+ "data": data[:, :dims].copy(),
+ }
+
+ pipeline.base.process_reduction = record_input
+ headers = ["P1", "P2", "P3"]
+ first_input = EmbeddingSet(
+ name="prot_t5",
+ data=np.zeros((3, 3), dtype=np.float32),
+ headers=headers,
+ )
+ changed_input = EmbeddingSet(
+ name="prot_t5",
+ data=np.full((3, 3), 7.0, dtype=np.float32),
+ headers=headers,
+ )
+
+ pipeline._run_reductions([first_input])
+ changed = pipeline._run_reductions([changed_input])[0]
+
+ assert len(inputs) == 2
+ np.testing.assert_array_equal(
+ changed["data"], np.full((3, 2), 7.0, dtype=np.float32)
+ )
diff --git a/openspec/changes/fix-notebook-projection-cache/README.md b/openspec/changes/fix-notebook-projection-cache/README.md
index b1f655dd..57a52366 100644
--- a/openspec/changes/fix-notebook-projection-cache/README.md
+++ b/openspec/changes/fix-notebook-projection-cache/README.md
@@ -1,3 +1,3 @@
# fix-notebook-projection-cache
-Invalidate cached notebook projections when input data changes.
+Keep retained Preparation-notebook intermediates aligned with the selected input.
diff --git a/openspec/changes/fix-notebook-projection-cache/design.md b/openspec/changes/fix-notebook-projection-cache/design.md
index 01718763..b0f78d0f 100644
--- a/openspec/changes/fix-notebook-projection-cache/design.md
+++ b/openspec/changes/fix-notebook-projection-cache/design.md
@@ -1,19 +1,20 @@
## Context
-`ProtSpace_Preparation.ipynb` keeps `output/tmp` so expensive FASTA downloads, embeddings, and annotations can survive repeated Generate actions. `ReductionPipeline` also stores projections there. Its projection cache key includes the logical embedding name, method, dimensions, and reducer parameters, but not the embedding matrix or headers. The notebook reuses generic embedding names such as `prot_t5`, so a changed input can collide with a prior projection. The issue's desired notebook behavior is simpler than the CLI's reusable-cache behavior: Generate must recompute projections.
+`ProtSpace_Preparation.ipynb` keeps `output/tmp` so expensive FASTA downloads, embeddings, and annotations can survive repeated Generate actions. `ReductionPipeline` also stores projections there. Projection keys already include the logical embedding name, method, dimensions, and every reducer parameter, so the slider-only symptom in issue #338 is not reproduced by the current implementation. The reproducible collision is broader: the notebook stores every query as `sequences.fasta`, every model as `{embedder}.h5`, every annotation set as `all_annotations.parquet`, and projections under one shared directory. Changing datasets can therefore reuse a different query's FASTA, append disjoint proteins to an embedding file, retain an old embedding for a changed sequence with the same identifier, or return annotations for unrelated identifiers.
## Goals / Non-Goals
**Goals:**
- Guarantee that every Preparation-notebook Generate action reduces the current embedding data.
-- Preserve caching for the notebook's more expensive input, embedding, and annotation stages.
-- Cover changed input with an observable reducer-execution regression.
+- Preserve caching for compatible query, embedding, and annotation inputs.
+- Prevent query, embedding, and annotation cache reuse across incompatible inputs.
+- Cover changed queries, disjoint inputs, same-ID sequence changes, annotation identifiers, and projection refresh with focused regressions.
**Non-Goals:**
- Redesign projection cache identity for CLI users.
-- Disable every notebook cache or alter query/embedding/annotation refresh semantics.
+- Disable every notebook cache or redesign backend resume semantics.
- Change reducer parameters, projection naming, bundle layout, or output paths.
## Decisions
@@ -24,25 +25,41 @@ The notebook will construct `PipelineConfig` with `refetch_stages=frozenset({"pr
This uses the pipeline's public configuration contract and keeps cache lifecycle in one place.
+Reducer-parameter changes already select a distinct projection cache key. Explicit projection refresh remains a notebook-level correctness guarantee and also protects same-name inputs whose matrices differ.
+
**Alternative: delete `proj_*.npz` files before each run.** Rejected because it duplicates cache naming/lifecycle knowledge in the notebook and introduces an unnecessary destructive filesystem operation.
-**Alternative: hash all embedding bytes and headers in the core cache key.** Rejected for this issue because it broadens CLI cache semantics and adds hashing cost to all callers. The notebook explicitly wants fresh projections on Generate, so selecting the existing refetch stage is both clearer and narrower.
+**Alternative: hash all embedding bytes and headers in the core projection key.** Rejected because it broadens CLI cache semantics and adds hashing cost to all callers. The notebook explicitly wants fresh projections on Generate, so selecting the existing refetch stage is clearer and narrower.
+
+### Partition retained notebook caches by their owning input
+
+UniProt query FASTA paths are derived from a short SHA-256 digest of the exact query text. Once an input file is available, the notebook derives its intermediate directory from a streaming SHA-256 digest of that file's bytes. Query and uploaded FASTA inputs therefore place embedding, annotation, and projection intermediates under a content-owned directory; H5 inputs use the same rule directly.
+
+This keeps byte-identical inputs reusable while separating changed queries, disjoint FASTA files, and same-identifier sequences whose residues changed. The helper functions live beside the existing pipeline cache logic, and the notebook supplies the resulting directory through the existing `PipelineConfig.intermediate_dir` contract.
+
+**Alternative: teach each embedding backend to reconcile per-sequence hashes inside H5.** Rejected because both backends already implement resumable H5 writes and changing that format would broaden this notebook-scoped fix.
+
+### Validate annotation identifiers before reuse
+
+`ReductionPipeline._fetch_annotations` compares the cached and requested identifier multisets before considering cached columns. A mismatch rebuilds the annotation cache for the current headers instead of passing incompatible rows into the bundle merge. Exact-identifier caches retain the existing incremental column/source behavior.
### Exercise actual cache behavior in the regression
-The regression will use the real `ReductionPipeline._run_reductions` cache path with a deterministic fake reducer. It will run two same-name embedding sets with different data through a configuration that requests projection refresh, then assert the reducer sees both inputs and the second result reflects the second input.
+The projection regression uses a normally constructed `ReductionPipeline` and substitutes only the reducer call. It runs two same-name embedding sets with different data through the notebook's configured projection refresh, then asserts the reducer sees both inputs and the second result reflects the second input. Additional focused tests assert cache paths differ for query changes, disjoint FASTA inputs, and same-ID changed sequences, and that annotation identifiers are validated before reuse.
The notebook artifact will also be validated as a parseable notebook with parseable code cells, following existing notebook verification practice.
## Risks / Trade-offs
- **Projection reruns take longer even when nothing changed.** → This is the explicit notebook correctness contract; expensive embedding and annotation intermediates remain cached.
+- **Hashing an input file adds one sequential read per Generate action.** → Example and uploaded inputs are already read for processing, and the bounded cost avoids far more expensive incompatible embedding reuse.
+- **Old shared cache files remain under `output/tmp`.** → New input-owned paths ignore them; no destructive migration is required.
- **The regression could test pipeline behavior without proving notebook wiring.** → Verification will additionally inspect the executed notebook configuration path and validate all notebook code cells.
- **A future pipeline refetch API rename could break the notebook.** → The focused pipeline regression and notebook configuration verification make that failure visible.
## Migration Plan
-No data migration is required. Existing projection cache files may remain in `output/tmp`; the notebook will stop reading them during Generate. Rollback is a one-line notebook configuration revert.
+No data migration is required. Existing shared FASTA, embedding, annotation, and projection files may remain in `output/tmp`; the notebook uses new query- and input-owned subpaths and stops reading incompatible shared entries. Rollback restores the shared cache paths and removes the explicit projection refresh.
## Open Questions
diff --git a/openspec/changes/fix-notebook-projection-cache/proposal.md b/openspec/changes/fix-notebook-projection-cache/proposal.md
index 9e2a677e..baeaceb9 100644
--- a/openspec/changes/fix-notebook-projection-cache/proposal.md
+++ b/openspec/changes/fix-notebook-projection-cache/proposal.md
@@ -1,18 +1,20 @@
## Why
-The Preparation notebook keeps one intermediate directory across Generate runs, but projection cache identity does not include the input embeddings. A later run can therefore rebundle stale coordinates when its input changes while the embedding name, method, and reducer parameters remain the same.
+Issue #338 reports stale projections after changing a dimensionality-reduction slider. The existing projection key already includes all reducer parameters, so that exact symptom is not reproduced by the current code. Auditing the same retained-cache flow exposed a separate reproducible problem: the Preparation notebook shares query FASTA, embedding, annotation, and projection caches across unrelated inputs. A later Generate action can therefore use stale or unioned upstream data when the selected query, FASTA, sequence content, or H5 input changes.
## What Changes
- Make every Generate action in `ProtSpace_Preparation.ipynb` explicitly recompute dimensionality-reduction projections.
-- Continue retaining the notebook's expensive query, embedding, and annotation intermediates; only projection reuse changes.
-- Add regression coverage proving an explicitly refreshed projection does not reuse coordinates from changed input data.
+- Partition retained query FASTA files by query text and other intermediates by input-file content so only compatible inputs share cache entries.
+- Validate annotation-cache identifiers before reuse.
+- Continue retaining compatible query, embedding, and annotation intermediates.
+- Add focused regression coverage for changed queries, disjoint FASTA inputs, same-ID sequence changes, cross-dataset annotations, and explicitly refreshed projections.
## Capabilities
### New Capabilities
-- `notebook-projection-cache-safety`: Defines how the Preparation notebook treats cached projections across Generate actions.
+- `notebook-projection-cache-safety`: Defines how the Preparation notebook owns retained query, embedding, annotation, and projection intermediates across Generate actions.
### Modified Capabilities
@@ -21,5 +23,6 @@ None.
## Impact
- Affected notebook: `apps/protspace/notebooks/ProtSpace_Preparation.ipynb`.
-- Affected tests: Python pipeline regression coverage for notebook-equivalent projection refresh behavior.
+- Affected pipeline helper: annotation cache validation and content-addressed notebook cache paths in `apps/protspace/src/protspace/data/processors/pipeline.py`.
+- Affected tests: focused Python pipeline regressions using normal pipeline construction.
- No CLI defaults, bundle format, public Python API, or dependencies change.
diff --git a/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
index 510ec89f..4c393d60 100644
--- a/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
+++ b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
@@ -2,7 +2,7 @@
### Requirement: Preparation notebook Generate actions use current projection inputs
-The Preparation notebook SHALL recompute dimensionality-reduction projections on every Generate action and SHALL NOT read cached projection coordinates from an earlier action. This projection refresh SHALL NOT disable caching for other intermediate stages.
+The Preparation notebook SHALL recompute dimensionality-reduction projections on every Generate action and SHALL NOT read cached projection coordinates from an earlier action. This projection refresh SHALL NOT disable compatible caching for other intermediate stages.
#### Scenario: Reducer parameters change between Generate actions
@@ -16,8 +16,40 @@ The Preparation notebook SHALL recompute dimensionality-reduction projections on
- **THEN** the reducer runs against the current embedding matrix
- **AND** cached coordinates from the earlier input are not used
-#### Scenario: Non-projection intermediates remain reusable
+#### Scenario: Compatible non-projection intermediates remain reusable
- **WHEN** the notebook requests fresh projections
- **THEN** only the projection stage is explicitly refreshed
-- **AND** retained query, embedding, and annotation intermediates remain eligible for their existing cache behavior
+- **AND** retained query, embedding, and annotation intermediates remain eligible for reuse when their cache identity matches the current input
+
+### Requirement: Preparation notebook caches are owned by their inputs
+
+The Preparation notebook SHALL partition retained query FASTA files by query text and SHALL partition embedding, annotation, and projection intermediates by the content of the selected input file.
+
+#### Scenario: UniProt query changes between Generate actions
+
+- **WHEN** a user generates from one UniProt query and then selects a different query
+- **THEN** the second action SHALL NOT reuse the first query's downloaded FASTA
+
+#### Scenario: Disjoint FASTA input replaces the current input
+
+- **WHEN** a user generates embeddings from one FASTA file and then selects a disjoint FASTA file
+- **THEN** the second action SHALL use an embedding cache owned by the second FASTA content
+- **AND** the downloaded bundle SHALL NOT contain the union of both inputs
+
+#### Scenario: Sequence changes without changing its identifier
+
+- **WHEN** a FASTA sequence changes while its identifier and selected embedder remain unchanged
+- **THEN** the changed FASTA content SHALL select a different embedding cache
+- **AND** the sequence SHALL be embedded from its current residues
+
+### Requirement: Annotation cache reuse validates identifiers
+
+The reduction pipeline SHALL reuse a retained annotation cache only when its identifier multiset matches the identifiers requested by the current run.
+
+#### Scenario: Input identifiers change between runs
+
+- **WHEN** a retained annotation cache contains identifiers from an earlier input
+- **AND** the current run requests a different identifier multiset
+- **THEN** annotations SHALL be fetched for the current identifiers
+- **AND** incompatible cached rows SHALL NOT be returned as the current metadata
diff --git a/openspec/changes/fix-notebook-projection-cache/tasks.md b/openspec/changes/fix-notebook-projection-cache/tasks.md
index 0e860fd6..355695c6 100644
--- a/openspec/changes/fix-notebook-projection-cache/tasks.md
+++ b/openspec/changes/fix-notebook-projection-cache/tasks.md
@@ -1,20 +1,25 @@
## 1. Regression coverage
-- [x] 1.1 Add the smallest pipeline regression that changes same-name input embeddings across retained-cache runs and asserts projection refresh processes the second input.
-- [x] 1.2 Run the regression before implementation and record the expected stale-cache failure.
+- [x] 1.1 Integrate the same-name changed-embedding projection regression into the normal pipeline suite and construct `ReductionPipeline` through its initializer.
+- [x] 1.2 Add focused regressions for query changes, disjoint FASTA inputs, same-ID sequence changes, and annotation identifier mismatches.
+- [x] 1.3 Run the new cache-identity regressions before implementation and record the expected failures.
## 2. Notebook implementation
- [x] 2.1 Configure `ProtSpace_Preparation.ipynb` to explicitly refresh only the projection stage on every Generate action.
-- [x] 2.2 Keep query, embedding, and annotation cache wiring unchanged.
+- [x] 2.2 Partition cached query FASTA files by query text.
+- [x] 2.3 Partition retained embedding, annotation, and projection intermediates by selected input-file content.
+- [x] 2.4 Validate cached annotation identifiers before reuse and preserve incremental reuse for matching inputs.
## 3. Focused verification
- [x] 3.1 Run the regression after implementation and observe it pass.
- [x] 3.2 Validate the notebook with `nbformat` and compile every code cell after removing Colab magics.
- [x] 3.3 Verify the original two-run reproduction returns coordinates from the changed input and invokes the reducer twice.
+- [x] 3.4 Run the consolidated pipeline regressions and the full non-slow Python suite.
## 4. Repository gates
- [x] 4.1 Run affected Python tests and Ruff checks.
-- [x] 4.2 Run `pnpm precommit` before commit and push.
+- [x] 4.2 Run `openspec validate fix-notebook-projection-cache --strict`.
+- [x] 4.3 Run `pnpm precommit` before commit and push.
From 82b6dbb25c93cf3eff1ed1234491ff4331985c92 Mon Sep 17 00:00:00 2001
From: Florin Senoner <23100806+FlorinSenoner@users.noreply.github.com>
Date: Wed, 5 Aug 2026 14:45:20 +0200
Subject: [PATCH 3/5] fix(notebook): isolate backends and publish fasta
atomically
---
apps/protspace/CLAUDE.md | 3 +-
.../notebooks/ProtSpace_Preparation.ipynb | 5 +-
.../src/protspace/data/loaders/query.py | 68 ++++++++++++-----
.../src/protspace/data/processors/pipeline.py | 5 ++
apps/protspace/tests/test_backend_switch.py | 66 +++++++++++++++-
apps/protspace/tests/test_pipeline_utils.py | 24 ++++++
apps/protspace/tests/test_query.py | 76 +++++++++++++++++++
.../fix-notebook-projection-cache/design.md | 22 +++++-
.../fix-notebook-projection-cache/proposal.md | 7 +-
.../notebook-projection-cache-safety/spec.md | 26 ++++++-
.../fix-notebook-projection-cache/tasks.md | 3 +
11 files changed, 276 insertions(+), 29 deletions(-)
create mode 100644 apps/protspace/tests/test_query.py
diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md
index 53b95df1..30c58316 100644
--- a/apps/protspace/CLAUDE.md
+++ b/apps/protspace/CLAUDE.md
@@ -277,9 +277,10 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
| `test_annotation_select.py` | Annotation selection: suitability filter (cardinality/numeric/id-like exclusion), `auto` vs explicit-list label building (explicit names bypass the heuristic), missing-value dropping |
| `test_annotation_validity.py` | `AnnotationValidityStatistic`: silhouette/DBI/CH scored per annotation on `ctx.coords`, embedding vs. projection `space_kind`, missing-value exclusion, single-category no-op, id-canonical subsample determinism |
| `test_biocentral_embedder.py` | Biocentral API client, embedding flow |
-| `test_backend_switch.py` | Embedding backend switch: `resolve_default_backend` (Colab+GPU→local), `embed_fasta` local/biocentral dispatch (short key vs resolved name), `protspace embed --backend` CLI wiring + enum validation + non-positive batch_size rejection |
+| `test_backend_switch.py` | Embedding backend switch: notebook cache ownership/reuse, `resolve_default_backend` (Colab+GPU→local), `embed_fasta` local/biocentral dispatch (short key vs resolved name), `protspace embed --backend` CLI wiring + enum validation + non-positive batch_size rejection |
| `test_local_embedder.py` | Local embedding backend: checkpoint resolution (12 short keys, Synthyra ESM-C), per-family preprocessing/residue pooling, `/`-in-header guard, LocalEmbedConfig validation, empty-output guard, esm2_8m end-to-end + resume (slow) |
| `test_fasta.py` | FASTA parsing, edge cases, CSV annotation loading |
+| `test_query.py` | UniProt query FASTA download validation and atomic cache publication |
| `test_biocentral_retriever.py` | Biocentral prediction retriever (TMbed parsing, per-sequence) |
| `test_taxonomy_annotation_retriever.py` | Taxonomy via UniProt Taxonomy API (mocked + integration) |
| `test_config_validation.py` | DimensionReductionConfig parameter validation |
diff --git a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
index e82a5789..ae8285dd 100644
--- a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
+++ b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
@@ -57,6 +57,7 @@
" PipelineConfig,\n",
" ReducerParams,\n",
" ReductionPipeline,\n",
+ " _embedding_cache_path,\n",
" _input_cache_dir,\n",
" parse_methods_arg,\n",
" _query_fasta_cache_path,\n",
@@ -626,7 +627,7 @@
" fasta_path, emb_name,\n",
" backend=backend,\n",
" embed_config=_emb_cfg,\n",
- " embedding_cache=cache_dir / f\"{emb_name}.h5\",\n",
+ " embedding_cache=_embedding_cache_path(cache_dir, emb_name, backend),\n",
" )\n",
" emb_set.fasta_path = fasta_path\n",
" embedding_sets.append(emb_set)\n",
@@ -649,7 +650,7 @@
" fasta_path, emb_name,\n",
" backend=backend,\n",
" embed_config=_emb_cfg,\n",
- " embedding_cache=cache_dir / f\"{emb_name}.h5\",\n",
+ " embedding_cache=_embedding_cache_path(cache_dir, emb_name, backend),\n",
" )\n",
" emb_set.fasta_path = fasta_path\n",
" embedding_sets.append(emb_set)\n",
diff --git a/apps/protspace/src/protspace/data/loaders/query.py b/apps/protspace/src/protspace/data/loaders/query.py
index 7c30119e..6084a735 100644
--- a/apps/protspace/src/protspace/data/loaders/query.py
+++ b/apps/protspace/src/protspace/data/loaders/query.py
@@ -33,43 +33,68 @@ def query_uniprot(
base_url = "https://rest.uniprot.org/uniprotkb/stream"
params = {"compressed": "true", "format": "fasta", "query": query}
+ temp_gz_file: Path | None = None
+ staged_path: Path | None = None
+ extracted_path: Path | None = None
+ completed = False
try:
response = requests.get(base_url, params=params, stream=True)
response.raise_for_status()
# Download to temporary compressed file
- temp_file = tempfile.NamedTemporaryFile(
- mode="wb", suffix=".fasta.gz", delete=False
- )
- temp_gz_file = Path(temp_file.name)
-
total_size = int(response.headers.get("content-length", 0))
- with tqdm(
- total=total_size, unit="B", unit_scale=True, desc="Downloading FASTA"
- ) as pbar:
- for chunk in response.iter_content(chunk_size=8192):
- if chunk:
- temp_file.write(chunk)
- pbar.update(len(chunk))
- temp_file.close()
+ with tempfile.NamedTemporaryFile(
+ mode="wb", suffix=".fasta.gz", delete=False
+ ) as temp_file:
+ temp_gz_file = Path(temp_file.name)
+ with tqdm(
+ total=total_size,
+ unit="B",
+ unit_scale=True,
+ desc="Downloading FASTA",
+ ) as pbar:
+ for chunk in response.iter_content(chunk_size=8192):
+ if chunk:
+ temp_file.write(chunk)
+ pbar.update(len(chunk))
# Extract identifiers from compressed FASTA
identifiers = _extract_identifiers_gz(temp_gz_file)
# Extract FASTA to final location
- if save_to:
- extracted_path = save_to
- extracted_path.parent.mkdir(parents=True, exist_ok=True)
+ if save_to is not None:
+ save_to = Path(save_to)
+ save_to.parent.mkdir(parents=True, exist_ok=True)
+ with tempfile.NamedTemporaryFile(
+ mode="w",
+ prefix=f".{save_to.name}.",
+ suffix=".tmp",
+ dir=save_to.parent,
+ delete=False,
+ ) as staged_file:
+ staged_path = Path(staged_file.name)
+ extracted_path = staged_path
else:
extracted_path = temp_gz_file.with_suffix("")
with gzip.open(temp_gz_file, "rt") as gz_file:
content = gz_file.read()
with open(extracted_path, "w") as out:
- out.write(content)
+ written = out.write(content)
+ if written != len(content):
+ raise OSError("Incomplete FASTA extraction")
+
+ extracted_identifiers = extract_identifiers_from_fasta(extracted_path)
+ if extracted_identifiers != identifiers:
+ raise ValueError("Extracted FASTA identifiers do not match the download")
+
+ if save_to is not None:
+ staged_path.replace(save_to)
+ staged_path = None
+ extracted_path = save_to
- temp_gz_file.unlink(missing_ok=True)
+ completed = True
logger.info(f"Downloaded and extracted {len(identifiers)} sequences")
return identifiers, extracted_path
@@ -80,6 +105,13 @@ def query_uniprot(
except Exception as e:
logger.error(f"Error processing FASTA: {e}")
raise
+ finally:
+ if temp_gz_file is not None:
+ temp_gz_file.unlink(missing_ok=True)
+ if staged_path is not None:
+ staged_path.unlink(missing_ok=True)
+ if not completed and save_to is None and extracted_path is not None:
+ extracted_path.unlink(missing_ok=True)
def extract_identifiers_from_fasta(fasta_path: Path) -> list[str]:
diff --git a/apps/protspace/src/protspace/data/processors/pipeline.py b/apps/protspace/src/protspace/data/processors/pipeline.py
index 6c963289..d697e3ae 100644
--- a/apps/protspace/src/protspace/data/processors/pipeline.py
+++ b/apps/protspace/src/protspace/data/processors/pipeline.py
@@ -97,6 +97,11 @@ def _input_cache_dir(cache_root: Path, input_path: Path) -> Path:
return cache_root / "inputs" / digest.hexdigest()[:12]
+def _embedding_cache_path(cache_dir: Path, embedder: str, backend: str) -> Path:
+ """Return the H5 path owned by one input, model, and producing backend."""
+ return cache_dir / f"{backend}-{embedder}.h5"
+
+
# Valid override parameter names (from ReducerParams fields)
_VALID_OVERRIDE_KEYS = {f.name for f in fields(ReducerParams)}
# Field types for coercion
diff --git a/apps/protspace/tests/test_backend_switch.py b/apps/protspace/tests/test_backend_switch.py
index c00899a3..79d3b94f 100644
--- a/apps/protspace/tests/test_backend_switch.py
+++ b/apps/protspace/tests/test_backend_switch.py
@@ -21,7 +21,7 @@
from protspace.data.loaders.fasta import embed_fasta
-def _fake_embed(captured):
+def _fake_embed(captured, fill_value=1.0):
"""A stand-in for ``embed_sequences`` that records its args and writes a
minimal valid HDF5 so the surrounding load_h5 machinery still works."""
@@ -29,7 +29,10 @@ def fake(sequences, embedder, h5_path, embed_config=None):
with h5py.File(h5_path, "a") as f:
for pid in sequences:
if pid not in f:
- f.create_dataset(pid, data=np.ones(4, dtype=np.float32))
+ f.create_dataset(
+ pid,
+ data=np.full(4, fill_value, dtype=np.float32),
+ )
captured["embedder"] = embedder
captured["ids"] = list(sequences)
captured["config"] = embed_config
@@ -121,6 +124,65 @@ def test_embed_fasta_unknown_backend_raises(tmp_path):
embed_fasta(fasta, "prot_t5", backend="nope", embedding_cache=tmp_path / "e.h5")
+def test_notebook_cache_switches_embedding_producer(tmp_path, monkeypatch):
+ from protspace.data.processors.pipeline import _embedding_cache_path
+
+ fasta = tmp_path / "s.fasta"
+ fasta.write_text(">P12345\nMKVLAAG\n")
+ local_capture = {}
+ biocentral_capture = {}
+ monkeypatch.setattr(
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed(local_capture, fill_value=1.0),
+ )
+ monkeypatch.setattr(
+ "protspace.data.embedding.biocentral.embed_sequences",
+ _fake_embed(biocentral_capture, fill_value=2.0),
+ )
+
+ embed_fasta(
+ fasta,
+ "prot_t5",
+ backend="local",
+ embedding_cache=_embedding_cache_path(tmp_path, "prot_t5", "local"),
+ )
+ result = embed_fasta(
+ fasta,
+ "prot_t5",
+ backend="biocentral",
+ embedding_cache=_embedding_cache_path(tmp_path, "prot_t5", "biocentral"),
+ )
+
+ assert biocentral_capture["ids"] == ["P12345"]
+ assert result.data.tolist() == [[2.0, 2.0, 2.0, 2.0]]
+
+
+def test_notebook_cache_reuses_same_embedding_producer(tmp_path, monkeypatch):
+ from protspace.data.processors.pipeline import _embedding_cache_path
+
+ fasta = tmp_path / "s.fasta"
+ fasta.write_text(">P12345\nMKVLAAG\n")
+ cache = _embedding_cache_path(tmp_path, "prot_t5", "local")
+ monkeypatch.setattr(
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed({}, fill_value=1.0),
+ )
+ embed_fasta(fasta, "prot_t5", backend="local", embedding_cache=cache)
+ monkeypatch.setattr(
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed({}, fill_value=2.0),
+ )
+
+ result = embed_fasta(
+ fasta,
+ "prot_t5",
+ backend="local",
+ embedding_cache=_embedding_cache_path(tmp_path, "prot_t5", "local"),
+ )
+
+ assert result.data.tolist() == [[1.0, 1.0, 1.0, 1.0]]
+
+
# ---------------------------------------------------------------------------
# CLI wiring
# ---------------------------------------------------------------------------
diff --git a/apps/protspace/tests/test_pipeline_utils.py b/apps/protspace/tests/test_pipeline_utils.py
index c2053830..d078f068 100644
--- a/apps/protspace/tests/test_pipeline_utils.py
+++ b/apps/protspace/tests/test_pipeline_utils.py
@@ -696,6 +696,13 @@ def test_query_cache_path_changes_with_query(self, tmp_path):
assert globin != phosphatase
assert globin.parent == phosphatase.parent == tmp_path / "queries"
+ def test_query_cache_path_is_reused_for_same_query(self, tmp_path):
+ query = "(family:globin) AND (reviewed:true)"
+
+ assert pipeline_module._query_fasta_cache_path(
+ tmp_path, query
+ ) == pipeline_module._query_fasta_cache_path(tmp_path, query)
+
def test_input_cache_dir_changes_for_disjoint_fasta_inputs(self, tmp_path):
first = tmp_path / "first.fasta"
second = tmp_path / "second.fasta"
@@ -727,6 +734,23 @@ def test_input_cache_dir_is_reused_for_identical_content(self, tmp_path):
tmp_path, first
) == pipeline_module._input_cache_dir(tmp_path, renamed)
+ def test_embedding_cache_path_changes_with_backend(self, tmp_path):
+ cache_dir = tmp_path / "inputs" / "content-key"
+
+ local = pipeline_module._embedding_cache_path(cache_dir, "prot_t5", "local")
+ biocentral = pipeline_module._embedding_cache_path(
+ cache_dir, "prot_t5", "biocentral"
+ )
+
+ assert local != biocentral
+
+ def test_embedding_cache_path_is_reused_for_same_backend(self, tmp_path):
+ cache_dir = tmp_path / "inputs" / "content-key"
+
+ assert pipeline_module._embedding_cache_path(
+ cache_dir, "prot_t5", "local"
+ ) == pipeline_module._embedding_cache_path(cache_dir, "prot_t5", "local")
+
# ---------------------------------------------------------------------------
# Annotation cache identity
diff --git a/apps/protspace/tests/test_query.py b/apps/protspace/tests/test_query.py
new file mode 100644
index 00000000..145e6582
--- /dev/null
+++ b/apps/protspace/tests/test_query.py
@@ -0,0 +1,76 @@
+"""Tests for UniProt query FASTA downloads and publication."""
+
+import builtins
+import gzip
+from pathlib import Path
+
+import pytest
+
+from protspace.data.loaders import query as query_module
+
+
+class _Response:
+ headers: dict[str, str] = {}
+
+ def __init__(self, content: bytes):
+ self.content = content
+
+ def raise_for_status(self) -> None:
+ pass
+
+ def iter_content(self, chunk_size: int):
+ yield self.content
+
+
+class _InterruptingWriter:
+ def __init__(self, wrapped):
+ self.wrapped = wrapped
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, *args):
+ return self.wrapped.__exit__(*args)
+
+ def write(self, content: str):
+ self.wrapped.write(content[:10])
+ self.wrapped.flush()
+ raise RuntimeError("interrupted extraction")
+
+
+def _mock_download(monkeypatch, fasta: str) -> None:
+ response = _Response(gzip.compress(fasta.encode()))
+ monkeypatch.setattr(query_module.requests, "get", lambda *args, **kwargs: response)
+
+
+def test_query_uniprot_does_not_publish_partial_fasta(tmp_path, monkeypatch):
+ target = tmp_path / "query.fasta"
+ _mock_download(monkeypatch, ">P1\nAAAA\n>P2\nCCCC\n")
+ real_open = builtins.open
+
+ def interrupt_cache_write(file, mode="r", *args, **kwargs):
+ opened = real_open(file, mode, *args, **kwargs)
+ if "w" in mode and Path(file).parent == tmp_path:
+ return _InterruptingWriter(opened)
+ return opened
+
+ monkeypatch.setattr(builtins, "open", interrupt_cache_write)
+
+ with pytest.raises(RuntimeError, match="interrupted extraction"):
+ query_module.query_uniprot("family:globin", save_to=target)
+
+ assert not target.exists()
+ assert list(tmp_path.iterdir()) == []
+
+
+def test_query_uniprot_atomically_publishes_complete_fasta(tmp_path, monkeypatch):
+ target = tmp_path / "query.fasta"
+ fasta = ">sp|P1|ONE Protein one\nAAAA\n>P2 Protein two\nCCCC\n"
+ _mock_download(monkeypatch, fasta)
+
+ identifiers, path = query_module.query_uniprot("family:globin", save_to=target)
+
+ assert identifiers == ["P1", "P2"]
+ assert path == target
+ assert target.read_text() == fasta
+ assert list(tmp_path.iterdir()) == [target]
diff --git a/openspec/changes/fix-notebook-projection-cache/design.md b/openspec/changes/fix-notebook-projection-cache/design.md
index b0f78d0f..3ee677ce 100644
--- a/openspec/changes/fix-notebook-projection-cache/design.md
+++ b/openspec/changes/fix-notebook-projection-cache/design.md
@@ -1,6 +1,6 @@
## Context
-`ProtSpace_Preparation.ipynb` keeps `output/tmp` so expensive FASTA downloads, embeddings, and annotations can survive repeated Generate actions. `ReductionPipeline` also stores projections there. Projection keys already include the logical embedding name, method, dimensions, and every reducer parameter, so the slider-only symptom in issue #338 is not reproduced by the current implementation. The reproducible collision is broader: the notebook stores every query as `sequences.fasta`, every model as `{embedder}.h5`, every annotation set as `all_annotations.parquet`, and projections under one shared directory. Changing datasets can therefore reuse a different query's FASTA, append disjoint proteins to an embedding file, retain an old embedding for a changed sequence with the same identifier, or return annotations for unrelated identifiers.
+`ProtSpace_Preparation.ipynb` keeps `output/tmp` so expensive FASTA downloads, embeddings, and annotations can survive repeated Generate actions. `ReductionPipeline` also stores projections there. Projection keys already include the logical embedding name, method, dimensions, and every reducer parameter, so the slider-only symptom in issue #338 is not reproduced by the current implementation. The reproducible collision is broader: the notebook originally shared every query FASTA, model H5, annotation set, and projection directory. Input-content partitioning separates datasets, but an H5 still needs producer ownership because Local and Biocentral both resume by identifier, and a query FASTA must not appear at its final cache path until extraction completes.
## Goals / Non-Goals
@@ -9,6 +9,8 @@
- Guarantee that every Preparation-notebook Generate action reduces the current embedding data.
- Preserve caching for compatible query, embedding, and annotation inputs.
- Prevent query, embedding, and annotation cache reuse across incompatible inputs.
+- Prevent embedding reuse across producing backends while preserving reuse within one backend and model.
+- Make a query FASTA visible as a cache hit only after complete, validated extraction.
- Cover changed queries, disjoint inputs, same-ID sequence changes, annotation identifiers, and projection refresh with focused regressions.
**Non-Goals:**
@@ -39,6 +41,20 @@ This keeps byte-identical inputs reusable while separating changed queries, disj
**Alternative: teach each embedding backend to reconcile per-sequence hashes inside H5.** Rejected because both backends already implement resumable H5 writes and changing that format would broaden this notebook-scoped fix.
+### Include the embedding producer in H5 ownership
+
+Within an input-content directory, the notebook names each embedding H5 with the resolved backend and selected model. The input digest still owns the sequences, the model name still owns the requested representation, and the backend namespace prevents Local-produced identifiers from satisfying Biocentral's resume check or vice versa. Repeating the same input, backend, and model selects the same H5 and preserves the intended resume behavior.
+
+The notebook constructs fixed default backend configurations. Their batch sizes affect scheduling rather than vector identity, so no additional configuration hash is introduced.
+
+**Alternative: store and validate producer metadata inside every H5.** Rejected because producer-specific paths close the notebook collision without changing the shared H5 format or backend APIs.
+
+### Publish query FASTA caches atomically
+
+`query_uniprot` extracts a downloaded gzip into a temporary sibling of the requested cache file. It parses that staged FASTA and requires its ordered identifiers to match those read from the compressed download. Only then does it replace the final path atomically. A `finally` cleanup removes the compressed download and any incomplete staged output, so interruption cannot leave a nonempty final-path artifact for the next Generate action to accept.
+
+**Alternative: persist a separate completion marker.** Rejected because same-directory atomic replacement makes final-path existence the completion signal without a two-file consistency problem.
+
### Validate annotation identifiers before reuse
`ReductionPipeline._fetch_annotations` compares the cached and requested identifier multisets before considering cached columns. A mismatch rebuilds the annotation cache for the current headers instead of passing incompatible rows into the bundle merge. Exact-identifier caches retain the existing incremental column/source behavior.
@@ -53,13 +69,15 @@ The notebook artifact will also be validated as a parseable notebook with parsea
- **Projection reruns take longer even when nothing changed.** → This is the explicit notebook correctness contract; expensive embedding and annotation intermediates remain cached.
- **Hashing an input file adds one sequential read per Generate action.** → Example and uploaded inputs are already read for processing, and the bounded cost avoids far more expensive incompatible embedding reuse.
+- **Backend-qualified H5 names leave prior unqualified files unused.** → They remain recoverable but are intentionally ignored because their producer cannot be proven.
+- **An interruption can leave the previous complete query FASTA in place.** → Atomic replacement preserves that known-complete artifact; incomplete staged output is removed and never published.
- **Old shared cache files remain under `output/tmp`.** → New input-owned paths ignore them; no destructive migration is required.
- **The regression could test pipeline behavior without proving notebook wiring.** → Verification will additionally inspect the executed notebook configuration path and validate all notebook code cells.
- **A future pipeline refetch API rename could break the notebook.** → The focused pipeline regression and notebook configuration verification make that failure visible.
## Migration Plan
-No data migration is required. Existing shared FASTA, embedding, annotation, and projection files may remain in `output/tmp`; the notebook uses new query- and input-owned subpaths and stops reading incompatible shared entries. Rollback restores the shared cache paths and removes the explicit projection refresh.
+No data migration is required. Existing shared FASTA, annotation, and projection files plus backend-unqualified embedding H5 files may remain in `output/tmp`; the notebook uses query-, input-, and producer-owned paths and stops reading entries whose ownership cannot be proven. Rollback restores the shared cache paths and removes the explicit projection refresh.
## Open Questions
diff --git a/openspec/changes/fix-notebook-projection-cache/proposal.md b/openspec/changes/fix-notebook-projection-cache/proposal.md
index baeaceb9..d1f7e2bb 100644
--- a/openspec/changes/fix-notebook-projection-cache/proposal.md
+++ b/openspec/changes/fix-notebook-projection-cache/proposal.md
@@ -5,10 +5,11 @@ Issue #338 reports stale projections after changing a dimensionality-reduction s
## What Changes
- Make every Generate action in `ProtSpace_Preparation.ipynb` explicitly recompute dimensionality-reduction projections.
-- Partition retained query FASTA files by query text and other intermediates by input-file content so only compatible inputs share cache entries.
+- Partition retained query FASTA files by query text, publish them atomically, and partition other intermediates by input-file content so only compatible inputs share cache entries.
+- Partition embedding H5 files by producing backend as well as input content and model.
- Validate annotation-cache identifiers before reuse.
- Continue retaining compatible query, embedding, and annotation intermediates.
-- Add focused regression coverage for changed queries, disjoint FASTA inputs, same-ID sequence changes, cross-dataset annotations, and explicitly refreshed projections.
+- Add focused regression coverage for changed queries, interrupted FASTA extraction, backend switches and reuse, disjoint FASTA inputs, same-ID sequence changes, cross-dataset annotations, and explicitly refreshed projections.
## Capabilities
@@ -23,6 +24,6 @@ None.
## Impact
- Affected notebook: `apps/protspace/notebooks/ProtSpace_Preparation.ipynb`.
-- Affected pipeline helper: annotation cache validation and content-addressed notebook cache paths in `apps/protspace/src/protspace/data/processors/pipeline.py`.
+- Affected loaders/helpers: atomic query FASTA publication plus annotation validation and content-/producer-addressed notebook cache paths.
- Affected tests: focused Python pipeline regressions using normal pipeline construction.
- No CLI defaults, bundle format, public Python API, or dependencies change.
diff --git a/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
index 4c393d60..ddacdc6b 100644
--- a/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
+++ b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
@@ -24,7 +24,7 @@ The Preparation notebook SHALL recompute dimensionality-reduction projections on
### Requirement: Preparation notebook caches are owned by their inputs
-The Preparation notebook SHALL partition retained query FASTA files by query text and SHALL partition embedding, annotation, and projection intermediates by the content of the selected input file.
+The Preparation notebook SHALL partition retained query FASTA files by query text and SHALL publish them only after validated extraction completes. It SHALL partition embedding, annotation, and projection intermediates by the content of the selected input file, and embedding H5 files SHALL additionally be owned by their producing backend and model.
#### Scenario: UniProt query changes between Generate actions
@@ -43,6 +43,30 @@ The Preparation notebook SHALL partition retained query FASTA files by query tex
- **THEN** the changed FASTA content SHALL select a different embedding cache
- **AND** the sequence SHALL be embedded from its current residues
+#### Scenario: Embedding backend changes for the same input and model
+
+- **WHEN** a user generates an embedding with one backend and then selects the other backend for the same input and model
+- **THEN** the second backend SHALL use a different embedding H5 cache
+- **AND** identifiers produced by the first backend SHALL NOT satisfy the second backend's resume check
+
+#### Scenario: Embedding backend remains unchanged
+
+- **WHEN** a user repeats Generate with the same input, backend, and model
+- **THEN** the notebook SHALL select the same embedding H5 cache
+- **AND** the backend's existing resume behavior SHALL remain available
+
+#### Scenario: Query FASTA extraction is interrupted
+
+- **WHEN** query FASTA extraction fails after writing part of its output
+- **THEN** the query-addressed final cache path SHALL NOT expose those incomplete bytes
+- **AND** incomplete temporary output SHALL be removed
+
+#### Scenario: Query FASTA extraction completes
+
+- **WHEN** the extracted FASTA identifiers match the downloaded query result
+- **THEN** the complete FASTA SHALL be atomically published at the query-addressed cache path
+- **AND** a later Generate action for that query MAY reuse it
+
### Requirement: Annotation cache reuse validates identifiers
The reduction pipeline SHALL reuse a retained annotation cache only when its identifier multiset matches the identifiers requested by the current run.
diff --git a/openspec/changes/fix-notebook-projection-cache/tasks.md b/openspec/changes/fix-notebook-projection-cache/tasks.md
index 355695c6..8c7919da 100644
--- a/openspec/changes/fix-notebook-projection-cache/tasks.md
+++ b/openspec/changes/fix-notebook-projection-cache/tasks.md
@@ -3,6 +3,7 @@
- [x] 1.1 Integrate the same-name changed-embedding projection regression into the normal pipeline suite and construct `ReductionPipeline` through its initializer.
- [x] 1.2 Add focused regressions for query changes, disjoint FASTA inputs, same-ID sequence changes, and annotation identifier mismatches.
- [x] 1.3 Run the new cache-identity regressions before implementation and record the expected failures.
+- [x] 1.4 Add RED regressions for Local/Biocentral cache ownership, same-backend reuse, and interrupted query FASTA publication.
## 2. Notebook implementation
@@ -10,6 +11,8 @@
- [x] 2.2 Partition cached query FASTA files by query text.
- [x] 2.3 Partition retained embedding, annotation, and projection intermediates by selected input-file content.
- [x] 2.4 Validate cached annotation identifiers before reuse and preserve incremental reuse for matching inputs.
+- [x] 2.5 Scope embedding H5 paths by producing backend while retaining same-backend/model reuse.
+- [x] 2.6 Stage, validate, and atomically publish query FASTA cache files, cleaning incomplete artifacts.
## 3. Focused verification
From af4ffcaf8d807a2877228364c0dfceec0378602a Mon Sep 17 00:00:00 2001
From: tsenoner
Date: Thu, 6 Aug 2026 11:45:38 +0200
Subject: [PATCH 4/5] refactor(query): drop dead short-write guard, narrow test
patch
- Remove the `written != len(content)` guard in query_uniprot; TextIOWrapper.write always returns len(s), so the OSError was unreachable. A real truncation still surfaces as an OSError from the close-time flush.
- Patch `open` on the query module instead of `builtins` in the interrupted-extraction test, so the fake only intercepts opens made by query.py.
Co-Authored-By: Claude Opus 5 (1M context)
Claude-Session: https://claude.ai/code/session_016qoU16kDQxz6U3H2UWbbm2
---
apps/protspace/src/protspace/data/loaders/query.py | 4 +---
apps/protspace/tests/test_query.py | 2 +-
2 files changed, 2 insertions(+), 4 deletions(-)
diff --git a/apps/protspace/src/protspace/data/loaders/query.py b/apps/protspace/src/protspace/data/loaders/query.py
index 6084a735..79dc7044 100644
--- a/apps/protspace/src/protspace/data/loaders/query.py
+++ b/apps/protspace/src/protspace/data/loaders/query.py
@@ -81,9 +81,7 @@ def query_uniprot(
with gzip.open(temp_gz_file, "rt") as gz_file:
content = gz_file.read()
with open(extracted_path, "w") as out:
- written = out.write(content)
- if written != len(content):
- raise OSError("Incomplete FASTA extraction")
+ out.write(content)
extracted_identifiers = extract_identifiers_from_fasta(extracted_path)
if extracted_identifiers != identifiers:
diff --git a/apps/protspace/tests/test_query.py b/apps/protspace/tests/test_query.py
index 145e6582..68bfcff5 100644
--- a/apps/protspace/tests/test_query.py
+++ b/apps/protspace/tests/test_query.py
@@ -54,7 +54,7 @@ def interrupt_cache_write(file, mode="r", *args, **kwargs):
return _InterruptingWriter(opened)
return opened
- monkeypatch.setattr(builtins, "open", interrupt_cache_write)
+ monkeypatch.setattr(query_module, "open", interrupt_cache_write, raising=False)
with pytest.raises(RuntimeError, match="interrupted extraction"):
query_module.query_uniprot("family:globin", save_to=target)
From 55837338ca191db09bd497c3c69f8eff193503dd Mon Sep 17 00:00:00 2001
From: Florin Senoner <23100806+FlorinSenoner@users.noreply.github.com>
Date: Thu, 6 Aug 2026 14:17:55 +0200
Subject: [PATCH 5/5] fix(protspace): preserve cache compatibility
---
apps/protspace/docs/cli.md | 3 +-
.../src/protspace/data/loaders/query.py | 4 ++
.../src/protspace/data/processors/pipeline.py | 35 +++++----
apps/protspace/tests/test_pipeline_utils.py | 72 ++++++++++++++++---
apps/protspace/tests/test_query.py | 15 ++++
.../fix-notebook-projection-cache/design.md | 5 +-
.../notebook-projection-cache-safety/spec.md | 14 +++-
.../fix-notebook-projection-cache/tasks.md | 4 ++
8 files changed, 121 insertions(+), 31 deletions(-)
diff --git a/apps/protspace/docs/cli.md b/apps/protspace/docs/cli.md
index 1b79014e..b75e4a69 100644
--- a/apps/protspace/docs/cli.md
+++ b/apps/protspace/docs/cli.md
@@ -348,11 +348,12 @@ With `--keep-tmp` (default), all intermediate results are cached in `{output}/tm
| ----------- | ---- | -------------- |
| FASTA sequences | `sequences.fasta` | Skip UniProt query download |
| Embeddings | `{embedder}.h5` | Skip already-embedded proteins |
-| Annotations | `all_annotations.parquet` | Fetch only missing annotation sources |
+| Annotations | `all_annotations.parquet` | Fetch missing sources when the cache covers every requested identifier; rebuild when requested identifiers are absent |
| Similarity matrix | `similarity_matrix.npy` | Skip MMseqs2 recomputation |
| DR projections | `proj_{name}_{method}_{hash}.npz` | Skip dimensionality reduction |
- Annotation cache always includes scores regardless of `--no-scores`
+- An annotation cache may cover more proteins than the current run; those extra rows are filtered later. If any requested identifier is absent, annotations are rebuilt for the current input and the cache is replaced.
- DR projection caches are keyed by embedding name, method, dimensions, and all parameters — changing any parameter creates a new cache entry
- Use `--refetch all` to bypass all caches, or `--refetch ` selectively (e.g., `--refetch ted,biocentral`)
diff --git a/apps/protspace/src/protspace/data/loaders/query.py b/apps/protspace/src/protspace/data/loaders/query.py
index 79dc7044..fda66f57 100644
--- a/apps/protspace/src/protspace/data/loaders/query.py
+++ b/apps/protspace/src/protspace/data/loaders/query.py
@@ -6,6 +6,7 @@
import gzip
import logging
+import os
import tempfile
from pathlib import Path
@@ -88,6 +89,9 @@ def query_uniprot(
raise ValueError("Extracted FASTA identifiers do not match the download")
if save_to is not None:
+ current_umask = os.umask(0)
+ os.umask(current_umask)
+ staged_path.chmod(0o666 & ~current_umask)
staged_path.replace(save_to)
staged_path = None
extracted_path = save_to
diff --git a/apps/protspace/src/protspace/data/processors/pipeline.py b/apps/protspace/src/protspace/data/processors/pipeline.py
index d697e3ae..1432b639 100644
--- a/apps/protspace/src/protspace/data/processors/pipeline.py
+++ b/apps/protspace/src/protspace/data/processors/pipeline.py
@@ -414,6 +414,15 @@ def _fetch_annotations(
intermediate_dir.mkdir(parents=True, exist_ok=True)
cache_path = intermediate_dir / "all_annotations.parquet"
+ def fetch_current_annotations() -> pd.DataFrame:
+ api_df = ProteinAnnotationManager(
+ headers=headers,
+ annotations=annotations_list,
+ output_path=cache_path,
+ sequences=sequences,
+ ).to_pd()
+ return self._merge_csv(api_df, csv_df)
+
if cache_path.exists():
cached_df = pd.read_parquet(cache_path)
cached_identifiers = (
@@ -422,19 +431,15 @@ def _fetch_annotations(
else Counter()
)
requested_identifiers = Counter(map(str, headers))
+ missing_identifiers = requested_identifiers - cached_identifiers
- if cached_identifiers != requested_identifiers:
- logger.info(
- "Annotation cache input changed; fetching annotations "
- "for the current identifiers"
+ if missing_identifiers:
+ logger.warning(
+ "Annotation cache is missing %d requested identifier(s); "
+ "fetching annotations for the current identifiers",
+ sum(missing_identifiers.values()),
)
- api_df = ProteinAnnotationManager(
- headers=headers,
- annotations=annotations_list,
- output_path=cache_path,
- sequences=sequences,
- ).to_pd()
- return self._merge_csv(api_df, csv_df)
+ return fetch_current_annotations()
cached_annotations = set(cached_df.columns) - {"identifier"}
@@ -516,13 +521,7 @@ def _fetch_annotations(
).to_pd()
return self._merge_csv(api_df, csv_df)
else:
- api_df = ProteinAnnotationManager(
- headers=headers,
- annotations=annotations_list,
- output_path=cache_path,
- sequences=sequences,
- ).to_pd()
- return self._merge_csv(api_df, csv_df)
+ return fetch_current_annotations()
else:
api_df = ProteinAnnotationManager(
headers=headers,
diff --git a/apps/protspace/tests/test_pipeline_utils.py b/apps/protspace/tests/test_pipeline_utils.py
index d078f068..40a20131 100644
--- a/apps/protspace/tests/test_pipeline_utils.py
+++ b/apps/protspace/tests/test_pipeline_utils.py
@@ -37,18 +37,37 @@ def _preparation_notebook_projection_refetch_stages() -> frozenset[str]:
for cell in notebook["cells"]
if cell["cell_type"] == "code"
)
- generate_source = next(source for source in code_sources if "def _on_gen" in source)
+ generate_source = next(
+ (source for source in code_sources if "def _on_gen" in source), None
+ )
+ if generate_source is None:
+ pytest.fail("Generate callback not found in the Preparation notebook")
tree = ast.parse(generate_source)
config_call = next(
- node
- for node in ast.walk(tree)
- if isinstance(node, ast.Call)
- and isinstance(node.func, ast.Name)
- and node.func.id == "PipelineConfig"
+ (
+ node
+ for node in ast.walk(tree)
+ if isinstance(node, ast.Call)
+ and isinstance(node.func, ast.Name)
+ and node.func.id == "PipelineConfig"
+ ),
+ None,
)
+ if config_call is None:
+ pytest.fail("PipelineConfig(...) not found in the notebook Generate callback")
refetch_keyword = next(
- keyword for keyword in config_call.keywords if keyword.arg == "refetch_stages"
+ (
+ keyword
+ for keyword in config_call.keywords
+ if keyword.arg == "refetch_stages"
+ ),
+ None,
)
+ if refetch_keyword is None:
+ pytest.fail(
+ "PipelineConfig(...) with refetch_stages not found in the notebook "
+ "Generate callback"
+ )
expression = ast.Expression(refetch_keyword.value)
return eval(
compile(expression, filename=str(notebook_path), mode="eval"),
@@ -836,19 +855,56 @@ def unexpected_fetch(_manager):
assert result["identifier"].tolist() == ["P1", "P2"]
+def test_annotation_cache_superset_is_reused_without_truncation(tmp_path, monkeypatch):
+ from protspace.data.annotations.manager import ProteinAnnotationManager
+
+ cache_path = tmp_path / "cache" / "all_annotations.parquet"
+ cache_path.parent.mkdir()
+ cached = pd.DataFrame(
+ {
+ "identifier": ["P1", "P2", "P3"],
+ "protein_name": ["one", "two", "three"],
+ "gene_name": ["gene-one", "gene-two", "gene-three"],
+ "uniprot_kb_id": ["id-one", "id-two", "id-three"],
+ }
+ )
+ cached.to_parquet(cache_path)
+ pipeline = ReductionPipeline(
+ PipelineConfig(
+ methods=[],
+ output_path=tmp_path / "output.parquetbundle",
+ keep_tmp=True,
+ intermediate_dir=cache_path.parent,
+ annotations=["protein_name"],
+ )
+ )
+
+ def unexpected_fetch(_manager):
+ pytest.fail("a cache covering every requested identifier should be reused")
+
+ monkeypatch.setattr(ProteinAnnotationManager, "to_pd", unexpected_fetch)
+
+ result = pipeline._fetch_annotations(["P2", "P1"])
+
+ assert result["identifier"].tolist() == ["P1", "P2", "P3"]
+ pd.testing.assert_frame_equal(pd.read_parquet(cache_path), cached)
+
+
# ---------------------------------------------------------------------------
# Preparation notebook projection refresh
# ---------------------------------------------------------------------------
def test_notebook_refreshes_same_name_changed_input_through_pipeline(tmp_path):
+ refetch_stages = _preparation_notebook_projection_refetch_stages()
+ assert refetch_stages == frozenset({"projections"})
config = PipelineConfig(
methods=parse_methods_arg(["umap2"]),
output_path=tmp_path / "output" / "data.parquetbundle",
keep_tmp=True,
intermediate_dir=tmp_path / "output" / "tmp",
annotations=None,
- refetch_stages=_preparation_notebook_projection_refetch_stages(),
+ refetch_stages=refetch_stages,
)
config.intermediate_dir.mkdir(parents=True)
pipeline = ReductionPipeline(config)
diff --git a/apps/protspace/tests/test_query.py b/apps/protspace/tests/test_query.py
index 68bfcff5..1608f6a6 100644
--- a/apps/protspace/tests/test_query.py
+++ b/apps/protspace/tests/test_query.py
@@ -2,6 +2,8 @@
import builtins
import gzip
+import os
+import stat
from pathlib import Path
import pytest
@@ -74,3 +76,16 @@ def test_query_uniprot_atomically_publishes_complete_fasta(tmp_path, monkeypatch
assert path == target
assert target.read_text() == fasta
assert list(tmp_path.iterdir()) == [target]
+
+
+def test_query_uniprot_publishes_fasta_with_process_umask(tmp_path, monkeypatch):
+ target = tmp_path / "query.fasta"
+ _mock_download(monkeypatch, ">P1\nAAAA\n")
+ previous_umask = os.umask(0o027)
+
+ try:
+ query_module.query_uniprot("family:globin", save_to=target)
+ finally:
+ os.umask(previous_umask)
+
+ assert stat.S_IMODE(target.stat().st_mode) == 0o640
diff --git a/openspec/changes/fix-notebook-projection-cache/design.md b/openspec/changes/fix-notebook-projection-cache/design.md
index 3ee677ce..2fa82aab 100644
--- a/openspec/changes/fix-notebook-projection-cache/design.md
+++ b/openspec/changes/fix-notebook-projection-cache/design.md
@@ -53,11 +53,13 @@ The notebook constructs fixed default backend configurations. Their batch sizes
`query_uniprot` extracts a downloaded gzip into a temporary sibling of the requested cache file. It parses that staged FASTA and requires its ordered identifiers to match those read from the compressed download. Only then does it replace the final path atomically. A `finally` cleanup removes the compressed download and any incomplete staged output, so interruption cannot leave a nonempty final-path artifact for the next Generate action to accept.
+Before publication, the staged file receives the permissions that a normal new file would receive under the process umask. Atomic replacement therefore does not make the retained FASTA less accessible than the direct-write behavior it replaces.
+
**Alternative: persist a separate completion marker.** Rejected because same-directory atomic replacement makes final-path existence the completion signal without a two-file consistency problem.
### Validate annotation identifiers before reuse
-`ReductionPipeline._fetch_annotations` compares the cached and requested identifier multisets before considering cached columns. A mismatch rebuilds the annotation cache for the current headers instead of passing incompatible rows into the bundle merge. Exact-identifier caches retain the existing incremental column/source behavior.
+`ReductionPipeline._fetch_annotations` verifies that the cached identifier multiset covers every requested identifier before considering cached columns. Missing requested identifiers rebuild the annotation cache for the current headers instead of passing incompatible rows into the bundle merge. A cached superset remains reusable because the pipeline's later identifier merge drops rows outside the current input; this preserves the existing subset-run behavior and avoids replacing a larger cache with a smaller one.
### Exercise actual cache behavior in the regression
@@ -69,6 +71,7 @@ The notebook artifact will also be validated as a parseable notebook with parsea
- **Projection reruns take longer even when nothing changed.** → This is the explicit notebook correctness contract; expensive embedding and annotation intermediates remain cached.
- **Hashing an input file adds one sequential read per Generate action.** → Example and uploaded inputs are already read for processing, and the bounded cost avoids far more expensive incompatible embedding reuse.
+- **Any FASTA content change selects a new embedding cache and re-embeds the complete file.** → This deliberately gives same-identifier sequence changes correct ownership without redesigning the shared H5 format around per-sequence hashes; incremental per-sequence invalidation remains outside this notebook-scoped change.
- **Backend-qualified H5 names leave prior unqualified files unused.** → They remain recoverable but are intentionally ignored because their producer cannot be proven.
- **An interruption can leave the previous complete query FASTA in place.** → Atomic replacement preserves that known-complete artifact; incomplete staged output is removed and never published.
- **Old shared cache files remain under `output/tmp`.** → New input-owned paths ignore them; no destructive migration is required.
diff --git a/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
index ddacdc6b..9d272ea6 100644
--- a/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
+++ b/openspec/changes/fix-notebook-projection-cache/specs/notebook-projection-cache-safety/spec.md
@@ -65,15 +65,23 @@ The Preparation notebook SHALL partition retained query FASTA files by query tex
- **WHEN** the extracted FASTA identifiers match the downloaded query result
- **THEN** the complete FASTA SHALL be atomically published at the query-addressed cache path
+- **AND** the published file SHALL use normal new-file permissions under the process umask
- **AND** a later Generate action for that query MAY reuse it
### Requirement: Annotation cache reuse validates identifiers
-The reduction pipeline SHALL reuse a retained annotation cache only when its identifier multiset matches the identifiers requested by the current run.
+The reduction pipeline SHALL reuse a retained annotation cache only when its identifier multiset contains every identifier requested by the current run. The cache MAY contain identifiers outside the current request.
-#### Scenario: Input identifiers change between runs
+#### Scenario: Requested identifiers are missing from the cache
- **WHEN** a retained annotation cache contains identifiers from an earlier input
-- **AND** the current run requests a different identifier multiset
+- **AND** the current run requests one or more identifiers absent from that cache
- **THEN** annotations SHALL be fetched for the current identifiers
- **AND** incompatible cached rows SHALL NOT be returned as the current metadata
+
+#### Scenario: Cache contains a superset of requested identifiers
+
+- **WHEN** a retained annotation cache contains every identifier requested by the current run
+- **AND** it also contains identifiers outside the current request
+- **THEN** the retained cache SHALL remain eligible for reuse
+- **AND** the larger retained cache SHALL NOT be replaced by a subset-only fetch
diff --git a/openspec/changes/fix-notebook-projection-cache/tasks.md b/openspec/changes/fix-notebook-projection-cache/tasks.md
index 8c7919da..ae74b4d1 100644
--- a/openspec/changes/fix-notebook-projection-cache/tasks.md
+++ b/openspec/changes/fix-notebook-projection-cache/tasks.md
@@ -4,6 +4,7 @@
- [x] 1.2 Add focused regressions for query changes, disjoint FASTA inputs, same-ID sequence changes, and annotation identifier mismatches.
- [x] 1.3 Run the new cache-identity regressions before implementation and record the expected failures.
- [x] 1.4 Add RED regressions for Local/Biocentral cache ownership, same-backend reuse, and interrupted query FASTA publication.
+- [x] 1.5 Add RED regressions for annotation-cache superset reuse and published FASTA permissions.
## 2. Notebook implementation
@@ -13,6 +14,8 @@
- [x] 2.4 Validate cached annotation identifiers before reuse and preserve incremental reuse for matching inputs.
- [x] 2.5 Scope embedding H5 paths by producing backend while retaining same-backend/model reuse.
- [x] 2.6 Stage, validate, and atomically publish query FASTA cache files, cleaning incomplete artifacts.
+- [x] 2.7 Preserve normal umask-derived permissions when atomically publishing query FASTA files.
+- [x] 2.8 Reuse annotation caches that cover all requested identifiers without truncating cached supersets.
## 3. Focused verification
@@ -20,6 +23,7 @@
- [x] 3.2 Validate the notebook with `nbformat` and compile every code cell after removing Colab magics.
- [x] 3.3 Verify the original two-run reproduction returns coordinates from the changed input and invokes the reducer twice.
- [x] 3.4 Run the consolidated pipeline regressions and the full non-slow Python suite.
+- [x] 3.5 Make notebook projection-refetch wiring failures explicit in the focused regression.
## 4. Repository gates