diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md
index d9c3fbec..eabf697d 100644
--- a/apps/protspace/CLAUDE.md
+++ b/apps/protspace/CLAUDE.md
@@ -268,24 +268,25 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
| File | What it covers |
|------|---------------|
-| `test_annotation_manager.py` | Annotation fetch, merge, cache, configuration, evidence parsing |
+| `test_annotation_manager.py` | Annotation fetch, merge, cache, configuration, evidence parsing; per-identifier reuse (a source is fetched only for the identifiers the cache lacks, taxonomy only for unseen organisms, rows outside the run are kept, a failed fill-in caches nothing) |
| `test_transformer.py` | Annotation transformers (field normalization, EC names) |
| `test_reducers.py` | All 6 DR methods: shapes, finite output, float16, config validation |
| `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, projection cache identity (a changed, reordered or grown matrix under one embedding name misses; an unchanged rerun hits; `--refetch projections` always recomputes), annotation cache fill-in wiring, 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, auto-cluster self-validity (filed under the membership column, gated on it, and equal to driving `AnnotationValidityStatistic` directly so an out-of-band re-score cannot drift), 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) |
| `test_stats_bundle.py` | Optional 5th (statistics) bundle part round-trip |
| `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, completeness gate (reads the .h5, not a counter), `/`-in-header rejection |
-| `test_embed_completeness.py` | Shared embed contract (`data/embedding/store.py`): `expected = requested - skipped`, skip-vs-fail, skip reporting, resume-covered runs, FASTA coverage direction + identifier normalisation |
-| `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_local_embedder.py` | Local embedding backend: checkpoint resolution (12 short keys, Synthyra ESM-C), the notebook-gating sets pinned to the registry each constrains (`COLAB_OVERSIZED`→`LOCAL_CHECKPOINTS`, `BIOCENTRAL_INVALID`→`ALL_SHORT_KEYS`), per-family preprocessing/residue pooling, `/`-in-header guard, LocalEmbedConfig validation, over-length + OOM skips reported not failed, non-skip shortfall fails, esm2_8m end-to-end + resume (slow) |
+| `test_biocentral_embedder.py` | Biocentral API client, embedding flow, completeness gate (reads the .h5, not a counter), `/`-in-header rejection, producer stamping, and a local-written cache refused before any API call |
+| `test_embed_completeness.py` | Shared embed contract (`data/embedding/store.py`): `expected = requested - skipped`, skip-vs-fail, skip reporting, resume-covered runs, FASTA coverage direction + identifier normalisation; producer ownership (another backend or model is refused with the remedies and the file left byte-identical, an unstamped file is adopted then owned) and residue identity (a changed sequence is outstanding again and replaces its vector + digest, an unchanged one resumes, a digest-less protein is trusted, digests read in one file open) |
+| `test_backend_switch.py` | Embedding backend switch: `embed_fasta` refuses another backend's cache and resumes its own, and returns only the requested FASTA's proteins; `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: producer/digest stamping, refusing a Biocentral cache before a checkpoint loads, re-embedding a changed sequence; checkpoint resolution (12 short keys, Synthyra ESM-C), the notebook-gating sets pinned to the registry each constrains (`COLAB_OVERSIZED`→`LOCAL_CHECKPOINTS`, `BIOCENTRAL_INVALID`→`ALL_SHORT_KEYS`), per-family preprocessing/residue pooling, `/`-in-header guard, LocalEmbedConfig validation, over-length + OOM skips reported not failed, non-skip shortfall fails, esm2_8m end-to-end + resume (slow) |
| `test_fasta.py` | FASTA parsing, edge cases, CSV annotation loading |
+| `test_query.py` | UniProt query FASTA download: a truncated download is never published, atomic cache publication, umask-derived permissions, and a retained FASTA owned by its query text (`prepare -q A` then `-q B` in one output directory) |
| `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 |
@@ -303,6 +304,7 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
| `test_display_decode.py` | Display-side decoding of encoded values, multi-hit rendering, gated-off passthrough |
| `test_toxprot_demo.py` | Signal-peptide bound parsing, mature-FASTA stripping, bundle post-processing (column filter/reorder) |
| `test_bundle_overlay.py` | Round-trip replacement of the annotations part of a bundle |
+| `test_atomic_publication.py` | `data/io/atomic.py`: staged rename keeps the previous content on failure, and a published file (bundle, statistics parquet, retained FASTA) carries the process umask rather than `mkstemp`'s owner-only mode |
| `test_classification.py` | Query/reference rules: id-prefix and case-insensitive `where` substring, query-over-reference precedence, empty-match and missing-column errors |
| `test_bundle_version.py` | `format_version=2` stamped into the annotations parquet |
| `test_uniprot_parser_encoding.py` | UniProtEntry free-text emit points percent-encode reserved chars |
@@ -310,7 +312,7 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
| `test_cli_no_frontend.py` | CLI imports without the optional `frontend` extra (plotly, dash) |
| `test_cli_no_similarity.py` | `-s/--similarity` without the optional `similarity` extra: up-front CLI guard (before any load/embed), loader `ImportError` backstop, `EMBEDDER_MODELS` pinned to the embedder registry |
| `test_docs_extras_sync.py` | `README.md` (PyPI) and `docs/guide/python-cli.md` (protspace.app) hold the same extras section; the guide's embedder shortcut list matches `EMBEDDER_MODELS` |
-| `test_notebooks.py` | Colab notebooks: cell magics only on line 1, every code cell compiles after IPython transformation, cell ids present for `nbformat >= 4.5`, `except ImportError` fallback sets equal the package constants they stand in for — read structurally off the guarded import, so a fallback that is missing, emptied or written in an unrecognised shape fails instead of matching nothing |
+| `test_notebooks.py` | Colab notebooks: cell magics only on line 1, every code cell compiles after IPython transformation, cell ids present for `nbformat >= 4.5`, no notebook imports a private `protspace` name (cell 1 installs the *released* package, so a private name added this release breaks setup until the next one), each Generate action names its bundle distinctly, `except ImportError` fallback sets equal the package constants they stand in for — read structurally off the guarded import, so a fallback that is missing, emptied or written in an unrecognised shape fails instead of matching nothing |
| `test_encoding_e2e.py` | Backend end-to-end round-trip proof for v2 annotation encoding |
| `test_scores_ted.py` | `--no-scores` strips TED domains |
diff --git a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
index b901a723..3baae6c0 100644
--- a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
+++ b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
@@ -49,6 +49,7 @@
"# instead — and cover every import below, not just protspace's: h5py, pandas and\n",
"# tqdm are the preloaded ones, so they are the likeliest to be the casualty.\n",
"try:\n",
+ " import hashlib\n",
" import urllib.request\n",
" from pathlib import Path\n",
"\n",
@@ -82,7 +83,7 @@
" f\"Import failed after install ({_exc}).\\n\"\n",
" \"This usually means an upgraded package needs a fresh interpreter: \"\n",
" \"Runtime > Restart session, then run this cell again.\"\n",
- " ) from _exc"
+ " ) from _exc\n"
]
},
{
@@ -518,7 +519,7 @@
" # reading, and that is also when a group *is* visible.\n",
" notes = []\n",
" if \"PCA\" in sel:\n",
- " notes.append(\"PCA has no parameters\")\n",
+ " notes.append(\"PCA has no parameters, so these sliders do not change it\")\n",
" if \"MDS\" in sel:\n",
" notes.append(\n",
" \"MDS runs with this notebook's fixed defaults\"\n",
@@ -652,7 +653,10 @@
" out_dir.mkdir(exist_ok=True)\n",
" cache_dir = out_dir / \"tmp\"\n",
" cache_dir.mkdir(exist_ok=True)\n",
- " output_path = out_dir / \"data.parquetbundle\"\n",
+ " # Distinct per run: a fixed name downloads as \"data (1).parquetbundle\"\n",
+ " # beside the previous one, and opening the older file looks exactly like\n",
+ " # a projection that never refreshed (issue #338).\n",
+ " output_path = out_dir / f\"protspace_{_time.strftime('%Y%m%d-%H%M%S')}.parquetbundle\"\n",
"\n",
" step_html = widgets.HTML(value=\"Loading embeddings...\")\n",
" display(step_html)\n",
@@ -661,59 +665,52 @@
" t0 = _time.time()\n",
" embedding_sets = []\n",
"\n",
- " if inp[\"type\"] == \"query\":\n",
+ " if inp[\"type\"] in (\"query\", \"fasta\"):\n",
" gated = _embedders_and_backend()\n",
" if gated is None:\n",
" return\n",
" embs, backend, _emb_cfg = gated\n",
- " step_html.value = \"Fetching sequences from UniProt...\"\n",
- " fasta_cache = cache_dir / \"sequences.fasta\"\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",
- " )\n",
- " headers = extract_identifiers_from_fasta(fasta_cache)\n",
- " fasta_path = fasta_cache\n",
- " print(f\"Reusing cached sequences from output/tmp ({len(headers):,} sequences) — delete output/tmp to re-run the query\")\n",
+ " if inp[\"type\"] == \"query\":\n",
+ " step_html.value = \"Fetching sequences from UniProt...\"\n",
+ " # Addressed by the query text: one shared name would hand a\n",
+ " # later query the previous one's sequences.\n",
+ " _q = hashlib.sha256(inp[\"query\"].encode()).hexdigest()[:12]\n",
+ " fasta_cache = cache_dir / \"queries\" / f\"{_q}.fasta\"\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",
+ " )\n",
+ " headers = extract_identifiers_from_fasta(fasta_cache)\n",
+ " fasta_path = fasta_cache\n",
+ " print(f\"Reusing cached sequences from output/tmp ({len(headers):,} sequences) — delete output/tmp to re-run the query\")\n",
+ " else:\n",
+ " headers, fasta_path = query_uniprot(inp[\"query\"], save_to=fasta_cache)\n",
+ " if not headers:\n",
+ " print(f\"No sequences found for query: {inp['query']}\")\n",
+ " return\n",
" else:\n",
- " headers, fasta_path = query_uniprot(inp[\"query\"], save_to=fasta_cache)\n",
- " if not headers:\n",
- " print(f\"No sequences found for query: {inp['query']}\")\n",
- " return\n",
+ " fasta_path = Path(inp[\"path\"])\n",
" for emb_name in embs:\n",
" step_html.value = f\"Computing {emb_name} embeddings ({backend})...\"\n",
" emb_set = embed_fasta(\n",
" fasta_path, emb_name,\n",
" backend=backend,\n",
" embed_config=_emb_cfg,\n",
- " embedding_cache=cache_dir / f\"{emb_name}.h5\",\n",
+ " # Backend in the name so switching the toggle keeps both\n",
+ " # caches; the store refuses to mix them either way.\n",
+ " embedding_cache=cache_dir / f\"{backend}-{emb_name}.h5\",\n",
" )\n",
" emb_set.fasta_path = fasta_path\n",
" embedding_sets.append(emb_set)\n",
- " elif inp[\"type\"] == \"fasta\":\n",
- " gated = _embedders_and_backend()\n",
- " if gated is None:\n",
- " return\n",
- " embs, backend, _emb_cfg = gated\n",
- " for emb_name in embs:\n",
- " step_html.value = f\"Computing {emb_name} embeddings ({backend})...\"\n",
- " emb_set = embed_fasta(\n",
- " Path(inp[\"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",
- " embedding_sets.append(emb_set)\n",
" else:\n",
" h5_path = Path(inp[\"path\"])\n",
" name_override = inp.get(\"name\")\n",
" emb_set = load_h5([h5_path], name_override=name_override)\n",
" embedding_sets.append(emb_set)\n",
"\n",
- " n_proteins = len(embedding_sets[0].headers)\n",
- "\n",
- " # Build pipeline with caching enabled\n",
+ " # Cached intermediates are keyed by what they were computed from\n",
+ " # (matrix, sequence, query), so reuse here is always reuse of this\n",
+ " # input's own work.\n",
" reducer_params = ReducerParams(\n",
" n_neighbors=pw[\"n_neighbors\"].value,\n",
" min_dist=pw[\"min_dist\"].value,\n",
@@ -734,15 +731,20 @@
" )\n",
" pipeline = ReductionPipeline(config)\n",
"\n",
+ " # Embedders can skip different proteins (an OOM on the larger model,\n",
+ " # say), and bundling pairs projection rows with one header list, so\n",
+ " # align every set to the shared identifiers as ReductionPipeline.run() does.\n",
+ " all_headers = pipeline._validate_headers(embedding_sets)\n",
+ " n_proteins = len(all_headers)\n",
+ "\n",
" # Step 2: Annotations (cached after first run)\n",
" step_html.value = \"Fetching annotations...\"\n",
- " metadata = pipeline._fetch_annotations(embedding_sets[0].headers)\n",
+ " metadata = pipeline._fetch_annotations(all_headers, embedding_sets)\n",
"\n",
" # Step 3: Dimensionality reduction\n",
" step_html.value = \"Reducing dimensions...\"\n",
"\n",
" # Build full metadata\n",
- " all_headers = embedding_sets[0].headers\n",
" full_metadata = _pd.DataFrame({\"identifier\": all_headers})\n",
" if len(metadata.columns) > 1:\n",
" metadata = metadata.astype(str)\n",
@@ -790,7 +792,7 @@
" print(f\"Processed {n_proteins} proteins with {len(method_specs)} method(s)\")\n",
" print(f\"\\nBundle written to {output_path.resolve()} (inside this Colab runtime).\")\n",
" print(\"Downloading it to your computer now, then open it at https://protspace.app/explore\")\n",
- " print(\"(Download blocked by the browser? Take it from the Colab Files pane: output/data.parquetbundle)\")\n",
+ " print(f\"(Download blocked by the browser? Take it from the Colab Files pane: output/{output_path.name})\")\n",
" files.download(str(output_path))\n",
"\n",
" except Exception as e:\n",
@@ -839,7 +841,7 @@
"display(widgets.VBox([\n",
" _panel,\n",
" gen_btn,\n",
- " widgets.HTML(\"The bundle is written to output/data.parquetbundle in this \"\n",
+ " widgets.HTML(\"
Each run writes a timestamped bundle to output/ in this \"\n",
" \"runtime, then downloaded to your computer.
\"),\n",
" gen_out,\n",
"]))\n",
diff --git a/apps/protspace/src/protspace/cli/prepare.py b/apps/protspace/src/protspace/cli/prepare.py
index bcea8564..260f07cc 100644
--- a/apps/protspace/src/protspace/cli/prepare.py
+++ b/apps/protspace/src/protspace/cli/prepare.py
@@ -420,10 +420,6 @@ def prepare(
# --- Build embedding sets ---
from protspace.data.loaders import EmbeddingSet, load_h5
from protspace.data.loaders.h5 import EMBEDDING_EXTENSIONS
- from protspace.data.loaders.query import (
- extract_identifiers_from_fasta,
- query_uniprot,
- )
embed_config = build_embed_config(backend, batch_size, max_length)
embedding_sets: list[EmbeddingSet] = []
@@ -431,21 +427,9 @@ def prepare(
try:
if query:
- fasta_save = cache_dir / "sequences.fasta" if cache_dir else None
- if (
- fasta_save
- and fasta_save.exists()
- and fasta_save.stat().st_size > 0
- and "query" not in refetch_stages
- ):
- headers = extract_identifiers_from_fasta(fasta_save)
- logger.warning(
- "Using cached FASTA (%s sequences)",
- f"{len(headers):,}",
- )
- fasta_path = fasta_save
- else:
- headers, fasta_path = query_uniprot(query, save_to=fasta_save)
+ from protspace.data.loaders.query import resolve_query_fasta
+
+ headers, fasta_path = resolve_query_fasta(query, cache_dir, refetch_stages)
if not headers:
raise typer.BadParameter(f"No sequences for query: '{query}'")
diff --git a/apps/protspace/src/protspace/cli/stats.py b/apps/protspace/src/protspace/cli/stats.py
index f3f438cb..131063e7 100644
--- a/apps/protspace/src/protspace/cli/stats.py
+++ b/apps/protspace/src/protspace/cli/stats.py
@@ -43,9 +43,10 @@ def _atomic_write_table(table, path: Path) -> None:
"""
import pyarrow.parquet as pq
- tmp = path.with_name(path.name + ".tmp")
- pq.write_table(table, str(tmp))
- tmp.replace(path)
+ from protspace.data.io.atomic import staged_write
+
+ with staged_write(path) as staged:
+ pq.write_table(table, str(staged))
def _load_reductions(
diff --git a/apps/protspace/src/protspace/data/annotations/manager.py b/apps/protspace/src/protspace/data/annotations/manager.py
index 04ed9920..56b25c89 100644
--- a/apps/protspace/src/protspace/data/annotations/manager.py
+++ b/apps/protspace/src/protspace/data/annotations/manager.py
@@ -41,6 +41,7 @@
UniProtRetriever,
)
from protspace.data.annotations.transformers.transformer import AnnotationTransformer
+from protspace.data.io.atomic import staged_write
from protspace.data.io.fasta import count_residues
from protspace.data.io.formatters import DataFormatter
from protspace.data.io.writers import AnnotationWriter
@@ -66,6 +67,26 @@ def resolve_fasta_sequence_length(
return str(residues) if residues > 0 else length
+def uncached_headers(headers: list[str], cached_data: pd.DataFrame | None) -> list[str]:
+ """Requested identifiers *cached_data* holds no row for, in request order.
+
+ One rule, one implementation: the pipeline decides from it whether a cache
+ may serve a run at all, and this manager decides from it which identifiers
+ each source is fetched for. Two copies would let the pipeline serve a frame
+ the manager knows is short.
+
+ The identifier column is the frame's first, which is how the cache is
+ written.
+ """
+ if cached_data is None or cached_data.empty:
+ # No cache holds no rows, so every requested identifier is uncached.
+ # Answering "none missing" here would let a caller serve an empty frame
+ # as a complete one.
+ return list(headers)
+ cached_ids = set(cached_data[cached_data.columns[0]].astype(str))
+ return [h for h in headers if str(h) not in cached_ids]
+
+
class ProteinAnnotationManager:
"""Orchestrator for protein annotation extraction workflow."""
@@ -143,6 +164,17 @@ def to_pd(self) -> pd.DataFrame:
# Track which annotation sources failed
failed_sources = []
+ # Identifiers this run wants that the cache has no row for. Every source
+ # served from that cache is fetched for exactly these, so "added a few
+ # sequences" costs a few lookups rather than a full refetch.
+ fill_in = uncached_headers(self.headers, self.cached_data)
+
+ def filled_in(cached_source, fetch):
+ """Cached annotations plus a fetch for the identifiers they lack."""
+ if not cached_source or not fill_in:
+ return cached_source
+ return list(cached_source) + list(fetch(fill_in))
+
# Extract cached annotations by source if available
cached_uniprot = (
self._extract_cached_source(UNIPROT_ANNOTATIONS)
@@ -175,28 +207,47 @@ def to_pd(self) -> pd.DataFrame:
uniprot_annotations = (
self._fetch_uniprot(failed_sources)
if self.sources_to_fetch["uniprot"]
- else cached_uniprot
+ else filled_in(
+ cached_uniprot,
+ lambda headers: self._fetch_uniprot(failed_sources, headers),
+ )
)
uniprot_annotations = self._fill_missing_fasta_lengths(uniprot_annotations)
+ # One call either way: `cached_taxonomy` is None whenever taxonomy is
+ # being fetched outright, and the helper treats that as "nothing cached".
taxonomy_annotations = (
- self._fetch_taxonomy(uniprot_annotations, failed_sources)
- if self.sources_to_fetch["taxonomy"]
+ self._fetch_taxonomy(
+ uniprot_annotations, failed_sources, cached=cached_taxonomy
+ )
+ if self.sources_to_fetch["taxonomy"] or (cached_taxonomy and fill_in)
else cached_taxonomy
)
interpro_annotations = (
self._fetch_interpro(uniprot_annotations, failed_sources)
if self.sources_to_fetch["interpro"]
- else cached_interpro
+ else filled_in(
+ cached_interpro,
+ lambda headers: self._fetch_interpro(
+ uniprot_annotations, failed_sources, headers
+ ),
+ )
)
ted_annotations = (
self._fetch_ted(failed_sources)
if self.sources_to_fetch.get("ted")
- else cached_ted
+ else filled_in(
+ cached_ted, lambda headers: self._fetch_ted(failed_sources, headers)
+ )
)
biocentral_annotations = (
self._fetch_biocentral(uniprot_annotations, failed_sources)
if self.sources_to_fetch.get("biocentral")
- else cached_biocentral
+ else filled_in(
+ cached_biocentral,
+ lambda headers: self._fetch_biocentral(
+ uniprot_annotations, failed_sources, headers
+ ),
+ )
)
# Report failed sources
@@ -242,6 +293,30 @@ def to_pd(self) -> pd.DataFrame:
return df
+ def _with_retained_rows(self, df: pd.DataFrame) -> pd.DataFrame:
+ """Append cached rows for identifiers outside this run, when they fit.
+
+ A run for part of a dataset would otherwise replace the cache with its
+ own rows, so alternating between two subsets refetches both forever. The
+ rows only fit when this run produced the columns the cache already had:
+ a row missing a column is an empty value nothing can tell from a real
+ absence, which is the same trap incomplete sources are kept out for.
+ """
+ if self.cached_data is None or self.cached_data.empty:
+ return df
+ identifier_col = df.columns[0]
+ cached = self.cached_data.rename(
+ columns={self.cached_data.columns[0]: identifier_col}
+ )
+ if set(cached.columns) != set(df.columns):
+ return df
+ retained = cached[
+ ~cached[identifier_col].astype(str).isin(df[identifier_col].astype(str))
+ ]
+ if retained.empty:
+ return df
+ return pd.concat([df, retained[df.columns]], ignore_index=True)
+
def _uncacheable_sources(self) -> set[str]:
"""Sources that must stay out of the cache this run.
@@ -291,7 +366,7 @@ def _write_cache_and_frame(
drop = self._incomplete_columns()
df = DataFormatter.to_dataframe(proteins)
if not drop:
- self._write_cache(df)
+ self._write_cache(self._with_retained_rows(df))
return df
incomplete = ", ".join(sorted(self.incomplete_sources))
@@ -331,7 +406,11 @@ def _write_cache(self, df: pd.DataFrame) -> None:
"""Persist *df* as the annotation cache, stamped with the current semantics."""
df = df.copy()
df.attrs.update(annotation_cache_version_attrs())
- df.to_parquet(self.output_path, index=False)
+ # Staged: with retained rows folded in, this frame is a superset holding
+ # rows for identifiers no other file has, so a half-written cache loses
+ # data rather than costing one refetch.
+ with staged_write(self.output_path) as staged:
+ df.to_parquet(staged, index=False)
def _fill_missing_fasta_lengths(
self, proteins: list[ProteinAnnotations]
@@ -382,11 +461,14 @@ def _fill_missing_fasta_lengths(
for protein in result
]
- def _fetch_uniprot(self, failed_sources: list) -> list[ProteinAnnotations]:
- """Fetch UniProt annotations."""
+ def _fetch_uniprot(
+ self, failed_sources: list, headers: list[str] | None = None
+ ) -> list[ProteinAnnotations]:
+ """Fetch UniProt annotations for *headers* (default: the whole run)."""
+ headers = self.headers if headers is None else headers
try:
retriever = UniProtRetriever(
- headers=self.headers,
+ headers=headers,
annotations=self.config.uniprot_annotations,
)
annotations = retriever.fetch_annotations()
@@ -398,7 +480,7 @@ def _fetch_uniprot(self, failed_sources: list) -> list[ProteinAnnotations]:
ProteinAnnotations(
identifier=header, annotations={TAXONOMY_LOOKUP_ANNOTATION: ""}
)
- for header in self.headers
+ for header in headers
]
# Read outside the try: a problem reading the failure counter must not be
@@ -408,19 +490,28 @@ def _fetch_uniprot(self, failed_sources: list) -> list[ProteinAnnotations]:
return annotations
def _fetch_taxonomy(
- self, uniprot_annotations: list[ProteinAnnotations], failed_sources: list
+ self,
+ uniprot_annotations: list[ProteinAnnotations],
+ failed_sources: list,
+ cached: dict | None = None,
) -> dict:
- """Fetch taxonomy annotations if requested."""
+ """Fetch taxonomy for organisms *cached* does not already cover.
+
+ Taxonomy is keyed by organism rather than by identifier, so a protein
+ the cache never saw usually needs no lookup at all: its organism is
+ almost always one the cache already resolved.
+ """
if not self.config.taxonomy_annotations:
return {}
+ cached = cached or {}
try:
# Extract unique taxonomy IDs
taxon_counts = self._get_taxon_counts(uniprot_annotations)
- unique_taxons = list(taxon_counts.keys())
+ unique_taxons = [t for t in taxon_counts if t not in cached]
if not unique_taxons:
- return {}
+ return dict(cached)
retriever = TaxonomyRetriever(
taxon_ids=unique_taxons, annotations=self.config.taxonomy_annotations
@@ -428,12 +519,16 @@ def _fetch_taxonomy(
annotations = retriever.fetch_annotations()
if retriever.failed_batch_count:
self.incomplete_sources.add("taxonomy")
- return annotations
+ return {**cached, **annotations}
except Exception as e:
self.incomplete_sources.add("taxonomy")
failed_sources.append(f"Taxonomy ({str(e)})")
logger.warning(f"Failed to retrieve Taxonomy annotations: {e}")
- return {}
+ # What was already resolved survives the failure. A fill-in run
+ # reaches here to look up one unseen organism, and discarding
+ # *cached* would blank the taxonomy columns of every protein in the
+ # run over a lookup that only concerned the new ones.
+ return dict(cached)
def _build_sequence_map(
self, uniprot_annotations: list[ProteinAnnotations]
@@ -450,17 +545,21 @@ def _build_sequence_map(
return sequences
def _fetch_interpro(
- self, uniprot_annotations: list[ProteinAnnotations], failed_sources: list
+ self,
+ uniprot_annotations: list[ProteinAnnotations],
+ failed_sources: list,
+ headers: list[str] | None = None,
) -> list[ProteinAnnotations]:
- """Fetch InterPro annotations if requested."""
+ """Fetch InterPro annotations for *headers* (default: the whole run)."""
if not self.config.interpro_annotations:
return []
+ headers = self.headers if headers is None else headers
try:
sequences = self._build_sequence_map(uniprot_annotations)
retriever = InterProRetriever(
- headers=self.headers,
+ headers=headers,
annotations=self.config.interpro_annotations,
sequences=sequences,
)
@@ -475,17 +574,21 @@ def _fetch_interpro(
return []
def _fetch_biocentral(
- self, uniprot_annotations: list[ProteinAnnotations], failed_sources: list
+ self,
+ uniprot_annotations: list[ProteinAnnotations],
+ failed_sources: list,
+ headers: list[str] | None = None,
) -> list[ProteinAnnotations]:
- """Fetch Biocentral prediction annotations if requested."""
+ """Fetch Biocentral predictions for *headers* (default: the whole run)."""
if not self.config.biocentral_annotations:
return []
+ headers = self.headers if headers is None else headers
try:
sequences = self._build_sequence_map(uniprot_annotations)
retriever = BiocentralPredictionRetriever(
- headers=self.headers,
+ headers=headers,
annotations=self.config.biocentral_annotations,
sequences=sequences,
)
@@ -499,14 +602,17 @@ def _fetch_biocentral(
logger.warning(f"Failed to retrieve Biocentral predictions: {e}")
return []
- def _fetch_ted(self, failed_sources: list) -> list[ProteinAnnotations]:
- """Fetch TED domain annotations if requested."""
+ def _fetch_ted(
+ self, failed_sources: list, headers: list[str] | None = None
+ ) -> list[ProteinAnnotations]:
+ """Fetch TED domain annotations for *headers* (default: the whole run)."""
if not self.config.ted_annotations:
return []
+ headers = self.headers if headers is None else headers
try:
retriever = TedRetriever(
- headers=self.headers,
+ headers=headers,
annotations=self.config.ted_annotations,
)
annotations = retriever.fetch_annotations()
@@ -555,22 +661,21 @@ def _extract_cached_source(
if not available:
return []
- # Convert DataFrame to ProteinAnnotations format
- result = []
+ # Column-wise rather than `iterrows`: the cache retains rows for
+ # identifiers outside the run, so this walks the whole frame on every
+ # fill-in run, and `iterrows` builds a Series per row (and upcasts a
+ # mixed-dtype row to one common dtype on the way).
identifier_col = self.cached_data.columns[0] # First column is identifier
+ identifiers = self.cached_data[identifier_col].tolist()
+ columns = {a: self.cached_data[a].tolist() for a in available}
- for _, row in self.cached_data.iterrows():
- annotations_dict = {}
- for annotation in available:
- annotations_dict[annotation] = row[annotation]
-
- result.append(
- ProteinAnnotations(
- identifier=row[identifier_col], annotations=annotations_dict
- )
+ return [
+ ProteinAnnotations(
+ identifier=identifier,
+ annotations={a: values[i] for a, values in columns.items()},
)
-
- return result
+ for i, identifier in enumerate(identifiers)
+ ]
def _extract_cached_taxonomy(self, taxonomy_annotations: list[str]) -> dict:
"""
@@ -596,19 +701,22 @@ def _extract_cached_taxonomy(self, taxonomy_annotations: list[str]) -> dict:
# Convert to taxonomy format: {organism_id: {"annotations": {annotation: value}}}
taxonomy_dict = {}
+ # Column-wise for the same reason as `_extract_cached_source`: the cache
+ # is a superset of the run and `iterrows` costs a Series per row.
+ organism_ids = self.cached_data[TAXONOMY_LOOKUP_ANNOTATION].tolist()
+ columns = {a: self.cached_data[a].tolist() for a in available}
+
# Group by organism_id
- for _, row in self.cached_data.iterrows():
- organism_id = row[TAXONOMY_LOOKUP_ANNOTATION]
+ for i, organism_id in enumerate(organism_ids):
if pd.isna(organism_id) or organism_id == "":
continue
try:
org_id = int(organism_id)
if org_id not in taxonomy_dict:
- annotations_dict = {}
- for annotation in available:
- annotations_dict[annotation] = row[annotation]
- taxonomy_dict[org_id] = {"annotations": annotations_dict}
+ taxonomy_dict[org_id] = {
+ "annotations": {a: values[i] for a, values in columns.items()}
+ }
except (ValueError, TypeError):
pass
diff --git a/apps/protspace/src/protspace/data/embedding/biocentral.py b/apps/protspace/src/protspace/data/embedding/biocentral.py
index 0ccb30b4..21de0c12 100644
--- a/apps/protspace/src/protspace/data/embedding/biocentral.py
+++ b/apps/protspace/src/protspace/data/embedding/biocentral.py
@@ -12,9 +12,12 @@
from biocentral_api import BiocentralAPI, CommonEmbedder, batched
from tqdm import tqdm
-# Re-exported: the HDF5 layer moved to `store` so neither backend owns it, but
-# local.py, cli/annotate.py and existing importers still reach it from here.
+# Re-exported: the HDF5 layer moved to `store` so neither backend owns it. All
+# but `load_existing_ids` are used below; they stay importable from here for
+# out-of-repo callers that predate the move, which is the only reason
+# `load_existing_ids` outlived its last in-repo caller (`begin_run` replaced it).
from protspace.data.embedding.store import ( # noqa: F401
+ begin_run,
finish_run,
load_existing_ids,
save_embeddings,
@@ -148,16 +151,9 @@ def embed_sequences(
# Reject HDF5-hostile identifiers before spending a single API call on them.
validate_headers(sequences)
- # Resume: skip already-embedded sequences
- existing_ids = load_existing_ids(h5_path)
- if existing_ids:
- logger.info("Found %d existing embeddings in %s", len(existing_ids), h5_path)
- remaining = {k: v for k, v in sequences.items() if k not in existing_ids}
- logger.info(
- "Remaining sequences to embed: %d (skipped %d)",
- len(remaining),
- len(sequences) - len(remaining),
- )
+ # Resume: claim the file for this backend and model, and drop the sequences
+ # it already holds a current vector for (see store.begin_run).
+ remaining = begin_run(h5_path, sequences, backend="biocentral", model=embedder)
if not remaining:
logger.info(
@@ -239,7 +235,13 @@ def embed_sequences(
seq = unique_seqs[rep_id]
for pid in seq_to_ids[seq]:
expanded[pid] = emb
- save_embeddings(h5_path, expanded)
+ save_embeddings(
+ h5_path,
+ expanded,
+ sequences=remaining,
+ backend="biocentral",
+ model=embedder,
+ )
missing_reps = batch_seqs.keys() - emb_dict.keys()
if missing_reps:
@@ -279,6 +281,9 @@ def embed_sequences(
return finish_run(
h5_path,
remaining,
+ # Scoped to this run's own work: begin_run already proved the rest
+ # current, and re-reading their digests doubles the scan at 570K.
+ sequences=remaining,
context=f"{failed_batches} of {len(api_batches)} batch(es) failed",
retry_hint="Check the Biocentral server status and rerun.",
)
diff --git a/apps/protspace/src/protspace/data/embedding/local.py b/apps/protspace/src/protspace/data/embedding/local.py
index c194dd35..3e3459be 100644
--- a/apps/protspace/src/protspace/data/embedding/local.py
+++ b/apps/protspace/src/protspace/data/embedding/local.py
@@ -34,8 +34,8 @@
from tqdm import tqdm
from protspace.data.embedding.store import (
+ begin_run,
finish_run,
- load_existing_ids,
save_embeddings,
validate_headers,
)
@@ -283,10 +283,9 @@ def embed_sequences(
)
h5_path = Path(h5_path)
- existing = load_existing_ids(h5_path)
- remaining = {k: v for k, v in sequences.items() if k not in existing}
- if existing:
- logger.info("Resuming: %d already embedded in %s", len(existing), h5_path)
+ # Claims the file for this backend and model, and drops the proteins it
+ # already holds a current vector for (see store.begin_run).
+ remaining = begin_run(h5_path, sequences, backend="local", model=embedder)
# Sequences a capability limit puts out of reach: recorded as skipped rather
# than failed, so they are reported and named but do not fail the run.
@@ -322,7 +321,13 @@ def embed_sequences(
vecs = _embed_batch(
processed, mod_type, model, tokenizer, device, cfg.max_length
)
- save_embeddings(h5_path, dict(zip(batch_ids, vecs, strict=True)))
+ save_embeddings(
+ h5_path,
+ dict(zip(batch_ids, vecs, strict=True)),
+ sequences=remaining,
+ backend="local",
+ model=embedder,
+ )
i += bs
pbar.update(len(batch_ids))
except torch.cuda.OutOfMemoryError:
@@ -352,6 +357,11 @@ def embed_sequences(
h5_path,
outstanding,
skipped=skipped,
+ # Only this run's own work needs its digests checked, and at Swiss-Prot
+ # scale that is the difference between one scan and two: everything else
+ # was proven current by begin_run, and a skipped sequence is on disk only
+ # under residues that are no longer this protein's.
+ sequences=remaining,
retry_hint=(
f"Raise --max-length (currently {cfg.max_length}) or use "
f"--backend biocentral, which has no length cap."
diff --git a/apps/protspace/src/protspace/data/embedding/store.py b/apps/protspace/src/protspace/data/embedding/store.py
index c342b0ad..c6b23a3b 100644
--- a/apps/protspace/src/protspace/data/embedding/store.py
+++ b/apps/protspace/src/protspace/data/embedding/store.py
@@ -7,6 +7,7 @@
from __future__ import annotations
+import hashlib
import logging
from collections.abc import Collection, Iterable, Mapping
from pathlib import Path
@@ -19,21 +20,180 @@
# Identifiers named in a message before it elides the rest.
_PREVIEW = 5
+# Who produced the file, on the root, and which residues each vector was computed
+# from, on its dataset. In the file rather than in its name because the name is a
+# caller's choice and the contract belongs to the file: `protspace embed -o
+# mine.h5` gets the same protection a managed cache does.
+_BACKEND_ATTR = "protspace_backend"
+_MODEL_ATTR = "protspace_model"
+_DIGEST_ATTR = "protspace_sequence_sha256"
+
+# Enough SHA-256 that two residue strings colliding is implausible, short enough
+# that the attribute stays cheap across 570K proteins.
+_DIGEST_CHARS = 16
+
+
+def sequence_digest(sequence: str) -> str:
+ """Return the residue digest stored alongside a protein's vector."""
+ return hashlib.sha256(sequence.encode()).hexdigest()[:_DIGEST_CHARS]
+
+
+def _current_ids(f: h5py.File, sequences: Mapping[str, str]) -> set[str]:
+ """Identifiers in the open file whose vector matches the residues in hand.
+
+ One pass over one open file: the digest lives in each dataset's attributes,
+ and this runs per protein at up to 570K of them, so reopening the file for
+ each would replace a seconds-long resume with 570K file opens.
+
+ A protein carrying no digest predates them and is trusted — refusing it would
+ force a full re-embed of every existing cache on upgrade.
+ """
+ current: set[str] = set()
+ for pid, sequence in sequences.items():
+ dataset = f.get(pid)
+ if dataset is None:
+ continue
+ stored = dataset.attrs.get(_DIGEST_ATTR)
+ if stored is None or str(stored) == sequence_digest(sequence):
+ current.add(pid)
+ return current
+
def load_existing_ids(h5_path: Path) -> set[str]:
- """Return the set of dataset keys already present in *h5_path*."""
+ """Every dataset key in *h5_path*, without checking residue identity.
+
+ Superseded by :func:`covered_ids`, which it delegates to, and kept as the
+ name out-of-repo callers already import. Resume goes through
+ :func:`begin_run` instead: a key on its own is no evidence the vector under
+ it was computed from the residues this run holds.
+ """
+ return covered_ids(h5_path)
+
+
+def covered_ids(h5_path: Path, sequences: Mapping[str, str] | None = None) -> set[str]:
+ """Identifiers *h5_path* holds a usable vector for.
+
+ Without *sequences* that is every dataset key. With them the answer is scoped
+ to the identifiers in hand and excludes any whose stored residue digest
+ disagrees: a protein whose sequence changed is on disk under residues that
+ are no longer the ones being embedded, so counting it as covered would let a
+ re-embed that never landed pass for a complete run.
+ """
+ h5_path = Path(h5_path)
if not h5_path.exists():
return set()
with h5py.File(h5_path, "r") as f:
- return set(f.keys())
+ if sequences is None:
+ return set(f.keys())
+ return _current_ids(f, sequences)
+
+
+def begin_run(
+ h5_path: Path,
+ sequences: Mapping[str, str],
+ *,
+ backend: str,
+ model: str,
+) -> dict[str, str]:
+ """Claim *h5_path* for *backend*/*model* and return what is left to embed.
+
+ Resume matches on identifier alone, which is evidence of a reusable vector
+ only once the file is known to have been written by this producer and each
+ identifier still carries the residues its vector was computed from. Another
+ producer's file is refused rather than extended -- both backends resume the
+ same way, so mixing them is silent -- and a protein whose residues changed is
+ outstanding work again.
+
+ A file recording no producer predates the stamps: it is adopted and stamped,
+ because refusing it would force a full re-embed of every existing cache.
+ """
+ h5_path = Path(h5_path)
+ if not h5_path.exists():
+ return dict(sequences)
+
+ # Read-only first, so a file this run has no claim to is left exactly as it
+ # was -- not extended, not truncated by an append-mode open.
+ with h5py.File(h5_path, "r") as f:
+ recorded_backend = f.attrs.get(_BACKEND_ATTR)
+ recorded_model = f.attrs.get(_MODEL_ATTR)
+ unstamped = recorded_backend is None and recorded_model is None
+ if not unstamped and (str(recorded_backend), str(recorded_model)) != (
+ backend,
+ model,
+ ):
+ raise ValueError(
+ f"{h5_path} holds embeddings produced by the "
+ f"{recorded_backend} backend with model {recorded_model}, but "
+ f"this run is {backend} with {model}. Resuming would mix two "
+ f"models' vectors into one dataset. Select that backend and "
+ f"model, choose another output path, or refetch the embeddings "
+ f"(--refetch embed)."
+ )
+ current = _current_ids(f, sequences)
+ if unstamped:
+ logger.info(
+ "Adopting %s: it records no producer, so this run stamps it as "
+ "%s/%s. Use --refetch embed if it was produced by something else.",
+ h5_path,
+ backend,
+ model,
+ )
+ with h5py.File(h5_path, "a") as f:
+ f.attrs[_BACKEND_ATTR] = backend
+ f.attrs[_MODEL_ATTR] = model
-def save_embeddings(h5_path: Path, embeddings: dict[str, np.ndarray]) -> None:
- """Append embeddings to an HDF5 file (one dataset per protein)."""
+ remaining = {k: v for k, v in sequences.items() if k not in current}
+ resumed = len(sequences) - len(remaining)
+ if resumed:
+ # Logged here rather than in each backend: both resume through this one
+ # call, so the count and its wording stay the same whoever is embedding.
+ logger.info(
+ "Resuming %s: %d of %d already embedded, %d to go.",
+ h5_path,
+ resumed,
+ len(sequences),
+ len(remaining),
+ )
+ return remaining
+
+
+def save_embeddings(
+ h5_path: Path,
+ embeddings: dict[str, np.ndarray],
+ *,
+ sequences: Mapping[str, str] | None = None,
+ backend: str | None = None,
+ model: str | None = None,
+) -> None:
+ """Append embeddings to an HDF5 file (one dataset per protein).
+
+ An identifier already present is left alone -- that is what makes resume
+ cheap -- unless *sequences* shows its vector was computed from different
+ residues, in which case dataset and digest are both replaced.
+
+ *backend* and *model* stamp a file that carries no producer yet. The stamp
+ rides along with the vectors rather than living in a separate step, so an
+ interrupted run's partial file is owned too, while a run that wrote nothing
+ leaves no stamped-but-empty cache behind for the next run to resume from.
+ """
with h5py.File(h5_path, "a") as f:
+ if backend is not None and f.attrs.get(_BACKEND_ATTR) is None:
+ f.attrs[_BACKEND_ATTR] = backend
+ if model is not None and f.attrs.get(_MODEL_ATTR) is None:
+ f.attrs[_MODEL_ATTR] = model
+
for protein_id, emb in embeddings.items():
- if protein_id not in f:
- f.create_dataset(protein_id, data=emb.astype(np.float32))
+ sequence = sequences.get(protein_id) if sequences else None
+ digest = sequence_digest(sequence) if sequence is not None else None
+ if protein_id in f:
+ stored = f[protein_id].attrs.get(_DIGEST_ATTR)
+ if digest is None or str(stored) == digest:
+ continue
+ del f[protein_id]
+ dataset = f.create_dataset(protein_id, data=emb.astype(np.float32))
+ if digest is not None:
+ dataset.attrs[_DIGEST_ATTR] = digest
def validate_headers(ids: Iterable[str]) -> None:
@@ -63,6 +223,7 @@ def finish_run(
requested: Collection[str],
*,
skipped: Mapping[str, str] | None = None,
+ sequences: Mapping[str, str] | None = None,
context: str = "",
retry_hint: str = "",
) -> Path:
@@ -77,13 +238,17 @@ def finish_run(
skips identifiers already present, so a counter can claim sequences the file
does not hold.
+ *sequences* makes that read residue-aware: a protein that was outstanding
+ because its sequence changed is still on disk under its old residues, so
+ without them its failed re-embed reads as covered.
+
*context* is backend detail for the failure message (e.g. how many batches
failed); *retry_hint* is the closing advice when nothing was produced.
"""
skipped = dict(skipped or {})
requested_ids = set(requested)
expected = requested_ids - set(skipped)
- on_disk = load_existing_ids(h5_path)
+ on_disk = covered_ids(h5_path, sequences)
missing = expected - on_disk
embedded = len(expected) - len(missing)
diff --git a/apps/protspace/src/protspace/data/io/atomic.py b/apps/protspace/src/protspace/data/io/atomic.py
new file mode 100644
index 00000000..91f8a035
--- /dev/null
+++ b/apps/protspace/src/protspace/data/io/atomic.py
@@ -0,0 +1,53 @@
+"""Publishing a file so it appears complete, or not at all.
+
+Every path here writes a sibling of the destination and renames it into place:
+rename within one filesystem is atomic, so a crash, a Ctrl-C or a full disk
+leaves the previous content rather than a half-written file. That matters wherever
+the existence of a file is itself a signal — a retained cache entry the next run
+will trust, or the bundle the user asked to overwrite in place.
+
+The staging file is created with a plain ``open`` rather than ``mkstemp``, so the
+published file carries the permissions the process umask gives any new file.
+``mkstemp`` creates owner-only, and renaming that into place hands the user a
+mode a direct write would never have produced.
+"""
+
+from __future__ import annotations
+
+import os
+import uuid
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+
+
+@contextmanager
+def staged_write(path: Path) -> Iterator[Path]:
+ """Yield a staging path beside *path*, published on a clean exit.
+
+ The caller writes to the yielded path with whatever writer it has (a plain
+ ``open``, ``pyarrow.parquet.write_table``, ...). Leaving the block normally
+ renames it onto *path*; leaving it by exception removes it.
+ """
+ path = Path(path)
+ path.parent.mkdir(parents=True, exist_ok=True)
+ # Short random part: the staging name is the destination's plus this, and a
+ # full uuid4 hex adds 38 characters to a name the user chose -- enough to
+ # push a long bundle name past the filesystem's 255-byte limit on a write
+ # that used to work.
+ staged = path.with_name(f".{path.name}.{uuid.uuid4().hex[:8]}.tmp")
+ try:
+ yield staged
+ os.replace(staged, path)
+ except BaseException:
+ staged.unlink(missing_ok=True)
+ raise
+
+
+def atomic_write_bytes(path: Path, data: bytes) -> None:
+ """Write *data* to *path* atomically, flushed to disk before publication."""
+ with staged_write(path) as staged:
+ with open(staged, "wb") as f:
+ f.write(data)
+ f.flush()
+ os.fsync(f.fileno())
diff --git a/apps/protspace/src/protspace/data/io/bundle.py b/apps/protspace/src/protspace/data/io/bundle.py
index 2c17504e..2b199ae4 100644
--- a/apps/protspace/src/protspace/data/io/bundle.py
+++ b/apps/protspace/src/protspace/data/io/bundle.py
@@ -14,7 +14,6 @@
import io
import json
import logging
-import os
import tempfile
from pathlib import Path
@@ -22,6 +21,7 @@
import pyarrow.parquet as pq
from protspace.data.annotations.encoding import stamp_format_version
+from protspace.data.io.atomic import atomic_write_bytes
logger = logging.getLogger(__name__)
@@ -62,28 +62,6 @@ def _table_to_parquet_bytes(table: pa.Table) -> bytes:
return buf.getvalue()
-def _atomic_write_bytes(path: Path, data: bytes) -> None:
- """Write ``data`` to ``path`` atomically (temp file + ``os.replace``).
-
- The destination is never left truncated or partial on interrupt — it keeps
- the old bytes until the rename completes, then atomically becomes the full
- new bytes. Critical for the in-place overwrite workflow that ``transfer``
- documents (``-b results.parquetbundle -o results.parquetbundle``): a Ctrl+C
- or crash mid-write can no longer destroy the user's bundle.
- """
- path.parent.mkdir(parents=True, exist_ok=True)
- fd, tmp = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp")
- try:
- with os.fdopen(fd, "wb") as f:
- f.write(data)
- f.flush()
- os.fsync(f.fileno())
- os.replace(tmp, path)
- except BaseException:
- Path(tmp).unlink(missing_ok=True)
- raise
-
-
def _check_no_delimiter(part_bytes: bytes) -> None:
"""Guard: a serialized part must not contain the bundle delimiter.
@@ -195,7 +173,7 @@ def write_bundle(
_check_no_delimiter(stats_bytes)
buf.write(stats_bytes)
- _atomic_write_bytes(bundle_path, buf.getvalue())
+ atomic_write_bytes(bundle_path, buf.getvalue())
logger.info(f"Saved bundled output to: {bundle_path}")
@@ -219,7 +197,7 @@ def replace_settings_in_bundle(
new_parts.append(statistics)
new_content = PARQUET_BUNDLE_DELIMITER.join(new_parts)
- _atomic_write_bytes(output_path, new_content)
+ atomic_write_bytes(output_path, new_content)
def replace_annotations_in_bundle(
@@ -255,7 +233,7 @@ def replace_annotations_in_bundle(
if statistics is not None:
new_parts.append(statistics)
- _atomic_write_bytes(output_path, PARQUET_BUNDLE_DELIMITER.join(new_parts))
+ atomic_write_bytes(output_path, PARQUET_BUNDLE_DELIMITER.join(new_parts))
logger.info(f"Wrote bundle with updated annotations to: {output_path}")
diff --git a/apps/protspace/src/protspace/data/loaders/__init__.py b/apps/protspace/src/protspace/data/loaders/__init__.py
index e864a9fd..b310b7c5 100644
--- a/apps/protspace/src/protspace/data/loaders/__init__.py
+++ b/apps/protspace/src/protspace/data/loaders/__init__.py
@@ -14,7 +14,9 @@
)
from protspace.data.loaders.query import (
extract_identifiers_from_fasta,
+ query_cache_path,
query_uniprot,
+ resolve_query_fasta,
)
from protspace.data.loaders.similarity import compute_similarity
@@ -28,6 +30,8 @@
"extract_identifiers_from_fasta",
"load_h5",
"parse_identifier",
+ "query_cache_path",
"query_uniprot",
+ "resolve_query_fasta",
"split_h5_spec",
]
diff --git a/apps/protspace/src/protspace/data/loaders/fasta.py b/apps/protspace/src/protspace/data/loaders/fasta.py
index bcbb48ef..c9f7c8b0 100644
--- a/apps/protspace/src/protspace/data/loaders/fasta.py
+++ b/apps/protspace/src/protspace/data/loaders/fasta.py
@@ -6,7 +6,8 @@
from __future__ import annotations
import logging
-from collections.abc import Iterable
+from collections.abc import Collection, Iterable
+from dataclasses import replace
from pathlib import Path
from typing import TYPE_CHECKING
@@ -97,7 +98,43 @@ def embed_fasta(
with h5py.File(h5_path, "a") as f:
f.attrs["model_name"] = embedder
- return load_h5([h5_path], name_override=embedder)
+ return _restrict_to(load_h5([h5_path], name_override=embedder), sequences)
+
+
+def _restrict_to(embedding_set: EmbeddingSet, wanted: Collection[str]) -> EmbeddingSet:
+ """Drop rows for proteins outside *wanted*, keeping the loaded row order.
+
+ The embedding cache legitimately accumulates proteins across inputs — that is
+ what makes resume work — so what a run is about is decided by the FASTA in
+ hand rather than by everything the cache has ever held. Returning the
+ accumulation unions unrelated datasets into one bundle.
+
+ Identifiers go through ``parse_identifier`` because a cache written elsewhere
+ may be keyed ``sp|P12345|NAME`` while the FASTA-derived keys are parsed.
+ """
+ requested = set(wanted)
+ keep = [
+ i
+ for i, header in enumerate(embedding_set.headers)
+ if parse_identifier(header) in requested
+ ]
+ if len(keep) == len(embedding_set.headers):
+ return embedding_set
+
+ logger.info(
+ "Returning %d of %d cached protein(s): the rest belong to other inputs "
+ "sharing this embedding cache.",
+ len(keep),
+ len(embedding_set.headers),
+ )
+ # A new set rather than an edit in place: "restrict to" reads as a query,
+ # and a caller still holding the loaded set would otherwise find its rows
+ # silently gone.
+ return replace(
+ embedding_set,
+ data=embedding_set.data[keep],
+ headers=[embedding_set.headers[i] for i in keep],
+ )
def check_fasta_coverage(
diff --git a/apps/protspace/src/protspace/data/loaders/query.py b/apps/protspace/src/protspace/data/loaders/query.py
index 7c30119e..9de4d021 100644
--- a/apps/protspace/src/protspace/data/loaders/query.py
+++ b/apps/protspace/src/protspace/data/loaders/query.py
@@ -5,16 +5,54 @@
"""
import gzip
+import hashlib
import logging
+import shutil
import tempfile
from pathlib import Path
import requests
from tqdm import tqdm
+from protspace.data.io.atomic import staged_write
+
logger = logging.getLogger(__name__)
+def query_cache_path(cache_dir: Path, query: str) -> Path:
+ """Return the retained FASTA path owned by one exact query text.
+
+ One shared file name per output directory hands a later run the previous
+ query's sequences: the file exists and parses, so nothing downstream can
+ tell it apart from the one this query would have produced.
+ """
+ digest = hashlib.sha256(query.encode()).hexdigest()[:12]
+ return cache_dir / "queries" / f"{digest}.fasta"
+
+
+def resolve_query_fasta(
+ query: str, cache_dir: Path | None, refetch_stages: frozenset[str]
+) -> tuple[list[str], Path]:
+ """Return ``(headers, fasta_path)`` for *query*, reusing only its own FASTA.
+
+ Path, reuse rule and download live together because they are one decision:
+ whether the retained FASTA on disk is the one *query* would produce. A caller
+ that owned the rule separately would have to be changed in step with the
+ path, and the two would drift.
+ """
+ fasta_save = query_cache_path(cache_dir, query) if cache_dir else None
+ if (
+ fasta_save
+ and fasta_save.exists()
+ and fasta_save.stat().st_size > 0
+ and "query" not in refetch_stages
+ ):
+ headers = extract_identifiers_from_fasta(fasta_save)
+ logger.warning("Using cached FASTA (%s sequences)", f"{len(headers):,}")
+ return headers, fasta_save
+ return query_uniprot(query, save_to=fasta_save)
+
+
def query_uniprot(
query: str,
*,
@@ -33,46 +71,38 @@ def query_uniprot(
base_url = "https://rest.uniprot.org/uniprotkb/stream"
params = {"compressed": "true", "format": "fasta", "query": query}
+ temp_gz_file: Path | None = None
try:
response = requests.get(base_url, params=params, stream=True)
response.raise_for_status()
# Download to temporary compressed file
- temp_file = tempfile.NamedTemporaryFile(
+ total_size = int(response.headers.get("content-length", 0))
+ with tempfile.NamedTemporaryFile(
mode="wb", suffix=".fasta.gz", delete=False
- )
- temp_gz_file = Path(temp_file.name)
+ ) 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))
+
+ # Published by rename either way: a retained FASTA's existence is the next
+ # run's cache hit, so it may not appear until the whole stream has been
+ # decompressed. Without *save_to* the extraction is the caller's own file.
+ extracted = Path(save_to) if save_to else temp_gz_file.with_suffix("")
+ with staged_write(extracted) as staged:
+ identifiers = _extract_fasta(temp_gz_file, staged)
- 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()
-
- # 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)
- 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)
-
- temp_gz_file.unlink(missing_ok=True)
logger.info(f"Downloaded and extracted {len(identifiers)} sequences")
-
- return identifiers, extracted_path
+ return identifiers, extracted
except requests.RequestException as e:
logger.error(f"Error downloading FASTA: {e}")
@@ -80,27 +110,28 @@ 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)
-def extract_identifiers_from_fasta(fasta_path: Path) -> list[str]:
- """Extract protein identifiers from an uncompressed FASTA file."""
- from protspace.data.loaders.h5 import parse_identifier
+def _extract_fasta(gz_path: Path, target: Path) -> list[str]:
+ """Decompress *gz_path* into *target* and return its identifiers.
- identifiers = []
- with open(fasta_path) as f:
- for line in f:
- if line.startswith(">"):
- raw = line[1:].strip().split()[0]
- identifiers.append(parse_identifier(raw))
- return identifiers
+ Streamed rather than read whole: a broad query decompresses to gigabytes. A
+ truncated or corrupt download raises here, before anything is published.
+ """
+ with gzip.open(gz_path, "rt") as gz_file, open(target, "w") as out:
+ shutil.copyfileobj(gz_file, out)
+ return extract_identifiers_from_fasta(target)
-def _extract_identifiers_gz(fasta_gz_path: Path) -> list[str]:
- """Extract protein identifiers from a gzipped FASTA file."""
+def extract_identifiers_from_fasta(fasta_path: Path) -> list[str]:
+ """Extract protein identifiers from an uncompressed FASTA file."""
from protspace.data.loaders.h5 import parse_identifier
identifiers = []
- with gzip.open(fasta_gz_path, "rt") as f:
+ with open(fasta_path) as f:
for line in f:
if line.startswith(">"):
raw = line[1:].strip().split()[0]
diff --git a/apps/protspace/src/protspace/data/processors/pipeline.py b/apps/protspace/src/protspace/data/processors/pipeline.py
index 7644b74a..3115e7ba 100644
--- a/apps/protspace/src/protspace/data/processors/pipeline.py
+++ b/apps/protspace/src/protspace/data/processors/pipeline.py
@@ -15,6 +15,7 @@
import numpy as np
import pandas as pd
+from protspace.data.io.atomic import staged_write
from protspace.data.loaders import EmbeddingSet
from protspace.data.loaders.embedding_set import (
format_param_suffix,
@@ -118,6 +119,24 @@ class PipelineConfig:
reducer_params: ReducerParams = field(default_factory=ReducerParams)
+def _embedding_fingerprint(emb_set: EmbeddingSet) -> str:
+ """Digest exactly what the reducer will be handed: identifiers and matrix.
+
+ The embedding name says where numbers came from, not which numbers they are:
+ a resumed embedding cache, a re-embedded input, a narrower intersection and a
+ reordered input all keep the name. Coordinates are stored as bare rows and
+ paired positionally with the current identifiers on load, so the identifier
+ order belongs in the digest too -- reusing a projection across a reorder
+ relabels every point.
+ """
+ data = np.ascontiguousarray(emb_set.data)
+ digest = hashlib.sha256()
+ digest.update("\0".join(emb_set.headers).encode())
+ digest.update(f"{data.dtype}{data.shape}".encode())
+ digest.update(memoryview(data).cast("B"))
+ return digest.hexdigest()[:16]
+
+
# Valid override parameter names (from ReducerParams fields)
_VALID_OVERRIDE_KEYS = {f.name for f in fields(ReducerParams)}
# Field types for coercion
@@ -334,10 +353,14 @@ def _extract_sequences(embedding_sets: list[EmbeddingSet]) -> dict[str, str]:
"""Extract protein sequences from FASTA files referenced by embedding sets."""
from protspace.data.loaders.fasta import parse_fasta_normalized
+ # One FASTA typically backs every embedder's set, so parse each file once.
+ fasta_paths = dict.fromkeys(
+ Path(emb_set.fasta_path) for emb_set in embedding_sets if emb_set.fasta_path
+ )
sequences = {}
- for emb_set in embedding_sets:
- if emb_set.fasta_path and Path(emb_set.fasta_path).exists():
- sequences.update(parse_fasta_normalized(Path(emb_set.fasta_path)))
+ for fasta_path in fasta_paths:
+ if fasta_path.exists():
+ sequences.update(parse_fasta_normalized(fasta_path))
return sequences
def _validate_headers(self, embedding_sets: list[EmbeddingSet]) -> list[str]:
@@ -416,6 +439,7 @@ def _fetch_annotations(
from protspace.data.annotations.manager import (
ProteinAnnotationManager,
resolve_fasta_sequence_length,
+ uncached_headers,
)
# Extract sequences from FASTA files (if available) to avoid re-fetching
@@ -458,6 +482,18 @@ def _fetch_annotations(
if cache_path.exists():
cached_df = pd.read_parquet(cache_path)
+ missing_identifiers = uncached_headers(headers, cached_df)
+ if missing_identifiers:
+ # Rows the cache has no entry for. The manager fetches each
+ # source for exactly these and reuses cached values for the
+ # rest, so the cache is filled in rather than rebuilt.
+ logger.warning(
+ "Annotation cache lacks %d of %d requested identifier(s); "
+ "fetching annotations for them",
+ len(missing_identifiers),
+ len(headers),
+ )
+
# Repair at the cache-read boundary, which dominates every path
# that reuses a stored column, then persist so it stays a
# one-time cost rather than a rewrite on every resumed run.
@@ -466,7 +502,8 @@ def _fetch_annotations(
"Rewrote legacy 'unclassified' TED domain labels in the "
"cached annotations to TED's '-'."
)
- cached_df.to_parquet(cache_path, index=False)
+ with staged_write(cache_path) as staged:
+ cached_df.to_parquet(staged, index=False)
cached_annotations = set(cached_df.columns) - {"identifier"}
if annotations_list is None:
@@ -502,7 +539,12 @@ def _fetch_annotations(
missing = required - cached_annotations
- if not missing and not refetching_annotations and not refresh_columns:
+ if (
+ not missing
+ and not missing_identifiers
+ and not refetching_annotations
+ and not refresh_columns
+ ):
logger.warning("Using cached annotations")
if annotations_list:
cols = ["identifier"] + [
@@ -704,6 +746,8 @@ def _projection_cache_path(
method: str,
dims: int,
effective_params: dict[str, Any] | None = None,
+ *,
+ fingerprint: str,
) -> Path | None:
cache_dir = self.config.intermediate_dir
if not cache_dir or not self.config.keep_tmp:
@@ -713,6 +757,7 @@ def _projection_cache_path(
"method": method,
"dims": dims,
"params": effective_params or asdict(self.config.reducer_params),
+ "fingerprint": fingerprint,
}
key_json = json.dumps(key_dict, sort_keys=True, default=str)
h = hashlib.sha256(key_json.encode()).hexdigest()[:12]
@@ -725,9 +770,11 @@ def _load_cached_projection(
dims: int,
effective_params: dict[str, Any] | None = None,
param_suffix: str = "",
+ *,
+ fingerprint: str,
) -> dict[str, Any] | None:
path = self._projection_cache_path(
- embedding_name, method, dims, effective_params
+ embedding_name, method, dims, effective_params, fingerprint=fingerprint
)
if (
path is None
@@ -757,15 +804,21 @@ def _save_projection_cache(
dims: int,
reduction: dict,
effective_params: dict[str, Any] | None = None,
+ *,
+ fingerprint: str,
) -> None:
path = self._projection_cache_path(
- embedding_name, method, dims, effective_params
+ embedding_name, method, dims, effective_params, fingerprint=fingerprint
)
if path is None:
return
- np.savez(
- path, data=reduction["data"], info=np.array(json.dumps(reduction["info"]))
- )
+ # Staged: `_load_cached_projection` trusts this entry on `exists()` alone,
+ # so a half-written zip would make every later run fail to load it.
+ # Written through a handle because `np.savez` appends `.npz` to a path.
+ with staged_write(path) as staged, open(staged, "wb") as fh:
+ np.savez(
+ fh, data=reduction["data"], info=np.array(json.dumps(reduction["info"]))
+ )
# --- Dimensionality reduction ---
@@ -793,9 +846,20 @@ def add(reduction: dict[str, Any]) -> None:
all_reductions.append(reduction)
for emb_set in embedding_sets:
+ # Once per set, not per method: the digest is a full pass over the
+ # matrix (~0.9 s for Swiss-Prot) and every method sees the same one.
+ # Not at all when nothing will be cached -- `_projection_cache_path`
+ # returns None then, so the digest would be a full scan of a 2 GB
+ # matrix computed for a key nobody looks up.
+ fingerprint = (
+ _embedding_fingerprint(emb_set)
+ if self.config.keep_tmp and self.config.intermediate_dir
+ else ""
+ )
+
if emb_set.precomputed:
cached = self._load_cached_projection(
- emb_set.name, MDS_NAME, 2, global_params
+ emb_set.name, MDS_NAME, 2, global_params, fingerprint=fingerprint
)
if cached:
add(cached)
@@ -809,7 +873,12 @@ def add(reduction: dict[str, Any]) -> None:
reduction["name"] = format_projection_name(emb_set.name, MDS_NAME, 2)
add(reduction)
self._save_projection_cache(
- emb_set.name, MDS_NAME, 2, reduction, global_params
+ emb_set.name,
+ MDS_NAME,
+ 2,
+ reduction,
+ global_params,
+ fingerprint=fingerprint,
)
computed_count += 1
continue
@@ -828,7 +897,12 @@ def add(reduction: dict[str, Any]) -> None:
param_suffix = disambiguation_suffix(spec, method_counts)
cached = self._load_cached_projection(
- emb_set.name, method, dims, effective_params, param_suffix
+ emb_set.name,
+ method,
+ dims,
+ effective_params,
+ param_suffix,
+ fingerprint=fingerprint,
)
if cached:
add(cached)
@@ -847,7 +921,12 @@ def add(reduction: dict[str, Any]) -> None:
)
add(reduction)
self._save_projection_cache(
- emb_set.name, method, dims, reduction, effective_params
+ emb_set.name,
+ method,
+ dims,
+ reduction,
+ effective_params,
+ fingerprint=fingerprint,
)
computed_count += 1
diff --git a/apps/protspace/tests/test_annotation_manager.py b/apps/protspace/tests/test_annotation_manager.py
index e5c5342f..f3e58d1d 100644
--- a/apps/protspace/tests/test_annotation_manager.py
+++ b/apps/protspace/tests/test_annotation_manager.py
@@ -1488,3 +1488,153 @@ def test_lost_batch_leaves_an_existing_cache_untouched(
).to_pd()
assert pd.read_parquet(cache_path)["length"].tolist() == ["110"]
+
+
+class TestPerIdentifierReuse:
+ """A cache covering part of a run must be filled in, not thrown away.
+
+ Adding sequences to an existing run is the routine case, and re-fetching
+ every source for every identifier costs hours at Swiss-Prot scale. Reuse is
+ therefore decided per identifier as well as per column.
+ """
+
+ NO_FETCHING = {
+ "uniprot": False,
+ "taxonomy": False,
+ "interpro": False,
+ "ted": False,
+ "biocentral": False,
+ }
+
+ @staticmethod
+ def _cached(identifiers, organism="9606"):
+ return pd.DataFrame(
+ {
+ "identifier": identifiers,
+ "organism_id": [organism] * len(identifiers),
+ "length": [f"{10 * (i + 1)}" for i in range(len(identifiers))],
+ }
+ )
+
+ @staticmethod
+ def _fetched(mock_retriever, annotations):
+ mock_retriever.return_value.failed_batch_count = 0
+ mock_retriever.return_value.fetch_annotations.return_value = annotations
+
+ @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+ def test_only_the_missing_identifiers_are_fetched(self, mock_uniprot):
+ self._fetched(
+ mock_uniprot,
+ [ProteinAnnotations(identifier="P3", annotations={"length": "30"})],
+ )
+
+ result = ProteinAnnotationExtractor(
+ headers=["P1", "P2", "P3"],
+ annotations=["length"],
+ cached_data=self._cached(["P1", "P2"]),
+ sources_to_fetch=dict(self.NO_FETCHING),
+ ).to_pd()
+
+ assert mock_uniprot.call_args.kwargs["headers"] == ["P3"]
+ assert dict(zip(result["identifier"], result["length"], strict=True)) == {
+ "P1": "10",
+ "P2": "20",
+ "P3": "30",
+ }
+
+ @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+ def test_a_cache_covering_the_run_fetches_nothing(self, mock_uniprot):
+ result = ProteinAnnotationExtractor(
+ headers=["P1", "P2"],
+ annotations=["length"],
+ cached_data=self._cached(["P1", "P2", "P3"]),
+ sources_to_fetch=dict(self.NO_FETCHING),
+ ).to_pd()
+
+ mock_uniprot.assert_not_called()
+ assert set(result["identifier"]) >= {"P1", "P2"}
+
+ @patch("src.protspace.data.annotations.manager.TaxonomyRetriever")
+ @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+ def test_taxonomy_is_looked_up_only_for_organisms_the_cache_lacks(
+ self, mock_uniprot, mock_taxonomy
+ ):
+ cached = self._cached(["P1"])
+ cached["genus"] = ["Homo"]
+ self._fetched(
+ mock_uniprot,
+ [
+ ProteinAnnotations(
+ identifier="P2", annotations={"organism_id": "9606", "length": "30"}
+ ),
+ ProteinAnnotations(
+ identifier="P3",
+ annotations={"organism_id": "10090", "length": "40"},
+ ),
+ ],
+ )
+ mock_taxonomy.return_value.fetch_annotations.return_value = {
+ 10090: {"annotations": {"genus": "Mus"}}
+ }
+
+ result = ProteinAnnotationExtractor(
+ headers=["P1", "P2", "P3"],
+ annotations=["length", "genus"],
+ cached_data=cached,
+ sources_to_fetch=dict(self.NO_FETCHING),
+ ).to_pd()
+
+ # 9606 is already in the cache; only the unseen organism is looked up.
+ assert mock_taxonomy.call_args.kwargs["taxon_ids"] == [10090]
+ assert dict(zip(result["identifier"], result["genus"], strict=True)) == {
+ "P1": "Homo",
+ "P2": "Homo",
+ "P3": "Mus",
+ }
+
+ @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+ def test_rows_outside_the_run_survive_a_fetch(self, mock_uniprot, tmp_path):
+ self._fetched(
+ mock_uniprot,
+ [
+ ProteinAnnotations(
+ identifier="P1", annotations={"organism_id": "9606", "length": "10"}
+ )
+ ],
+ )
+ cache_path = tmp_path / "all_annotations.parquet"
+
+ ProteinAnnotationExtractor(
+ headers=["P1"],
+ annotations=["length"],
+ output_path=cache_path,
+ cached_data=self._cached(["P1", "P2", "P3"]),
+ ).to_pd()
+
+ assert set(pd.read_parquet(cache_path)["identifier"]) == {"P1", "P2", "P3"}
+
+ @patch("src.protspace.data.annotations.manager.UniProtRetriever")
+ def test_a_failed_fill_in_does_not_cache_empty_values(self, mock_uniprot, tmp_path):
+ mock_uniprot.return_value.fetch_annotations.side_effect = RuntimeError(
+ "offline"
+ )
+ cache_path = tmp_path / "all_annotations.parquet"
+ cached = self._cached(["P1", "P2"])
+ cached.to_parquet(cache_path, index=False)
+
+ result = ProteinAnnotationExtractor(
+ headers=["P1", "P2", "P3"],
+ annotations=["length"],
+ output_path=cache_path,
+ cached_data=cached,
+ sources_to_fetch=dict(self.NO_FETCHING),
+ ).to_pd()
+
+ # The run still returns what it had, but P3's empty length must not be
+ # cached: the next run would read it as "this protein has no length".
+ assert dict(zip(result["identifier"], result["length"], strict=True)) == {
+ "P1": "10",
+ "P2": "20",
+ "P3": "",
+ }
+ assert set(pd.read_parquet(cache_path)["identifier"]) == {"P1", "P2"}
diff --git a/apps/protspace/tests/test_atomic_publication.py b/apps/protspace/tests/test_atomic_publication.py
new file mode 100644
index 00000000..7a167a09
--- /dev/null
+++ b/apps/protspace/tests/test_atomic_publication.py
@@ -0,0 +1,84 @@
+"""Files this package publishes appear complete, and readable as usual.
+
+Staging into a private temp file and renaming it is what makes an interrupted
+write harmless, but `mkstemp` creates that file owner-only — so publishing by
+rename quietly handed the user a mode a direct write would never have produced.
+"""
+
+import os
+import stat
+
+import pyarrow as pa
+import pytest
+
+from protspace.data.io.atomic import atomic_write_bytes, staged_write
+
+
+@pytest.fixture
+def permissive_umask():
+ previous = os.umask(0o022)
+ yield
+ os.umask(previous)
+
+
+def _mode(path):
+ return stat.S_IMODE(path.stat().st_mode)
+
+
+def test_staged_write_publishes_with_the_process_umask(tmp_path, permissive_umask):
+ target = tmp_path / "published.txt"
+
+ with staged_write(target) as staged:
+ staged.write_text("content")
+
+ assert target.read_text() == "content"
+ assert _mode(target) == 0o644
+
+
+def test_staged_write_leaves_the_previous_content_on_failure(tmp_path):
+ target = tmp_path / "published.txt"
+ target.write_text("original")
+
+ with pytest.raises(RuntimeError, match="interrupted"):
+ with staged_write(target) as staged:
+ staged.write_text("half")
+ raise RuntimeError("interrupted")
+
+ assert target.read_text() == "original"
+ assert list(tmp_path.iterdir()) == [target]
+
+
+def test_atomic_write_bytes_publishes_with_the_process_umask(
+ tmp_path, permissive_umask
+):
+ target = tmp_path / "data.bin"
+
+ atomic_write_bytes(target, b"payload")
+
+ assert target.read_bytes() == b"payload"
+ assert _mode(target) == 0o644
+
+
+def test_a_bundle_is_not_owner_only(tmp_path, permissive_umask):
+ from protspace.data.io.bundle import write_bundle
+
+ bundle_path = tmp_path / "data.parquetbundle"
+ tables = [
+ pa.table({"identifier": ["P1"]}),
+ pa.table({"projection_name": ["PCA 2"]}),
+ pa.table({"identifier": ["P1"], "x": [1.0]}),
+ ]
+
+ write_bundle(tables, bundle_path)
+
+ assert _mode(bundle_path) == 0o644
+
+
+def test_a_rewritten_statistics_table_is_not_owner_only(tmp_path, permissive_umask):
+ from protspace.cli.stats import _atomic_write_table
+
+ target = tmp_path / "statistics.parquet"
+
+ _atomic_write_table(pa.table({"metric": ["silhouette"], "value": [0.5]}), target)
+
+ assert _mode(target) == 0o644
diff --git a/apps/protspace/tests/test_backend_switch.py b/apps/protspace/tests/test_backend_switch.py
index 6d6cadb0..148f2734 100644
--- a/apps/protspace/tests/test_backend_switch.py
+++ b/apps/protspace/tests/test_backend_switch.py
@@ -17,20 +17,28 @@
from typer.testing import CliRunner
from protspace.cli.app import app
-from protspace.data.embedding import local
+from protspace.data.embedding import local, store
from protspace.data.embedding.biocentral import resolve_embedder
from protspace.data.loaders.fasta import embed_fasta
-def _fake_embed(captured):
+def _fake_embed(captured, fill_value=1.0, backend="biocentral"):
"""A stand-in for ``embed_sequences`` that records its args and writes a
- minimal valid HDF5 so the surrounding load_h5 machinery still works."""
+ minimal valid HDF5.
+
+ It writes through the shared store, so the file carries the producer stamps
+ and residue digests a real run would leave -- a fake that wrote datasets
+ directly would make every cache look like a legacy one.
+ """
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))
+ store.save_embeddings(
+ Path(h5_path),
+ {pid: np.full(4, fill_value, dtype=np.float32) for pid in sequences},
+ sequences=sequences,
+ backend=backend,
+ model=embedder,
+ )
captured["embedder"] = embedder
captured["ids"] = list(sequences)
captured["config"] = embed_config
@@ -75,7 +83,8 @@ def test_embed_fasta_local_passes_short_key_and_remapped_ids(tmp_path, monkeypat
fasta.write_text(">sp|P12345|SOME_NAME\nMKVLAAG\n")
captured = {}
monkeypatch.setattr(
- "protspace.data.embedding.local.embed_sequences", _fake_embed(captured)
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed(captured, backend="local"),
)
result = embed_fasta(
@@ -122,6 +131,67 @@ def test_embed_fasta_unknown_backend_raises(tmp_path):
embed_fasta(fasta, "prot_t5", backend="nope", embedding_cache=tmp_path / "e.h5")
+def test_embed_fasta_refuses_a_cache_the_other_backend_wrote(tmp_path, monkeypatch):
+ """Both backends resume by identifier, so without ownership the second run
+ reuses the first backend's vectors and the two models land in one dataset."""
+ fasta = tmp_path / "s.fasta"
+ fasta.write_text(">P12345\nMKVLAAG\n")
+ cache = tmp_path / "prot_t5.h5"
+ monkeypatch.setattr(
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed({}, backend="local"),
+ )
+ embed_fasta(fasta, "prot_t5", backend="local", embedding_cache=cache)
+
+ # The real Biocentral backend: the refusal lands before any API call, so this
+ # needs no stub to stay off the network.
+ with pytest.raises(ValueError, match="--refetch embed"):
+ embed_fasta(fasta, "prot_t5", backend="biocentral", embedding_cache=cache)
+
+
+def test_embed_fasta_resumes_its_own_backends_cache(tmp_path, monkeypatch):
+ """The other half of ownership: the same producer must still resume."""
+ fasta = tmp_path / "s.fasta"
+ fasta.write_text(">P12345\nMKVLAAG\n")
+ cache = tmp_path / "prot_t5.h5"
+
+ for fill in (1.0, 2.0):
+ monkeypatch.setattr(
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed({}, fill_value=fill, backend="local"),
+ )
+ result = embed_fasta(fasta, "prot_t5", backend="local", embedding_cache=cache)
+
+ # The second run writes 2.0, so 1.0 is proof the first run's vector was kept.
+ assert result.data.tolist() == [[1.0] * 4]
+
+
+def test_embed_fasta_returns_only_the_fastas_proteins(tmp_path, monkeypatch):
+ """A cache shared by successive inputs accumulates every protein it has ever
+ embedded; returning the accumulation unions unrelated datasets into one
+ bundle."""
+ cache = tmp_path / "prot_t5.h5"
+ store.save_embeddings(
+ cache,
+ {"P99999": np.zeros(4, dtype=np.float32)},
+ sequences={"P99999": "MMMM"},
+ backend="local",
+ model="prot_t5",
+ )
+ fasta = tmp_path / "s.fasta"
+ fasta.write_text(">P12345\nMKVLAAG\n")
+ monkeypatch.setattr(
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed({}, backend="local"),
+ )
+
+ result = embed_fasta(fasta, "prot_t5", backend="local", embedding_cache=cache)
+
+ assert result.headers == ["P12345"]
+ # ...and the cache keeps what it already had, so the next run still resumes.
+ assert store.load_existing_ids(cache) == {"P12345", "P99999"}
+
+
# ---------------------------------------------------------------------------
# CLI wiring
# ---------------------------------------------------------------------------
@@ -132,7 +202,8 @@ def test_embed_cli_backend_local_dispatches_to_local(tmp_path, monkeypatch):
fasta.write_text(">P12345\nMKVLAAG\n")
captured = {}
monkeypatch.setattr(
- "protspace.data.embedding.local.embed_sequences", _fake_embed(captured)
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed(captured, backend="local"),
)
result = CliRunner().invoke(
@@ -164,7 +235,8 @@ def test_embed_cli_rejects_nonpositive_batch_size(tmp_path, monkeypatch):
# Guard against regressions: even if validation were skipped, don't let a
# real model load / hang.
monkeypatch.setattr(
- "protspace.data.embedding.local.embed_sequences", _fake_embed({})
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed({}, backend="local"),
)
result = CliRunner().invoke(
@@ -319,7 +391,8 @@ def test_embed_cli_wires_max_length_to_local_config(tmp_path, monkeypatch):
fasta.write_text(">P12345\nMKVLAAG\n")
captured = {}
monkeypatch.setattr(
- "protspace.data.embedding.local.embed_sequences", _fake_embed(captured)
+ "protspace.data.embedding.local.embed_sequences",
+ _fake_embed(captured, backend="local"),
)
result = CliRunner().invoke(
diff --git a/apps/protspace/tests/test_biocentral_embedder.py b/apps/protspace/tests/test_biocentral_embedder.py
index d016f354..6925c3a7 100644
--- a/apps/protspace/tests/test_biocentral_embedder.py
+++ b/apps/protspace/tests/test_biocentral_embedder.py
@@ -434,14 +434,54 @@ def test_gate_reads_the_file_not_a_counter(self, monkeypatch, tmp_path):
real_save = bc.save_embeddings
dropped = sorted(seqs)[0]
- def lossy_save(path, embeddings):
- real_save(path, {k: v for k, v in embeddings.items() if k != dropped})
+ def lossy_save(path, embeddings, **kwargs):
+ real_save(
+ path, {k: v for k, v in embeddings.items() if k != dropped}, **kwargs
+ )
monkeypatch.setattr(bc, "save_embeddings", lossy_save)
with pytest.raises(ValueError, match="Embedding incomplete"):
bc_mod.embed_sequences(seqs, "m", h5_path, embed_config=bc.EmbedConfig(2))
+ def test_a_run_stamps_its_producer(self, monkeypatch, tmp_path):
+ """Both backends stamp through the same shared store, so the file records
+ the model this backend was handed -- the resolved name, not a short key."""
+ import h5py
+
+ seqs, h5_path, _, bc = self._run(
+ monkeypatch, tmp_path, to_dict=lambda s: self._embeddings(s)
+ )
+ bc.embed_sequences(seqs, "m", h5_path, embed_config=bc.EmbedConfig(2))
+
+ with h5py.File(h5_path, "r") as f:
+ assert f.attrs["protspace_backend"] == "biocentral"
+ assert f.attrs["protspace_model"] == "m"
+
+ def test_a_local_cache_is_refused_before_any_api_call(self, monkeypatch, tmp_path):
+ """A Local-written vector satisfies this backend's resume check, so
+ without ownership the two models are silently mixed in one dataset."""
+ from src.protspace.data.embedding import biocentral as bc
+ from src.protspace.data.embedding import store
+
+ h5_path = tmp_path / "out.h5"
+ store.save_embeddings(
+ h5_path,
+ {"P0000": np.zeros(3, dtype=np.float32)},
+ sequences={"P0000": "AAAA"},
+ backend="local",
+ model="prot_t5",
+ )
+ called = []
+ monkeypatch.setattr(
+ bc, "BiocentralAPI", lambda **kw: called.append(1) or self._fake_api()
+ )
+
+ with pytest.raises(ValueError, match="--refetch embed"):
+ bc.embed_sequences({"P0000": "AAAA"}, "m", h5_path)
+
+ assert not called, "must refuse before connecting to the API"
+
def test_rerun_embeds_only_what_is_missing(self, monkeypatch, tmp_path):
"""A failed run must leave the pipeline able to converge on a retry."""
import h5py
diff --git a/apps/protspace/tests/test_bundle_overlay.py b/apps/protspace/tests/test_bundle_overlay.py
index 0b90bc6f..2139394c 100644
--- a/apps/protspace/tests/test_bundle_overlay.py
+++ b/apps/protspace/tests/test_bundle_overlay.py
@@ -89,7 +89,7 @@ def test_in_place_overwrite_works_and_leaves_no_temp(tmp_path):
def test_failed_replace_preserves_original_in_place(tmp_path, monkeypatch):
# If the rename is interrupted, the original bundle must survive intact
# (atomic write) rather than being left truncated.
- import protspace.data.io.bundle as bundle_mod
+ import protspace.data.io.atomic as atomic_mod
path = tmp_path / "b.parquetbundle"
write_bundle(_tables(), path)
@@ -101,7 +101,7 @@ def test_failed_replace_preserves_original_in_place(tmp_path, monkeypatch):
def boom(*args, **kwargs):
raise OSError("simulated interrupt before rename")
- monkeypatch.setattr(bundle_mod.os, "replace", boom)
+ monkeypatch.setattr(atomic_mod.os, "replace", boom)
with pytest.raises(OSError):
replace_annotations_in_bundle(path, path, new_annotations)
assert path.read_bytes() == original # untouched
diff --git a/apps/protspace/tests/test_embed_completeness.py b/apps/protspace/tests/test_embed_completeness.py
index edb44404..f1b3353a 100644
--- a/apps/protspace/tests/test_embed_completeness.py
+++ b/apps/protspace/tests/test_embed_completeness.py
@@ -22,6 +22,20 @@ def _write(h5_path: Path, ids) -> None:
f.create_dataset(pid, data=np.zeros(4, dtype=np.float32))
+def _save(
+ h5_path: Path, sequences, *, backend="local", model="prot_t5", fill=1.0
+) -> Path:
+ """Write one vector per identifier in *sequences*, stamped with its producer."""
+ store.save_embeddings(
+ h5_path,
+ {pid: np.full(4, fill, dtype=np.float32) for pid in sequences},
+ sequences=sequences,
+ backend=backend,
+ model=model,
+ )
+ return h5_path
+
+
class TestFinishRun:
def test_complete_run_succeeds(self, tmp_path):
h5 = tmp_path / "o.h5"
@@ -97,6 +111,143 @@ def test_message_cannot_be_mistaken_for_a_service_outage(self, tmp_path):
assert not [p for p in patterns if p in str(exc.value).lower()]
+class TestProducerOwnership:
+ """A cache belongs to the backend and model that wrote it.
+
+ Both backends resume by identifier alone, so without a recorded producer a
+ Local-written vector satisfies a Biocentral run's resume check and the two
+ models end up mixed in one dataset.
+ """
+
+ def test_another_backend_is_refused_and_named_with_the_remedies(self, tmp_path):
+ h5 = _save(tmp_path / "local-prot_t5.h5", {"a": "MKV"})
+ with pytest.raises(ValueError) as exc:
+ store.begin_run(
+ h5,
+ {"a": "MKV"},
+ backend="biocentral",
+ model="Rostlab/prot_t5_xl_uniref50",
+ )
+ msg = str(exc.value)
+ assert str(h5) in msg
+ assert "local" in msg and "prot_t5" in msg
+ # The three ways forward, or the message is a dead end.
+ assert "--refetch embed" in msg
+ assert "backend" in msg and "path" in msg
+
+ def test_another_model_on_the_same_backend_is_refused(self, tmp_path):
+ """Two models' vectors are as unmixable as two backends'."""
+ h5 = _save(tmp_path / "c.h5", {"a": "MKV"})
+ with pytest.raises(ValueError, match="prot_t5"):
+ store.begin_run(h5, {"a": "MKV"}, backend="local", model="esm2_8m")
+
+ def test_a_refused_file_is_left_untouched(self, tmp_path):
+ h5 = _save(tmp_path / "c.h5", {"a": "MKV"})
+ before = h5.read_bytes()
+ with pytest.raises(ValueError):
+ store.begin_run(h5, {"b": "MKW"}, backend="biocentral", model="m")
+ assert h5.exists(), "a refused cache must not be deleted"
+ assert h5.read_bytes() == before, "a refused cache must not be extended"
+
+ def test_the_same_producer_resumes(self, tmp_path):
+ h5 = _save(tmp_path / "c.h5", {"a": "MKV"})
+ outstanding = store.begin_run(
+ h5, {"a": "MKV", "b": "MKW"}, backend="local", model="prot_t5"
+ )
+ assert outstanding == {"b": "MKW"}
+
+ def test_a_file_predating_producers_is_adopted_and_reported(self, tmp_path, caplog):
+ """Refusing legacy files would force a full re-embed of every existing
+ cache on upgrade, so they are adopted -- audibly."""
+ h5 = tmp_path / "legacy.h5"
+ _write(h5, ["a"])
+ with caplog.at_level("INFO"):
+ outstanding = store.begin_run(
+ h5, {"a": "MKV"}, backend="local", model="prot_t5"
+ )
+ assert outstanding == {}
+ assert "local" in caplog.text and "prot_t5" in caplog.text
+ with h5py.File(h5, "r") as f:
+ assert f.attrs["protspace_backend"] == "local"
+ assert f.attrs["protspace_model"] == "prot_t5"
+
+ def test_an_adopted_file_is_owned_from_then_on(self, tmp_path):
+ h5 = tmp_path / "legacy.h5"
+ _write(h5, ["a"])
+ store.begin_run(h5, {"a": "MKV"}, backend="local", model="prot_t5")
+ with pytest.raises(ValueError, match="--refetch embed"):
+ store.begin_run(h5, {"a": "MKV"}, backend="biocentral", model="prot_t5")
+
+
+class TestSequenceIdentity:
+ """A vector belongs to the residues it was computed from."""
+
+ def test_a_changed_sequence_is_outstanding_again(self, tmp_path):
+ h5 = tmp_path / "c.h5"
+ _save(h5, {"a": "MKV", "b": "MKW"})
+ outstanding = store.begin_run(
+ h5, {"a": "MKV", "b": "EDITED"}, backend="local", model="prot_t5"
+ )
+ assert outstanding == {"b": "EDITED"}
+
+ def test_re_embedding_replaces_the_vector_and_the_digest(self, tmp_path):
+ """save_embeddings skips identifiers already present, so without this the
+ re-embed is computed and then thrown away."""
+ h5 = tmp_path / "c.h5"
+ _save(h5, {"a": "MKV"}, fill=1.0)
+ _save(h5, {"a": "EDITED"}, fill=2.0)
+ with h5py.File(h5, "r") as f:
+ assert f["a"][:].tolist() == [2.0] * 4
+ assert f["a"].attrs["protspace_sequence_sha256"] == store.sequence_digest(
+ "EDITED"
+ )
+
+ def test_an_unchanged_sequence_keeps_its_vector(self, tmp_path):
+ h5 = tmp_path / "c.h5"
+ _save(h5, {"a": "MKV"}, fill=1.0)
+ _save(h5, {"a": "MKV"}, fill=2.0)
+ with h5py.File(h5, "r") as f:
+ assert f["a"][:].tolist() == [1.0] * 4
+
+ def test_a_protein_without_a_digest_is_trusted(self, tmp_path):
+ h5 = tmp_path / "legacy.h5"
+ _write(h5, ["a"])
+ assert (
+ store.begin_run(h5, {"a": "ANYTHING"}, backend="local", model="prot_t5")
+ == {}
+ )
+
+ def test_finish_run_fails_when_a_stale_protein_was_not_re_embedded(self, tmp_path):
+ """Its old vector is still on disk under its old residues, so a presence
+ check alone reports the run complete."""
+ h5 = tmp_path / "c.h5"
+ _save(h5, {"a": "MKV", "b": "MKW"})
+ with pytest.raises(ValueError, match="Embedding incomplete"):
+ store.finish_run(h5, ["a", "b"], sequences={"a": "EDITED", "b": "MKW"})
+
+ def test_finish_run_accepts_a_stale_protein_that_was_re_embedded(self, tmp_path):
+ h5 = tmp_path / "c.h5"
+ _save(h5, {"a": "MKV"}, fill=1.0)
+ _save(h5, {"a": "EDITED"}, fill=2.0)
+ assert store.finish_run(h5, ["a"], sequences={"a": "EDITED"}) == h5
+
+ def test_resume_reads_the_digests_in_one_pass(self, tmp_path, monkeypatch):
+ """The digest is per protein at up to 570K of them: one open for the file,
+ not one per protein."""
+ h5 = tmp_path / "c.h5"
+ sequences = {f"p{i}": "MKV" for i in range(50)}
+ _save(h5, sequences)
+
+ opens = []
+ real_file = h5py.File
+ monkeypatch.setattr(
+ h5py, "File", lambda *a, **kw: opens.append(1) or real_file(*a, **kw)
+ )
+ store.begin_run(h5, sequences, backend="local", model="prot_t5")
+
+ assert len(opens) == 1, f"{len(opens)} opens for {len(sequences)} proteins"
+
+
class TestValidateHeaders:
def test_rejects_slash(self):
with pytest.raises(ValueError, match="invalid for HDF5 dataset names"):
diff --git a/apps/protspace/tests/test_local_embedder.py b/apps/protspace/tests/test_local_embedder.py
index 3bc79606..53ce835f 100644
--- a/apps/protspace/tests/test_local_embedder.py
+++ b/apps/protspace/tests/test_local_embedder.py
@@ -10,7 +10,7 @@
import numpy as np
import pytest
-from protspace.data.embedding import biocentral, local
+from protspace.data.embedding import biocentral, local, store
from protspace.data.embedding.biocentral import ALL_SHORT_KEYS
# ---------------------------------------------------------------------------
@@ -219,17 +219,22 @@ def test_embed_sequences_resumes_and_skips_existing(tmp_path):
# ---------------------------------------------------------------------------
-def _stub_model(monkeypatch, *, oom_ids=()):
+def _stub_model(monkeypatch, *, oom_ids=(), fill=0.0, loaded=None):
"""Replace model loading and inference so the contract can be tested without
downloading a checkpoint."""
import torch
- monkeypatch.setattr(local, "setup_model", lambda ckpt, mt: (None, None, "cpu"))
+ def setup(ckpt, mt):
+ if loaded is not None:
+ loaded.append(ckpt)
+ return (None, None, "cpu")
+
+ monkeypatch.setattr(local, "setup_model", setup)
def fake_embed_batch(processed, mod_type, model, tokenizer, device, max_length):
if len(processed) == 1 and processed[0] in oom_ids:
raise torch.cuda.OutOfMemoryError("stub OOM")
- return [np.zeros(4, dtype=np.float32) for _ in processed]
+ return [np.full(4, fill, dtype=np.float32) for _ in processed]
monkeypatch.setattr(local, "_embed_batch", fake_embed_batch)
@@ -282,6 +287,62 @@ def test_oom_at_batch_size_one_is_skipped(tmp_path, monkeypatch):
assert set(f.keys()) == {"ok"}
+# ---------------------------------------------------------------------------
+# Cache ownership: the shared store's contract reached through this backend
+# ---------------------------------------------------------------------------
+
+
+def test_local_run_stamps_its_producer_and_digests(tmp_path, monkeypatch):
+ """The stamps have to be written by the run, not by a caller that remembers
+ to -- `protspace embed -o mine.h5` gets the same protection as a cache."""
+ _stub_model(monkeypatch)
+ out = tmp_path / "emb.h5"
+
+ local.embed_sequences({"a": "MKVL"}, "esm2_8m", out)
+
+ with h5py.File(out, "r") as f:
+ assert f.attrs["protspace_backend"] == "local"
+ assert f.attrs["protspace_model"] == "esm2_8m" # the id this backend takes
+ assert f["a"].attrs["protspace_sequence_sha256"] == store.sequence_digest(
+ "MKVL"
+ )
+
+
+def test_local_run_refuses_a_biocentral_cache_before_loading_a_model(
+ tmp_path, monkeypatch
+):
+ out = tmp_path / "emb.h5"
+ store.save_embeddings(
+ out,
+ {"a": np.zeros(4, dtype=np.float32)},
+ sequences={"a": "MKVL"},
+ backend="biocentral",
+ model="facebook/esm2_t6_8M_UR50D",
+ )
+ loaded: list[str] = []
+ _stub_model(monkeypatch, loaded=loaded)
+
+ with pytest.raises(ValueError, match="--refetch embed"):
+ local.embed_sequences({"a": "MKVL"}, "esm2_8m", out)
+
+ assert not loaded, "must refuse before paying to load a checkpoint"
+
+
+def test_local_run_re_embeds_a_changed_sequence(tmp_path, monkeypatch):
+ """Resume matches on identifier alone, so an edited sequence otherwise keeps
+ the vector of the residues it used to have."""
+ out = tmp_path / "emb.h5"
+ _stub_model(monkeypatch, fill=1.0)
+ local.embed_sequences({"a": "MKVL", "b": "MKVA"}, "esm2_8m", out)
+
+ _stub_model(monkeypatch, fill=2.0)
+ local.embed_sequences({"a": "MKVL", "b": "EDITED"}, "esm2_8m", out)
+
+ with h5py.File(out, "r") as f:
+ assert f["a"][:].tolist() == [1.0] * 4 # untouched
+ assert f["b"][:].tolist() == [2.0] * 4 # re-embedded
+
+
def test_shortfall_that_is_not_a_skip_still_fails(tmp_path, monkeypatch):
"""Everything absent from the .h5 that was NOT deliberately skipped is a
failure -- this is what the local backend used to miss entirely."""
@@ -290,7 +351,9 @@ def test_shortfall_that_is_not_a_skip_still_fails(tmp_path, monkeypatch):
monkeypatch.setattr(
local,
"save_embeddings",
- lambda p, e: real_save(p, {k: v for k, v in e.items() if k != "dropped"}),
+ lambda p, e, **kw: real_save(
+ p, {k: v for k, v in e.items() if k != "dropped"}, **kw
+ ),
)
with pytest.raises(ValueError, match="Embedding incomplete"):
diff --git a/apps/protspace/tests/test_notebooks.py b/apps/protspace/tests/test_notebooks.py
index ff7bcd1b..ed2fd4e5 100644
--- a/apps/protspace/tests/test_notebooks.py
+++ b/apps/protspace/tests/test_notebooks.py
@@ -199,6 +199,67 @@ def test_notebook_fallback_sets_match_the_package(path: Path):
)
+@pytest.mark.parametrize("path", NOTEBOOKS, ids=lambda p: p.name)
+def test_no_notebook_imports_a_private_protspace_name(path: Path):
+ """Cell 1 installs the *released* protspace while this notebook is main's.
+
+ A name added this release does not exist in the release the cell installs,
+ so importing one breaks setup for every reader until the next release — and
+ the failure surfaces as "restart the session", which cannot fix it. Public
+ names are the contract that survives the lag; private ones are not, and a
+ duplicated fallback copy is a second source of truth that drifts.
+ """
+ transform = pytest.importorskip(
+ "IPython.core.inputtransformer2",
+ reason="IPython is a dev-group dependency (via jupyter)",
+ ).TransformerManager()
+
+ private = []
+ for index, source in _code_cells(path):
+ for node in ast.walk(ast.parse(transform.transform_cell(source))):
+ if not isinstance(node, ast.ImportFrom) or not node.module:
+ continue
+ if not node.module.startswith("protspace"):
+ continue
+ private += [
+ f"cell {index}: {node.module}.{alias.name}"
+ for alias in node.names
+ if alias.name.startswith("_")
+ ]
+ assert not private, (
+ f"{path.name} imports private protspace names: {', '.join(private)}. "
+ "Use a public name, or inline the few lines the notebook needs."
+ )
+
+
+def test_preparation_names_each_bundle_distinctly():
+ """A fixed download name is what issue #338 reads as a stale projection.
+
+ Two Generate actions land as `data.parquetbundle` and `data (1).parquetbundle`,
+ and opening the first shows the first run's coordinates. Structural rather
+ than a substring match: the point is that the name carries something from
+ this run, not that it is spelled any particular way.
+ """
+ transform = pytest.importorskip(
+ "IPython.core.inputtransformer2",
+ reason="IPython is a dev-group dependency (via jupyter)",
+ ).TransformerManager()
+
+ values = [
+ value
+ for _, source in _code_cells(NOTEBOOK_DIR / "ProtSpace_Preparation.ipynb")
+ for node in ast.walk(ast.parse(transform.transform_cell(source)))
+ for name, value in _assignments(node)
+ if name == "output_path"
+ ]
+
+ assert values, "the Generate callback assigns no output_path"
+ assert all(
+ any(isinstance(part, ast.FormattedValue) for part in ast.walk(value))
+ for value in values
+ ), "every bundle name is a constant, so two runs download the same file name"
+
+
@pytest.mark.parametrize("path", NOTEBOOKS, ids=lambda p: p.name)
def test_cells_carry_ids_when_the_format_requires_them(path: Path):
"""nbformat >= 4.5 requires cell ids.
diff --git a/apps/protspace/tests/test_pipeline_utils.py b/apps/protspace/tests/test_pipeline_utils.py
index d8b2c054..cf4b7c93 100644
--- a/apps/protspace/tests/test_pipeline_utils.py
+++ b/apps/protspace/tests/test_pipeline_utils.py
@@ -1186,3 +1186,266 @@ def boom(data, method, dims):
assert id(pipeline.base.config) == original_config_id, (
"base.config reference should be the original dict, not a replacement"
)
+
+
+# ---------------------------------------------------------------------------
+# Annotation cache identity
+# ---------------------------------------------------------------------------
+
+
+def _write_annotation_cache(cache_dir, identifiers, value):
+ cached = pd.DataFrame(
+ {
+ "identifier": identifiers,
+ "protein_name": [value] * len(identifiers),
+ "gene_name": [value] * len(identifiers),
+ "uniprot_kb_id": [value] * len(identifiers),
+ }
+ )
+ cached.to_parquet(cache_dir / "all_annotations.parquet")
+ return cached
+
+
+def test_a_cache_missing_identifiers_is_filled_in_not_discarded(tmp_path, monkeypatch):
+ """The cache still reaches the manager, which fetches only what it lacks.
+
+ Discarding it instead refetches every source for every identifier, which is
+ hours at Swiss-Prot scale for one added protein.
+ """
+ from protspace.data.annotations.manager import ProteinAnnotationManager
+
+ _write_annotation_cache(tmp_path, ["P1", "P2"], "cached")
+ captured = {}
+
+ def fill_in(manager):
+ captured["headers"] = manager.headers
+ captured["cached_data"] = manager.cached_data
+ return pd.DataFrame(
+ {
+ "identifier": ["P1", "P2", "P3"],
+ "protein_name": ["cached", "cached", "new"],
+ }
+ )
+
+ monkeypatch.setattr(ProteinAnnotationManager, "to_pd", fill_in)
+
+ result = _cache_pipeline(tmp_path, annotations=["protein_name"])._fetch_annotations(
+ ["P1", "P2", "P3"]
+ )
+
+ assert captured["headers"] == ["P1", "P2", "P3"]
+ assert captured["cached_data"]["identifier"].tolist() == ["P1", "P2"]
+ assert result["identifier"].tolist() == ["P1", "P2", "P3"]
+
+
+def test_a_failed_source_while_filling_in_leaves_the_cache_alone(tmp_path, monkeypatch):
+ """A source that failed for the new identifiers must not reach the cache.
+
+ Its columns would be empty for them and indistinguishable from a real
+ absence. Leaving the cache untouched is cheap now that the next run fills in
+ only the identifiers it lacks rather than rebuilding everything.
+ """
+ from protspace.data.annotations.retrievers.interpro_retriever import (
+ InterProRetriever,
+ )
+ from protspace.data.annotations.retrievers.uniprot_retriever import (
+ ProteinAnnotations,
+ UniProtRetriever,
+ )
+
+ cache_path = tmp_path / "all_annotations.parquet"
+ pd.DataFrame(
+ {"identifier": ["OLD1"], "gene_name": ["OLD"], "pfam": ["PF00001"]}
+ ).to_parquet(cache_path, index=False)
+
+ monkeypatch.setattr(
+ UniProtRetriever,
+ "fetch_annotations",
+ lambda _self: [
+ ProteinAnnotations(
+ identifier="P01308",
+ annotations={"gene_name": "INS", "organism_id": "9606"},
+ )
+ ],
+ )
+
+ def interpro_down(_self):
+ raise RuntimeError("InterPro unavailable")
+
+ monkeypatch.setattr(InterProRetriever, "fetch_annotations", interpro_down)
+
+ result = _cache_pipeline(
+ tmp_path, annotations=["gene_name", "pfam"]
+ )._fetch_annotations(["P01308"])
+
+ assert result.set_index("identifier").loc["P01308", "gene_name"] == "INS"
+ assert pd.read_parquet(cache_path)["identifier"].tolist() == ["OLD1"]
+
+
+@pytest.mark.parametrize(
+ "cached_ids", [["P1", "P2"], ["P1", "P2", "P3"]], ids=["exact", "superset"]
+)
+def test_annotation_cache_covering_the_request_is_reused(
+ tmp_path, monkeypatch, cached_ids
+):
+ from protspace.data.annotations.manager import ProteinAnnotationManager
+
+ cached = _write_annotation_cache(tmp_path, cached_ids, "cached")
+
+ def unexpected_fetch(_manager):
+ pytest.fail("a cache covering every requested identifier should be reused")
+
+ monkeypatch.setattr(ProteinAnnotationManager, "to_pd", unexpected_fetch)
+
+ result = _cache_pipeline(tmp_path, annotations=["protein_name"])._fetch_annotations(
+ ["P2", "P1"]
+ )
+
+ assert result["identifier"].tolist() == cached_ids
+ pd.testing.assert_frame_equal(
+ pd.read_parquet(tmp_path / "all_annotations.parquet"), cached
+ )
+
+
+# ---------------------------------------------------------------------------
+# Projection cache identity
+# ---------------------------------------------------------------------------
+
+
+def _recording_pipeline(tmp_path, **overrides):
+ """A pipeline whose reducer records what it was handed and slices it."""
+ pipeline = ReductionPipeline(
+ PipelineConfig(
+ methods=parse_methods_arg(["umap2"]),
+ output_path=None,
+ keep_tmp=True,
+ intermediate_dir=tmp_path,
+ **overrides,
+ )
+ )
+ reduced = []
+
+ def record_input(data, method, dims):
+ reduced.append(data.copy())
+ return {
+ "name": f"{method}{dims}",
+ "dimensions": dims,
+ "info": {},
+ "data": data[:, :dims].copy(),
+ }
+
+ pipeline.base.process_reduction = record_input
+ return pipeline, reduced
+
+
+def test_projection_cache_misses_a_changed_matrix_under_one_name(tmp_path):
+ pipeline, reduced = _recording_pipeline(tmp_path)
+ headers = ["P1", "P2", "P3"]
+
+ pipeline._run_reductions(
+ [_make_es("prot_t5", headers, data=np.zeros((3, 3), dtype=np.float32))]
+ )
+ changed = pipeline._run_reductions(
+ [_make_es("prot_t5", headers, data=np.full((3, 3), 7.0, dtype=np.float32))]
+ )[0]
+
+ assert len(reduced) == 2
+ np.testing.assert_array_equal(
+ changed["data"], np.full((3, 2), 7.0, dtype=np.float32)
+ )
+
+
+def test_projection_cache_misses_reordered_identifiers(tmp_path):
+ pipeline, reduced = _recording_pipeline(tmp_path)
+ headers = ["P1", "P2", "P3"]
+ data = np.array([[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]], dtype=np.float32)
+
+ pipeline._run_reductions([_make_es("prot_t5", headers, data=data)])
+ reordered = pipeline._run_reductions(
+ [_make_es("prot_t5", headers[::-1], data=data[::-1])]
+ )[0]
+
+ assert len(reduced) == 2
+ # Coordinates are paired positionally with the current identifiers, so a
+ # cache hit here would put P3's row on P1.
+ np.testing.assert_array_equal(reordered["data"], data[::-1][:, :2])
+
+
+def test_projection_cache_misses_a_grown_input(tmp_path):
+ pipeline, reduced = _recording_pipeline(tmp_path)
+ rng = np.random.default_rng(0)
+
+ pipeline._run_reductions(
+ [
+ _make_es(
+ "prot_t5",
+ [f"P{i}" for i in range(3)],
+ data=rng.normal(size=(3, 4)).astype(np.float32),
+ )
+ ]
+ )
+ grown = pipeline._run_reductions(
+ [
+ _make_es(
+ "prot_t5",
+ [f"P{i}" for i in range(5)],
+ data=rng.normal(size=(5, 4)).astype(np.float32),
+ )
+ ]
+ )[0]
+
+ assert len(reduced) == 2
+ assert grown["data"].shape[0] == 5
+
+
+def test_projection_cache_hits_an_unchanged_rerun(tmp_path):
+ pipeline, reduced = _recording_pipeline(tmp_path)
+ headers = ["P1", "P2", "P3"]
+ data = np.full((3, 3), 4.0, dtype=np.float32)
+
+ first = pipeline._run_reductions([_make_es("prot_t5", headers, data=data)])[0]
+ second = pipeline._run_reductions([_make_es("prot_t5", headers, data=data)])[0]
+
+ assert len(reduced) == 1
+ np.testing.assert_array_equal(second["data"], first["data"])
+
+
+def test_projection_refetch_recomputes_an_unchanged_rerun(tmp_path):
+ pipeline, reduced = _recording_pipeline(
+ tmp_path, refetch_stages=frozenset({"projections"})
+ )
+ headers = ["P1", "P2", "P3"]
+ data = np.full((3, 3), 4.0, dtype=np.float32)
+
+ pipeline._run_reductions([_make_es("prot_t5", headers, data=data)])
+ pipeline._run_reductions([_make_es("prot_t5", headers, data=data)])
+
+ assert len(reduced) == 2
+
+
+def test_precomputed_projection_cache_misses_a_changed_matrix(tmp_path):
+ pipeline, reduced = _recording_pipeline(tmp_path)
+ headers = ["P1", "P2", "P3"]
+
+ pipeline._run_reductions(
+ [
+ _make_es(
+ "MMseqs2", headers, data=np.zeros((3, 3), np.float32), precomputed=True
+ )
+ ]
+ )
+ changed = pipeline._run_reductions(
+ [
+ _make_es(
+ "MMseqs2",
+ headers,
+ data=np.full((3, 3), 9.0, np.float32),
+ precomputed=True,
+ )
+ ]
+ )[0]
+
+ assert len(reduced) == 2
+ np.testing.assert_array_equal(
+ changed["data"], np.full((3, 2), 9.0, dtype=np.float32)
+ )
diff --git a/apps/protspace/tests/test_query.py b/apps/protspace/tests/test_query.py
new file mode 100644
index 00000000..15cf434f
--- /dev/null
+++ b/apps/protspace/tests/test_query.py
@@ -0,0 +1,127 @@
+"""Tests for UniProt query FASTA downloads and publication."""
+
+import gzip
+import os
+import random
+import stat
+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
+
+
+def _mock_download(monkeypatch, fasta: str, *, truncate: bool = False) -> None:
+ payload = gzip.compress(fasta.encode())
+ if truncate:
+ payload = payload[: len(payload) // 2]
+ response = _Response(payload)
+ monkeypatch.setattr(query_module.requests, "get", lambda *args, **kwargs: response)
+
+
+def test_query_uniprot_does_not_publish_a_truncated_download(tmp_path, monkeypatch):
+ target = tmp_path / "query.fasta"
+ # Incompressible residues, so half the gzip stream still decompresses to
+ # something: extraction fails after writing part of its output.
+ residues = random.Random(0).choices("ACDEFGHIKLMNPQRSTVWY", k=200_000)
+ fasta = "".join(
+ f">P{i}\n{''.join(residues[i * 1000 : (i + 1) * 1000])}\n" for i in range(200)
+ )
+ _mock_download(monkeypatch, fasta, truncate=True)
+
+ with pytest.raises(EOFError):
+ query_module.query_uniprot("family:globin", save_to=target)
+
+ 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]
+
+
+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
+
+
+# ---------------------------------------------------------------------------
+# Retained query FASTA ownership
+# ---------------------------------------------------------------------------
+
+
+def _recording_download(monkeypatch, fasta=">P1\nAAAA\n"):
+ """Record every query that reaches query_uniprot, and write its FASTA."""
+ downloaded = []
+
+ def fake_query_uniprot(query, *, save_to=None):
+ downloaded.append(query)
+ save_to.parent.mkdir(parents=True, exist_ok=True)
+ save_to.write_text(fasta)
+ return ["P1"], save_to
+
+ monkeypatch.setattr(query_module, "query_uniprot", fake_query_uniprot)
+ return downloaded
+
+
+def test_a_second_query_does_not_reuse_the_first_query_fasta(tmp_path, monkeypatch):
+ downloaded = _recording_download(monkeypatch)
+
+ _, first = query_module.resolve_query_fasta("family:globin", tmp_path, frozenset())
+ _, second = query_module.resolve_query_fasta(
+ "family:phosphatase", tmp_path, frozenset()
+ )
+
+ assert downloaded == ["family:globin", "family:phosphatase"]
+ assert first != second
+
+
+def test_the_same_query_reuses_its_retained_fasta(tmp_path, monkeypatch):
+ downloaded = _recording_download(monkeypatch)
+ query = "family:globin"
+
+ _, first = query_module.resolve_query_fasta(query, tmp_path, frozenset())
+ headers, again = query_module.resolve_query_fasta(query, tmp_path, frozenset())
+
+ assert downloaded == [query]
+ assert again == first
+ assert headers == ["P1"]
+
+
+def test_refetch_query_downloads_again(tmp_path, monkeypatch):
+ downloaded = _recording_download(monkeypatch)
+ query = "family:globin"
+
+ query_module.resolve_query_fasta(query, tmp_path, frozenset())
+ query_module.resolve_query_fasta(query, tmp_path, frozenset({"query"}))
+
+ assert downloaded == [query, query]
diff --git a/docs/guide/fetching-and-caching.md b/docs/guide/fetching-and-caching.md
index 58702d4b..5bf48c75 100644
--- a/docs/guide/fetching-and-caching.md
+++ b/docs/guide/fetching-and-caching.md
@@ -39,28 +39,45 @@ With `--keep-tmp` (**the default**), everything expensive lands in `{output}/tmp
| Cached item | File | Saves you |
| ----------------- | --------------------------------------- | ------------------------------ |
-| FASTA sequences | `sequences.fasta` | re-downloading a UniProt query |
+| FASTA sequences | `queries/{query hash}.fasta` | re-downloading a UniProt query |
| Embeddings | `{embedder}.h5` | re-embedding proteins |
| Annotations | `all_annotations.parquet` | re-querying the APIs |
| Similarity matrix | `similarity_matrix.npy` | re-running MMseqs2 |
| DR projections | `proj_{name}_{method}{dims}_{hash}.npz` | recomputing UMAP/PaCMAP/… |
-Embeddings are cached per protein, so adding sequences to an existing run only embeds the new ones.
-Projections are keyed by a hash of their parameters, so changing a slider computes one new file and
-leaves the others alone.
+Every entry is owned by what produced it, so reuse can only ever be reuse of your own work:
+
+- **Query FASTA** by the exact query text, so a second query in the same `-o` downloads its own
+ sequences.
+- **Embeddings** per protein _and_ per residue: a protein whose sequence changed under an unchanged
+ identifier is embedded again, and the file records which backend and model wrote it (see
+ [below](#embeddings-belong-to-one-backend-and-model)). Adding sequences to an existing run still
+ only embeds the new ones.
+- **Projections** by the embedding matrix, the identifier order, the method, the dimensions and
+ every reducer parameter. Changing a slider computes one new file and leaves the others alone;
+ re-running an input that changed under the same name recomputes rather than returning the earlier
+ coordinates.
+- **Annotations** per identifier and per column — the next section.
## How the annotation cache decides
This is the part worth understanding, because it explains most "why didn't it refetch?" questions.
-The cache is judged **by column, not by row**. On each run ProtSpace compares the columns you asked
-for against the columns already in `all_annotations.parquet`:
+The cache is judged **by column and by row**. ProtSpace compares what you asked for against what
+`all_annotations.parquet` holds:
-- **Every column present** → the cache is used as-is, and no API is called.
+- **Every column present, every protein present** → the cache is used as-is, and no API is called.
- **Some column missing** → only the sources owning the missing columns are queried; cached columns
from other sources are reused.
+- **Some protein missing** → each source is queried for exactly those proteins, and cached values
+ serve the rest. Taxonomy is looked up only for organisms the cache has not resolved before.
+
+So asking for a new annotation is cheap, asking for the same ones again is free, and adding a
+handful of proteins to a large run costs a handful of lookups rather than a full refetch.
-So asking for a new annotation is cheap, and asking for the same ones again is free.
+A cache holding _more_ proteins than the current run is fine: the extra rows are filtered out of the
+bundle, and a run for part of a dataset keeps them rather than replacing the cache with its own
+subset.
The consequence of column-level granularity is that an **empty value is not a signal**. A protein
with an empty `ec` may have no EC number, may not be in UniProt at all, or may be a custom
@@ -78,7 +95,8 @@ ProtSpace therefore **caches only the sources that completed**:
- Sources that succeeded are still cached, so one flaky API does not throw away an expensive UniProt
fetch.
- If leaving it out would mean overwriting an existing cache with _fewer_ columns, the existing
- cache is kept untouched instead.
+ cache is kept untouched instead — unless that cache covers other proteins, in which case the
+ sources that completed replace it.
Either way the run still returns everything it did retrieve — your bundle is built, and the message
says which source was short.
@@ -94,6 +112,16 @@ thousands of sequential requests, and without retries a single blip would be nea
fetched one request per protein (TED) use a smaller retry budget, so a full outage does not multiply
the backoff by the number of proteins.
+## Embeddings belong to one backend and model
+
+An HDF5 records the backend and model that produced it. Both backends resume by identifier, so
+without that record a run with `--backend local` would resume from vectors the Biocentral API wrote
+and silently mix two embedding spaces in one dataset. A run that points at another producer's file
+stops and names your options: select that backend, choose another output, or `--refetch embed`.
+
+Files written before this existed carry no record; they are adopted, stamped and reported the first
+time a run resumes from them.
+
## Forcing a refresh
`--refetch` recomputes specific stages, comma-separated:
@@ -137,7 +165,8 @@ add `--refetch annotations`.
UniProt. If a message said a source was incomplete, re-run: that source was deliberately not cached.
**A re-run refetches more than expected** — a source that did not complete on the previous run is
-not in the cache, by design.
+not in the cache, by design. A run whose input includes proteins the cache does not cover also
+rebuilds every annotation for that input.
**Nothing is being cached** — check that `--keep-tmp` is on (it is by default) and that `{output}/`
is writable.
diff --git a/docs/guide/python-cli.md b/docs/guide/python-cli.md
index 2b91d0dc..f7bdf23d 100644
--- a/docs/guide/python-cli.md
+++ b/docs/guide/python-cli.md
@@ -350,16 +350,21 @@ python -c "import h5py; print(dict(h5py.File('file.h5','r').attrs))"
With `--keep-tmp` (the default), intermediate results are cached in `{output}/tmp/` and reused on
subsequent runs:
-| Cached item | File | Reuse behavior |
-| ----------------- | --------------------------------------- | ------------------------------- |
-| FASTA sequences | `sequences.fasta` | Skip the UniProt query download |
-| Embeddings | `{embedder}.h5` | Skip already-embedded proteins |
-| Annotations | `all_annotations.parquet` | Fetch missing or stale columns |
-| Similarity matrix | `similarity_matrix.npy` | Skip MMseqs2 recomputation |
-| DR projections | `proj_{name}_{method}{dims}_{hash}.npz` | Skip dimensionality reduction |
+| Cached item | File | Reuse behavior |
+| ----------------- | --------------------------------------- | ----------------------------------------------------- |
+| FASTA sequences | `queries/{query hash}.fasta` | Skip the download for that query |
+| Embeddings | `{embedder}.h5` | Skip proteins already embedded from the same residues |
+| Annotations | `all_annotations.parquet` | Fetch missing or stale columns, and missing proteins |
+| Similarity matrix | `similarity_matrix.npy` | Skip MMseqs2 recomputation |
+| DR projections | `proj_{name}_{method}{dims}_{hash}.npz` | Skip dimensionality reduction |
The annotation cache always stores scores; `--no-scores` strips them from the output afterwards.
+The annotation cache is read per column and per protein: a source is queried only for the proteins
+whose values the cache cannot supply, and a cache covering more proteins than the current run keeps
+those extra rows. An embedding HDF5 records the backend and model that wrote it, and a run that
+points at another producer's file stops rather than mixing two embedding spaces.
+
If a source could not be fully retrieved, its columns are **left out of the cache**: a partly empty
column is indistinguishable from one where those proteins genuinely have no entry, so caching it
would make every later run reuse the gaps instead of refetching. Sources that did complete are
diff --git a/openspec/changes/fix-cache-ownership/.openspec.yaml b/openspec/changes/fix-cache-ownership/.openspec.yaml
new file mode 100644
index 00000000..5849c2db
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/.openspec.yaml
@@ -0,0 +1,2 @@
+schema: spec-driven
+created: 2026-08-01
diff --git a/openspec/changes/fix-cache-ownership/README.md b/openspec/changes/fix-cache-ownership/README.md
new file mode 100644
index 00000000..bf68bda0
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/README.md
@@ -0,0 +1,4 @@
+# fix-cache-ownership
+
+Make every retained intermediate — projection, embedding, query FASTA, annotation row —
+owned by what it was computed from, in the shared layer rather than in one notebook.
diff --git a/openspec/changes/fix-cache-ownership/design.md b/openspec/changes/fix-cache-ownership/design.md
new file mode 100644
index 00000000..339ca4c3
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/design.md
@@ -0,0 +1,89 @@
+## Context
+
+`{output}/tmp` is shared by every run that writes to the same `-o`, and the notebook's `output/tmp` is shared by every Generate action. Cache entries in it are addressed by _names_ — the embedding name for a projection, the model name for an HDF5, `sequences.fasta` for a query — while what makes an entry reusable is the _data_ behind that name. Every defect in the proposal is that gap, and each was reproduced against the CLI layout, not only the notebook.
+
+The first iteration closed the gap for the notebook by addressing directories with a digest of the input file. That is the same idea applied one layer too high: it re-reads multi-gigabyte inputs on every Generate action, forfeits per-protein embedding resume on any byte change, and leaves `protspace prepare` broken.
+
+## Goals / Non-Goals
+
+**Goals:**
+
+- Make each retained artifact's identity include what it was computed from, in the shared layer, so the CLI, the notebook, and the hosted prep service all get it.
+- Keep resumability: per protein for embeddings, per identifier and column for annotations, per parameter set for projections.
+- Keep retained caches from earlier versions usable without a forced recompute.
+- Let the notebook shrink back to configuration plus display.
+
+**Non-Goals:**
+
+- Redesigning the bundle format, projection naming, or output paths.
+- Making the notebook call `ReductionPipeline.run()` (it needs per-stage progress; tracked separately).
+- Garbage-collecting cache entries that ownership makes unreachable.
+
+## Decisions
+
+### Projection identity includes the matrix and the identifier order
+
+`_projection_cache_path` adds a fingerprint to its key: a SHA-256 over the identifier list, the dtype, the shape, and the matrix bytes, computed once per embedding set per run and reused across that set's methods.
+
+Measured at 2.7 GB/s, so ~0.9 s for a 573K × 1024 float32 matrix (a few seconds on a Colab CPU), against reducer runtimes of minutes at that size. Including the identifiers is what makes a reordered or grown input miss: coordinates are stored as bare rows and paired positionally with the current identifiers on load, so an order change silently relabels every point.
+
+The notebook then drops `refetch_stages={"projections"}`: with data identity in the key, returning to an earlier slider value is a correct cache hit rather than a stale one.
+
+**Alternative: keep the notebook's blanket refetch.** Rejected: it fixes one caller, leaves `prepare -o` stale, and makes the notebook write `proj_*.npz` files nothing ever reads.
+
+**Alternative: hash a sample of the matrix.** Rejected: a cheap fingerprint that can collide re-introduces exactly the silent staleness being removed, and the full hash is already negligible next to the reducers.
+
+### Embedding identity is stamped in the HDF5, not in the file name
+
+The shared store writes two root attributes (`protspace_backend`, `protspace_model`) and one per-protein attribute (`protspace_sequence_sha256`, the first 16 hex characters of the residue digest). Resume rejects a file stamped by another producer with a `ValueError` naming the remedies, and treats a protein whose digest differs from its current residues as outstanding, replacing that dataset.
+
+Attributes rather than file names because the file name is a caller's choice and the contract belongs to the file: `protspace embed -o mine.h5` gets the same protection as a managed cache. Cost measured at 50K proteins: writing +0.6 s (against hours of embedding), reading digests on resume +1.6 s per 50K (~19 s for Swiss-Prot, against a full `load_h5` of the same file).
+
+Legacy files carry neither attribute. A file with no producer is adopted and stamped, with a log line; a protein with no digest is trusted. Refusing them instead would force a full re-embed of every existing cache on upgrade, which for a Biocentral-sized run costs hours to protect against a mix that ownership now prevents going forward.
+
+The CLI keeps its `tmp/{model}.h5` naming, so existing caches stay valid and a backend switch is an explicit error. The notebook prefixes the backend into the name, because switching backends there is a toggle in the panel rather than a new command, and both files are worth keeping.
+
+**Alternative: per-sequence hashes in a side file.** Rejected: two files that can disagree, where the attribute cannot.
+
+### An embedding load returns the requested proteins
+
+`embed_fasta` restricts the returned set to the FASTA's identifiers. The shared cache legitimately accumulates proteins across inputs — that is what makes resume work — so the loader, not the cache, decides what a run is about.
+
+### Annotation reuse is per identifier as well as per column
+
+The manager already decides per source whether to fetch or to read the cache. That decision becomes per source _and_ per identifier: for a source served from the cache, identifiers the cache does not hold are fetched and merged with the cached rows. Taxonomy is keyed by organism rather than identifier, so it fills in only the organisms its cached lookup lacks.
+
+The cache write keeps rows for identifiers outside the current run when the run's columns are the columns the cache already had, so a subset run no longer shrinks a superset cache. When the columns differ, the existing behaviour stands — the frame for this run is what gets written.
+
+The incomplete-source rules are unchanged and continue to decide what may be written: a source that did not complete for the identifiers being filled in does not reach the cache as empty values.
+
+**Alternative: keep the full rebuild.** Rejected: it is hours of re-fetching for one added protein at Swiss-Prot scale, and it replaces a large cache with the current run's rows.
+
+### One staged-rename helper, respecting the umask
+
+`data/io/atomic.py` provides the write-then-rename helper the bundle writer, the `stats` table rewrites, and the retained query FASTA all use. It creates its temporary file with a normal `open`, so the published file gets the process umask rather than the owner-only mode `mkstemp` gives — which is what a plain write always did, and what the bundle writer has silently not been doing.
+
+### The notebook stops carrying cache logic
+
+With identity in the shared layer, the notebook keeps only what is genuinely a notebook policy: the backend-prefixed HDF5 name and the query-addressed FASTA path, both written inline from `hashlib` and an f-string. No private helper is imported from the package, so the window where the notebook on `main` runs against the previously released package cannot break it.
+
+### Issue #338 is addressed where the reporter sees it
+
+Each Generate action names its bundle `protspace_.parquetbundle` and prints the name, so two downloads are told apart. The parameter section states which methods each control applies to, so a slider that PCA ignores is visibly a slider PCA ignores.
+
+## Risks / Trade-offs
+
+- **A backend switch against an existing cache now fails.** → It silently returned the other backend's vectors before; the message names the three ways forward.
+- **Fingerprinting adds one pass over the matrix per run.** → Sub-second at Swiss-Prot scale, against reducers that take minutes.
+- **Reading residue digests adds to resume at scale.** → Seconds against an embedding run's hours, and `load_h5` reads the same file anyway.
+- **Legacy files are trusted rather than refused.** → Documented, with `--refetch embed` as the remedy; new writes are stamped from now on.
+- **A shared annotation cache keeps rows from other inputs.** → Bounded by what the user ran in that output directory, and the pipeline's identifier merge already drops rows outside the current run.
+- **Projection and embedding caches recompute once after upgrade.** → Only where their identity was previously unproven.
+
+## Migration Plan
+
+None required. Existing `proj_*.npz` entries miss on their new key and are recomputed once; existing HDF5 files are adopted and stamped on first resume; existing `sequences.fasta` files are ignored in favour of the query-addressed path and can be deleted.
+
+## Open Questions
+
+None.
diff --git a/openspec/changes/fix-cache-ownership/proposal.md b/openspec/changes/fix-cache-ownership/proposal.md
new file mode 100644
index 00000000..efeb851d
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/proposal.md
@@ -0,0 +1,44 @@
+## Why
+
+Issue #338 reports stale projections after changing a dimensionality-reduction slider. The projection key has always included every reducer parameter, so that symptom is not a parameter-cache bug; the report is reproduced instead by PCA (the first projection the app shows) taking none of those parameters, and by every Generate action downloading the same file name.
+
+Auditing the same retained-cache flow found real staleness underneath, and it is **not** notebook-specific. Every one of these was reproduced against the CLI's own cache layout (`{output}/tmp`, one file per model name):
+
+- A changed matrix under the same embedding name returns the earlier run's coordinates; a reordered input pairs cached coordinates with the wrong proteins; a grown input yields a projection with fewer rows than proteins.
+- A disjoint FASTA embedded into the same cache returns the union of both inputs.
+- A run with the other `--backend` resumes from the file the first backend wrote and returns its vectors.
+- A sequence edited under an unchanged identifier keeps the vector of the previous residues.
+- `prepare -q A -o out` followed by `-q B -o out` reuses A's sequences.
+- Requesting annotations for identifiers a cache does not hold re-fetches every source for every identifier, which is hours at Swiss-Prot scale for one added protein.
+
+The first iteration of this change addressed these by giving the notebook content-addressed cache directories. That hid the defects for one caller while leaving them in place for `protspace prepare`, cost a full re-read of the input on every Generate action, and discarded per-protein embedding resume whenever a single byte of the input changed.
+
+## What Changes
+
+- Include the embedding matrix and identifier order in the projection cache key, so cached coordinates are reused only for the data that produced them, for every caller.
+- Record the producing backend, the resolved model, and a per-protein residue digest in the embedding HDF5; refuse to resume from another producer's file, and re-embed proteins whose residues changed.
+- Return only the requested FASTA's proteins from an embedding load, rather than everything the shared cache accumulated.
+- Address a retained query FASTA by its query text in the CLI as well as the notebook.
+- Fetch each annotation source only for the identifiers the cache cannot supply, reusing cached rows for the rest, and keep a cached superset's rows when writing.
+- Publish bundles, rewritten statistics/annotation tables, and retained FASTA files through one staged-rename helper that respects the process umask.
+- Simplify the notebook accordingly: no content-addressed directories, no projection refetch override, and no imports that a released package may not have yet. Give each Generate action a distinguishable bundle name, and state which methods the parameter controls apply to.
+
+## Capabilities
+
+### New Capabilities
+
+- `intermediate-cache-ownership`: What a retained projection, embedding, query FASTA, or annotation row is owned by, and when it may be reused.
+- `atomic-file-publication`: How a published file becomes visible, and with what permissions.
+- `notebook-generate-output-identity`: How a Generate action's output and parameter scope are made legible.
+
+### Modified Capabilities
+
+- `annotation-cache-semantics`: Reuse becomes per identifier as well as per column; the incomplete-source rules continue to apply to the identifiers being filled in.
+- `embed-completeness`: Unchanged in what counts as complete; resume now additionally rejects another producer's file and re-embeds changed residues.
+
+## Impact
+
+- Affected code: `data/processors/pipeline.py`, `data/embedding/store.py` and both backends, `data/loaders/fasta.py`, `data/loaders/query.py`, `data/annotations/manager.py`, `cli/prepare.py`, `cli/stats.py`, `data/io/bundle.py`, `ProtSpace_Preparation.ipynb`.
+- Behaviour changes for CLI users: a backend switch against an existing cache now fails with guidance instead of silently mixing vectors; projection and embedding caches recompute once where their identity was previously unproven.
+- Retained caches from earlier versions stay readable: files without a producer or residue digest are adopted and stamped.
+- No bundle format, no public Python API signature, and no dependency changes.
diff --git a/openspec/changes/fix-cache-ownership/specs/atomic-file-publication/spec.md b/openspec/changes/fix-cache-ownership/specs/atomic-file-publication/spec.md
new file mode 100644
index 00000000..67a42020
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/specs/atomic-file-publication/spec.md
@@ -0,0 +1,46 @@
+## ADDED Requirements
+
+### Requirement: A file this package publishes appears complete or not at all
+
+A file written for a user or a later run to read SHALL be staged beside its
+destination and renamed into place, so an interrupted write leaves either the
+previous content or nothing. This covers the bundle, the statistics and
+annotation tables `stats` rewrites in place, and a retained query FASTA. A
+half-written file at the final path is indistinguishable from a complete one:
+the bundle overwrite workflow documents `-b results.parquetbundle -o
+results.parquetbundle`, and a retained FASTA's existence is what the next run
+reads as a cache hit.
+
+#### Scenario: A write is interrupted
+
+- **WHEN** an interruption or error occurs while writing a published file
+- **THEN** the destination keeps its previous content, or stays absent if there was none
+- **AND** the staged temporary file is removed
+
+#### Scenario: A download is truncated
+
+- **WHEN** a query download ends before its compressed stream is complete
+- **THEN** decompression fails and nothing is published at the retained path
+- **AND** a previously retained complete FASTA is left in place
+
+#### Scenario: A write completes
+
+- **WHEN** the staged file is written in full
+- **THEN** it replaces the destination in one rename
+
+### Requirement: A published file carries ordinary new-file permissions
+
+A file published by staging and rename SHALL carry the permissions a normal new
+file would receive under the process umask. Private temporary files are created
+owner-only, so publishing one by rename hands the user a file their own group
+and other tooling cannot read, which a direct write would never have done.
+
+#### Scenario: A bundle is written under a permissive umask
+
+- **WHEN** a bundle is written by a process whose umask allows group or other read
+- **THEN** the published bundle is readable accordingly rather than owner-only
+
+#### Scenario: A retained FASTA is published under a restrictive umask
+
+- **WHEN** a retained query FASTA is published by a process with a restrictive umask
+- **THEN** the published file respects that umask
diff --git a/openspec/changes/fix-cache-ownership/specs/intermediate-cache-ownership/spec.md b/openspec/changes/fix-cache-ownership/specs/intermediate-cache-ownership/spec.md
new file mode 100644
index 00000000..b9194de8
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/specs/intermediate-cache-ownership/spec.md
@@ -0,0 +1,148 @@
+## ADDED Requirements
+
+### Requirement: A cached projection belongs to the data it was computed from
+
+Cached projection coordinates SHALL be reused only when the embedding matrix and
+its identifier order are the same as the run that produced them, in addition to
+the existing method, dimension count, and reducer parameters. The logical
+embedding name is not evidence of data identity: a resumed embedding cache, a
+re-embedded input, a different intersection, or a reordered input all keep the
+name while changing the matrix.
+
+#### Scenario: Input data changes without changing its logical embedding name
+
+- **WHEN** an embedding set carries the same name, method, and reducer parameters as an earlier run but a different matrix
+- **THEN** the reducer runs against the current matrix
+- **AND** the earlier coordinates are not reused
+
+#### Scenario: The same proteins arrive in a different order
+
+- **WHEN** an embedding set holds the same matrix rows as an earlier run under a different identifier order
+- **THEN** the reducer runs again rather than pairing the cached coordinates with the reordered identifiers
+
+#### Scenario: An input grows between runs
+
+- **WHEN** proteins are added to an input that keeps its embedding name
+- **THEN** the projection is recomputed for every protein in the current run
+- **AND** no projection is emitted whose row count differs from the current identifier count
+
+#### Scenario: Nothing about the input changed
+
+- **WHEN** a run repeats with the same matrix, identifier order, method, dimensions, and parameters
+- **THEN** the cached coordinates are reused without running the reducer
+
+#### Scenario: Projections are explicitly refreshed
+
+- **WHEN** a run requests the `projections` refetch stage
+- **THEN** cached coordinates are ignored regardless of data identity
+
+### Requirement: An embedding cache is owned by the backend and model that produced it
+
+An embedding HDF5 SHALL record the backend and resolved model that produced it,
+and a run SHALL refuse to resume from a file another producer wrote. Both
+backends resume by identifier alone, so without this a Local-produced vector
+satisfies a Biocentral run's resume check and the two are silently mixed in one
+dataset.
+
+#### Scenario: A different backend resumes from the same file
+
+- **WHEN** a run embeds into an HDF5 recorded as produced by the other backend
+- **THEN** the run fails with an error naming the file, the recorded producer, and the remedies (select that backend, choose another output, or refetch the embeddings)
+- **AND** the existing file is left untouched rather than extended or deleted
+
+#### Scenario: The same producer resumes
+
+- **WHEN** a run embeds into an HDF5 recorded as produced by the same backend and model
+- **THEN** the existing identifiers are reused and only outstanding sequences are embedded
+
+#### Scenario: A file predates producer tracking
+
+- **WHEN** a run resumes from an HDF5 that records no producer
+- **THEN** the run adopts the file, records the current producer, and reports that it did so
+
+### Requirement: A cached embedding belongs to the sequence it was computed from
+
+An embedding HDF5 SHALL record, per protein, a digest of the residues the vector
+was computed from, and a run SHALL re-embed a protein whose current residues
+differ from that digest. Resume matches on identifier alone, so an identifier
+whose sequence changed otherwise keeps a vector of the previous residues.
+
+#### Scenario: A sequence changes while its identifier does not
+
+- **WHEN** a run embeds a FASTA whose residues for an already-embedded identifier changed
+- **THEN** that protein is embedded again and its stored vector and digest are replaced
+- **AND** proteins whose residues are unchanged are not embedded again
+
+#### Scenario: A protein predates sequence digests
+
+- **WHEN** a stored protein carries no residue digest
+- **THEN** it is reused as before and the run does not fail
+
+### Requirement: An embedding load returns only the proteins that were requested
+
+Embedding a FASTA SHALL return only that FASTA's proteins, even when the
+embedding cache holds more. A cache shared by successive inputs accumulates
+every protein it has ever embedded, and returning the accumulation silently
+unions unrelated datasets into one bundle.
+
+#### Scenario: A disjoint input reuses the same embedding cache
+
+- **WHEN** a FASTA is embedded into a cache that already holds a disjoint input's proteins
+- **THEN** the returned embedding set holds only the current FASTA's proteins
+- **AND** the retained cache keeps the proteins it already had
+
+### Requirement: A retained query FASTA is owned by its query text
+
+A retained query FASTA SHALL be addressed by the query that produced it, for
+every caller that retains one. A single shared file name reuses one query's
+sequences for a different query whenever both write to the same output
+directory.
+
+#### Scenario: A second query reuses the same output directory
+
+- **WHEN** a run downloads sequences for one query and a later run requests a different query with the same retained cache directory
+- **THEN** the second run downloads its own sequences instead of reusing the first query's FASTA
+
+#### Scenario: The same query runs again
+
+- **WHEN** a run repeats a query whose retained FASTA is present and complete
+- **THEN** the retained FASTA is reused without downloading it again
+
+### Requirement: Annotation reuse is decided per identifier and per source
+
+Annotation retrieval SHALL fetch a source only for the identifiers whose values
+the retained cache cannot supply, and SHALL reuse cached values for the rest.
+Requesting annotations for identifiers a cache does not hold is the routine
+"added sequences to an existing run" case, and re-fetching every source for
+every identifier can cost hours at Swiss-Prot scale.
+
+#### Scenario: The cache covers some of the requested identifiers
+
+- **WHEN** a run requests annotations for identifiers the retained cache only partly covers
+- **THEN** each source is fetched only for the identifiers missing from the cache
+- **AND** cached values supply the identifiers the cache already holds
+- **AND** the returned annotations cover every requested identifier
+
+#### Scenario: Taxonomy is reused for organisms already cached
+
+- **WHEN** identifiers missing from the cache resolve to organisms the cached taxonomy already covers
+- **THEN** no taxonomy lookup is made for those organisms
+- **AND** organisms the cache does not cover are looked up
+
+#### Scenario: The cache holds none of the requested identifiers
+
+- **WHEN** a retained cache describes an entirely different input
+- **THEN** every source is fetched for the current identifiers
+- **AND** the cached rows are not returned as the current run's annotations
+
+#### Scenario: A cached superset survives a run that fetches
+
+- **WHEN** a run fetches annotations for identifiers missing from a cache that also holds identifiers outside the request
+- **AND** the run's columns are the columns the cache already holds
+- **THEN** the retained cache keeps its other identifiers' rows alongside the newly fetched ones
+
+#### Scenario: A source fails while filling in missing identifiers
+
+- **WHEN** a source does not complete for the identifiers being filled in
+- **THEN** that source's columns are not written to the cache as empty values for those identifiers
+- **AND** the run still returns everything it did retrieve
diff --git a/openspec/changes/fix-cache-ownership/specs/notebook-generate-output-identity/spec.md b/openspec/changes/fix-cache-ownership/specs/notebook-generate-output-identity/spec.md
new file mode 100644
index 00000000..aedb6c48
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/specs/notebook-generate-output-identity/spec.md
@@ -0,0 +1,28 @@
+## ADDED Requirements
+
+### Requirement: Each Generate action produces a distinguishable bundle
+
+The Preparation notebook SHALL name each generated bundle so that bundles from
+different Generate actions are told apart after download. A fixed name makes a
+browser file the user opens ambiguous — a second download lands beside the first
+as a copy, and opening the earlier file shows the earlier coordinates, which
+reads as a stale projection.
+
+#### Scenario: A second Generate action downloads another bundle
+
+- **WHEN** a user runs Generate twice in one session
+- **THEN** the second bundle's file name differs from the first
+- **AND** the notebook reports the name it wrote
+
+### Requirement: The panel states which parameters affect which methods
+
+The Preparation notebook SHALL state that the reducer parameter controls apply
+only to the methods that consume them. PCA and MDS take no neighbourhood or
+perplexity parameter, so changing a slider and regenerating returns a
+bit-identical PCA view — the default first projection — which reads as a cache
+that ignored the change.
+
+#### Scenario: A user changes a parameter that the selected method ignores
+
+- **WHEN** the parameter controls are shown
+- **THEN** the panel names the methods each parameter group applies to
diff --git a/openspec/changes/fix-cache-ownership/tasks.md b/openspec/changes/fix-cache-ownership/tasks.md
new file mode 100644
index 00000000..937d8a08
--- /dev/null
+++ b/openspec/changes/fix-cache-ownership/tasks.md
@@ -0,0 +1,52 @@
+## 1. Earlier iteration (notebook-scoped, already on the branch)
+
+- [x] 1.1 Atomic, query-addressed retained FASTA publication with umask-derived permissions.
+- [x] 1.2 Annotation-cache identifier validation before reuse.
+- [x] 1.3 Align embedding sets with `_validate_headers` in the notebook Generate path.
+- [x] 1.4 Consolidated cache regressions in the normal pipeline suite.
+
+## 2. Projection identity
+
+- [x] 2.1 Add a failing regression for a changed matrix, a reordered input, and a grown input under one embedding name.
+- [x] 2.2 Fingerprint each embedding set once per run and include it in the projection cache key, covering the precomputed-MDS branch.
+- [x] 2.3 Drop `refetch_stages={"projections"}` from the notebook and pin that a repeated identical run is a cache hit.
+
+## 3. Embedding identity
+
+- [x] 3.1 Add failing regressions for a cross-backend resume, a changed sequence under an unchanged identifier, a legacy unstamped file, and a disjoint FASTA reusing one cache.
+- [x] 3.2 Record `protspace_backend` / `protspace_model` root attributes and the per-protein residue digest in the shared store.
+- [x] 3.3 Refuse to resume from another producer's file, naming the remedies; adopt and stamp an unstamped file.
+- [x] 3.4 Re-embed proteins whose residues no longer match their digest.
+- [x] 3.5 Return only the requested FASTA's proteins from `embed_fasta`.
+
+## 4. Query FASTA identity
+
+- [x] 4.1 Add a failing regression for `-q A` then `-q B` sharing one output directory.
+- [x] 4.2 Address the retained FASTA by query text in `cli/prepare.py`, and keep the notebook's path inline.
+
+## 5. Annotation identity
+
+- [x] 5.1 Add failing regressions for a partially covering cache, taxonomy reuse by organism, a superset cache surviving a fetch, and a source failing while filling in.
+- [x] 5.2 Fetch each cached source only for the identifiers the cache lacks, merging with cached rows.
+- [x] 5.3 Fill in taxonomy only for organisms the cached lookup does not cover.
+- [x] 5.4 Keep rows outside the current run when the run's columns match the cache's.
+- [x] 5.5 Remove the pipeline's full-rebuild path now that the manager fills in per identifier.
+
+## 6. Atomic publication
+
+- [x] 6.1 Add a failing regression for bundle permissions under a permissive umask.
+- [x] 6.2 Add `data/io/atomic.py` and route the bundle writer, the `stats` rewrites, and the retained FASTA through it.
+
+## 7. Notebook and issue #338
+
+- [x] 7.1 Remove the content-addressed directories, the private imports and their fallbacks, keeping the backend-prefixed HDF5 name and the query-addressed FASTA path inline.
+- [x] 7.2 Name each Generate action's bundle distinctly and report the name.
+- [x] 7.3 State which methods the parameter controls apply to.
+
+## 8. Documentation and gates
+
+- [x] 8.1 Update `docs/guide/fetching-and-caching.md` and `docs/guide/python-cli.md` for per-identifier annotation reuse, projection identity, producer ownership, and the query-addressed FASTA.
+- [x] 8.2 Update the test table and the caching notes in `apps/protspace/CLAUDE.md`.
+- [x] 8.3 Open the follow-up issue for anything deliberately left out, and link it from the PR.
+- [x] 8.4 Run the non-slow Python suite, Ruff, `openspec validate fix-cache-ownership --strict`, and `pnpm precommit`.
+- [x] 8.5 Re-run the seven reproductions from the proposal and record the result.