Claude-Session: https://claude.ai/code/session_01P7zAyQDRnfDCw3o7c3eUAn
---
.../notebooks/ProtSpace_Preparation.ipynb | 61 ++++------
.../src/protspace/data/processors/pipeline.py | 20 ----
apps/protspace/tests/test_backend_switch.py | 4 +-
apps/protspace/tests/test_notebooks.py | 105 +++++++-----------
apps/protspace/tests/test_pipeline_utils.py | 35 ------
5 files changed, 62 insertions(+), 163 deletions(-)
diff --git a/apps/protspace/notebooks/ProtSpace_Preparation.ipynb b/apps/protspace/notebooks/ProtSpace_Preparation.ipynb
index ed9f2742..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,34 +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\n",
- "\n",
- "# Cache-path helpers newer than the *released* protspace this cell installs,\n",
- "# which lags this notebook (served from main) by up to one release. Fall back\n",
- "# to identical copies rather than failing setup with a misleading \"restart the\n",
- "# session\". Kept honest by test_preparation_cache_helper_fallbacks_match_the_package.\n",
- "try:\n",
- " from protspace.data.processors.pipeline import (\n",
- " _embedding_cache_path,\n",
- " _input_cache_dir,\n",
- " _query_fasta_cache_path,\n",
- " )\n",
- "except ImportError: # protspace <= 4.12.2\n",
- " import hashlib as _hashlib\n",
- "\n",
- " def _query_fasta_cache_path(cache_root, query):\n",
- " digest = _hashlib.sha256(query.encode()).hexdigest()[:12]\n",
- " return cache_root / \"queries\" / f\"{digest}.fasta\"\n",
- "\n",
- " def _input_cache_dir(cache_root, input_path):\n",
- " with input_path.open(\"rb\") as source:\n",
- " digest = _hashlib.file_digest(source, \"sha256\").hexdigest()[:12]\n",
- " cache_dir = cache_root / \"inputs\" / digest\n",
- " cache_dir.mkdir(parents=True, exist_ok=True)\n",
- " return cache_dir\n",
- "\n",
- " def _embedding_cache_path(cache_dir, embedder, backend):\n",
- " return cache_dir / f\"{backend}-{embedder}.h5\""
+ " ) from _exc\n"
]
},
{
@@ -545,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",
@@ -677,9 +651,12 @@
"\n",
" out_dir = Path(\"output\")\n",
" out_dir.mkdir(exist_ok=True)\n",
- " cache_root = out_dir / \"tmp\"\n",
- " cache_root.mkdir(exist_ok=True)\n",
- " output_path = out_dir / \"data.parquetbundle\"\n",
+ " cache_dir = out_dir / \"tmp\"\n",
+ " cache_dir.mkdir(exist_ok=True)\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",
@@ -695,7 +672,10 @@
" embs, backend, _emb_cfg = gated\n",
" if inp[\"type\"] == \"query\":\n",
" step_html.value = \"Fetching sequences from UniProt...\"\n",
- " fasta_cache = _query_fasta_cache_path(cache_root, inp[\"query\"])\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",
@@ -710,25 +690,27 @@
" return\n",
" else:\n",
" fasta_path = Path(inp[\"path\"])\n",
- " cache_dir = _input_cache_dir(cache_root, fasta_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=_embedding_cache_path(cache_dir, emb_name, backend),\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",
" else:\n",
" h5_path = Path(inp[\"path\"])\n",
- " cache_dir = _input_cache_dir(cache_root, h5_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",
- " # Build pipeline with non-projection 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",
@@ -743,7 +725,6 @@
" bundled=True,\n",
" keep_tmp=True,\n",
" intermediate_dir=cache_dir,\n",
- " refetch_stages=frozenset({\"projections\"}),\n",
" annotations=ann,\n",
" reducer_params=reducer_params,\n",
" stats=compute_stats_cb.value,\n",
@@ -811,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",
@@ -860,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/data/processors/pipeline.py b/apps/protspace/src/protspace/data/processors/pipeline.py
index b2e957b8..a8398e35 100644
--- a/apps/protspace/src/protspace/data/processors/pipeline.py
+++ b/apps/protspace/src/protspace/data/processors/pipeline.py
@@ -118,21 +118,6 @@ class PipelineConfig:
reducer_params: ReducerParams = field(default_factory=ReducerParams)
-def _query_fasta_cache_path(cache_root: Path, query: str) -> Path:
- """Return the retained FASTA path owned by one exact UniProt query."""
- digest = hashlib.sha256(query.encode()).hexdigest()[:12]
- return cache_root / "queries" / f"{digest}.fasta"
-
-
-def _input_cache_dir(cache_root: Path, input_path: Path) -> Path:
- """Create and return the intermediate directory owned by one input file."""
- with input_path.open("rb") as source:
- digest = hashlib.file_digest(source, "sha256").hexdigest()[:12]
- cache_dir = cache_root / "inputs" / digest
- cache_dir.mkdir(parents=True, exist_ok=True)
- return cache_dir
-
-
def _embedding_fingerprint(emb_set: EmbeddingSet) -> str:
"""Digest exactly what the reducer will be handed: identifiers and matrix.
@@ -151,11 +136,6 @@ def _embedding_fingerprint(emb_set: EmbeddingSet) -> str:
return digest.hexdigest()[:16]
-def _embedding_cache_path(cache_dir: Path, embedder: str, backend: str) -> Path:
- """Return the H5 path owned by one input, model, and producing backend."""
- return cache_dir / f"{backend}-{embedder}.h5"
-
-
# Valid override parameter names (from ReducerParams fields)
_VALID_OVERRIDE_KEYS = {f.name for f in fields(ReducerParams)}
# Field types for coercion
diff --git a/apps/protspace/tests/test_backend_switch.py b/apps/protspace/tests/test_backend_switch.py
index dff61b12..541c6277 100644
--- a/apps/protspace/tests/test_backend_switch.py
+++ b/apps/protspace/tests/test_backend_switch.py
@@ -133,8 +133,6 @@ def test_embed_fasta_unknown_backend_raises(tmp_path):
def test_notebook_embedding_cache_is_owned_by_its_backend(
tmp_path, monkeypatch, second_backend, expected
):
- from protspace.data.processors.pipeline import _embedding_cache_path
-
fasta = tmp_path / "s.fasta"
fasta.write_text(">P12345\nMKVLAAG\n")
@@ -147,7 +145,7 @@ def embed(backend, fill_value):
fasta,
"prot_t5",
backend=backend,
- embedding_cache=_embedding_cache_path(tmp_path, "prot_t5", backend),
+ embedding_cache=tmp_path / f"{backend}-prot_t5.h5",
)
embed("local", 1.0)
diff --git a/apps/protspace/tests/test_notebooks.py b/apps/protspace/tests/test_notebooks.py
index f2caa0b7..ed2fd4e5 100644
--- a/apps/protspace/tests/test_notebooks.py
+++ b/apps/protspace/tests/test_notebooks.py
@@ -199,90 +199,65 @@ def test_notebook_fallback_sets_match_the_package(path: Path):
)
-_CACHE_HELPERS = frozenset(
- {"_embedding_cache_path", "_input_cache_dir", "_query_fasta_cache_path"}
-)
-
-
-def test_preparation_cache_helper_fallbacks_match_the_package(tmp_path: Path):
- """The cache-path helpers' `except ImportError` copies behave like the package.
-
- Cell 1 installs the *released* protspace while the notebook is served from
- `main`, so a helper added alongside a notebook change is missing on first
- run. Imported unguarded, that takes the whole setup cell down with a
- "restart the session" message that cannot help. The guarded copies are what
- runs during that lag, so they are executed here and compared, not trusted.
+@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.
"""
- from protspace.data.processors import pipeline
-
transform = pytest.importorskip(
"IPython.core.inputtransformer2",
reason="IPython is a dev-group dependency (via jupyter)",
).TransformerManager()
- handlers = [
- handler
- for _, source in _code_cells(NOTEBOOK_DIR / "ProtSpace_Preparation.ipynb")
- for node in ast.walk(ast.parse(transform.transform_cell(source)))
- if isinstance(node, ast.Try)
- and _CACHE_HELPERS
- <= {
- alias.asname or alias.name
- for stmt in node.body
- for sub in ast.walk(stmt)
- if isinstance(sub, ast.ImportFrom)
- for alias in sub.names
- }
- for handler in node.handlers
- ]
- assert len(handlers) == 1, (
- "expected exactly one `except ImportError` guarding the cache-path helpers"
- )
- fallback: dict = {}
- exec( # noqa: S102 - the notebook's own source, to test what it runs
- compile(ast.Module(body=handlers[0].body, type_ignores=[]), "cell1", "exec"),
- fallback,
- )
- assert _CACHE_HELPERS <= set(fallback), "a helper has no fallback definition"
-
- fasta = tmp_path / "input.fasta"
- fasta.write_text(">P1\nAAAA\n")
- query = "(family:globin) AND (reviewed:true)"
-
- assert fallback["_query_fasta_cache_path"](
- tmp_path, query
- ) == pipeline._query_fasta_cache_path(tmp_path, query)
- assert fallback["_input_cache_dir"](tmp_path, fasta) == pipeline._input_cache_dir(
- tmp_path, fasta
+ 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."
)
- assert fallback["_embedding_cache_path"](
- tmp_path, "prot_t5", "local"
- ) == pipeline._embedding_cache_path(tmp_path, "prot_t5", "local")
-def test_preparation_generate_refreshes_only_projections():
- """Generate must recompute projections while keeping the other caches.
+def test_preparation_names_each_bundle_distinctly():
+ """A fixed download name is what issue #338 reads as a stale projection.
- The pipeline half of this contract (a `projections` refetch reduces the
- current matrix) is pinned in test_pipeline_utils.py; this pins the notebook
- actually asking for it, and for nothing broader.
+ 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()
- stages = [
- _literal_set(keyword.value)
+ values = [
+ value
for _, source in _code_cells(NOTEBOOK_DIR / "ProtSpace_Preparation.ipynb")
for node in ast.walk(ast.parse(transform.transform_cell(source)))
- if isinstance(node, ast.Call)
- and isinstance(node.func, ast.Name)
- and node.func.id == "PipelineConfig"
- for keyword in node.keywords
- if keyword.arg == "refetch_stages"
+ for name, value in _assignments(node)
+ if name == "output_path"
]
- assert stages == [frozenset({"projections"})]
+
+ 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)
diff --git a/apps/protspace/tests/test_pipeline_utils.py b/apps/protspace/tests/test_pipeline_utils.py
index 500873f8..cf4b7c93 100644
--- a/apps/protspace/tests/test_pipeline_utils.py
+++ b/apps/protspace/tests/test_pipeline_utils.py
@@ -18,9 +18,7 @@
MethodSpec,
PipelineConfig,
ReductionPipeline,
- _input_cache_dir,
_migrate_legacy_ted_labels,
- _query_fasta_cache_path,
_run_with_overridden_config,
disambiguation_suffix,
parse_method_spec,
@@ -1190,39 +1188,6 @@ def boom(data, method, dims):
)
-# ---------------------------------------------------------------------------
-# Preparation notebook cache identity
-# ---------------------------------------------------------------------------
-
-
-def test_query_cache_path_is_owned_by_the_query_text(tmp_path):
- globin = "(family:globin) AND (reviewed:true)"
- phosphatase = "(family:phosphatase) AND (reviewed:true)"
-
- assert _query_fasta_cache_path(tmp_path, globin) == _query_fasta_cache_path(
- tmp_path, globin
- )
- assert _query_fasta_cache_path(tmp_path, globin) != _query_fasta_cache_path(
- tmp_path, phosphatase
- )
- assert _query_fasta_cache_path(tmp_path, globin).parent == tmp_path / "queries"
-
-
-def test_input_cache_dir_is_owned_by_file_content(tmp_path):
- fasta = tmp_path / "input.fasta"
- renamed = tmp_path / "renamed.fasta"
- fasta.write_text(">P1\nAAAA\n")
- renamed.write_text(">P1\nAAAA\n")
- original = _input_cache_dir(tmp_path, fasta)
-
- assert original.is_dir()
- assert _input_cache_dir(tmp_path, renamed) == original
-
- # Same identifier, changed residues: must not share embeddings.
- fasta.write_text(">P1\nCCCC\n")
- assert _input_cache_dir(tmp_path, fasta) != original
-
-
# ---------------------------------------------------------------------------
# Annotation cache identity
# ---------------------------------------------------------------------------
From 07cedb80f662dc3dd5b1023c36a8e753102dac33 Mon Sep 17 00:00:00 2001
From: tsenoner
Date: Fri, 18 Sep 2026 15:18:18 +0200
Subject: [PATCH 12/16] docs: describe cache ownership as the code now decides
it
Annotation reuse is per identifier as well as per column; a retained query
FASTA is addressed by its query; an embedding HDF5 records the backend and
model that wrote it; projections are keyed by the matrix and identifier
order.
Co-Authored-By: Claude Opus 5 (1M context)
Claude-Session: https://claude.ai/code/session_01P7zAyQDRnfDCw3o7c3eUAn
---
apps/protspace/CLAUDE.md | 9 +++---
docs/guide/fetching-and-caching.md | 47 ++++++++++++++++++++++--------
docs/guide/python-cli.md | 23 ++++++++-------
3 files changed, 52 insertions(+), 27 deletions(-)
diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md
index 29888a94..55aa86ce 100644
--- a/apps/protspace/CLAUDE.md
+++ b/apps/protspace/CLAUDE.md
@@ -268,13 +268,13 @@ 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, notebook query/input cache identity, annotation-cache identifier coverage (a partial rebuild replaces a cache for other proteins), `projections` refetch on a changed same-name input, 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) |
@@ -286,7 +286,7 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
| `test_backend_switch.py` | Embedding backend switch: notebook cache ownership/reuse, `resolve_default_backend` (Colab+GPU→local), `embed_fasta` local/biocentral dispatch (short key vs resolved name), `protspace embed --backend` CLI wiring + enum validation + non-positive batch_size rejection |
| `test_local_embedder.py` | Local embedding backend: checkpoint resolution (12 short keys, Synthyra ESM-C), 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 |
+| `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 |
@@ -304,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 |
@@ -311,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`, Preparation Generate requests a `projections`-only refetch, the Preparation cache-path helper fallbacks executed and compared against the package, `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/docs/guide/fetching-and-caching.md b/docs/guide/fetching-and-caching.md
index 0a153bda..5bf48c75 100644
--- a/docs/guide/fetching-and-caching.md
+++ b/docs/guide/fetching-and-caching.md
@@ -39,32 +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.
-First, the cache must **cover your proteins**. If any requested identifier has no row in
-`all_annotations.parquet`, annotations are rebuilt for the current input and the cache is replaced. A
-cache covering _more_ proteins than the run is fine — the extra rows are filtered out later.
-
-Beyond that, the cache is judged **by column, not by row**. ProtSpace compares the columns you asked
-for against the columns already cached:
+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
@@ -99,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:
diff --git a/docs/guide/python-cli.md b/docs/guide/python-cli.md
index 17d74d52..f7bdf23d 100644
--- a/docs/guide/python-cli.md
+++ b/docs/guide/python-cli.md
@@ -350,26 +350,27 @@ 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 may cover more proteins than the current run; the extra rows are filtered out
-later. If any requested identifier is absent from it, annotations are rebuilt for the current input
-and the cache is replaced.
+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
still cached, so one flaky API does not cost an expensive UniProt fetch — unless leaving the failed
source out would overwrite an existing cache with fewer columns, in which case the cache is kept
-untouched (a cache rebuilt because it lacked requested proteins is replaced regardless). Either way the run still returns everything it did retrieve, and the next run fetches
+untouched. Either way the run still returns everything it did retrieve, and the next run fetches
the rest. Transient HTTP failures are retried with backoff first, so this is reserved for a source
that is genuinely unavailable. Use `--refetch annotations` to rewrite the cache regardless — that
is the repair path for a cache already holding empty values. See
From 825e3a1b02ae04ec790c2f6ac6c424b060668fb4 Mon Sep 17 00:00:00 2001
From: tsenoner
Date: Fri, 18 Sep 2026 15:16:56 +0200
Subject: [PATCH 13/16] feat(embed): stamp embedding identity into the HDF5
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
An embedding cache was addressed by its file name alone, so both backends
resumed by identifier: a Local-produced vector satisfied a Biocentral run's
resume check and the two models were silently mixed in one dataset, an
identifier whose sequence changed kept the vector of residues it no longer
has, and `embed_fasta` returned everything the shared cache had ever
accumulated rather than the FASTA it was asked about.
The shared store now records who produced a file (`protspace_backend`,
`protspace_model`) and, per protein, which residues its vector was computed
from (`protspace_sequence_sha256`). Resuming from another producer's file
raises, naming the recorded producer and the three remedies, and leaves that
file untouched; a file recording no producer is adopted and stamped, and a
protein carrying no digest is trusted, so existing caches keep working. A
protein whose residues changed is outstanding work again: it is embedded
again, its dataset and digest are replaced, and `finish_run` reads the
digests so a re-embed that never landed cannot pass for a complete run.
Both backends go through `store.begin_run` / `save_embeddings`, so neither
can drift, and the digests are read in one pass over one open file — at
570K proteins a per-protein reopen would replace a seconds-long resume with
570K file opens.
Co-Authored-By: Claude Opus 5 (1M context)
Claude-Session: https://claude.ai/code/session_01P7zAyQDRnfDCw3o7c3eUAn
---
.../protspace/data/embedding/biocentral.py | 25 ++-
.../src/protspace/data/embedding/local.py | 25 ++-
.../src/protspace/data/embedding/store.py | 159 +++++++++++++++++-
.../src/protspace/data/loaders/fasta.py | 35 +++-
apps/protspace/tests/test_backend_switch.py | 110 ++++++++----
.../tests/test_biocentral_embedder.py | 44 ++++-
.../tests/test_embed_completeness.py | 158 +++++++++++++++++
apps/protspace/tests/test_local_embedder.py | 73 +++++++-
8 files changed, 566 insertions(+), 63 deletions(-)
diff --git a/apps/protspace/src/protspace/data/embedding/biocentral.py b/apps/protspace/src/protspace/data/embedding/biocentral.py
index 0ccb30b4..52cec836 100644
--- a/apps/protspace/src/protspace/data/embedding/biocentral.py
+++ b/apps/protspace/src/protspace/data/embedding/biocentral.py
@@ -15,6 +15,7 @@
# 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.
from protspace.data.embedding.store import ( # noqa: F401
+ begin_run,
finish_run,
load_existing_ids,
save_embeddings,
@@ -148,15 +149,16 @@ 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}
+ # 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)
+ resumed = len(sequences) - len(remaining)
+ if resumed:
+ logger.info("Found %d existing embeddings in %s", resumed, h5_path)
logger.info(
"Remaining sequences to embed: %d (skipped %d)",
len(remaining),
- len(sequences) - len(remaining),
+ resumed,
)
if not remaining:
@@ -239,7 +241,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 +287,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..316f7463 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,12 @@ 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)
+ resumed = len(sequences) - len(remaining)
+ if resumed:
+ logger.info("Resuming: %d already embedded in %s", resumed, h5_path)
# 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 +324,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 +360,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..c9a0a639 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,162 @@
# 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*."""
+ 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 save_embeddings(h5_path: Path, embeddings: dict[str, np.ndarray]) -> None:
- """Append embeddings to an HDF5 file (one dataset per protein)."""
+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
+
+ return {k: v for k, v in sequences.items() if k not in current}
+
+
+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 +205,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 +220,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/loaders/fasta.py b/apps/protspace/src/protspace/data/loaders/fasta.py
index bcbb48ef..222fb2ea 100644
--- a/apps/protspace/src/protspace/data/loaders/fasta.py
+++ b/apps/protspace/src/protspace/data/loaders/fasta.py
@@ -6,7 +6,7 @@
from __future__ import annotations
import logging
-from collections.abc import Iterable
+from collections.abc import Collection, Iterable
from pathlib import Path
from typing import TYPE_CHECKING
@@ -97,7 +97,38 @@ 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),
+ )
+ embedding_set.data = embedding_set.data[keep]
+ embedding_set.headers = [embedding_set.headers[i] for i in keep]
+ return embedding_set
def check_fasta_coverage(
diff --git a/apps/protspace/tests/test_backend_switch.py b/apps/protspace/tests/test_backend_switch.py
index 541c6277..148f2734 100644
--- a/apps/protspace/tests/test_backend_switch.py
+++ b/apps/protspace/tests/test_backend_switch.py
@@ -17,23 +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, fill_value=1.0):
+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.full(4, fill_value, 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
@@ -78,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(
@@ -125,34 +131,65 @@ def test_embed_fasta_unknown_backend_raises(tmp_path):
embed_fasta(fasta, "prot_t5", backend="nope", embedding_cache=tmp_path / "e.h5")
-@pytest.mark.parametrize(
- ("second_backend", "expected"),
- [("biocentral", 2.0), ("local", 1.0)],
- ids=["switched-backend-embeds-again", "same-backend-resumes"],
-)
-def test_notebook_embedding_cache_is_owned_by_its_backend(
- tmp_path, monkeypatch, second_backend, expected
-):
+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 embed(backend, fill_value):
+
+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(
- f"protspace.data.embedding.{backend}.embed_sequences",
- _fake_embed({}, fill_value=fill_value),
- )
- return embed_fasta(
- fasta,
- "prot_t5",
- backend=backend,
- embedding_cache=tmp_path / f"{backend}-prot_t5.h5",
+ "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"),
+ )
- embed("local", 1.0)
- # The second producer writes 2.0, so only reusing the first H5 returns 1.0.
- result = embed(second_backend, 2.0)
+ result = embed_fasta(fasta, "prot_t5", backend="local", embedding_cache=cache)
- assert result.data.tolist() == [[expected] * 4]
+ 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"}
# ---------------------------------------------------------------------------
@@ -165,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(
@@ -197,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(
@@ -352,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_embed_completeness.py b/apps/protspace/tests/test_embed_completeness.py
index edb44404..065b6428 100644
--- a/apps/protspace/tests/test_embed_completeness.py
+++ b/apps/protspace/tests/test_embed_completeness.py
@@ -97,6 +97,164 @@ 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.
+ """
+
+ @staticmethod
+ def _owned(h5_path, *, backend="local", model="prot_t5", ids=("a",)):
+ store.save_embeddings(
+ h5_path,
+ {pid: np.zeros(4, dtype=np.float32) for pid in ids},
+ sequences=dict.fromkeys(ids, "MKV"),
+ backend=backend,
+ model=model,
+ )
+ return h5_path
+
+ def test_another_backend_is_refused_and_named_with_the_remedies(self, tmp_path):
+ h5 = self._owned(tmp_path / "local-prot_t5.h5")
+ 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 = self._owned(tmp_path / "c.h5")
+ 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 = self._owned(tmp_path / "c.h5")
+ 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 = self._owned(tmp_path / "c.h5", ids=("a",))
+ 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."""
+
+ @staticmethod
+ def _save(h5_path, sequences, fill=1.0):
+ store.save_embeddings(
+ h5_path,
+ {pid: np.full(4, fill, dtype=np.float32) for pid in sequences},
+ sequences=sequences,
+ backend="local",
+ model="prot_t5",
+ )
+
+ def test_a_changed_sequence_is_outstanding_again(self, tmp_path):
+ h5 = tmp_path / "c.h5"
+ self._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"
+ self._save(h5, {"a": "MKV"}, fill=1.0)
+ self._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"
+ self._save(h5, {"a": "MKV"}, fill=1.0)
+ self._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"
+ self._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"
+ self._save(h5, {"a": "MKV"}, fill=1.0)
+ self._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)}
+ self._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"):
From 840a896b6c1b23a6e049e33cb3c2409b0e37fd83 Mon Sep 17 00:00:00 2001
From: tsenoner
Date: Fri, 18 Sep 2026 15:21:24 +0200
Subject: [PATCH 14/16] docs: record the embedding-identity tests and tick the
change tasks
Co-Authored-By: Claude Opus 5 (1M context)
Claude-Session: https://claude.ai/code/session_01P7zAyQDRnfDCw3o7c3eUAn
---
apps/protspace/CLAUDE.md | 8 ++--
openspec/changes/fix-cache-ownership/tasks.md | 48 +++++++++----------
2 files changed, 28 insertions(+), 28 deletions(-)
diff --git a/apps/protspace/CLAUDE.md b/apps/protspace/CLAUDE.md
index 55aa86ce..eabf697d 100644
--- a/apps/protspace/CLAUDE.md
+++ b/apps/protspace/CLAUDE.md
@@ -281,10 +281,10 @@ For a live count run `uv run pytest tests/ --collect-only -q`.
| `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: notebook cache ownership/reuse, `resolve_default_backend` (Colab+GPU→local), `embed_fasta` local/biocentral dispatch (short key vs resolved name), `protspace embed --backend` CLI wiring + enum validation + non-positive batch_size rejection |
-| `test_local_embedder.py` | Local embedding backend: checkpoint resolution (12 short keys, Synthyra ESM-C), 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) |
diff --git a/openspec/changes/fix-cache-ownership/tasks.md b/openspec/changes/fix-cache-ownership/tasks.md
index df83a3ff..11337281 100644
--- a/openspec/changes/fix-cache-ownership/tasks.md
+++ b/openspec/changes/fix-cache-ownership/tasks.md
@@ -7,46 +7,46 @@
## 2. Projection identity
-- [ ] 2.1 Add a failing regression for a changed matrix, a reordered input, and a grown input under one embedding name.
-- [ ] 2.2 Fingerprint each embedding set once per run and include it in the projection cache key, covering the precomputed-MDS branch.
-- [ ] 2.3 Drop `refetch_stages={"projections"}` from the notebook and pin that a repeated identical run is a cache hit.
+- [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
-- [ ] 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.
-- [ ] 3.2 Record `protspace_backend` / `protspace_model` root attributes and the per-protein residue digest in the shared store.
-- [ ] 3.3 Refuse to resume from another producer's file, naming the remedies; adopt and stamp an unstamped file.
-- [ ] 3.4 Re-embed proteins whose residues no longer match their digest.
-- [ ] 3.5 Return only the requested FASTA's proteins from `embed_fasta`.
+- [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
-- [ ] 4.1 Add a failing regression for `-q A` then `-q B` sharing one output directory.
-- [ ] 4.2 Address the retained FASTA by query text in `cli/prepare.py`, and keep the notebook's path inline.
+- [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
-- [ ] 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.
-- [ ] 5.2 Fetch each cached source only for the identifiers the cache lacks, merging with cached rows.
-- [ ] 5.3 Fill in taxonomy only for organisms the cached lookup does not cover.
-- [ ] 5.4 Keep rows outside the current run when the run's columns match the cache's.
-- [ ] 5.5 Remove the pipeline's full-rebuild path now that the manager fills in per identifier.
+- [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
-- [ ] 6.1 Add a failing regression for bundle permissions under a permissive umask.
-- [ ] 6.2 Add `data/io/atomic.py` and route the bundle writer, the `stats` rewrites, and the retained FASTA through it.
+- [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
-- [ ] 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.
-- [ ] 7.2 Name each Generate action's bundle distinctly and report the name.
-- [ ] 7.3 State which methods the parameter controls apply to.
+- [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
-- [ ] 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.
-- [ ] 8.2 Update the test table and the caching notes in `apps/protspace/CLAUDE.md`.
-- [ ] 8.3 Open the follow-up issue for anything deliberately left out, and link it from the PR.
+- [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.
- [ ] 8.4 Run the non-slow Python suite, Ruff, `openspec validate fix-cache-ownership --strict`, and `pnpm precommit`.
-- [ ] 8.5 Re-run the seven reproductions from the proposal and record the result.
+- [x] 8.5 Re-run the seven reproductions from the proposal and record the result.
From 37ddb8a915bfecf16509465b3568f623f10ed326 Mon Sep 17 00:00:00 2001
From: tsenoner
Date: Fri, 18 Sep 2026 15:25:02 +0200
Subject: [PATCH 15/16] chore(openspec): tick the verification task
Co-Authored-By: Claude Opus 5 (1M context)
Claude-Session: https://claude.ai/code/session_01P7zAyQDRnfDCw3o7c3eUAn
---
openspec/changes/fix-cache-ownership/tasks.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/openspec/changes/fix-cache-ownership/tasks.md b/openspec/changes/fix-cache-ownership/tasks.md
index 11337281..937d8a08 100644
--- a/openspec/changes/fix-cache-ownership/tasks.md
+++ b/openspec/changes/fix-cache-ownership/tasks.md
@@ -48,5 +48,5 @@
- [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.
-- [ ] 8.4 Run the non-slow Python suite, Ruff, `openspec validate fix-cache-ownership --strict`, and `pnpm precommit`.
+- [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.
From d6a16f324a1ead8c759e3e443cc258378b7dad1e Mon Sep 17 00:00:00 2001
From: tsenoner
Date: Fri, 18 Sep 2026 18:09:54 +0200
Subject: [PATCH 16/16] fix(protspace): repair the cache-ownership work's own
gaps
A cleanup pass and a correctness pass over the changes this branch already
made.
Fixed:
- `_fetch_taxonomy` returned `{}` on any exception, discarding the cached
taxonomy it had just been handed. Only reachable since this branch started
calling it to fill in unseen organisms, where one failed lookup blanked the
taxonomy columns of every protein in the run.
- The annotation parquet, the legacy-TED rewrite and the projection `.npz`
wrote straight to their final path, while this branch made their existence
the next run's trust signal; an interrupted write left a half-file read as a
cache hit. All three now publish through `staged_write`.
- `staged_write` appended a full uuid4 hex -- 38 characters onto a name the
user chose -- so a long bundle name hit NAME_MAX. Truncated to 8.
- `_restrict_to` mutated the EmbeddingSet it was passed; it returns a copy.
- `_embedding_fingerprint` ran even when nothing would be cached, a full pass
over the matrix for a key nobody looks up.
- `_extract_cached_source` and `_extract_cached_taxonomy` walked the cache with
`iterrows()`, now hot on every "added an identifier" run.
Simplified:
- One `uncached_headers` rule, shared by the pipeline's cache-serve gate and
the manager's fill-in. They had separate copies keyed off different columns,
and the lifted rule reports every identifier uncached when no cache holds
rows, so an empty frame cannot serve as a complete one.
- `fingerprint` is a required keyword on the projection-cache methods: an empty
default silently restored the pre-fix key.
- The query FASTA's reuse rule moved next to the artifact it governs, as
`loaders.query.resolve_query_fasta`.
- Resume logging lives in `begin_run` rather than duplicated per backend.
Co-Authored-By: Claude Opus 5 (1M context)
Claude-Session: https://claude.ai/code/session_01P7zAyQDRnfDCw3o7c3eUAn
---
apps/protspace/src/protspace/cli/prepare.py | 27 +-----
.../src/protspace/data/annotations/manager.py | 90 ++++++++++++-------
.../protspace/data/embedding/biocentral.py | 14 +--
.../src/protspace/data/embedding/local.py | 3 -
.../src/protspace/data/embedding/store.py | 22 ++++-
.../protspace/src/protspace/data/io/atomic.py | 6 +-
.../src/protspace/data/loaders/__init__.py | 2 +
.../src/protspace/data/loaders/fasta.py | 12 ++-
.../src/protspace/data/loaders/query.py | 43 ++++++---
.../src/protspace/data/processors/pipeline.py | 65 +++++++++-----
.../tests/test_embed_completeness.py | 61 ++++++-------
apps/protspace/tests/test_query.py | 26 +++---
12 files changed, 211 insertions(+), 160 deletions(-)
diff --git a/apps/protspace/src/protspace/cli/prepare.py b/apps/protspace/src/protspace/cli/prepare.py
index d7cd5a24..260f07cc 100644
--- a/apps/protspace/src/protspace/cli/prepare.py
+++ b/apps/protspace/src/protspace/cli/prepare.py
@@ -237,29 +237,6 @@ def _parse_refetch(raw: str | None) -> frozenset[str]:
return frozenset(stages)
-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."""
- from protspace.data.loaders.query import (
- extract_identifiers_from_fasta,
- query_cache_path,
- query_uniprot,
- )
-
- 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 _embed_all(
embedders: list[str],
fasta_path: Path,
@@ -450,7 +427,9 @@ def prepare(
try:
if query:
- headers, fasta_path = _resolve_query_fasta(query, cache_dir, refetch_stages)
+ 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/data/annotations/manager.py b/apps/protspace/src/protspace/data/annotations/manager.py
index d923c50a..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."""
@@ -146,7 +167,7 @@ def to_pd(self) -> pd.DataFrame:
# 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 = self._uncached_headers()
+ 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."""
@@ -192,13 +213,13 @@ def filled_in(cached_source, fetch):
)
)
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"]
- else self._fetch_taxonomy(
+ self._fetch_taxonomy(
uniprot_annotations, failed_sources, cached=cached_taxonomy
)
- if cached_taxonomy and fill_in
+ if self.sources_to_fetch["taxonomy"] or (cached_taxonomy and fill_in)
else cached_taxonomy
)
interpro_annotations = (
@@ -272,13 +293,6 @@ def filled_in(cached_source, fetch):
return df
- def _uncached_headers(self) -> list[str]:
- """Requested identifiers the cache holds no row for, in request order."""
- if self.cached_data is None or self.cached_data.empty:
- return []
- cached_ids = set(self.cached_data[self.cached_data.columns[0]].astype(str))
- return [h for h in self.headers if str(h) not in cached_ids]
-
def _with_retained_rows(self, df: pd.DataFrame) -> pd.DataFrame:
"""Append cached rows for identifiers outside this run, when they fit.
@@ -392,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]
@@ -506,7 +524,11 @@ def _fetch_taxonomy(
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]
@@ -639,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:
"""
@@ -680,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 52cec836..21de0c12 100644
--- a/apps/protspace/src/protspace/data/embedding/biocentral.py
+++ b/apps/protspace/src/protspace/data/embedding/biocentral.py
@@ -12,8 +12,10 @@
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,
@@ -152,14 +154,6 @@ def embed_sequences(
# 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)
- resumed = len(sequences) - len(remaining)
- if resumed:
- logger.info("Found %d existing embeddings in %s", resumed, h5_path)
- logger.info(
- "Remaining sequences to embed: %d (skipped %d)",
- len(remaining),
- resumed,
- )
if not remaining:
logger.info(
diff --git a/apps/protspace/src/protspace/data/embedding/local.py b/apps/protspace/src/protspace/data/embedding/local.py
index 316f7463..3e3459be 100644
--- a/apps/protspace/src/protspace/data/embedding/local.py
+++ b/apps/protspace/src/protspace/data/embedding/local.py
@@ -286,9 +286,6 @@ def embed_sequences(
# 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)
- resumed = len(sequences) - len(remaining)
- if resumed:
- logger.info("Resuming: %d already embedded in %s", resumed, h5_path)
# 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.
diff --git a/apps/protspace/src/protspace/data/embedding/store.py b/apps/protspace/src/protspace/data/embedding/store.py
index c9a0a639..c6b23a3b 100644
--- a/apps/protspace/src/protspace/data/embedding/store.py
+++ b/apps/protspace/src/protspace/data/embedding/store.py
@@ -60,7 +60,13 @@ def _current_ids(f: h5py.File, sequences: Mapping[str, str]) -> set[str]:
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)
@@ -137,7 +143,19 @@ def begin_run(
f.attrs[_BACKEND_ATTR] = backend
f.attrs[_MODEL_ATTR] = model
- return {k: v for k, v in sequences.items() if k not in current}
+ 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(
diff --git a/apps/protspace/src/protspace/data/io/atomic.py b/apps/protspace/src/protspace/data/io/atomic.py
index cd236217..91f8a035 100644
--- a/apps/protspace/src/protspace/data/io/atomic.py
+++ b/apps/protspace/src/protspace/data/io/atomic.py
@@ -31,7 +31,11 @@ def staged_write(path: Path) -> Iterator[Path]:
"""
path = Path(path)
path.parent.mkdir(parents=True, exist_ok=True)
- staged = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
+ # 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)
diff --git a/apps/protspace/src/protspace/data/loaders/__init__.py b/apps/protspace/src/protspace/data/loaders/__init__.py
index 45b99b96..b310b7c5 100644
--- a/apps/protspace/src/protspace/data/loaders/__init__.py
+++ b/apps/protspace/src/protspace/data/loaders/__init__.py
@@ -16,6 +16,7 @@
extract_identifiers_from_fasta,
query_cache_path,
query_uniprot,
+ resolve_query_fasta,
)
from protspace.data.loaders.similarity import compute_similarity
@@ -31,5 +32,6 @@
"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 222fb2ea..c9f7c8b0 100644
--- a/apps/protspace/src/protspace/data/loaders/fasta.py
+++ b/apps/protspace/src/protspace/data/loaders/fasta.py
@@ -7,6 +7,7 @@
import logging
from collections.abc import Collection, Iterable
+from dataclasses import replace
from pathlib import Path
from typing import TYPE_CHECKING
@@ -126,9 +127,14 @@ def _restrict_to(embedding_set: EmbeddingSet, wanted: Collection[str]) -> Embedd
len(keep),
len(embedding_set.headers),
)
- embedding_set.data = embedding_set.data[keep]
- embedding_set.headers = [embedding_set.headers[i] for i in keep]
- return embedding_set
+ # 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 b0568d0b..9de4d021 100644
--- a/apps/protspace/src/protspace/data/loaders/query.py
+++ b/apps/protspace/src/protspace/data/loaders/query.py
@@ -30,6 +30,29 @@ def query_cache_path(cache_dir: Path, query: str) -> Path:
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,
*,
@@ -71,20 +94,12 @@ def query_uniprot(
temp_file.write(chunk)
pbar.update(len(chunk))
- if save_to is None:
- # Nothing is retained, so the extraction is the caller's own file.
- extracted = temp_gz_file.with_suffix("")
- try:
- identifiers = _extract_fasta(temp_gz_file, extracted)
- except BaseException:
- extracted.unlink(missing_ok=True)
- raise
- else:
- # A retained FASTA's existence is the next run's cache hit, so it may
- # not appear until the whole stream has been decompressed.
- extracted = Path(save_to)
- with staged_write(extracted) as staged:
- identifiers = _extract_fasta(temp_gz_file, staged)
+ # 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)
logger.info(f"Downloaded and extracted {len(identifiers)} sequences")
return identifiers, extracted
diff --git a/apps/protspace/src/protspace/data/processors/pipeline.py b/apps/protspace/src/protspace/data/processors/pipeline.py
index a8398e35..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,
@@ -438,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
@@ -478,24 +480,20 @@ def _fetch_annotations(
intermediate_dir.mkdir(parents=True, exist_ok=True)
cache_path = intermediate_dir / "all_annotations.parquet"
- cached_df = None
- missing_identifiers: set[str] = set()
if cache_path.exists():
cached_df = pd.read_parquet(cache_path)
- missing_identifiers = set(map(str, headers)).difference(
- map(str, cached_df.get("identifier", ()))
- )
+ 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 covers none of %d requested "
- "identifier(s); fetching annotations for them",
+ "Annotation cache lacks %d of %d requested identifier(s); "
+ "fetching annotations for them",
len(missing_identifiers),
+ len(headers),
)
- if cached_df is not None:
# 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.
@@ -504,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:
@@ -747,7 +746,8 @@ def _projection_cache_path(
method: str,
dims: int,
effective_params: dict[str, Any] | None = None,
- fingerprint: str = "",
+ *,
+ fingerprint: str,
) -> Path | None:
cache_dir = self.config.intermediate_dir
if not cache_dir or not self.config.keep_tmp:
@@ -770,10 +770,11 @@ def _load_cached_projection(
dims: int,
effective_params: dict[str, Any] | None = None,
param_suffix: str = "",
- fingerprint: str = "",
+ *,
+ fingerprint: str,
) -> dict[str, Any] | None:
path = self._projection_cache_path(
- embedding_name, method, dims, effective_params, fingerprint
+ embedding_name, method, dims, effective_params, fingerprint=fingerprint
)
if (
path is None
@@ -803,16 +804,21 @@ def _save_projection_cache(
dims: int,
reduction: dict,
effective_params: dict[str, Any] | None = None,
- fingerprint: str = "",
+ *,
+ fingerprint: str,
) -> None:
path = self._projection_cache_path(
- embedding_name, method, dims, effective_params, fingerprint
+ 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 ---
@@ -842,7 +848,14 @@ def add(reduction: dict[str, Any]) -> None:
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.
- fingerprint = _embedding_fingerprint(emb_set)
+ # 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(
@@ -860,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, fingerprint
+ emb_set.name,
+ MDS_NAME,
+ 2,
+ reduction,
+ global_params,
+ fingerprint=fingerprint,
)
computed_count += 1
continue
@@ -884,7 +902,7 @@ def add(reduction: dict[str, Any]) -> None:
dims,
effective_params,
param_suffix,
- fingerprint,
+ fingerprint=fingerprint,
)
if cached:
add(cached)
@@ -903,7 +921,12 @@ def add(reduction: dict[str, Any]) -> None:
)
add(reduction)
self._save_projection_cache(
- emb_set.name, method, dims, reduction, effective_params, fingerprint
+ emb_set.name,
+ method,
+ dims,
+ reduction,
+ effective_params,
+ fingerprint=fingerprint,
)
computed_count += 1
diff --git a/apps/protspace/tests/test_embed_completeness.py b/apps/protspace/tests/test_embed_completeness.py
index 065b6428..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"
@@ -105,19 +119,8 @@ class TestProducerOwnership:
models end up mixed in one dataset.
"""
- @staticmethod
- def _owned(h5_path, *, backend="local", model="prot_t5", ids=("a",)):
- store.save_embeddings(
- h5_path,
- {pid: np.zeros(4, dtype=np.float32) for pid in ids},
- sequences=dict.fromkeys(ids, "MKV"),
- backend=backend,
- model=model,
- )
- return h5_path
-
def test_another_backend_is_refused_and_named_with_the_remedies(self, tmp_path):
- h5 = self._owned(tmp_path / "local-prot_t5.h5")
+ h5 = _save(tmp_path / "local-prot_t5.h5", {"a": "MKV"})
with pytest.raises(ValueError) as exc:
store.begin_run(
h5,
@@ -134,12 +137,12 @@ def test_another_backend_is_refused_and_named_with_the_remedies(self, tmp_path):
def test_another_model_on_the_same_backend_is_refused(self, tmp_path):
"""Two models' vectors are as unmixable as two backends'."""
- h5 = self._owned(tmp_path / "c.h5")
+ 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 = self._owned(tmp_path / "c.h5")
+ 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")
@@ -147,7 +150,7 @@ def test_a_refused_file_is_left_untouched(self, tmp_path):
assert h5.read_bytes() == before, "a refused cache must not be extended"
def test_the_same_producer_resumes(self, tmp_path):
- h5 = self._owned(tmp_path / "c.h5", ids=("a",))
+ h5 = _save(tmp_path / "c.h5", {"a": "MKV"})
outstanding = store.begin_run(
h5, {"a": "MKV", "b": "MKW"}, backend="local", model="prot_t5"
)
@@ -179,19 +182,9 @@ def test_an_adopted_file_is_owned_from_then_on(self, tmp_path):
class TestSequenceIdentity:
"""A vector belongs to the residues it was computed from."""
- @staticmethod
- def _save(h5_path, sequences, fill=1.0):
- store.save_embeddings(
- h5_path,
- {pid: np.full(4, fill, dtype=np.float32) for pid in sequences},
- sequences=sequences,
- backend="local",
- model="prot_t5",
- )
-
def test_a_changed_sequence_is_outstanding_again(self, tmp_path):
h5 = tmp_path / "c.h5"
- self._save(h5, {"a": "MKV", "b": "MKW"})
+ _save(h5, {"a": "MKV", "b": "MKW"})
outstanding = store.begin_run(
h5, {"a": "MKV", "b": "EDITED"}, backend="local", model="prot_t5"
)
@@ -201,8 +194,8 @@ 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"
- self._save(h5, {"a": "MKV"}, fill=1.0)
- self._save(h5, {"a": "EDITED"}, fill=2.0)
+ _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(
@@ -211,8 +204,8 @@ def test_re_embedding_replaces_the_vector_and_the_digest(self, tmp_path):
def test_an_unchanged_sequence_keeps_its_vector(self, tmp_path):
h5 = tmp_path / "c.h5"
- self._save(h5, {"a": "MKV"}, fill=1.0)
- self._save(h5, {"a": "MKV"}, fill=2.0)
+ _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
@@ -228,14 +221,14 @@ def test_finish_run_fails_when_a_stale_protein_was_not_re_embedded(self, tmp_pat
"""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"
- self._save(h5, {"a": "MKV", "b": "MKW"})
+ _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"
- self._save(h5, {"a": "MKV"}, fill=1.0)
- self._save(h5, {"a": "EDITED"}, fill=2.0)
+ _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):
@@ -243,7 +236,7 @@ def test_resume_reads_the_digests_in_one_pass(self, tmp_path, monkeypatch):
not one per protein."""
h5 = tmp_path / "c.h5"
sequences = {f"p{i}": "MKV" for i in range(50)}
- self._save(h5, sequences)
+ _save(h5, sequences)
opens = []
real_file = h5py.File
diff --git a/apps/protspace/tests/test_query.py b/apps/protspace/tests/test_query.py
index d8c048ce..15cf434f 100644
--- a/apps/protspace/tests/test_query.py
+++ b/apps/protspace/tests/test_query.py
@@ -75,14 +75,12 @@ def test_query_uniprot_publishes_fasta_with_process_umask(tmp_path, monkeypatch)
# ---------------------------------------------------------------------------
-# Retained query FASTA ownership (CLI)
+# Retained query FASTA ownership
# ---------------------------------------------------------------------------
def _recording_download(monkeypatch, fasta=">P1\nAAAA\n"):
"""Record every query that reaches query_uniprot, and write its FASTA."""
- from protspace.cli import prepare as prepare_module
-
downloaded = []
def fake_query_uniprot(query, *, save_to=None):
@@ -92,16 +90,14 @@ def fake_query_uniprot(query, *, save_to=None):
return ["P1"], save_to
monkeypatch.setattr(query_module, "query_uniprot", fake_query_uniprot)
- return prepare_module, downloaded
+ return downloaded
def test_a_second_query_does_not_reuse_the_first_query_fasta(tmp_path, monkeypatch):
- prepare_module, downloaded = _recording_download(monkeypatch)
+ downloaded = _recording_download(monkeypatch)
- _, first = prepare_module._resolve_query_fasta(
- "family:globin", tmp_path, frozenset()
- )
- _, second = prepare_module._resolve_query_fasta(
+ _, first = query_module.resolve_query_fasta("family:globin", tmp_path, frozenset())
+ _, second = query_module.resolve_query_fasta(
"family:phosphatase", tmp_path, frozenset()
)
@@ -110,11 +106,11 @@ def test_a_second_query_does_not_reuse_the_first_query_fasta(tmp_path, monkeypat
def test_the_same_query_reuses_its_retained_fasta(tmp_path, monkeypatch):
- prepare_module, downloaded = _recording_download(monkeypatch)
+ downloaded = _recording_download(monkeypatch)
query = "family:globin"
- _, first = prepare_module._resolve_query_fasta(query, tmp_path, frozenset())
- headers, again = prepare_module._resolve_query_fasta(query, tmp_path, frozenset())
+ _, 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
@@ -122,10 +118,10 @@ def test_the_same_query_reuses_its_retained_fasta(tmp_path, monkeypatch):
def test_refetch_query_downloads_again(tmp_path, monkeypatch):
- prepare_module, downloaded = _recording_download(monkeypatch)
+ downloaded = _recording_download(monkeypatch)
query = "family:globin"
- prepare_module._resolve_query_fasta(query, tmp_path, frozenset())
- prepare_module._resolve_query_fasta(query, tmp_path, frozenset({"query"}))
+ query_module.resolve_query_fasta(query, tmp_path, frozenset())
+ query_module.resolve_query_fasta(query, tmp_path, frozenset({"query"}))
assert downloaded == [query, query]