From 149a4819dcd8335ce519afbf9a3f6a467ce69164 Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Tue, 11 Aug 2026 16:04:50 +0100 Subject: [PATCH 1/7] Add pinned PRISM2 slide embeddings --- CHANGELOG.md | 12 +- docs/MODELS.md | 22 ++ docs/MODEL_LICENSES.md | 8 +- docs/SPEC.md | 7 +- docs/adding_a_model.md | 3 + pyproject.toml | 13 + src/raw2features/embedders/fingerprint.py | 72 ++++- src/raw2features/embedders/registry.yaml | 48 +++ src/raw2features/slide_embedders/base.py | 3 + src/raw2features/slide_embedders/encoding.py | 4 + src/raw2features/slide_embedders/prism2.py | 311 +++++++++++++++++++ tests/test_prism2.py | 278 +++++++++++++++++ tests/test_registry_contract.py | 11 + tests/test_registry_revisions_network.py | 46 +++ 14 files changed, 827 insertions(+), 11 deletions(-) create mode 100644 src/raw2features/slide_embedders/prism2.py create mode 100644 tests/test_prism2.py diff --git a/CHANGELOG.md b/CHANGELOG.md index bcb6e80..70d2f73 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,11 +5,19 @@ Notable changes to raw2features, newest first. This project follows ## [0.2.1] - Unreleased -This focused model update adds KRONOS2 to the existing native multiplex path. The store -layout, geometry-only `grid_hash`, and brightfield execution path are unchanged. +This focused model update adds KRONOS2 to the existing native multiplex path and PRISM2 +to slide embedding. The store layout, geometry-only `grid_hash`, and brightfield +extraction path are unchanged. ### Models and multiplex inputs +- Add the gated PRISM2 base (`prism2`, 2560-d) and diagnostic + (`prism2_diagnostic`, 3072-d) slide representations. Both consume the 1280-d CLS + portion already preserved in raw2features' 2560-d `virchow2` patch output. The four + checkpoint shards, shard index, and Phi-3 construction/tokenizer assets are pinned and + SHA-256 verified in local runtime snapshots; the published CUDA BF16 path is retained + with validated Transformers 4.51.3 and FlashAttention 2.8.3 pins. Coordinates are not + used. - Add the gated `kronos2` native multiplex patch encoder: a pinned, SHA-verified marker-aware ViT-B/16 that produces a 768-dimensional CLS vector from ordered marker stacks. The released metadata contains 288 usable marker entries, 268 of them diff --git a/docs/MODELS.md b/docs/MODELS.md index ae8f519..d007c03 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -80,6 +80,8 @@ Each needs the matching patch encoder run first (`-f `). | `meanmax` | any patch encoder | 2 × patch dim | MIT | no | built-in (no weights) | | `titan` | `conch_v1_5` (768) | 768 | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/TITAN) · [GH](https://github.com/mahmoodlab/TITAN) · [paper](https://arxiv.org/abs/2411.19666) | | `prism` | `virchow` (2560) | 1280 | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/paige-ai/Prism) · [paper](https://arxiv.org/abs/2405.10254) | +| `prism2` | `virchow2` (2560; CLS 1280 consumed) | 2560 | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/paige-ai/Prism2) · [paper](https://doi.org/10.1038/s41591-026-04521-4) | +| `prism2_diagnostic` | `virchow2` (2560; CLS 1280 consumed) | 3072 | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/paige-ai/Prism2) · [paper](https://doi.org/10.1038/s41591-026-04521-4) | | `madeleine` | `conch` (512) | 512 | MIT / CC-BY-NC-ND ‡ | yes | [HF](https://huggingface.co/MahmoodLab/madeleine) · [GH](https://github.com/mahmoodlab/MADELEINE) · [paper](https://arxiv.org/abs/2408.02859) | | `feather_conch_v15` | `conch_v1_5` (768) | 512 | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/abmil.base.conch_v15.pc108-24k) · [GH](https://github.com/mahmoodlab/MIL-Lab) · [paper](https://arxiv.org/abs/2506.09022) | | `feather_uni_v2` | `uni2_h` (1536) | 512 | CC-BY-NC-ND-4.0 | yes | [HF](https://huggingface.co/MahmoodLab/abmil.base.uni_v2.pc108-24k) · [GH](https://github.com/mahmoodlab/MIL-Lab) · [paper](https://arxiv.org/abs/2506.09022) | @@ -92,6 +94,26 @@ Each needs the matching patch encoder run first (`-f `). TITAN runs on CONCH v1.5 patch features, **not** UNI: `raw2features embed slide.ome.zarr out -f conch_v1_5 --patch-size 512 -s titan`. +PRISM2 requires Virchow2 CLS tokens from foreground 224 px tiles at 0.5 µm/px. The +registered `virchow2` output is `[CLS, mean-patch]`, so the slide adapter takes its first +1280 values without recomputing patches. `prism2` is the authors' general base +representation; `prism2_diagnostic` is the Phi-3 hidden-state representation intended +for diagnosis-focused tasks. They may be requested separately or together: + +```bash +pip install "raw2features[prism2]" +raw2features embed slide.ome.zarr out -f virchow2 \ + -s prism2 -s prism2_diagnostic --mpp 0.5 --patch-size 224 +``` + +Both PRISM2 outputs require an accepted model gate, a CUDA GPU, and a `flash-attn` build +compatible with the installed Torch/CUDA stack. The `[prism2]` extra pins Transformers +4.51.3 because the released config records that version and its diagnostic helper is not +compatible with Transformers 4.57. It pins the raw2features-validated FlashAttention +2.8.3 build as well. The checkpoint is approximately 17.9 GB. +The base path still constructs the complete released model, so it has the same download +and device-memory requirements as the diagnostic path. + For standalone encoding, `slide-embed` selects the sole grid automatically. In a multi-grid store it can infer the grid from a slide encoder's required patch model when that model occurs once. Use `--grid ` when a model appears in several grids, or for diff --git a/docs/MODEL_LICENSES.md b/docs/MODEL_LICENSES.md index 2afbac8..a8a945b 100644 --- a/docs/MODEL_LICENSES.md +++ b/docs/MODEL_LICENSES.md @@ -54,7 +54,7 @@ table below lists each model's weights licence. | `plip` | MIT (project `setup.py`; no LICENSE file) | no | | `keep` | MIT | no | | `kronos` | CC-BY-NC-ND-4.0 (KRONOS, MahmoodLab; multiplex, non-commercial) | yes | -| `kronos2` | CC-BY-NC-ND-4.0 (KRONOS2, MahmoodLab; gated non-commercial academic terms) | yes (institutional approval) | +| `kronos2` | CC-BY-NC-ND-4.0 (KRONOS2, MahmoodLab; gated non-commercial academic terms) | yes | KRONOS2's optional novel-marker path also downloads `michiyasunaga/BioLinkBERT-large`, whose repository declares Apache-2.0. Its source, @@ -71,6 +71,7 @@ below. `mean` / `max` / `meanmax` are weightless pooling baselines (MIT). |---|---|---| | `titan` | CC-BY-NC-ND-4.0 | yes | | `prism` | CC-BY-NC-ND-4.0 | yes | +| `prism2` / `prism2_diagnostic` | CC-BY-NC-ND-4.0 | yes | | `madeleine` | MIT (HF card) / CC-BY-NC-ND-4.0 (GitHub `LICENSE`) - conflict, verify | yes | | `feather_conch_v15` / `feather_uni_v2` / `feather_uni` | CC-BY-NC-ND-4.0 | yes | | `gigapath_slide` / `gigapath_flash_slide` | Apache-2.0 | yes (terms acceptance) | @@ -79,6 +80,11 @@ below. `mean` / `max` / `meanmax` are weightless pooling baselines (MIT). Check each model's licence on its model card before use. +PRISM2's released checkpoint includes its Phi-3 decoder weights. The loader also uses +the pinned MIT-licensed `microsoft/Phi-3-mini-128k-instruct` config and tokenizer to +construct that decoder; their immutable revision and file hashes are recorded in the +slide-output fingerprint. + ## Quality control / preprocessing QC tools run before or alongside extraction; like the encoders, our **code is an MIT diff --git a/docs/SPEC.md b/docs/SPEC.md index a356886..d8c3fc3 100644 --- a/docs/SPEC.md +++ b/docs/SPEC.md @@ -216,8 +216,11 @@ or incomplete labels. Each `slide/` array carries an analogous `output_fingerprint`, mirrored in `slide_embeddings`. It includes the selected patch array's fingerprint, so replacing -patch features invalidates any derived slide vector, and records the effective fp16/fp32 -slide-forward precision for the resolved device. Fingerprint fields are optional in +patch features invalidates any derived slide vector, and records the effective +fp16/fp32/bf16 slide-forward precision for the resolved device. Sharded slide checkpoints +additionally record a `weights_manifest` binding the index and every shard; construction +dependencies such as a separately sourced config or tokenizer belong in the loader contract. +Fingerprint fields are optional in the 0.1 schema so older stores remain readable and exportable, but current completion checks never reconstruct them from legacy headers: requesting an unfingerprinted patch model recomputes it, and standalone slide encoding asks for that recomputation first. diff --git a/docs/adding_a_model.md b/docs/adding_a_model.md index f98497b..560b181 100644 --- a/docs/adding_a_model.md +++ b/docs/adding_a_model.md @@ -63,3 +63,6 @@ constructor semantics change, bump the patch or slide loader-contract version so finite arrays cannot be mistaken for current outputs. Composite loaders must identify and verify every component they can; any deliberately unpinned component must be marked `experimental` in the registry and documented as outside the pinning guarantee. +For a sharded slide checkpoint, set `weights_filename` / `weights_sha256` to its index +and provide `weights_manifest` with the SHA-256 of the index and every shard. The loader +must verify the whole manifest before constructing the model. diff --git a/pyproject.toml b/pyproject.toml index 1c7eff1..53fb9aa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -93,6 +93,8 @@ max = "raw2features.slide_embedders.pool:MaxPoolSlideEmbedder" meanmax = "raw2features.slide_embedders.pool:MeanMaxPoolSlideEmbedder" titan = "raw2features.slide_embedders.titan:TITANSlideEmbedder" prism = "raw2features.slide_embedders.prism:PrismSlideEmbedder" +prism2 = "raw2features.slide_embedders.prism2:Prism2SlideEmbedder" +prism2_diagnostic = "raw2features.slide_embedders.prism2:Prism2DiagnosticSlideEmbedder" madeleine = "raw2features.slide_embedders.madeleine:MadeleineSlideEmbedder" feather_conch_v15 = "raw2features.slide_embedders.feather:FeatherSlideEmbedder" feather_uni_v2 = "raw2features.slide_embedders.feather:FeatherSlideEmbedder" @@ -173,6 +175,17 @@ grandqc = ["raw2features[torch]", "segmentation-models-pytorch>=0.3", "timm>=1.0 # PRISM slide encoder (Paige): its trust_remote_code modeling imports environs + # sacremoses at module load (the BioGPT decoder path, unused for slide embeddings). Gated. prism = ["raw2features[torch,models]", "environs", "sacremoses"] +# PRISM2's released custom code requires FlashAttention for the perceiver and +# Transformers 4.51.3 for its Phi-3 image decoder. The model card specifies CUDA bf16; +# keep this gated 17.9 GB checkpoint opt-in rather than adding it to [all]. +prism2 = [ + "raw2features[torch,models]", + "torch>=2.3", + # The released config records 4.51.3 and its diagnostic helper calls a private + # Phi-3 method removed by 4.57; keep the forward on the validated upstream stack. + "transformers==4.51.3", + "flash-attn==2.8.3", +] # MADELEINE slide encoder (Mahmood Lab): loads via the authors' non-PyPI `madeleine` git # package (post-install, like conch): # pip install git+https://github.com/mahmoodlab/MADELEINE.git@419287dc60a57296d959840b893481019c4f0d21 diff --git a/src/raw2features/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index 2e84bf8..f5729f8 100644 --- a/src/raw2features/embedders/fingerprint.py +++ b/src/raw2features/embedders/fingerprint.py @@ -678,6 +678,13 @@ def expected_patch_outputs( def _effective_slide_checkpoint(spec) -> dict[str, Any]: repo = _hf_repo(spec.source) + if spec.family == "prism2": + return { + "repo": repo, + "filename": spec.weights_filename, + "files": deepcopy(spec.weights_manifest), + "mechanism": "pinned_sha256_verified_local_snapshot", + } if spec.family == "gigapath_slide": return { "repo": repo, @@ -706,6 +713,24 @@ def _effective_slide_checkpoint(spec) -> dict[str, Any]: } +def _prism2_phi3_contract() -> dict[str, Any]: + from raw2features.slide_embedders.prism2 import ( + PRISM2_PHI3_ARTIFACT_SHA256, + PRISM2_PHI3_LICENSE, + PRISM2_PHI3_REVISION, + PRISM2_PHI3_SOURCE, + ) + + return { + "repo": _hf_repo(PRISM2_PHI3_SOURCE), + "revision": PRISM2_PHI3_REVISION, + "artifacts": deepcopy(PRISM2_PHI3_ARTIFACT_SHA256), + "license": PRISM2_PHI3_LICENSE, + "mechanism": "pinned_sha256_verified_local_config_and_tokenizer", + "weights_source": "included_in_PRISM2_checkpoint", + } + + def _slide_constructor(spec) -> dict[str, Any]: contracts: dict[str, dict[str, Any]] = { "pool": { @@ -733,6 +758,33 @@ def _slide_constructor(spec) -> dict[str, Any]: "features_dtype": "float32", "batched": True, }, + "prism2": { + "entrypoint": "transformers.AutoModel.from_pretrained", + "input": "pinned_sha256_verified_local_snapshot", + "trust_remote_code": True, + "transformers_version": "4.51.3", + "flash_attn_version": "2.8.3", + "stored_patch_encoder": "virchow2", + "stored_patch_dim": 2560, + "model_input_projection": { + "operation": "slice", + "axis": 1, + "start": 0, + "stop": 1280, + "meaning": "Virchow2 CLS token", + }, + "model_context_dim": 1280, + "forward": ( + "get_diagnostic_embedding" + if spec.name == "prism2_diagnostic" + else "get_base_embedding" + ), + "features_dtype": "float32", + "batched": True, + "attention_mask": {"dtype": "int32", "single_slide_values": "ones"}, + "uses_coords": False, + "uses_patch_size_lv0": False, + }, "feather": { "entrypoint": "transformers.AutoModel.from_pretrained", "input": "pinned_local_snapshot", @@ -820,7 +872,10 @@ def _slide_constructor(spec) -> dict[str, Any]: }, }, } - return deepcopy(contracts.get(spec.family, {"entrypoint": spec.family})) + contract = deepcopy(contracts.get(spec.family, {"entrypoint": spec.family})) + if spec.family == "prism2": + contract["phi3_construction_dependency"] = _prism2_phi3_contract() + return contract def slide_output_dim(spec, patch_dim: int) -> int: @@ -836,6 +891,8 @@ def resolved_slide_amp(spec, device: str) -> str: # These loaders follow their model-card examples with fp16 CUDA autocast and # deliberately run fp32 on CPU/MPS. Other slide encoders currently run fp32. + if str(device).startswith("cuda") and spec.family == "prism2": + return "bf16" if str(device).startswith("cuda") and spec.family in { "gigapath_slide", "prism", @@ -862,6 +919,13 @@ def slide_output_fingerprint( "raw2features embed for that patch model before slide encoding." ) output_dim = slide_output_dim(spec, patch_dim) + checkpoint = { + "effective": _effective_slide_checkpoint(spec), + "weights_revision": spec.weights_revision, + "weights_sha256": spec.weights_sha256, + } + if spec.weights_manifest is not None: + checkpoint["weights_manifest"] = deepcopy(spec.weights_manifest) payload = { "kind": "slide_embedding", "model": spec.name, @@ -871,11 +935,7 @@ def slide_output_fingerprint( "source": spec.source, "constructor": _slide_constructor(spec), }, - "checkpoint": { - "effective": _effective_slide_checkpoint(spec), - "weights_revision": spec.weights_revision, - "weights_sha256": spec.weights_sha256, - }, + "checkpoint": checkpoint, "input": { "patch_model": patch_model, "patch_dim": int(patch_dim), diff --git a/src/raw2features/embedders/registry.yaml b/src/raw2features/embedders/registry.yaml index 4b71142..d6e8759 100644 --- a/src/raw2features/embedders/registry.yaml +++ b/src/raw2features/embedders/registry.yaml @@ -1340,6 +1340,54 @@ slide_encoders: [prism] extra (environs, sacremoses, imported by its remote code) + accepted paige-ai/Prism gate. Run with `-f virchow -s prism`. + # PRISM2 exposes distinct base and diagnosis-oriented representations over + # Virchow2 CLS tokens. raw2features' virchow2 array stores CLS first, followed by + # mean-patch pooling, so both adapters consume the first 1280 values losslessly. + prism2: + family: prism2 + source: "hf-hub:paige-ai/Prism2" + embedding_dim: 2560 + patch_encoder: virchow2 + patch_dim: 2560 + gated: true + license: "CC-BY-NC-ND-4.0" + transform_source_url: "https://huggingface.co/paige-ai/Prism2" + weights_revision: "450352d0ddc6b42b21ce20794ce0fbefe6b5a47a" + weights_sha256: "bda099bfdad33ee0e2fe49b6101feb58a0389d8c8a000cc6636c7a24b8c41bbd" + weights_filename: "model.safetensors.index.json" + weights_manifest: &prism2_weights_manifest + model.safetensors.index.json: "bda099bfdad33ee0e2fe49b6101feb58a0389d8c8a000cc6636c7a24b8c41bbd" + model-00001-of-00004.safetensors: "3eec1f46c9dab8ac0dc744ef6f7b2db19028b5dd3bc722499e25a802f9a95556" + model-00002-of-00004.safetensors: "8cba90ae79ea934770673e9461720b0bf37665e0ac00e86124d116d9dfa526a4" + model-00003-of-00004.safetensors: "f899aefe28c8111bbd6a14ef6e295925e1588c88b1741d06935015162f0d4dd5" + model-00004-of-00004.safetensors: "da79790c9a73cf1ce85867aaa57c47cfd8d69e13b57ac86252d24a2687c2998c" + doi: "10.1038/s41591-026-04521-4" + notes: >- + PRISM2 base perceiver embedding. Consumes the 1280-d CLS half of the existing + 2560-d virchow2 output and returns 2560-d. No coordinates. The four checkpoint + shards and the Phi-3 construction assets are pinned and SHA-verified. Needs the + [prism2] extra, CUDA, flash-attn, and accepted paige-ai/Prism2 access. + + prism2_diagnostic: + family: prism2 + source: "hf-hub:paige-ai/Prism2" + embedding_dim: 3072 + patch_encoder: virchow2 + patch_dim: 2560 + gated: true + license: "CC-BY-NC-ND-4.0" + transform_source_url: "https://huggingface.co/paige-ai/Prism2" + weights_revision: "450352d0ddc6b42b21ce20794ce0fbefe6b5a47a" + weights_sha256: "bda099bfdad33ee0e2fe49b6101feb58a0389d8c8a000cc6636c7a24b8c41bbd" + weights_filename: "model.safetensors.index.json" + weights_manifest: *prism2_weights_manifest + doi: "10.1038/s41591-026-04521-4" + notes: >- + PRISM2 diagnostic Phi-3 hidden-state embedding. Consumes the 1280-d CLS half of + the existing 2560-d virchow2 output and returns 3072-d. No coordinates. Uses the + same verified checkpoint and runtime requirements as prism2. Run separately or + request both names when both representations are needed. + # FEATHER (Mahmood Lab MIL-Lab, ICML 2025) - supervised ABMIL slide aggregators, one # per patch encoder. Pure-HF load; slide vector = forward_features pre-head pooling # (512-d). See slide_embedders/feather.py. diff --git a/src/raw2features/slide_embedders/base.py b/src/raw2features/slide_embedders/base.py index 98154a9..57b20c9 100644 --- a/src/raw2features/slide_embedders/base.py +++ b/src/raw2features/slide_embedders/base.py @@ -49,6 +49,9 @@ class SlideModelSpec: weights_revision: str | None = None # Exact artifact identified by ``weights_sha256``; None for weight-free pools. weights_filename: str | None = None + # Sharded checkpoints bind every downloaded file as well as the primary artifact. + # Single-file models leave this unset. + weights_manifest: dict[str, str] | None = None # Upstream architecture where one loader family serves several compatible sizes. # ``patch_dim`` supplies the constructor's input-channel width. architecture: str | None = None diff --git a/src/raw2features/slide_embedders/encoding.py b/src/raw2features/slide_embedders/encoding.py index def9a85..f88185e 100644 --- a/src/raw2features/slide_embedders/encoding.py +++ b/src/raw2features/slide_embedders/encoding.py @@ -367,6 +367,10 @@ def encode_slide_embedding( "computed_utc": now_utc_iso(), "raw2features_version": __version__, } + if slide_embedder.spec.weights_manifest is not None: + provenance["weights_manifest"] = dict( + slide_embedder.spec.weights_manifest + ) finally: slide_embedder.unload() diff --git a/src/raw2features/slide_embedders/prism2.py b/src/raw2features/slide_embedders/prism2.py new file mode 100644 index 0000000..c250a4a --- /dev/null +++ b/src/raw2features/slide_embedders/prism2.py @@ -0,0 +1,311 @@ +"""PRISM2 slide encoders (paige-ai/Prism2, CC-BY-NC-ND-4.0). + +The gated model consumes 1280-d Virchow2 CLS tokens from 224 px tissue tiles at +0.5 micrometres per pixel. raw2features already preserves those tokens as the first +half of its 2560-d ``virchow2`` output, so the adapter slices them losslessly instead +of requiring a duplicate patch extraction. The base and diagnostic representations +are separate outputs because the authors give them different downstream roles. +""" + +from __future__ import annotations + +from importlib import metadata as importlib_metadata +from pathlib import Path +from typing import Any + +import numpy as np + +from raw2features.core.plugins import register +from raw2features.embedders._hub import ( + download_pinned_hf_snapshot, + pinned_model_cache_dir, + verify_sha256, +) + +from .base import SlideEmbedder, SlideModelSpec + +PRISM2_SOURCE = "hf-hub:paige-ai/Prism2" +PRISM2_REVISION = "450352d0ddc6b42b21ce20794ce0fbefe6b5a47a" +PRISM2_TRANSFORMERS_VERSION = "4.51.3" +PRISM2_FLASH_ATTN_VERSION = "2.8.3" +PRISM2_CHECKPOINT_SHA256 = { + "model.safetensors.index.json": ( + "bda099bfdad33ee0e2fe49b6101feb58a0389d8c8a000cc6636c7a24b8c41bbd" + ), + "model-00001-of-00004.safetensors": ( + "3eec1f46c9dab8ac0dc744ef6f7b2db19028b5dd3bc722499e25a802f9a95556" + ), + "model-00002-of-00004.safetensors": ( + "8cba90ae79ea934770673e9461720b0bf37665e0ac00e86124d116d9dfa526a4" + ), + "model-00003-of-00004.safetensors": ( + "f899aefe28c8111bbd6a14ef6e295925e1588c88b1741d06935015162f0d4dd5" + ), + "model-00004-of-00004.safetensors": ( + "da79790c9a73cf1ce85867aaa57c47cfd8d69e13b57ac86252d24a2687c2998c" + ), +} +PRISM2_CODE_FILES = ( + "config.json", + "configuration_prism2.py", + "modeling_prism2.py", + "processing_prism2.py", +) +PRISM2_SNAPSHOT_ALLOW_PATTERNS = ( + *PRISM2_CODE_FILES, + *PRISM2_CHECKPOINT_SHA256, +) + +PRISM2_PHI3_SOURCE = "hf-hub:microsoft/Phi-3-mini-128k-instruct" +PRISM2_PHI3_REVISION = "f3c06aed622e14ca0abf5115094e4fc9a9948f36" +PRISM2_PHI3_LICENSE = "MIT" +PRISM2_PHI3_ARTIFACT_SHA256 = { + "added_tokens.json": ( + "f8e5a880c6c563a126d9efacb654e105fc4b66f21e5a540f04436e9487a955fb" + ), + "config.json": ("af70441b867458d18788c564ea69894adddaa4de47ab6534391604fc4893f03b"), + "special_tokens_map.json": ( + "810adc6e6c6ef2f56c285ef930d243358a3a9f05e36a01c5a10bafc6fac4609b" + ), + "tokenizer.json": ( + "072ab882d6c7192a42f78790945d16c064691321a73251a4b18f6a380f0fbe39" + ), + "tokenizer.model": ( + "9e556afd44213b6bd1be2b850ebbbd98f5481437a8021afaf58ee7fb1818d347" + ), + "tokenizer_config.json": ( + "aaa87217a0f61c684cdc8703d3d4030a1f5b1077183610b61a14d7f28addbb58" + ), +} + +_BASE_SPEC = SlideModelSpec( + name="prism2", + family="prism2", + source=PRISM2_SOURCE, + embedding_dim=2560, + patch_encoder="virchow2", + patch_dim=2560, + gated=True, + license="CC-BY-NC-ND-4.0", + transform_source_url="https://huggingface.co/paige-ai/Prism2", + doi="10.1038/s41591-026-04521-4", + weights_sha256=PRISM2_CHECKPOINT_SHA256["model.safetensors.index.json"], + weights_revision=PRISM2_REVISION, + weights_filename="model.safetensors.index.json", + weights_manifest=PRISM2_CHECKPOINT_SHA256, + notes=( + "Base perceiver embedding over the 1280-d CLS half of raw2features' " + "Virchow2 output. CUDA and flash-attn are required." + ), +) + +_DIAGNOSTIC_SPEC = SlideModelSpec( + name="prism2_diagnostic", + family="prism2", + source=PRISM2_SOURCE, + embedding_dim=3072, + patch_encoder="virchow2", + patch_dim=2560, + gated=True, + license="CC-BY-NC-ND-4.0", + transform_source_url="https://huggingface.co/paige-ai/Prism2", + doi="10.1038/s41591-026-04521-4", + weights_sha256=PRISM2_CHECKPOINT_SHA256["model.safetensors.index.json"], + weights_revision=PRISM2_REVISION, + weights_filename="model.safetensors.index.json", + weights_manifest=PRISM2_CHECKPOINT_SHA256, + notes=( + "Diagnostic Phi-3 hidden-state embedding over the 1280-d CLS half of " + "raw2features' Virchow2 output. CUDA and flash-attn are required." + ), +) + + +def _assert_snapshot_files( + snapshot: str, + expected: tuple[str, ...], + *, + what: str, +) -> None: + """Reject a stale managed directory before its custom code can execute.""" + + root = Path(snapshot) + allowed = {*expected, ".cache", "__pycache__"} + unexpected = sorted( + path.name for path in root.iterdir() if path.name not in allowed + ) + missing = sorted(name for name in expected if not (root / name).is_file()) + if unexpected or missing: + details = [] + if missing: + details.append(f"missing={missing}") + if unexpected: + details.append(f"unexpected={unexpected}") + joined = "; ".join(details) + raise ValueError(f"{what}: invalid pinned runtime snapshot ({joined})") + + +def _verified_snapshot( + source: str, + revision: str, + artifacts: dict[str, str], + *, + extra_files: tuple[str, ...] = (), + what: str, +) -> str: + expected = (*extra_files, *artifacts) + snapshot = download_pinned_hf_snapshot( + source, + revision, + allow_patterns=expected, + local_dir=pinned_model_cache_dir(source, revision), + ) + _assert_snapshot_files(snapshot, expected, what=what) + for filename, sha256 in artifacts.items(): + verify_sha256(str(Path(snapshot) / filename), sha256, what=f"{what}:{filename}") + return snapshot + + +def _require_prism2_runtime() -> None: + required = { + "transformers": PRISM2_TRANSFORMERS_VERSION, + "flash-attn": PRISM2_FLASH_ATTN_VERSION, + } + mismatches = [] + for distribution, expected in required.items(): + try: + installed = importlib_metadata.version(distribution) + except importlib_metadata.PackageNotFoundError: + installed = None + if installed != expected: + found = installed or "missing" + mismatches.append(f"{distribution}=={expected} (found {found})") + if mismatches: + raise RuntimeError( + "PRISM2 requires its validated runtime; install raw2features[prism2]: " + + ", ".join(mismatches) + ) + + +class _Prism2SlideEmbedder(SlideEmbedder): + def __init__(self, spec: SlideModelSpec) -> None: + super().__init__(spec) + self._model = None + self._device = "cpu" + + def load(self, device: str = "cuda", dtype=None) -> _Prism2SlideEmbedder: + if not str(device).startswith("cuda"): + raise ValueError("PRISM2 requires a CUDA GPU and flash-attn==2.8.3") + _require_prism2_runtime() + del dtype # the published path keeps fp32 weights and uses bf16 autocast + + if self.spec.weights_manifest != PRISM2_CHECKPOINT_SHA256: + raise ValueError( + f"{self.spec.name}: registry weights_manifest does not match the " + "pinned PRISM2 checkpoint" + ) + prism_snapshot = _verified_snapshot( + self.spec.source, + str(self.spec.weights_revision), + dict(self.spec.weights_manifest), + extra_files=PRISM2_CODE_FILES, + what=self.spec.name, + ) + phi3_snapshot = _verified_snapshot( + PRISM2_PHI3_SOURCE, + PRISM2_PHI3_REVISION, + PRISM2_PHI3_ARTIFACT_SHA256, + what=f"{self.spec.name}:Phi-3", + ) + + from transformers import AutoConfig, AutoModel + + config = AutoConfig.from_pretrained( + prism_snapshot, + trust_remote_code=True, + local_files_only=True, + ) + # Pinned upstream code constructs Phi-3 from this field. Point it at the + # verified local config/tokenizer snapshot instead of mutable Hub main. + config.text_decoder_model_id = phi3_snapshot + model = AutoModel.from_pretrained( + prism_snapshot, + config=config, + trust_remote_code=True, + local_files_only=True, + torch_dtype="auto", + ) + model.eval().to(device) + self._model = model + self._device = str(device) + return self + + def encode( + self, + features: np.ndarray, + coords: np.ndarray | None = None, # noqa: ARG002 - PRISM2 is non-spatial + patch_size_lv0: int | None = None, # noqa: ARG002 - fixed upstream tile scale + ) -> np.ndarray: + import torch + + if self._model is None: + raise RuntimeError("call load() before encode()") + if features.ndim != 2 or int(features.shape[1]) != 2560: + raise ValueError( + "PRISM2 requires raw2features virchow2 features shaped (N, 2560); " + f"got {tuple(features.shape)}" + ) + + # virchow2 is stored as [CLS(1280), mean-patch(1280)]. PRISM2 consumes CLS. + cls = np.ascontiguousarray(features[:, :1280], dtype=np.float32) + tile_embeddings = torch.from_numpy(cls).unsqueeze(0).to(self._device) + attention_mask = torch.ones( + tile_embeddings.shape[:2], + dtype=torch.int32, + device=self._device, + ) + autocast = ( + torch.autocast("cuda", torch.bfloat16) + if self._device.startswith("cuda") + else torch.autocast("cpu", enabled=False) + ) + forwards: dict[str, Any] = { + "prism2": self._model.get_base_embedding, + "prism2_diagnostic": self._model.get_diagnostic_embedding, + } + try: + forward = forwards[self.spec.name] + except KeyError as exc: # pragma: no cover - registry construction guards this + raise ValueError(f"unknown PRISM2 output {self.spec.name!r}") from exc + with autocast, torch.inference_mode(): + vector = forward( + tile_embeddings=tile_embeddings, + attention_mask=attention_mask, + ) + return vector.reshape(-1).float().cpu().numpy() + + def unload(self) -> None: + import torch + + if self._model is not None: + self._model.cpu() + del self._model + self._model = None + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +@register("slide_embedders", "prism2") +class Prism2SlideEmbedder(_Prism2SlideEmbedder): + """PRISM2's general 2560-dimensional base representation.""" + + def __init__(self) -> None: + super().__init__(_BASE_SPEC) + + +@register("slide_embedders", "prism2_diagnostic") +class Prism2DiagnosticSlideEmbedder(_Prism2SlideEmbedder): + """PRISM2's diagnosis-oriented 3072-dimensional representation.""" + + def __init__(self) -> None: + super().__init__(_DIAGNOSTIC_SPEC) diff --git a/tests/test_prism2.py b/tests/test_prism2.py new file mode 100644 index 0000000..534f4ee --- /dev/null +++ b/tests/test_prism2.py @@ -0,0 +1,278 @@ +"""PRISM2 registry, loader, and forward-contract tests.""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +from raw2features.embedders.fingerprint import ( + patch_output_fingerprint, + resolved_slide_amp, + slide_output_fingerprint, +) +from raw2features.embedders.model_registry import get_spec +from raw2features.slide_embedders.model_registry import get_slide_spec +from raw2features.slide_embedders.prism2 import ( + PRISM2_CHECKPOINT_SHA256, + PRISM2_CODE_FILES, + PRISM2_FLASH_ATTN_VERSION, + PRISM2_PHI3_ARTIFACT_SHA256, + PRISM2_PHI3_REVISION, + PRISM2_PHI3_SOURCE, + PRISM2_REVISION, + PRISM2_TRANSFORMERS_VERSION, + Prism2DiagnosticSlideEmbedder, + Prism2SlideEmbedder, + _assert_snapshot_files, + _require_prism2_runtime, +) + + +def test_prism2_registry_records_both_published_embeddings(): + base = get_slide_spec("prism2") + diagnostic = get_slide_spec("prism2_diagnostic") + + for spec in (base, diagnostic): + assert spec.family == "prism2" + assert spec.patch_encoder == "virchow2" + # Stored virchow2 is [CLS(1280), mean-patch(1280)]; the loader slices CLS. + assert spec.patch_dim == 2560 + assert spec.gated is True + assert spec.license == "CC-BY-NC-ND-4.0" + assert spec.weights_revision == PRISM2_REVISION + assert spec.weights_filename == "model.safetensors.index.json" + assert spec.weights_manifest == PRISM2_CHECKPOINT_SHA256 + assert ( + spec.weights_sha256 + == PRISM2_CHECKPOINT_SHA256["model.safetensors.index.json"] + ) + assert spec.doi == "10.1038/s41591-026-04521-4" + assert base.embedding_dim == 2560 + assert diagnostic.embedding_dim == 3072 + + +@pytest.mark.parametrize( + ("embedder_cls", "method", "dimension"), + [ + (Prism2SlideEmbedder, "base", 2560), + (Prism2DiagnosticSlideEmbedder, "diagnostic", 3072), + ], +) +def test_prism2_forward_slices_preserved_virchow2_cls(embedder_cls, method, dimension): + torch = pytest.importorskip("torch") + calls = [] + + class Model: + def _forward(self, name, tile_embeddings, attention_mask): + calls.append((name, tile_embeddings.cpu(), attention_mask.cpu())) + return torch.full((1, dimension), 3.0) + + def get_base_embedding(self, tile_embeddings, attention_mask): + return self._forward("base", tile_embeddings, attention_mask) + + def get_diagnostic_embedding(self, tile_embeddings, attention_mask): + return self._forward("diagnostic", tile_embeddings, attention_mask) + + features = np.empty((3, 2560), dtype=np.float32) + features[:, :1280] = np.arange(3, dtype=np.float32)[:, None] + features[:, 1280:] = 1000.0 + embedder = embedder_cls() + embedder._model = Model() + embedder._device = "cpu" + + output = embedder.encode(features) + + assert output.shape == (dimension,) + assert np.isfinite(output).all() + name, received, mask = calls.pop() + assert name == method + np.testing.assert_array_equal(received.numpy(), features[:, :1280][None]) + assert received.shape == (1, 3, 1280) + assert mask.dtype == torch.int32 + np.testing.assert_array_equal(mask.numpy(), np.ones((1, 3), dtype=np.int32)) + + +def test_prism2_rejects_an_incompatible_patch_array(): + pytest.importorskip("torch") + embedder = Prism2SlideEmbedder() + embedder._model = SimpleNamespace() + with pytest.raises(ValueError, match=r"virchow2 features shaped \(N, 2560\)"): + embedder.encode(np.zeros((4, 1280), dtype=np.float32)) + + +def test_prism2_load_uses_only_verified_local_snapshots(monkeypatch): + calls = [] + + def verified(source, revision, artifacts, *, extra_files=(), what): + calls.append((source, revision, dict(artifacts), tuple(extra_files), what)) + return "/verified/phi3" if "Phi-3" in what else "/verified/prism2" + + config = SimpleNamespace(text_decoder_model_id="mutable-upstream-name") + + class AutoConfig: + @staticmethod + def from_pretrained(path, **kwargs): + assert path == "/verified/prism2" + assert kwargs == {"trust_remote_code": True, "local_files_only": True} + return config + + class Model: + def eval(self): + return self + + def to(self, device): + assert device == "cuda:3" + return self + + class AutoModel: + @staticmethod + def from_pretrained(path, **kwargs): + assert path == "/verified/prism2" + assert kwargs["config"] is config + assert kwargs["trust_remote_code"] is True + assert kwargs["local_files_only"] is True + assert kwargs["torch_dtype"] == "auto" + return Model() + + import raw2features.slide_embedders.prism2 as module + + monkeypatch.setattr(module, "_verified_snapshot", verified) + monkeypatch.setattr(module, "_require_prism2_runtime", lambda: None) + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace(AutoConfig=AutoConfig, AutoModel=AutoModel), + ) + + loaded = Prism2SlideEmbedder().load(device="cuda:3") + + assert loaded._device == "cuda:3" + assert config.text_decoder_model_id == "/verified/phi3" + assert calls[0][1:] == ( + PRISM2_REVISION, + PRISM2_CHECKPOINT_SHA256, + PRISM2_CODE_FILES, + "prism2", + ) + assert calls[1][0:4] == ( + PRISM2_PHI3_SOURCE, + PRISM2_PHI3_REVISION, + PRISM2_PHI3_ARTIFACT_SHA256, + (), + ) + + +def test_prism2_rejects_cpu_before_downloading(monkeypatch): + import raw2features.slide_embedders.prism2 as module + + monkeypatch.setattr( + module, + "_verified_snapshot", + lambda *args, **kwargs: pytest.fail("download attempted"), + ) + with pytest.raises(ValueError, match="requires a CUDA GPU"): + Prism2SlideEmbedder().load(device="cpu") + + +def test_prism2_runtime_pin_fails_clearly(monkeypatch): + import raw2features.slide_embedders.prism2 as module + + versions = {"transformers": "4.57.6", "flash-attn": None} + + def version(distribution): + value = versions[distribution] + if value is None: + raise module.importlib_metadata.PackageNotFoundError(distribution) + return value + + monkeypatch.setattr(module.importlib_metadata, "version", version) + with pytest.raises( + RuntimeError, + match=( + r"transformers==4\.51\.3 \(found 4\.57\.6\).*" + r"flash-attn==2\.8\.3 \(found missing\)" + ), + ): + _require_prism2_runtime() + + +def test_prism2_snapshot_allowlist_fails_closed(tmp_path): + for filename in ("config.json", "model.safetensors.index.json"): + (tmp_path / filename).write_bytes(b"test") + _assert_snapshot_files( + str(tmp_path), + ("config.json", "model.safetensors.index.json"), + what="prism2", + ) + + (tmp_path / "stale.py").write_bytes(b"unexpected") + with pytest.raises(ValueError, match="unexpected=.*stale.py"): + _assert_snapshot_files( + str(tmp_path), + ("config.json", "model.safetensors.index.json"), + what="prism2", + ) + + +def test_prism2_fingerprint_binds_cls_projection_shards_and_phi3(): + patch = patch_output_fingerprint(get_spec("virchow2"), "fp16") + base_spec = get_slide_spec("prism2") + diagnostic_spec = get_slide_spec("prism2_diagnostic") + base = slide_output_fingerprint( + base_spec, + patch_model="virchow2", + patch_output_fingerprint=patch, + patch_dim=2560, + resolved_amp=resolved_slide_amp(base_spec, "cuda:0"), + ) + diagnostic = slide_output_fingerprint( + diagnostic_spec, + patch_model="virchow2", + patch_output_fingerprint=patch, + patch_dim=2560, + resolved_amp=resolved_slide_amp(diagnostic_spec, "cuda:0"), + ) + + assert base["payload"]["checkpoint"]["weights_manifest"] == ( + PRISM2_CHECKPOINT_SHA256 + ) + constructor = base["payload"]["loader"]["constructor"] + assert constructor["model_input_projection"] == { + "operation": "slice", + "axis": 1, + "start": 0, + "stop": 1280, + "meaning": "Virchow2 CLS token", + } + assert constructor["forward"] == "get_base_embedding" + assert constructor["transformers_version"] == PRISM2_TRANSFORMERS_VERSION + assert constructor["flash_attn_version"] == PRISM2_FLASH_ATTN_VERSION + phi3 = constructor["phi3_construction_dependency"] + assert phi3["repo"] == "microsoft/Phi-3-mini-128k-instruct" + assert phi3["revision"] == PRISM2_PHI3_REVISION + assert phi3["artifacts"] == PRISM2_PHI3_ARTIFACT_SHA256 + assert base["payload"]["output"]["resolved_amp"] == "bf16" + assert diagnostic["payload"]["loader"]["constructor"]["forward"] == ( + "get_diagnostic_embedding" + ) + assert diagnostic["payload"]["output"]["embedding_dim"] == 3072 + assert base["digest"] != diagnostic["digest"] + + +def test_prism2_fields_do_not_change_existing_prism_fingerprint_shape(): + patch = patch_output_fingerprint(get_spec("virchow"), "fp16") + prism = get_slide_spec("prism") + fingerprint = slide_output_fingerprint( + prism, + patch_model="virchow", + patch_output_fingerprint=patch, + patch_dim=2560, + resolved_amp="fp16", + ) + assert "weights_manifest" not in fingerprint["payload"]["checkpoint"] + assert fingerprint["payload"]["loader"]["constructor"]["forward"] == ( + "slide_representations" + ) diff --git a/tests/test_registry_contract.py b/tests/test_registry_contract.py index 2ee5227..a7535ad 100644 --- a/tests/test_registry_contract.py +++ b/tests/test_registry_contract.py @@ -80,6 +80,17 @@ def test_every_slide_spec_is_complete(): if spec.family != "pool": assert spec.doi and spec.doi.startswith("10."), f"{name}: needs a DOI" assert spec.weights_filename, f"{name}: needs the exact weights_filename" + if spec.weights_manifest is not None: + assert spec.weights_filename in spec.weights_manifest + assert spec.weights_manifest[spec.weights_filename] == ( + spec.weights_sha256 + ) + assert all( + len(filename) > 0 + and len(digest) == 64 + and digest == digest.lower() + for filename, digest in spec.weights_manifest.items() + ), f"{name}: invalid weights_manifest" def test_optional_fields_round_trip_through_loader(): diff --git a/tests/test_registry_revisions_network.py b/tests/test_registry_revisions_network.py index ab99c3e..93d8209 100644 --- a/tests/test_registry_revisions_network.py +++ b/tests/test_registry_revisions_network.py @@ -28,6 +28,11 @@ ) from raw2features.embedders.seal_embedder import _SEAL_REPO from raw2features.slide_embedders.model_registry import load_slide_registry +from raw2features.slide_embedders.prism2 import ( + PRISM2_PHI3_ARTIFACT_SHA256, + PRISM2_PHI3_REVISION, + PRISM2_PHI3_SOURCE, +) _FULL_COMMIT = re.compile(r"[0-9a-f]{40}") _SHA256 = re.compile(r"[0-9a-f]{64}") @@ -238,6 +243,47 @@ def test_new_model_artifact_checksum_matches_huggingface_metadata(pin: _HubArtif ) +@pytest.mark.network +def test_prism2_shard_manifest_matches_huggingface(): + """Every file used to construct the sharded checkpoint matches its pin.""" + + spec = load_slide_registry()["prism2"] + assert spec.weights_manifest + huggingface_hub = pytest.importorskip("huggingface_hub") + info = huggingface_hub.HfApi().model_info( + repo_id=_hf_repo(spec) or "", + revision=spec.weights_revision, + timeout=30, + files_metadata=True, + token=os.environ.get("HF_TOKEN") or False, + ) + siblings = {sibling.rfilename: sibling for sibling in (info.siblings or ())} + for filename, expected in spec.weights_manifest.items(): + assert filename in siblings, filename + lfs = siblings[filename].lfs + if lfs is None: + path = download_pinned_hf_file( + spec.source, + filename, + spec.weights_revision, + ) + verify_sha256(path, expected, what=f"prism2:{filename}") + continue + sha256 = lfs.get("sha256") if isinstance(lfs, Mapping) else lfs.sha256 + assert sha256 == expected, filename + + +@pytest.mark.network +def test_prism2_phi3_construction_artifacts_match_pins(): + for filename, expected in PRISM2_PHI3_ARTIFACT_SHA256.items(): + path = download_pinned_hf_file( + PRISM2_PHI3_SOURCE, + filename, + PRISM2_PHI3_REVISION, + ) + verify_sha256(path, expected, what=f"prism2:Phi-3:{filename}") + + @pytest.mark.network def test_hipt_official_git_lfs_pointer_matches_registry(): """The lightweight official pointer proves the 704 MB asset's immutable hash.""" From 1da9a8beac381b42ad769ded54233ad566590884 Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Tue, 11 Aug 2026 18:27:36 +0100 Subject: [PATCH 2/7] Share PRISM2 and KRONOS2 runtime --- CHANGELOG.md | 5 +- docs/MODELS.md | 8 ++- pyproject.toml | 8 +-- src/raw2features/embedders/fingerprint.py | 10 ++- src/raw2features/slide_embedders/prism2.py | 75 +++++++++++++++++++++- tests/test_prism2.py | 49 +++++++++++++- uv.lock | 27 +++++++- 7 files changed, 169 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 70d2f73..59f2d0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,8 +16,9 @@ extraction path are unchanged. portion already preserved in raw2features' 2560-d `virchow2` patch output. The four checkpoint shards, shard index, and Phi-3 construction/tokenizer assets are pinned and SHA-256 verified in local runtime snapshots; the published CUDA BF16 path is retained - with validated Transformers 4.51.3 and FlashAttention 2.8.3 pins. Coordinates are not - used. + with validated Transformers 4.56.0 and FlashAttention 2.8.3 pins. A fingerprinted + compatibility helper preserves the released Phi-3 causal-mask construction. Coordinates + are not used. - Add the gated `kronos2` native multiplex patch encoder: a pinned, SHA-verified marker-aware ViT-B/16 that produces a 768-dimensional CLS vector from ordered marker stacks. The released metadata contains 288 usable marker entries, 268 of them diff --git a/docs/MODELS.md b/docs/MODELS.md index d007c03..499ae51 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -108,9 +108,11 @@ raw2features embed slide.ome.zarr out -f virchow2 \ Both PRISM2 outputs require an accepted model gate, a CUDA GPU, and a `flash-attn` build compatible with the installed Torch/CUDA stack. The `[prism2]` extra pins Transformers -4.51.3 because the released config records that version and its diagnostic helper is not -compatible with Transformers 4.57. It pins the raw2features-validated FlashAttention -2.8.3 build as well. The checkpoint is approximately 17.9 GB. +4.56.0, shared with the KRONOS2 extra. PRISM2's diagnostic helper calls a Phi-3 mask +method removed after its released 4.51.3 stack, so raw2features restores that deterministic +mask construction and records the compatibility contract in the output fingerprint. The +extra also pins the raw2features-validated FlashAttention 2.8.3 build. The checkpoint is +approximately 17.9 GB. The base path still constructs the complete released model, so it has the same download and device-memory requirements as the diagnostic path. diff --git a/pyproject.toml b/pyproject.toml index 53fb9aa..6b31513 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -176,14 +176,14 @@ grandqc = ["raw2features[torch]", "segmentation-models-pytorch>=0.3", "timm>=1.0 # sacremoses at module load (the BioGPT decoder path, unused for slide embeddings). Gated. prism = ["raw2features[torch,models]", "environs", "sacremoses"] # PRISM2's released custom code requires FlashAttention for the perceiver and -# Transformers 4.51.3 for its Phi-3 image decoder. The model card specifies CUDA bf16; +# Transformers 4.56.0 for its Phi-3 image decoder. The model card specifies CUDA bf16; # keep this gated 17.9 GB checkpoint opt-in rather than adding it to [all]. prism2 = [ "raw2features[torch,models]", "torch>=2.3", - # The released config records 4.51.3 and its diagnostic helper calls a private - # Phi-3 method removed by 4.57; keep the forward on the validated upstream stack. - "transformers==4.51.3", + # Share KRONOS2's validated runtime. The loader restores the deterministic Phi-3 + # mask helper called by PRISM2 and records that compatibility contract. + "transformers==4.56.0", "flash-attn==2.8.3", ] # MADELEINE slide encoder (Mahmood Lab): loads via the authors' non-PyPI `madeleine` git diff --git a/src/raw2features/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index f5729f8..387d159 100644 --- a/src/raw2features/embedders/fingerprint.py +++ b/src/raw2features/embedders/fingerprint.py @@ -717,6 +717,7 @@ def _prism2_phi3_contract() -> dict[str, Any]: from raw2features.slide_embedders.prism2 import ( PRISM2_PHI3_ARTIFACT_SHA256, PRISM2_PHI3_LICENSE, + PRISM2_PHI3_MASK_COMPAT_VERSION, PRISM2_PHI3_REVISION, PRISM2_PHI3_SOURCE, ) @@ -728,6 +729,13 @@ def _prism2_phi3_contract() -> dict[str, Any]: "license": PRISM2_PHI3_LICENSE, "mechanism": "pinned_sha256_verified_local_config_and_tokenizer", "weights_source": "included_in_PRISM2_checkpoint", + "mask_compatibility": { + "version": PRISM2_PHI3_MASK_COMPAT_VERSION, + "restored_contract": ( + "Phi3Model._prepare_4d_causal_attention_mask_with_cache_position" + ), + "source_runtime": "transformers==4.51.3", + }, } @@ -762,7 +770,7 @@ def _slide_constructor(spec) -> dict[str, Any]: "entrypoint": "transformers.AutoModel.from_pretrained", "input": "pinned_sha256_verified_local_snapshot", "trust_remote_code": True, - "transformers_version": "4.51.3", + "transformers_version": "4.56.0", "flash_attn_version": "2.8.3", "stored_patch_encoder": "virchow2", "stored_patch_dim": 2560, diff --git a/src/raw2features/slide_embedders/prism2.py b/src/raw2features/slide_embedders/prism2.py index c250a4a..e9de437 100644 --- a/src/raw2features/slide_embedders/prism2.py +++ b/src/raw2features/slide_embedders/prism2.py @@ -26,8 +26,9 @@ PRISM2_SOURCE = "hf-hub:paige-ai/Prism2" PRISM2_REVISION = "450352d0ddc6b42b21ce20794ce0fbefe6b5a47a" -PRISM2_TRANSFORMERS_VERSION = "4.51.3" +PRISM2_TRANSFORMERS_VERSION = "4.56.0" PRISM2_FLASH_ATTN_VERSION = "2.8.3" +PRISM2_PHI3_MASK_COMPAT_VERSION = "phi3_4_51_3_causal_mask_v1" PRISM2_CHECKPOINT_SHA256 = { "model.safetensors.index.json": ( "bda099bfdad33ee0e2fe49b6101feb58a0389d8c8a000cc6636c7a24b8c41bbd" @@ -78,6 +79,76 @@ ), } + +def _prepare_phi3_4d_causal_mask( + attention_mask, + sequence_length: int, + target_length: int, + dtype, + device, + cache_position, + batch_size: int, + config, + past_key_values, +): + """Restore the Phi-3 mask contract called by the pinned PRISM2 code.""" + + # This is the mask algorithm expected by the pinned PRISM2 custom code. Its + # Phi3Model entry point was present in Transformers 4.51.3 (Apache-2.0) but + # removed before the shared 4.56.0 runtime. + import torch + from transformers.cache_utils import SlidingWindowCache + + if attention_mask is not None and attention_mask.dim() == 4: + return attention_mask + + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), + fill_value=min_dtype, + dtype=dtype, + device=device, + ) + target_positions = torch.arange(target_length, device=device) + diagonal_attend_mask = target_positions > cache_position.reshape(-1, 1) + if config.sliding_window is not None: + if ( + not isinstance(past_key_values, SlidingWindowCache) + or sequence_length > target_length + ): + sliding_attend_mask = target_positions <= ( + cache_position.reshape(-1, 1) - config.sliding_window + ) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] + mask_length = attention_mask.shape[-1] + padding_mask = ( + causal_mask[:, :, :, :mask_length] + + attention_mask[:, None, None, :].to(causal_mask.device) + ) == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[ + :, :, :, :mask_length + ].masked_fill(padding_mask, min_dtype) + return causal_mask + + +def _install_phi3_mask_compatibility(phi3_model_cls=None) -> None: + """Install the 4.51.3 Phi-3 helper removed before the shared 4.56 runtime.""" + + if phi3_model_cls is None: + from transformers.models.phi3.modeling_phi3 import Phi3Model + + phi3_model_cls = Phi3Model + phi3_model_cls._prepare_4d_causal_attention_mask_with_cache_position = staticmethod( + _prepare_phi3_4d_causal_mask + ) + + _BASE_SPEC = SlideModelSpec( name="prism2", family="prism2", @@ -220,6 +291,8 @@ def load(self, device: str = "cuda", dtype=None) -> _Prism2SlideEmbedder: from transformers import AutoConfig, AutoModel + _install_phi3_mask_compatibility() + config = AutoConfig.from_pretrained( prism_snapshot, trust_remote_code=True, diff --git a/tests/test_prism2.py b/tests/test_prism2.py index 534f4ee..bd1dc6f 100644 --- a/tests/test_prism2.py +++ b/tests/test_prism2.py @@ -20,6 +20,7 @@ PRISM2_CODE_FILES, PRISM2_FLASH_ATTN_VERSION, PRISM2_PHI3_ARTIFACT_SHA256, + PRISM2_PHI3_MASK_COMPAT_VERSION, PRISM2_PHI3_REVISION, PRISM2_PHI3_SOURCE, PRISM2_REVISION, @@ -27,6 +28,8 @@ Prism2DiagnosticSlideEmbedder, Prism2SlideEmbedder, _assert_snapshot_files, + _install_phi3_mask_compatibility, + _prepare_phi3_4d_causal_mask, _require_prism2_runtime, ) @@ -141,6 +144,7 @@ def from_pretrained(path, **kwargs): monkeypatch.setattr(module, "_verified_snapshot", verified) monkeypatch.setattr(module, "_require_prism2_runtime", lambda: None) + monkeypatch.setattr(module, "_install_phi3_mask_compatibility", lambda: None) monkeypatch.setitem( sys.modules, "transformers", @@ -192,13 +196,55 @@ def version(distribution): with pytest.raises( RuntimeError, match=( - r"transformers==4\.51\.3 \(found 4\.57\.6\).*" + r"transformers==4\.56\.0 \(found 4\.57\.6\).*" r"flash-attn==2\.8\.3 \(found missing\)" ), ): _require_prism2_runtime() +def test_prism2_restores_released_phi3_mask_contract(): + torch = pytest.importorskip("torch") + + class Phi3Model: + pass + + _install_phi3_mask_compatibility(Phi3Model) + helper = Phi3Model._prepare_4d_causal_attention_mask_with_cache_position + mask = helper( + attention_mask=torch.tensor([[1, 1, 0]], dtype=torch.int8), + sequence_length=3, + target_length=3, + dtype=torch.float32, + device=torch.device("cpu"), + cache_position=torch.arange(3), + batch_size=1, + config=SimpleNamespace(sliding_window=None), + past_key_values=None, + ) + minimum = torch.finfo(torch.float32).min + expected = torch.tensor( + [[[[0.0, minimum, minimum], [0.0, 0.0, minimum], [0.0, 0.0, minimum]]]] + ) + torch.testing.assert_close(mask, expected, rtol=0, atol=0) + + supplied = torch.zeros((1, 1, 2, 2), dtype=torch.float32) + assert ( + _prepare_phi3_4d_causal_mask( + supplied, + 2, + 2, + torch.float32, + torch.device("cpu"), + torch.arange(2), + 1, + SimpleNamespace(sliding_window=None), + None, + ) + is supplied + ) + + def test_prism2_snapshot_allowlist_fails_closed(tmp_path): for filename in ("config.json", "model.safetensors.index.json"): (tmp_path / filename).write_bytes(b"test") @@ -254,6 +300,7 @@ def test_prism2_fingerprint_binds_cls_projection_shards_and_phi3(): assert phi3["repo"] == "microsoft/Phi-3-mini-128k-instruct" assert phi3["revision"] == PRISM2_PHI3_REVISION assert phi3["artifacts"] == PRISM2_PHI3_ARTIFACT_SHA256 + assert phi3["mask_compatibility"]["version"] == (PRISM2_PHI3_MASK_COMPAT_VERSION) assert base["payload"]["output"]["resolved_amp"] == "bf16" assert diagnostic["payload"]["loader"]["constructor"]["forward"] == ( "get_diagnostic_embedding" diff --git a/uv.lock b/uv.lock index af8a22a..56dd062 100644 --- a/uv.lock +++ b/uv.lock @@ -748,6 +748,16 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/47/dd9a212ef6e343a6857485ffe25bba537304f1913bdbed446a23f7f592e1/filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258", size = 39812, upload-time = "2026-04-19T15:39:08.752Z" }, ] +[[package]] +name = "flash-attn" +version = "2.8.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "einops" }, + { name = "torch" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3b/b2/8d76c41ad7974ee264754709c22963447f7f8134613fd9ce80984ed0dab7/flash_attn-2.8.3.tar.gz", hash = "sha256:1e71dd64a9e0280e0447b8a0c2541bad4bf6ac65bdeaa2f90e51a9e57de0370d", size = 8447812, upload-time = "2025-08-15T08:28:12.911Z" } + [[package]] name = "fonttools" version = "4.63.0" @@ -3451,6 +3461,17 @@ prism = [ { name = "torchvision" }, { name = "transformers" }, ] +prism2 = [ + { name = "einops" }, + { name = "einops-exts" }, + { name = "flash-attn" }, + { name = "huggingface-hub" }, + { name = "safetensors" }, + { name = "timm" }, + { name = "torch" }, + { name = "torchvision" }, + { name = "transformers" }, +] s3 = [ { name = "s3fs" }, ] @@ -3525,6 +3546,7 @@ requires-dist = [ { name = "einops-exts", marker = "extra == 'models'", specifier = ">=0.0.4" }, { name = "environs", marker = "extra == 'prism'" }, { name = "fairscale", marker = "extra == 'musk'" }, + { name = "flash-attn", marker = "extra == 'prism2'", specifier = "==2.8.3" }, { name = "fsspec", extras = ["http"], marker = "extra == 'zarr'", specifier = ">=2024.0" }, { name = "gdown", marker = "extra == 'chief'", specifier = ">=5.2" }, { name = "gdown", marker = "extra == 'tangle'", specifier = ">=5.2" }, @@ -3560,6 +3582,7 @@ requires-dist = [ { name = "raw2features", extras = ["torch", "models"], marker = "extra == 'musk'" }, { name = "raw2features", extras = ["torch", "models"], marker = "extra == 'open-clip'" }, { name = "raw2features", extras = ["torch", "models"], marker = "extra == 'prism'" }, + { name = "raw2features", extras = ["torch", "models"], marker = "extra == 'prism2'" }, { name = "raw2features", extras = ["torch", "models"], marker = "extra == 'seal'" }, { name = "raw2features", extras = ["torch", "models"], marker = "extra == 'tangle'" }, { name = "raw2features", extras = ["zarr", "image", "torch", "models", "stain"], marker = "extra == 'all'" }, @@ -3577,17 +3600,19 @@ requires-dist = [ { name = "timm", marker = "extra == 'kronos2'", specifier = "==1.0.19" }, { name = "timm", marker = "extra == 'models'", specifier = ">=1.0" }, { name = "torch", marker = "extra == 'kronos2'", specifier = "==2.6.0" }, + { name = "torch", marker = "extra == 'prism2'", specifier = ">=2.3" }, { name = "torch", marker = "extra == 'torch'", specifier = ">=2.2" }, { name = "torchvision", marker = "extra == 'torch'", specifier = ">=0.17" }, { name = "tqdm", specifier = ">=4.66" }, { name = "transformers", marker = "extra == 'hibou'", specifier = "<5" }, { name = "transformers", marker = "extra == 'kronos2'", specifier = "==4.56.0" }, { name = "transformers", marker = "extra == 'models'", specifier = ">=4.40" }, + { name = "transformers", marker = "extra == 'prism2'", specifier = "==4.56.0" }, { name = "typer", specifier = ">=0.12" }, { name = "xformers", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux' and extra == 'kronos2'", specifier = "==0.0.29.post3" }, { name = "zarr", marker = "extra == 'zarr'", specifier = ">=3.1" }, ] -provides-extras = ["all", "zarr", "image", "stain", "torch", "models", "conch", "hibou", "kronos", "kronos2", "musk", "open-clip", "grandqc", "prism", "madeleine", "gigapath-slide", "chief", "tangle", "seal", "spatialdata", "h5", "tables", "s3", "benchmark"] +provides-extras = ["all", "zarr", "image", "stain", "torch", "models", "conch", "hibou", "kronos", "kronos2", "musk", "open-clip", "grandqc", "prism", "prism2", "madeleine", "gigapath-slide", "chief", "tangle", "seal", "spatialdata", "h5", "tables", "s3", "benchmark"] [package.metadata.requires-dev] dev = [ From 470a0868ad7a07bd0d019ca6075fc82de7740c9f Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Thu, 13 Aug 2026 13:16:58 +0100 Subject: [PATCH 3/7] Harden PRISM2 runtime provenance --- CHANGELOG.md | 12 +- docs/MODELS.md | 7 +- docs/SPEC.md | 2 + docs/adding_a_model.md | 2 + pyproject.toml | 9 +- src/raw2features/cli/embed_many.py | 50 +++ src/raw2features/cli/slide_embed.py | 12 +- src/raw2features/embedders/fingerprint.py | 42 ++- src/raw2features/pipeline/runner.py | 43 +++ src/raw2features/slide_embedders/encoding.py | 41 +++ .../slide_embedders/model_registry.py | 23 ++ src/raw2features/slide_embedders/prism2.py | 162 +++++++-- tests/test_embed_many.py | 69 ++++ tests/test_prism2.py | 317 ++++++++++++++++-- tests/test_registry_contract.py | 4 + tests/test_registry_revisions_network.py | 28 ++ uv.lock | 42 ++- 17 files changed, 785 insertions(+), 80 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 59f2d0c..f4e872a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,12 +13,12 @@ extraction path are unchanged. - Add the gated PRISM2 base (`prism2`, 2560-d) and diagnostic (`prism2_diagnostic`, 3072-d) slide representations. Both consume the 1280-d CLS - portion already preserved in raw2features' 2560-d `virchow2` patch output. The four - checkpoint shards, shard index, and Phi-3 construction/tokenizer assets are pinned and - SHA-256 verified in local runtime snapshots; the published CUDA BF16 path is retained - with validated Transformers 4.56.0 and FlashAttention 2.8.3 pins. A fingerprinted - compatibility helper preserves the released Phi-3 causal-mask construction. Coordinates - are not used. + portion already preserved in raw2features' 2560-d `virchow2` patch output. The custom + code, four checkpoint shards, shard index, and consumed Phi-3 construction assets are + pinned and SHA-256 verified in local runtime snapshots. The published CUDA BF16 path + is retained with validated Transformers 4.56.0 and FlashAttention 2.8.3 pins. A + fingerprinted helper scoped to PRISM2's verified module preserves the released Phi-3 + causal-mask construction. Coordinates are not used. - Add the gated `kronos2` native multiplex patch encoder: a pinned, SHA-verified marker-aware ViT-B/16 that produces a 768-dimensional CLS vector from ordered marker stacks. The released metadata contains 288 usable marker entries, 268 of them diff --git a/docs/MODELS.md b/docs/MODELS.md index 499ae51..347b568 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -102,6 +102,7 @@ for diagnosis-focused tasks. They may be requested separately or together: ```bash pip install "raw2features[prism2]" +MAX_JOBS=4 pip install "flash-attn==2.8.3" --no-build-isolation raw2features embed slide.ome.zarr out -f virchow2 \ -s prism2 -s prism2_diagnostic --mpp 0.5 --patch-size 224 ``` @@ -111,8 +112,10 @@ compatible with the installed Torch/CUDA stack. The `[prism2]` extra pins Transf 4.56.0, shared with the KRONOS2 extra. PRISM2's diagnostic helper calls a Phi-3 mask method removed after its released 4.51.3 stack, so raw2features restores that deterministic mask construction and records the compatibility contract in the output fingerprint. The -extra also pins the raw2features-validated FlashAttention 2.8.3 build. The checkpoint is -approximately 17.9 GB. +validated path requires FlashAttention 2.8.3. PyPI distributes it as source, so install +it after the extra with build isolation disabled as shown above, or install an +ABI-compatible 2.8.3 wheel. Increase `MAX_JOBS` only when the build host has sufficient +RAM. The checkpoint is approximately 17.9 GB. The base path still constructs the complete released model, so it has the same download and device-memory requirements as the diagnostic path. diff --git a/docs/SPEC.md b/docs/SPEC.md index d8c3fc3..a139b5f 100644 --- a/docs/SPEC.md +++ b/docs/SPEC.md @@ -220,6 +220,8 @@ patch features invalidates any derived slide vector, and records the effective fp16/fp32/bf16 slide-forward precision for the resolved device. Sharded slide checkpoints additionally record a `weights_manifest` binding the index and every shard; construction dependencies such as a separately sourced config or tokenizer belong in the loader contract. +When a slide loader executes pinned repository code, the consumed code/config files and +their SHA-256 digests are recorded there as well. Fingerprint fields are optional in the 0.1 schema so older stores remain readable and exportable, but current completion checks never reconstruct them from legacy headers: requesting an unfingerprinted patch diff --git a/docs/adding_a_model.md b/docs/adding_a_model.md index 560b181..5414562 100644 --- a/docs/adding_a_model.md +++ b/docs/adding_a_model.md @@ -66,3 +66,5 @@ and verify every component they can; any deliberately unpinned component must be For a sharded slide checkpoint, set `weights_filename` / `weights_sha256` to its index and provide `weights_manifest` with the SHA-256 of the index and every shard. The loader must verify the whole manifest before constructing the model. +If loading with `trust_remote_code=True`, pin and SHA-256 verify every executed code/config +file before construction and record that allowlist in the loader fingerprint. diff --git a/pyproject.toml b/pyproject.toml index 6b31513..14935c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -176,15 +176,18 @@ grandqc = ["raw2features[torch]", "segmentation-models-pytorch>=0.3", "timm>=1.0 # sacremoses at module load (the BioGPT decoder path, unused for slide embeddings). Gated. prism = ["raw2features[torch,models]", "environs", "sacremoses"] # PRISM2's released custom code requires FlashAttention for the perceiver and -# Transformers 4.56.0 for its Phi-3 image decoder. The model card specifies CUDA bf16; -# keep this gated 17.9 GB checkpoint opt-in rather than adding it to [all]. +# Transformers 4.56.0 for its Phi-3 image decoder. FlashAttention publishes only a +# source distribution on PyPI and requires a build against an already-installed +# Torch/CUDA stack, so install it separately with --no-build-isolation (see MODELS.md). +# Keep this gated 17.9 GB checkpoint opt-in rather than adding it to [all]. prism2 = [ "raw2features[torch,models]", "torch>=2.3", # Share KRONOS2's validated runtime. The loader restores the deterministic Phi-3 # mask helper called by PRISM2 and records that compatibility contract. "transformers==4.56.0", - "flash-attn==2.8.3", + "packaging>=23", + "ninja>=1.11", ] # MADELEINE slide encoder (Mahmood Lab): loads via the authors' non-PyPI `madeleine` git # package (post-install, like conch): diff --git a/src/raw2features/cli/embed_many.py b/src/raw2features/cli/embed_many.py index 34d6075..7b17a82 100644 --- a/src/raw2features/cli/embed_many.py +++ b/src/raw2features/cli/embed_many.py @@ -427,6 +427,18 @@ def embed_many( ) raise typer.Exit(1) from exc + try: + _preflight_slide_encoder_runtime( + shard, + out_dir, + cfg, + device_list, + force=force, + ) + except (RuntimeError, ValueError) as exc: + typer.echo(f"Error: slide-encoder runtime preflight failed: {exc}", err=True) + raise typer.Exit(1) from exc + t0 = time.time() try: if len(device_list) > 1: @@ -470,6 +482,44 @@ def _classify(summary: dict) -> str: return "skipped" if summary["status"] == "skipped" else "done" +def _preflight_slide_encoder_runtime( + shard: list[dict], + out_dir: str, + cfg: RunConfig, + devices: list[str], + *, + force: bool, +) -> None: + """Check a required slide runtime before warm patch models are loaded.""" + + if not cfg.slide_encoders: + return + from raw2features.slide_embedders.encoding import ( + slide_encoders_requiring_compute, + ) + from raw2features.slide_embedders.model_registry import ( + validate_slide_encoder_runtime, + ) + + required = set(cfg.slide_encoders) if force else set() + if not force: + inspection_device = devices[0] if devices else cfg.device + for row in shard: + slide_id = slide_id_from_path(row["path"]) + store_path = os.path.join(out_dir, f"{slide_id}.embeddings.zarr") + required.update( + slide_encoders_requiring_compute( + store_path, + cfg.slide_encoders, + device=inspection_device, + ) + ) + if len(required) == len(set(cfg.slide_encoders)): + break + ordered = [name for name in dict.fromkeys(cfg.slide_encoders) if name in required] + validate_slide_encoder_runtime(ordered, devices=devices) + + def _resolve_manifest_sources(rows: list[dict], slide_dir: str) -> list[dict]: """Resolve only relative local manifest paths against ``slide_dir`` in place.""" diff --git a/src/raw2features/cli/slide_embed.py b/src/raw2features/cli/slide_embed.py index dab837c..7c116a5 100644 --- a/src/raw2features/cli/slide_embed.py +++ b/src/raw2features/cli/slide_embed.py @@ -76,6 +76,7 @@ def slide_embed( ) from raw2features.slide_embedders.model_registry import ( validate_slide_encoder_names, + validate_slide_encoder_runtime, ) try: @@ -111,7 +112,7 @@ def slide_embed( grid=grid, patch_model=patch_model, ) - except (KeyError, ValueError) as exc: + except (KeyError, RuntimeError, ValueError) as exc: typer.echo(f"Error: {exc}", err=True) raise typer.Exit(1) from exc @@ -133,6 +134,15 @@ def slide_embed( ) continue + try: + validate_slide_encoder_runtime( + [slide_model_name], + devices=[device], + ) + except (RuntimeError, ValueError) as exc: + typer.echo(f"Error: {exc}", err=True) + raise typer.Exit(1) from exc + typer.echo( f"{slide_model_name} [{selected_grid}]: encoding from " f"'{selected_patch_model}' patch features …" diff --git a/src/raw2features/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index 387d159..de5d236 100644 --- a/src/raw2features/embedders/fingerprint.py +++ b/src/raw2features/embedders/fingerprint.py @@ -713,33 +713,54 @@ def _effective_slide_checkpoint(spec) -> dict[str, Any]: } -def _prism2_phi3_contract() -> dict[str, Any]: +def _prism2_phi3_contract(*, diagnostic: bool) -> dict[str, Any]: from raw2features.slide_embedders.prism2 import ( PRISM2_PHI3_ARTIFACT_SHA256, + PRISM2_PHI3_CONFIG_SHA256, PRISM2_PHI3_LICENSE, PRISM2_PHI3_MASK_COMPAT_VERSION, PRISM2_PHI3_REVISION, PRISM2_PHI3_SOURCE, ) - return { + contract = { "repo": _hf_repo(PRISM2_PHI3_SOURCE), "revision": PRISM2_PHI3_REVISION, - "artifacts": deepcopy(PRISM2_PHI3_ARTIFACT_SHA256), + "artifacts": deepcopy( + PRISM2_PHI3_ARTIFACT_SHA256 if diagnostic else PRISM2_PHI3_CONFIG_SHA256 + ), "license": PRISM2_PHI3_LICENSE, - "mechanism": "pinned_sha256_verified_local_config_and_tokenizer", + "mechanism": ( + "pinned_sha256_verified_local_config_and_tokenizer" + if diagnostic + else "pinned_sha256_verified_local_config" + ), "weights_source": "included_in_PRISM2_checkpoint", - "mask_compatibility": { + } + if diagnostic: + contract["mask_compatibility"] = { "version": PRISM2_PHI3_MASK_COMPAT_VERSION, "restored_contract": ( "Phi3Model._prepare_4d_causal_attention_mask_with_cache_position" ), "source_runtime": "transformers==4.51.3", - }, - } + "binding_scope": "verified_prism2_dynamic_module", + } + return contract def _slide_constructor(spec) -> dict[str, Any]: + prism2_snapshot_contract = None + if spec.family == "prism2": + from raw2features.slide_embedders.prism2 import ( + PRISM2_CODE_SHA256, + PRISM2_SNAPSHOT_ALLOW_PATTERNS, + ) + + prism2_snapshot_contract = { + "allow_patterns": list(PRISM2_SNAPSHOT_ALLOW_PATTERNS), + "custom_code_sha256": deepcopy(PRISM2_CODE_SHA256), + } contracts: dict[str, dict[str, Any]] = { "pool": { "operation": spec.name, @@ -769,6 +790,7 @@ def _slide_constructor(spec) -> dict[str, Any]: "prism2": { "entrypoint": "transformers.AutoModel.from_pretrained", "input": "pinned_sha256_verified_local_snapshot", + "snapshot": prism2_snapshot_contract, "trust_remote_code": True, "transformers_version": "4.56.0", "flash_attn_version": "2.8.3", @@ -882,7 +904,9 @@ def _slide_constructor(spec) -> dict[str, Any]: } contract = deepcopy(contracts.get(spec.family, {"entrypoint": spec.family})) if spec.family == "prism2": - contract["phi3_construction_dependency"] = _prism2_phi3_contract() + contract["phi3_construction_dependency"] = _prism2_phi3_contract( + diagnostic=spec.name == "prism2_diagnostic" + ) return contract @@ -899,7 +923,7 @@ def resolved_slide_amp(spec, device: str) -> str: # These loaders follow their model-card examples with fp16 CUDA autocast and # deliberately run fp32 on CPU/MPS. Other slide encoders currently run fp32. - if str(device).startswith("cuda") and spec.family == "prism2": + if spec.family == "prism2": return "bf16" if str(device).startswith("cuda") and spec.family in { "gigapath_slide", diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index 481dc97..036f2e8 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -1453,6 +1453,32 @@ def run_slide( else: models_to_do = list(cfg.models) + if cfg.slide_encoders: + from raw2features.slide_embedders.encoding import ( + slide_encoders_requiring_compute, + ) + from raw2features.slide_embedders.model_registry import ( + get_slide_spec, + validate_slide_encoder_runtime, + ) + + runtime_names = ( + list(cfg.slide_encoders) + if force + else slide_encoders_requiring_compute( + out_path, + cfg.slide_encoders, + device=cfg.device, + ) + ) + for name in cfg.slide_encoders: + if get_slide_spec(name).patch_encoder in models_to_do: + runtime_names.append(name) + validate_slide_encoder_runtime( + list(dict.fromkeys(runtime_names)), + devices=devices, + ) + qc_to_do = list(cfg.qc) thumbnail_to_do = bool(cfg.emit_thumbnail) thumb_meta = None @@ -2202,6 +2228,23 @@ def embed_slide( "grids": grids, } + if cfg.slide_encoders: + from raw2features.slide_embedders.encoding import ( + slide_encoders_requiring_compute, + ) + from raw2features.slide_embedders.model_registry import ( + validate_slide_encoder_runtime, + ) + + validate_slide_encoder_runtime( + slide_encoders_requiring_compute( + out_path, + cfg.slide_encoders, + device=cfg.device, + ), + devices=devices, + ) + started = time.time() results = [] for i, gc in enumerate(group_cfgs): diff --git a/src/raw2features/slide_embedders/encoding.py b/src/raw2features/slide_embedders/encoding.py index f88185e..da903d7 100644 --- a/src/raw2features/slide_embedders/encoding.py +++ b/src/raw2features/slide_embedders/encoding.py @@ -7,6 +7,7 @@ from __future__ import annotations +import os import warnings from dataclasses import dataclass @@ -195,6 +196,46 @@ def slide_embedding_is_complete( return bool(vector.size and np.isfinite(vector).all() and (vector != 0).any()) +def slide_encoders_requiring_compute( + store_path: str, + slide_models: list[str], + *, + device: str = "cpu", +) -> list[str]: + """Return requested slide encoders without a complete current store output.""" + + if not slide_models: + return [] + path = store_path.removeprefix("file://") + if not os.path.exists(path): + return list(slide_models) + + import zarr + + try: + root = zarr.open_group(path, mode="r", use_consolidated=False) + except Exception: # noqa: BLE001 - an unreadable store cannot prove completion + return list(slide_models) + + missing = [] + for slide_model in slide_models: + try: + _, group, patch_model = resolve_slide_grid(root, slide_model) + output_name = slide_output_key(group, slide_model, patch_model) + complete = slide_embedding_is_complete( + group, + slide_model, + patch_model=patch_model, + device=device, + output_name=output_name, + ) + except (KeyError, TypeError, ValueError): + complete = False + if not complete: + missing.append(slide_model) + return missing + + def resolve_slide_grid( root, slide_model: str, diff --git a/src/raw2features/slide_embedders/model_registry.py b/src/raw2features/slide_embedders/model_registry.py index b62e20b..c5d6b88 100644 --- a/src/raw2features/slide_embedders/model_registry.py +++ b/src/raw2features/slide_embedders/model_registry.py @@ -53,6 +53,29 @@ def validate_slide_encoder_names(names: list[str]) -> None: ) +def validate_slide_encoder_runtime( + names: list[str], + *, + devices: list[str] | None = None, +) -> None: + """Fail before patch work when a requested slide runtime is unavailable.""" + + specs = [get_slide_spec(name) for name in dict.fromkeys(names)] + if not any(spec.family == "prism2" for spec in specs): + return + + from raw2features.slide_embedders.prism2 import _require_prism2_runtime + + _require_prism2_runtime() + if devices is not None: + incompatible = [device for device in devices if not device.startswith("cuda")] + if incompatible: + raise ValueError( + "PRISM2 requires a CUDA GPU; resolved incompatible device(s): " + f"{incompatible}" + ) + + def build_slide_embedder(name: str) -> SlideEmbedder: from raw2features.core.plugins import get diff --git a/src/raw2features/slide_embedders/prism2.py b/src/raw2features/slide_embedders/prism2.py index e9de437..19c8766 100644 --- a/src/raw2features/slide_embedders/prism2.py +++ b/src/raw2features/slide_embedders/prism2.py @@ -9,6 +9,8 @@ from __future__ import annotations +import sys +import threading from importlib import metadata as importlib_metadata from pathlib import Path from typing import Any @@ -46,16 +48,24 @@ "da79790c9a73cf1ce85867aaa57c47cfd8d69e13b57ac86252d24a2687c2998c" ), } -PRISM2_CODE_FILES = ( - "config.json", - "configuration_prism2.py", - "modeling_prism2.py", - "processing_prism2.py", -) +PRISM2_CODE_SHA256 = { + "config.json": "2170566eeb216810e32f69293e07429694fc308e4290efb4ca2157dc7438f7a9", + "configuration_prism2.py": ( + "24874c747af4ad499db5ce4ac9e6113c0289fb937badce4a6ccf37b3e6f8ede4" + ), + "modeling_prism2.py": ( + "57546d28c94d4e8fd0fd29600f0b0190d32cc8fedf7c82ec85f525183f28669a" + ), + "processing_prism2.py": ( + "a3dcc7b466234fcc96b47b32f774f1ffbc7271b8af37b6352b13d537b51c9bc9" + ), +} +PRISM2_CODE_FILES = tuple(PRISM2_CODE_SHA256) PRISM2_SNAPSHOT_ALLOW_PATTERNS = ( *PRISM2_CODE_FILES, *PRISM2_CHECKPOINT_SHA256, ) +PRISM2_RUNTIME_SHA256 = {**PRISM2_CODE_SHA256, **PRISM2_CHECKPOINT_SHA256} PRISM2_PHI3_SOURCE = "hf-hub:microsoft/Phi-3-mini-128k-instruct" PRISM2_PHI3_REVISION = "f3c06aed622e14ca0abf5115094e4fc9a9948f36" @@ -78,6 +88,13 @@ "aaa87217a0f61c684cdc8703d3d4030a1f5b1077183610b61a14d7f28addbb58" ), } +PRISM2_PHI3_CONFIG_SHA256 = {"config.json": PRISM2_PHI3_ARTIFACT_SHA256["config.json"]} + +_VERIFIED_SNAPSHOT_CACHE: dict[ + tuple[str, tuple[tuple[str, str], ...]], + tuple[tuple[str, int, int, int, int, int], ...], +] = {} +_VERIFIED_SNAPSHOT_LOCK = threading.Lock() def _prepare_phi3_4d_causal_mask( @@ -137,18 +154,43 @@ def _prepare_phi3_4d_causal_mask( return causal_mask -def _install_phi3_mask_compatibility(phi3_model_cls=None) -> None: - """Install the 4.51.3 Phi-3 helper removed before the shared 4.56 runtime.""" - - if phi3_model_cls is None: - from transformers.models.phi3.modeling_phi3 import Phi3Model +class _Prism2Phi3MaskCompatibility: + """Private namespace for the Phi-3 helper used by pinned PRISM2 code.""" - phi3_model_cls = Phi3Model - phi3_model_cls._prepare_4d_causal_attention_mask_with_cache_position = staticmethod( + _prepare_4d_causal_attention_mask_with_cache_position = staticmethod( _prepare_phi3_4d_causal_mask ) +def _bind_phi3_mask_compatibility(model) -> None: + """Bind the helper only inside PRISM2's verified dynamic-code module.""" + + module_name = type(model).__module__ + module = sys.modules.get(module_name) + if module is None or not module_name.startswith("transformers_modules."): + raise RuntimeError( + "PRISM2 loaded outside its expected Transformers dynamic-code module; " + "refusing to install the Phi-3 compatibility contract" + ) + current = getattr(module, "Phi3Model", None) + if current is _Prism2Phi3MaskCompatibility: + return + if ( + current is None + or getattr(current, "__name__", None) != "Phi3Model" + or getattr(current, "__module__", None) + != "transformers.models.phi3.modeling_phi3" + ): + raise RuntimeError( + "PRISM2's verified module does not expose the expected Phi3Model import; " + "refusing to alter an unknown runtime" + ) + # modeling_prism2.py calls its imported Phi3Model name directly during the + # diagnostic forward. Rebind that module global rather than changing the shared + # Transformers class process-wide. + module.Phi3Model = _Prism2Phi3MaskCompatibility + + _BASE_SPEC = SlideModelSpec( name="prism2", family="prism2", @@ -196,12 +238,13 @@ def _assert_snapshot_files( snapshot: str, expected: tuple[str, ...], *, + allowed_files: tuple[str, ...] | None = None, what: str, ) -> None: """Reject a stale managed directory before its custom code can execute.""" root = Path(snapshot) - allowed = {*expected, ".cache", "__pycache__"} + allowed = {*(allowed_files or expected), ".cache", "__pycache__"} unexpected = sorted( path.name for path in root.iterdir() if path.name not in allowed ) @@ -221,23 +264,81 @@ def _verified_snapshot( revision: str, artifacts: dict[str, str], *, - extra_files: tuple[str, ...] = (), + allowed_files: tuple[str, ...] | None = None, what: str, ) -> str: - expected = (*extra_files, *artifacts) + expected = tuple(artifacts) snapshot = download_pinned_hf_snapshot( source, revision, allow_patterns=expected, local_dir=pinned_model_cache_dir(source, revision), ) - _assert_snapshot_files(snapshot, expected, what=what) - for filename, sha256 in artifacts.items(): - verify_sha256(str(Path(snapshot) / filename), sha256, what=f"{what}:{filename}") + _assert_snapshot_files( + snapshot, + expected, + allowed_files=allowed_files, + what=what, + ) + _verify_snapshot_artifacts_once(snapshot, artifacts, what=what) return snapshot +def _snapshot_signature( + snapshot: str, + artifacts: dict[str, str], +) -> tuple[tuple[str, int, int, int, int, int], ...]: + root = Path(snapshot) + signature = [] + for filename in sorted(artifacts): + stat = (root / filename).stat() + signature.append( + ( + filename, + int(stat.st_dev), + int(stat.st_ino), + int(stat.st_size), + int(stat.st_mtime_ns), + int(stat.st_ctime_ns), + ) + ) + return tuple(signature) + + +def _verify_snapshot_artifacts_once( + snapshot: str, + artifacts: dict[str, str], + *, + what: str, +) -> None: + """Hash an immutable snapshot once per process, unless any file changes.""" + + key = (str(Path(snapshot).resolve()), tuple(sorted(artifacts.items()))) + with _VERIFIED_SNAPSHOT_LOCK: + before = _snapshot_signature(snapshot, artifacts) + if _VERIFIED_SNAPSHOT_CACHE.get(key) == before: + return + for filename, sha256 in artifacts.items(): + verify_sha256( + str(Path(snapshot) / filename), + sha256, + what=f"{what}:{filename}", + ) + after = _snapshot_signature(snapshot, artifacts) + if after != before: + raise ValueError(f"{what}: pinned runtime snapshot changed while hashing") + _VERIFIED_SNAPSHOT_CACHE[key] = after + + def _require_prism2_runtime() -> None: + try: + from packaging.specifiers import SpecifierSet + except ImportError as exc: # pragma: no cover - declared by the optional extra + raise RuntimeError( + "PRISM2 requires its optional runtime; install raw2features[prism2], " + "then install flash-attn==2.8.3 with --no-build-isolation" + ) from exc + required = { "transformers": PRISM2_TRANSFORMERS_VERSION, "flash-attn": PRISM2_FLASH_ATTN_VERSION, @@ -248,12 +349,15 @@ def _require_prism2_runtime() -> None: installed = importlib_metadata.version(distribution) except importlib_metadata.PackageNotFoundError: installed = None - if installed != expected: + if installed is None or not SpecifierSet(f"=={expected}").contains( + installed, prereleases=True + ): found = installed or "missing" mismatches.append(f"{distribution}=={expected} (found {found})") if mismatches: raise RuntimeError( - "PRISM2 requires its validated runtime; install raw2features[prism2]: " + "PRISM2 requires its validated runtime. Install raw2features[prism2], " + "then install flash-attn==2.8.3 with --no-build-isolation. Mismatch: " + ", ".join(mismatches) ) @@ -278,21 +382,24 @@ def load(self, device: str = "cuda", dtype=None) -> _Prism2SlideEmbedder: prism_snapshot = _verified_snapshot( self.spec.source, str(self.spec.weights_revision), - dict(self.spec.weights_manifest), - extra_files=PRISM2_CODE_FILES, + PRISM2_RUNTIME_SHA256, what=self.spec.name, ) + phi3_artifacts = ( + PRISM2_PHI3_ARTIFACT_SHA256 + if self.spec.name == "prism2_diagnostic" + else PRISM2_PHI3_CONFIG_SHA256 + ) phi3_snapshot = _verified_snapshot( PRISM2_PHI3_SOURCE, PRISM2_PHI3_REVISION, - PRISM2_PHI3_ARTIFACT_SHA256, + phi3_artifacts, + allowed_files=tuple(PRISM2_PHI3_ARTIFACT_SHA256), what=f"{self.spec.name}:Phi-3", ) from transformers import AutoConfig, AutoModel - _install_phi3_mask_compatibility() - config = AutoConfig.from_pretrained( prism_snapshot, trust_remote_code=True, @@ -308,6 +415,8 @@ def load(self, device: str = "cuda", dtype=None) -> _Prism2SlideEmbedder: local_files_only=True, torch_dtype="auto", ) + if self.spec.name == "prism2_diagnostic": + _bind_phi3_mask_compatibility(model) model.eval().to(device) self._model = model self._device = str(device) @@ -361,7 +470,6 @@ def unload(self) -> None: import torch if self._model is not None: - self._model.cpu() del self._model self._model = None if torch.cuda.is_available(): diff --git a/tests/test_embed_many.py b/tests/test_embed_many.py index 4198f64..0c7ceef 100644 --- a/tests/test_embed_many.py +++ b/tests/test_embed_many.py @@ -78,6 +78,75 @@ def test_embed_many_validates_slide_encoders_before_discovery_or_model_load( assert "no slides found" not in result.output +def test_embed_many_runtime_preflight_runs_before_warm_model_load( + tmp_path, monkeypatch +): + import raw2features.cli.embed_many as em + + slides = tmp_path / "slides" + slides.mkdir() + (slides / "S.zarr").mkdir() + monkeypatch.setattr( + em, + "_preflight_slide_encoder_runtime", + lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("runtime missing")), + ) + monkeypatch.setattr( + em, + "load_embedders", + lambda *args, **kwargs: pytest.fail("models must not be loaded"), + ) + + result = CliRunner().invoke( + app, + [ + "embed-many", + str(slides), + str(tmp_path / "out"), + "-f", + "mock", + "-s", + "prism2", + "--device", + "cpu", + ], + ) + + assert result.exit_code == 1 + assert "runtime missing" in result.output + assert "models must not be loaded" not in result.output + + +def test_embed_many_runtime_preflight_ignores_complete_slide_outputs( + tmp_path, monkeypatch +): + import raw2features.slide_embedders.encoding as encoding + import raw2features.slide_embedders.model_registry as slide_registry + from raw2features.cli.embed_many import _preflight_slide_encoder_runtime + + checked = [] + monkeypatch.setattr( + encoding, + "slide_encoders_requiring_compute", + lambda *args, **kwargs: [], + ) + monkeypatch.setattr( + slide_registry, + "validate_slide_encoder_runtime", + lambda names, *, devices: checked.append((names, devices)), + ) + + _preflight_slide_encoder_runtime( + [{"path": str(tmp_path / "S.zarr")}], + str(tmp_path / "out"), + RunConfig(models=["mock"], slide_encoders=["prism2"]), + ["cpu"], + force=False, + ) + + assert checked == [([], ["cpu"])] + + @pytest.mark.parametrize( "rows", [ diff --git a/tests/test_prism2.py b/tests/test_prism2.py index bd1dc6f..351a4a6 100644 --- a/tests/test_prism2.py +++ b/tests/test_prism2.py @@ -3,34 +3,49 @@ from __future__ import annotations import sys +from pathlib import Path from types import SimpleNamespace import numpy as np import pytest +import zarr from raw2features.embedders.fingerprint import ( patch_output_fingerprint, resolved_slide_amp, + slide_output_dim, slide_output_fingerprint, ) from raw2features.embedders.model_registry import get_spec -from raw2features.slide_embedders.model_registry import get_slide_spec +from raw2features.slide_embedders.encoding import ( + slide_encoders_requiring_compute, + write_slide_embedding, +) +from raw2features.slide_embedders.model_registry import ( + get_slide_spec, + validate_slide_encoder_runtime, +) from raw2features.slide_embedders.prism2 import ( + _VERIFIED_SNAPSHOT_CACHE, PRISM2_CHECKPOINT_SHA256, - PRISM2_CODE_FILES, + PRISM2_CODE_SHA256, PRISM2_FLASH_ATTN_VERSION, PRISM2_PHI3_ARTIFACT_SHA256, + PRISM2_PHI3_CONFIG_SHA256, PRISM2_PHI3_MASK_COMPAT_VERSION, PRISM2_PHI3_REVISION, PRISM2_PHI3_SOURCE, PRISM2_REVISION, + PRISM2_RUNTIME_SHA256, + PRISM2_SNAPSHOT_ALLOW_PATTERNS, PRISM2_TRANSFORMERS_VERSION, Prism2DiagnosticSlideEmbedder, Prism2SlideEmbedder, _assert_snapshot_files, - _install_phi3_mask_compatibility, + _bind_phi3_mask_compatibility, _prepare_phi3_4d_causal_mask, _require_prism2_runtime, + _verified_snapshot, ) @@ -109,8 +124,10 @@ def test_prism2_rejects_an_incompatible_patch_array(): def test_prism2_load_uses_only_verified_local_snapshots(monkeypatch): calls = [] - def verified(source, revision, artifacts, *, extra_files=(), what): - calls.append((source, revision, dict(artifacts), tuple(extra_files), what)) + def verified(source, revision, artifacts, *, allowed_files=None, what): + calls.append( + (source, revision, dict(artifacts), tuple(allowed_files or ()), what) + ) return "/verified/phi3" if "Phi-3" in what else "/verified/prism2" config = SimpleNamespace(text_decoder_model_id="mutable-upstream-name") @@ -144,7 +161,6 @@ def from_pretrained(path, **kwargs): monkeypatch.setattr(module, "_verified_snapshot", verified) monkeypatch.setattr(module, "_require_prism2_runtime", lambda: None) - monkeypatch.setattr(module, "_install_phi3_mask_compatibility", lambda: None) monkeypatch.setitem( sys.modules, "transformers", @@ -157,16 +173,63 @@ def from_pretrained(path, **kwargs): assert config.text_decoder_model_id == "/verified/phi3" assert calls[0][1:] == ( PRISM2_REVISION, - PRISM2_CHECKPOINT_SHA256, - PRISM2_CODE_FILES, + PRISM2_RUNTIME_SHA256, + (), "prism2", ) - assert calls[1][0:4] == ( + assert calls[1] == ( PRISM2_PHI3_SOURCE, PRISM2_PHI3_REVISION, + PRISM2_PHI3_CONFIG_SHA256, + tuple(PRISM2_PHI3_ARTIFACT_SHA256), + "prism2:Phi-3", + ) + + +def test_prism2_diagnostic_load_binds_helper_and_full_phi3_assets(monkeypatch): + calls = [] + model = SimpleNamespace( + eval=lambda: model, + to=lambda device: model, + ) + + def verified(source, revision, artifacts, *, allowed_files=None, what): + calls.append((dict(artifacts), tuple(allowed_files or ()), what)) + return "/verified/phi3" if "Phi-3" in what else "/verified/prism2" + + class AutoConfig: + @staticmethod + def from_pretrained(path, **kwargs): + return SimpleNamespace(text_decoder_model_id="mutable") + + class AutoModel: + @staticmethod + def from_pretrained(path, **kwargs): + return model + + import raw2features.slide_embedders.prism2 as module + + monkeypatch.setattr(module, "_verified_snapshot", verified) + monkeypatch.setattr(module, "_require_prism2_runtime", lambda: None) + monkeypatch.setattr( + module, + "_bind_phi3_mask_compatibility", + lambda received: calls.append(("bind", received)), + ) + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace(AutoConfig=AutoConfig, AutoModel=AutoModel), + ) + + Prism2DiagnosticSlideEmbedder().load("cuda:0") + + assert calls[1] == ( PRISM2_PHI3_ARTIFACT_SHA256, - (), + tuple(PRISM2_PHI3_ARTIFACT_SHA256), + "prism2_diagnostic:Phi-3", ) + assert calls[2] == ("bind", model) def test_prism2_rejects_cpu_before_downloading(monkeypatch): @@ -196,6 +259,7 @@ def version(distribution): with pytest.raises( RuntimeError, match=( + r"install flash-attn==2\.8\.3 with --no-build-isolation.*" r"transformers==4\.56\.0 \(found 4\.57\.6\).*" r"flash-attn==2\.8\.3 \(found missing\)" ), @@ -203,15 +267,41 @@ def version(distribution): _require_prism2_runtime() +def test_prism2_runtime_pin_accepts_pep440_equivalent_versions(monkeypatch): + import raw2features.slide_embedders.prism2 as module + + versions = {"transformers": "4.56", "flash-attn": "2.8.3+cu128"} + monkeypatch.setattr( + module.importlib_metadata, + "version", + lambda distribution: versions[distribution], + ) + + _require_prism2_runtime() + + +def test_prism2_runtime_preflight_checks_packages_and_device(monkeypatch): + import raw2features.slide_embedders.prism2 as module + + calls = [] + monkeypatch.setattr(module, "_require_prism2_runtime", lambda: calls.append(True)) + + validate_slide_encoder_runtime(["mean"], devices=["cpu"]) + assert calls == [] + validate_slide_encoder_runtime(["prism2"], devices=["cuda:0"]) + assert calls == [True] + with pytest.raises(ValueError, match="requires a CUDA GPU"): + validate_slide_encoder_runtime( + ["prism2_diagnostic"], + devices=["cpu"], + ) + + def test_prism2_restores_released_phi3_mask_contract(): torch = pytest.importorskip("torch") + pytest.importorskip("transformers") - class Phi3Model: - pass - - _install_phi3_mask_compatibility(Phi3Model) - helper = Phi3Model._prepare_4d_causal_attention_mask_with_cache_position - mask = helper( + mask = _prepare_phi3_4d_causal_mask( attention_mask=torch.tensor([[1, 1, 0]], dtype=torch.int8), sequence_length=3, target_length=3, @@ -245,6 +335,40 @@ class Phi3Model: ) +def test_prism2_binds_phi3_helper_only_to_verified_dynamic_module(monkeypatch): + pytest.importorskip("transformers") + module_name = "transformers_modules.paige_prism2.modeling_prism2" + original = type( + "Phi3Model", + (), + {"__module__": "transformers.models.phi3.modeling_phi3"}, + ) + dynamic_module = SimpleNamespace(Phi3Model=original) + model_cls = type("Prism2Model", (), {"__module__": module_name}) + monkeypatch.setitem(sys.modules, module_name, dynamic_module) + + _bind_phi3_mask_compatibility(model_cls()) + + helper = dynamic_module.Phi3Model + assert helper is not original + assert callable(helper._prepare_4d_causal_attention_mask_with_cache_position) + from transformers.models.phi3.modeling_phi3 import Phi3Model + + assert not hasattr( + Phi3Model, "_prepare_4d_causal_attention_mask_with_cache_position" + ) + + +def test_prism2_rejects_unknown_dynamic_module_binding(monkeypatch): + module_name = "transformers_modules.paige_prism2.modeling_prism2" + dynamic_module = SimpleNamespace(Phi3Model=object) + model_cls = type("Prism2Model", (), {"__module__": module_name}) + monkeypatch.setitem(sys.modules, module_name, dynamic_module) + + with pytest.raises(RuntimeError, match="does not expose the expected Phi3Model"): + _bind_phi3_mask_compatibility(model_cls()) + + def test_prism2_snapshot_allowlist_fails_closed(tmp_path): for filename in ("config.json", "model.safetensors.index.json"): (tmp_path / filename).write_bytes(b"test") @@ -263,6 +387,76 @@ def test_prism2_snapshot_allowlist_fails_closed(tmp_path): ) +def test_prism2_runtime_code_is_sha_bound_and_cached(monkeypatch, tmp_path): + artifact = tmp_path / "modeling_prism2.py" + artifact.write_bytes(b"verified code") + calls = [] + + monkeypatch.setattr( + "raw2features.slide_embedders.prism2.download_pinned_hf_snapshot", + lambda *args, **kwargs: str(tmp_path), + ) + + def verify(path, expected, *, what): + calls.append((Path(path).name, expected, what)) + + monkeypatch.setattr( + "raw2features.slide_embedders.prism2.verify_sha256", + verify, + ) + _VERIFIED_SNAPSHOT_CACHE.clear() + artifacts = {"modeling_prism2.py": "a" * 64} + + _verified_snapshot("hf-hub:test/model", "a" * 40, artifacts, what="test") + _verified_snapshot("hf-hub:test/model", "a" * 40, artifacts, what="test") + assert len(calls) == 1 + + artifact.write_bytes(b"changed code") + _verified_snapshot("hf-hub:test/model", "a" * 40, artifacts, what="test") + assert len(calls) == 2 + + +def test_prism2_rejects_bad_runtime_digest_before_custom_code(monkeypatch, tmp_path): + for filename in PRISM2_SNAPSHOT_ALLOW_PATTERNS: + (tmp_path / filename).write_bytes(b"untrusted") + events = [] + + monkeypatch.setattr( + "raw2features.slide_embedders.prism2.download_pinned_hf_snapshot", + lambda *args, **kwargs: str(tmp_path), + ) + monkeypatch.setattr( + "raw2features.slide_embedders.prism2._require_prism2_runtime", + lambda: None, + ) + + def reject(*args, **kwargs): + events.append("verify") + raise ValueError("digest mismatch") + + monkeypatch.setattr( + "raw2features.slide_embedders.prism2.verify_sha256", + reject, + ) + _VERIFIED_SNAPSHOT_CACHE.clear() + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + AutoConfig=SimpleNamespace( + from_pretrained=lambda *a, **k: pytest.fail("custom code executed") + ), + AutoModel=SimpleNamespace( + from_pretrained=lambda *a, **k: pytest.fail("custom code executed") + ), + ), + ) + + with pytest.raises(ValueError, match="digest mismatch"): + Prism2SlideEmbedder().load("cuda:0") + assert events == ["verify"] + + def test_prism2_fingerprint_binds_cls_projection_shards_and_phi3(): patch = patch_output_fingerprint(get_spec("virchow2"), "fp16") base_spec = get_slide_spec("prism2") @@ -299,16 +493,99 @@ def test_prism2_fingerprint_binds_cls_projection_shards_and_phi3(): phi3 = constructor["phi3_construction_dependency"] assert phi3["repo"] == "microsoft/Phi-3-mini-128k-instruct" assert phi3["revision"] == PRISM2_PHI3_REVISION - assert phi3["artifacts"] == PRISM2_PHI3_ARTIFACT_SHA256 - assert phi3["mask_compatibility"]["version"] == (PRISM2_PHI3_MASK_COMPAT_VERSION) + assert phi3["artifacts"] == PRISM2_PHI3_CONFIG_SHA256 + assert "mask_compatibility" not in phi3 + assert constructor["snapshot"] == { + "allow_patterns": list(PRISM2_SNAPSHOT_ALLOW_PATTERNS), + "custom_code_sha256": PRISM2_CODE_SHA256, + } assert base["payload"]["output"]["resolved_amp"] == "bf16" - assert diagnostic["payload"]["loader"]["constructor"]["forward"] == ( - "get_diagnostic_embedding" + diagnostic_constructor = diagnostic["payload"]["loader"]["constructor"] + assert diagnostic_constructor["forward"] == "get_diagnostic_embedding" + diagnostic_phi3 = diagnostic_constructor["phi3_construction_dependency"] + assert diagnostic_phi3["artifacts"] == PRISM2_PHI3_ARTIFACT_SHA256 + assert diagnostic_phi3["mask_compatibility"]["version"] == ( + PRISM2_PHI3_MASK_COMPAT_VERSION + ) + assert diagnostic_phi3["mask_compatibility"]["binding_scope"] == ( + "verified_prism2_dynamic_module" ) assert diagnostic["payload"]["output"]["embedding_dim"] == 3072 assert base["digest"] != diagnostic["digest"] +def test_prism2_amp_contract_is_device_independent(): + for name in ("prism2", "prism2_diagnostic"): + spec = get_slide_spec(name) + assert resolved_slide_amp(spec, "cpu") == "bf16" + assert resolved_slide_amp(spec, "cuda:0") == "bf16" + + +def test_prism2_complete_output_can_be_recognized_without_gpu_runtime(tmp_path): + path = str(tmp_path / "complete.embeddings.zarr") + root = zarr.open_group(path, mode="w", zarr_format=2) + grids = root.create_group("grids") + group = grids.create_group("mpp0.5_px224") + coords = group.create_array("coords", shape=(2, 2), dtype="int32") + coords[:] = np.asarray([[0, 0], [224, 0]], dtype=np.int32) + + patch_spec = get_spec("virchow2") + patch_fingerprint = patch_output_fingerprint(patch_spec, "fp16") + features = group.create_group("features").create_array( + "virchow2", + shape=(2, patch_spec.embedding_dim), + chunks=(2, patch_spec.embedding_dim), + dtype="float32", + ) + features[:] = 1.0 + features.attrs.update( + { + "role": "features", + "model": "virchow2", + "output_fingerprint": patch_fingerprint, + } + ) + group.attrs["raw2features"] = { + "schema_version": "0.1", + "models": { + "virchow2": { + "embedding_dim": patch_spec.embedding_dim, + "output_fingerprint": patch_fingerprint, + } + }, + "patching": {"level0_patch": 224}, + } + root.attrs["raw2features"] = { + "schema_version": "0.1", + "grids": {"mpp0.5_px224": {}}, + } + + slide_spec = get_slide_spec("prism2") + slide_fingerprint = slide_output_fingerprint( + slide_spec, + patch_model="virchow2", + patch_output_fingerprint=patch_fingerprint, + patch_dim=patch_spec.embedding_dim, + resolved_amp="bf16", + ) + write_slide_embedding( + group, + "prism2", + np.ones(slide_output_dim(slide_spec, patch_spec.embedding_dim)), + { + "patch_encoder": "virchow2", + "embedding_dim": slide_spec.embedding_dim, + "output_fingerprint": slide_fingerprint, + }, + ) + + assert slide_encoders_requiring_compute(path, ["prism2"], device="cpu") == [] + del group["slide"]["prism2"] + assert slide_encoders_requiring_compute(path, ["prism2"], device="cpu") == [ + "prism2" + ] + + def test_prism2_fields_do_not_change_existing_prism_fingerprint_shape(): patch = patch_output_fingerprint(get_spec("virchow"), "fp16") prism = get_slide_spec("prism") diff --git a/tests/test_registry_contract.py b/tests/test_registry_contract.py index a7535ad..4a04b4d 100644 --- a/tests/test_registry_contract.py +++ b/tests/test_registry_contract.py @@ -80,6 +80,10 @@ def test_every_slide_spec_is_complete(): if spec.family != "pool": assert spec.doi and spec.doi.startswith("10."), f"{name}: needs a DOI" assert spec.weights_filename, f"{name}: needs the exact weights_filename" + if spec.weights_filename.endswith(".safetensors.index.json"): + assert spec.weights_manifest is not None, ( + f"{name}: a sharded safetensors checkpoint needs a weights_manifest" + ) if spec.weights_manifest is not None: assert spec.weights_filename in spec.weights_manifest assert spec.weights_manifest[spec.weights_filename] == ( diff --git a/tests/test_registry_revisions_network.py b/tests/test_registry_revisions_network.py index 93d8209..68fd2be 100644 --- a/tests/test_registry_revisions_network.py +++ b/tests/test_registry_revisions_network.py @@ -11,6 +11,7 @@ from __future__ import annotations +import json import os import re import urllib.request @@ -29,9 +30,12 @@ from raw2features.embedders.seal_embedder import _SEAL_REPO from raw2features.slide_embedders.model_registry import load_slide_registry from raw2features.slide_embedders.prism2 import ( + PRISM2_CODE_SHA256, PRISM2_PHI3_ARTIFACT_SHA256, PRISM2_PHI3_REVISION, PRISM2_PHI3_SOURCE, + PRISM2_REVISION, + PRISM2_SOURCE, ) _FULL_COMMIT = re.compile(r"[0-9a-f]{40}") @@ -262,6 +266,10 @@ def test_prism2_shard_manifest_matches_huggingface(): assert filename in siblings, filename lfs = siblings[filename].lfs if lfs is None: + assert not filename.endswith(".safetensors"), ( + f"prism2:{filename}: missing Git LFS digest; refusing to download " + "a multi-GB shard in a network pin test" + ) path = download_pinned_hf_file( spec.source, filename, @@ -272,6 +280,26 @@ def test_prism2_shard_manifest_matches_huggingface(): sha256 = lfs.get("sha256") if isinstance(lfs, Mapping) else lfs.sha256 assert sha256 == expected, filename + index_path = download_pinned_hf_file( + spec.source, + spec.weights_filename, + spec.weights_revision, + ) + with open(index_path, encoding="utf-8") as handle: + index = json.load(handle) + indexed_shards = set(index["weight_map"].values()) + recorded_shards = set(spec.weights_manifest) - {spec.weights_filename} + assert indexed_shards == recorded_shards + + +@pytest.mark.network +def test_prism2_custom_code_matches_pinned_sha256(): + """Every Python/config file executed by trust_remote_code matches its pin.""" + + for filename, expected in PRISM2_CODE_SHA256.items(): + path = download_pinned_hf_file(PRISM2_SOURCE, filename, PRISM2_REVISION) + verify_sha256(path, expected, what=f"prism2:custom-code:{filename}") + @pytest.mark.network def test_prism2_phi3_construction_artifacts_match_pins(): diff --git a/uv.lock b/uv.lock index 56dd062..e0e4739 100644 --- a/uv.lock +++ b/uv.lock @@ -748,16 +748,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/81/47/dd9a212ef6e343a6857485ffe25bba537304f1913bdbed446a23f7f592e1/filelock-3.29.0-py3-none-any.whl", hash = "sha256:96f5f6344709aa1572bbf631c640e4ebeeb519e08da902c39a001882f30ac258", size = 39812, upload-time = "2026-04-19T15:39:08.752Z" }, ] -[[package]] -name = "flash-attn" -version = "2.8.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "einops" }, - { name = "torch" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/3b/b2/8d76c41ad7974ee264754709c22963447f7f8134613fd9ce80984ed0dab7/flash_attn-2.8.3.tar.gz", hash = "sha256:1e71dd64a9e0280e0447b8a0c2541bad4bf6ac65bdeaa2f90e51a9e57de0370d", size = 8447812, upload-time = "2025-08-15T08:28:12.911Z" } - [[package]] name = "fonttools" version = "4.63.0" @@ -1926,6 +1916,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bb/ba/659fb680390f648e28e256659640e44fbe02b03e12cccdd20b14827e16f3/ngff_zarr-0.38.0-py3-none-any.whl", hash = "sha256:2f07ca50ad3b912938e4777ce2de9b8f3350f370bf73702143018933f73a0e4a", size = 195036, upload-time = "2026-07-09T20:27:27.59Z" }, ] +[[package]] +name = "ninja" +version = "1.13.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/43/73/79a0b22fc731989c708068427579e840a6cf4e937fe7ae5c5d0b7356ac22/ninja-1.13.0.tar.gz", hash = "sha256:4a40ce995ded54d9dc24f8ea37ff3bf62ad192b547f6c7126e7e25045e76f978", size = 242558, upload-time = "2025-08-11T15:10:19.421Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/74/d02409ed2aa865e051b7edda22ad416a39d81a84980f544f8de717cab133/ninja-1.13.0-py3-none-macosx_10_9_universal2.whl", hash = "sha256:fa2a8bfc62e31b08f83127d1613d10821775a0eb334197154c4d6067b7068ff1", size = 310125, upload-time = "2025-08-11T15:09:50.971Z" }, + { url = "https://files.pythonhosted.org/packages/8e/de/6e1cd6b84b412ac1ef327b76f0641aeb5dcc01e9d3f9eee0286d0c34fd93/ninja-1.13.0-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3d00c692fb717fd511abeb44b8c5d00340c36938c12d6538ba989fe764e79630", size = 177467, upload-time = "2025-08-11T15:09:52.767Z" }, + { url = "https://files.pythonhosted.org/packages/c8/83/49320fb6e58ae3c079381e333575fdbcf1cca3506ee160a2dcce775046fa/ninja-1.13.0-py3-none-manylinux2014_i686.manylinux_2_17_i686.whl", hash = "sha256:be7f478ff9f96a128b599a964fc60a6a87b9fa332ee1bd44fa243ac88d50291c", size = 187834, upload-time = "2025-08-11T15:09:54.115Z" }, + { url = "https://files.pythonhosted.org/packages/56/c7/ba22748fb59f7f896b609cd3e568d28a0a367a6d953c24c461fe04fc4433/ninja-1.13.0-py3-none-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:60056592cf495e9a6a4bea3cd178903056ecb0943e4de45a2ea825edb6dc8d3e", size = 202736, upload-time = "2025-08-11T15:09:55.745Z" }, + { url = "https://files.pythonhosted.org/packages/79/22/d1de07632b78ac8e6b785f41fa9aad7a978ec8c0a1bf15772def36d77aac/ninja-1.13.0-py3-none-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:1c97223cdda0417f414bf864cfb73b72d8777e57ebb279c5f6de368de0062988", size = 179034, upload-time = "2025-08-11T15:09:57.394Z" }, + { url = "https://files.pythonhosted.org/packages/ed/de/0e6edf44d6a04dabd0318a519125ed0415ce437ad5a1ec9b9be03d9048cf/ninja-1.13.0-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fb46acf6b93b8dd0322adc3a4945452a4e774b75b91293bafcc7b7f8e6517dfa", size = 180716, upload-time = "2025-08-11T15:09:58.696Z" }, + { url = "https://files.pythonhosted.org/packages/54/28/938b562f9057aaa4d6bfbeaa05e81899a47aebb3ba6751e36c027a7f5ff7/ninja-1.13.0-py3-none-manylinux_2_28_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:4be9c1b082d244b1ad7ef41eb8ab088aae8c109a9f3f0b3e56a252d3e00f42c1", size = 146843, upload-time = "2025-08-11T15:10:00.046Z" }, + { url = "https://files.pythonhosted.org/packages/2a/fb/d06a3838de4f8ab866e44ee52a797b5491df823901c54943b2adb0389fbb/ninja-1.13.0-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:6739d3352073341ad284246f81339a384eec091d9851a886dfa5b00a6d48b3e2", size = 154402, upload-time = "2025-08-11T15:10:01.657Z" }, + { url = "https://files.pythonhosted.org/packages/31/bf/0d7808af695ceddc763cf251b84a9892cd7f51622dc8b4c89d5012779f06/ninja-1.13.0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:11be2d22027bde06f14c343f01d31446747dbb51e72d00decca2eb99be911e2f", size = 552388, upload-time = "2025-08-11T15:10:03.349Z" }, + { url = "https://files.pythonhosted.org/packages/9d/70/c99d0c2c809f992752453cce312848abb3b1607e56d4cd1b6cded317351a/ninja-1.13.0-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:aa45b4037b313c2f698bc13306239b8b93b4680eb47e287773156ac9e9304714", size = 472501, upload-time = "2025-08-11T15:10:04.735Z" }, + { url = "https://files.pythonhosted.org/packages/9f/43/c217b1153f0e499652f5e0766da8523ce3480f0a951039c7af115e224d55/ninja-1.13.0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:5f8e1e8a1a30835eeb51db05cf5a67151ad37542f5a4af2a438e9490915e5b72", size = 638280, upload-time = "2025-08-11T15:10:06.512Z" }, + { url = "https://files.pythonhosted.org/packages/8c/45/9151bba2c8d0ae2b6260f71696330590de5850e5574b7b5694dce6023e20/ninja-1.13.0-py3-none-musllinux_1_2_ppc64le.whl", hash = "sha256:3d7d7779d12cb20c6d054c61b702139fd23a7a964ec8f2c823f1ab1b084150db", size = 642420, upload-time = "2025-08-11T15:10:08.35Z" }, + { url = "https://files.pythonhosted.org/packages/3c/fb/95752eb635bb8ad27d101d71bef15bc63049de23f299e312878fc21cb2da/ninja-1.13.0-py3-none-musllinux_1_2_riscv64.whl", hash = "sha256:d741a5e6754e0bda767e3274a0f0deeef4807f1fec6c0d7921a0244018926ae5", size = 585106, upload-time = "2025-08-11T15:10:09.818Z" }, + { url = "https://files.pythonhosted.org/packages/c1/31/aa56a1a286703800c0cbe39fb4e82811c277772dc8cd084f442dd8e2938a/ninja-1.13.0-py3-none-musllinux_1_2_s390x.whl", hash = "sha256:e8bad11f8a00b64137e9b315b137d8bb6cbf3086fbdc43bf1f90fd33324d2e96", size = 707138, upload-time = "2025-08-11T15:10:11.366Z" }, + { url = "https://files.pythonhosted.org/packages/34/6f/5f5a54a1041af945130abdb2b8529cbef0cdcbbf9bcf3f4195378319d29a/ninja-1.13.0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:b4f2a072db3c0f944c32793e91532d8948d20d9ab83da9c0c7c15b5768072200", size = 581758, upload-time = "2025-08-11T15:10:13.295Z" }, + { url = "https://files.pythonhosted.org/packages/95/97/51359c77527d45943fe7a94d00a3843b81162e6c4244b3579fe8fc54cb9c/ninja-1.13.0-py3-none-win32.whl", hash = "sha256:8cfbb80b4a53456ae8a39f90ae3d7a2129f45ea164f43fadfa15dc38c4aef1c9", size = 267201, upload-time = "2025-08-11T15:10:15.158Z" }, + { url = "https://files.pythonhosted.org/packages/29/45/c0adfbfb0b5895aa18cec400c535b4f7ff3e52536e0403602fc1a23f7de9/ninja-1.13.0-py3-none-win_amd64.whl", hash = "sha256:fb8ee8719f8af47fed145cced4a85f0755dd55d45b2bddaf7431fa89803c5f3e", size = 309975, upload-time = "2025-08-11T15:10:16.697Z" }, + { url = "https://files.pythonhosted.org/packages/df/93/a7b983643d1253bb223234b5b226e69de6cda02b76cdca7770f684b795f5/ninja-1.13.0-py3-none-win_arm64.whl", hash = "sha256:3c0b40b1f0bba764644385319028650087b4c1b18cdfa6f45cb39a3669b81aa9", size = 290806, upload-time = "2025-08-11T15:10:18.018Z" }, +] + [[package]] name = "numba" version = "0.65.1" @@ -3464,8 +3480,9 @@ prism = [ prism2 = [ { name = "einops" }, { name = "einops-exts" }, - { name = "flash-attn" }, { name = "huggingface-hub" }, + { name = "ninja" }, + { name = "packaging" }, { name = "safetensors" }, { name = "timm" }, { name = "torch" }, @@ -3546,7 +3563,6 @@ requires-dist = [ { name = "einops-exts", marker = "extra == 'models'", specifier = ">=0.0.4" }, { name = "environs", marker = "extra == 'prism'" }, { name = "fairscale", marker = "extra == 'musk'" }, - { name = "flash-attn", marker = "extra == 'prism2'", specifier = "==2.8.3" }, { name = "fsspec", extras = ["http"], marker = "extra == 'zarr'", specifier = ">=2024.0" }, { name = "gdown", marker = "extra == 'chief'", specifier = ">=5.2" }, { name = "gdown", marker = "extra == 'tangle'", specifier = ">=5.2" }, @@ -3557,6 +3573,7 @@ requires-dist = [ { name = "imagecodecs", marker = "extra == 'kronos'" }, { name = "jsonschema", specifier = ">=4.18" }, { name = "ngff-zarr", marker = "extra == 'zarr'", specifier = ">=0.13" }, + { name = "ninja", marker = "extra == 'prism2'", specifier = ">=1.11" }, { name = "numcodecs", marker = "extra == 'zarr'", specifier = ">=0.15" }, { name = "numpy", specifier = ">=2.0" }, { name = "ome-zarr", marker = "extra == 'zarr'", specifier = ">=0.10" }, @@ -3564,6 +3581,7 @@ requires-dist = [ { name = "open-clip-torch", marker = "extra == 'musk'" }, { name = "open-clip-torch", marker = "extra == 'open-clip'", specifier = ">=3.0" }, { name = "opencv-python-headless", marker = "extra == 'image'", specifier = ">=4.9" }, + { name = "packaging", marker = "extra == 'prism2'", specifier = ">=23" }, { name = "pandas", marker = "extra == 'kronos2'", specifier = ">=2.2.2" }, { name = "pandas", marker = "extra == 'tables'", specifier = ">=2.2" }, { name = "peft", marker = "extra == 'seal'" }, From 0eda9f1b32380d7ac79f4fffc4bcd306030a6851 Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Fri, 14 Aug 2026 12:50:58 +0100 Subject: [PATCH 4/7] Refine PRISM2 runtime preflight --- docs/MODELS.md | 5 +- pyproject.toml | 2 + src/raw2features/cli/embed.py | 27 ++- src/raw2features/cli/embed_many.py | 63 ++---- src/raw2features/cli/slide_embed.py | 76 ++++--- src/raw2features/pipeline/runner.py | 119 +++++++---- src/raw2features/slide_embedders/encoding.py | 64 ++---- src/raw2features/slide_embedders/prism2.py | 29 ++- tests/test_embed_many.py | 38 +--- tests/test_prism2.py | 201 ++++++++++++++----- tests/test_registry_contract.py | 4 +- tests/test_release_metadata.py | 10 + tests/test_slide_embedders.py | 69 +++++++ uv.lock | 16 ++ 14 files changed, 455 insertions(+), 268 deletions(-) diff --git a/docs/MODELS.md b/docs/MODELS.md index 347b568..39bd99a 100644 --- a/docs/MODELS.md +++ b/docs/MODELS.md @@ -114,8 +114,9 @@ method removed after its released 4.51.3 stack, so raw2features restores that de mask construction and records the compatibility contract in the output fingerprint. The validated path requires FlashAttention 2.8.3. PyPI distributes it as source, so install it after the extra with build isolation disabled as shown above, or install an -ABI-compatible 2.8.3 wheel. Increase `MAX_JOBS` only when the build host has sufficient -RAM. The checkpoint is approximately 17.9 GB. +ABI-compatible 2.8.3 wheel. The extra supplies the source build's packaging, Ninja, +setuptools and wheel prerequisites. Increase `MAX_JOBS` only when the build host has +sufficient RAM. The checkpoint is approximately 17.9 GB. The base path still constructs the complete released model, so it has the same download and device-memory requirements as the diagnostic path. diff --git a/pyproject.toml b/pyproject.toml index 14935c8..c686667 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -188,6 +188,8 @@ prism2 = [ "transformers==4.56.0", "packaging>=23", "ninja>=1.11", + "setuptools>=64", + "wheel>=0.41", ] # MADELEINE slide encoder (Mahmood Lab): loads via the authors' non-PyPI `madeleine` git # package (post-install, like conch): diff --git a/src/raw2features/cli/embed.py b/src/raw2features/cli/embed.py index 7f51df2..4e97f5a 100644 --- a/src/raw2features/cli/embed.py +++ b/src/raw2features/cli/embed.py @@ -8,6 +8,7 @@ import typer from raw2features.core.provenance import sanitize_argv +from raw2features.core.uris import redact_uri_credentials from raw2features.pipeline.runner import RunConfig, embed_slide from raw2features.viz import DEFAULT_THUMBNAIL_MPP @@ -330,16 +331,20 @@ def embed( qc_artifact_mpp=qc_artifact_mpp, output_zarr_format=output_zarr_format, ) - summary = embed_slide( - slide, - out_dir, - cfg, - requested_mpp=mpp, - requested_patch_px=patch_size, - geometry_config=geometry_config, - receipts_dir=receipts_dir, - cli=sanitize_argv(sys.argv), - force=force, - ) + try: + summary = embed_slide( + slide, + out_dir, + cfg, + requested_mpp=mpp, + requested_patch_px=patch_size, + geometry_config=geometry_config, + receipts_dir=receipts_dir, + cli=sanitize_argv(sys.argv), + force=force, + ) + except (KeyError, RuntimeError, ValueError) as exc: + typer.echo(redact_uri_credentials(f"Error: {exc}"), err=True) + raise typer.Exit(1) from exc for key, value in summary.items(): typer.echo(f"{key}: {value}") diff --git a/src/raw2features/cli/embed_many.py b/src/raw2features/cli/embed_many.py index 7b17a82..9500144 100644 --- a/src/raw2features/cli/embed_many.py +++ b/src/raw2features/cli/embed_many.py @@ -391,6 +391,19 @@ def embed_many( cli = sanitize_argv(sys.argv) device_list = cfg.device_list() + try: + from raw2features.slide_embedders.model_registry import ( + validate_slide_encoder_runtime, + ) + + # embed-many loads patch models once before visiting individual stores. Its + # slide-parallel workers run slide encoders on these resolved devices, so + # validate that actual runtime now rather than after an expensive warm load. + validate_slide_encoder_runtime(cfg.slide_encoders, devices=device_list) + except (RuntimeError, ValueError) as exc: + typer.echo(f"Error: slide-encoder runtime preflight failed: {exc}", err=True) + raise typer.Exit(1) from exc + # Validate every positional panel before a warm worker loads model weights. This is # especially important for one --channel-names-file applied to a whole cohort: a # count or metadata conflict must not surface only after an expensive gated load. @@ -427,18 +440,6 @@ def embed_many( ) raise typer.Exit(1) from exc - try: - _preflight_slide_encoder_runtime( - shard, - out_dir, - cfg, - device_list, - force=force, - ) - except (RuntimeError, ValueError) as exc: - typer.echo(f"Error: slide-encoder runtime preflight failed: {exc}", err=True) - raise typer.Exit(1) from exc - t0 = time.time() try: if len(device_list) > 1: @@ -482,44 +483,6 @@ def _classify(summary: dict) -> str: return "skipped" if summary["status"] == "skipped" else "done" -def _preflight_slide_encoder_runtime( - shard: list[dict], - out_dir: str, - cfg: RunConfig, - devices: list[str], - *, - force: bool, -) -> None: - """Check a required slide runtime before warm patch models are loaded.""" - - if not cfg.slide_encoders: - return - from raw2features.slide_embedders.encoding import ( - slide_encoders_requiring_compute, - ) - from raw2features.slide_embedders.model_registry import ( - validate_slide_encoder_runtime, - ) - - required = set(cfg.slide_encoders) if force else set() - if not force: - inspection_device = devices[0] if devices else cfg.device - for row in shard: - slide_id = slide_id_from_path(row["path"]) - store_path = os.path.join(out_dir, f"{slide_id}.embeddings.zarr") - required.update( - slide_encoders_requiring_compute( - store_path, - cfg.slide_encoders, - device=inspection_device, - ) - ) - if len(required) == len(set(cfg.slide_encoders)): - break - ordered = [name for name in dict.fromkeys(cfg.slide_encoders) if name in required] - validate_slide_encoder_runtime(ordered, devices=devices) - - def _resolve_manifest_sources(rows: list[dict], slide_dir: str) -> list[dict]: """Resolve only relative local manifest paths against ``slide_dir`` in place.""" diff --git a/src/raw2features/cli/slide_embed.py b/src/raw2features/cli/slide_embed.py index 7c116a5..63c80cf 100644 --- a/src/raw2features/cli/slide_embed.py +++ b/src/raw2features/cli/slide_embed.py @@ -102,47 +102,65 @@ def slide_embed( # use_consolidated=False: the store was consolidated by `embed`, so its # consolidated metadata predates any slide/ group we add here. Reading the # live metadata avoids stale-key KeyErrors; we re-consolidate at the end. - root = zarr.open_group(path, mode="r+", use_consolidated=False) - - for slide_model_name in slide_encoder: - try: + plans = [] + try: + root = zarr.open_group(path, mode="r+", use_consolidated=False) + for slide_model_name in slide_encoder: selected_grid, group, selected_patch_model = resolve_slide_grid( root, slide_model_name, grid=grid, patch_model=patch_model, ) - except (KeyError, RuntimeError, ValueError) as exc: - typer.echo(f"Error: {exc}", err=True) - raise typer.Exit(1) from exc + # Strategy-derived multiplex pools are patch-qualified so several marker + # recipes can coexist on the same grid without replacing one another. + output_name = slide_output_key( + group, slide_model_name, selected_patch_model + ) + complete = not force and slide_embedding_is_complete( + group, + slide_model_name, + patch_model=selected_patch_model, + device=device, + output_name=output_name, + ) + plans.append( + ( + slide_model_name, + selected_grid, + group, + selected_patch_model, + output_name, + complete, + ) + ) + except (KeyError, OSError, RuntimeError, TypeError, ValueError) as exc: + typer.echo(f"Error: {exc}", err=True) + raise typer.Exit(1) from exc - # Strategy-derived multiplex pools are patch-qualified so several marker - # recipes can coexist on the same grid without replacing one another. - output_name = slide_output_key( - group, slide_model_name, selected_patch_model + try: + validate_slide_encoder_runtime( + [plan[0] for plan in plans if not plan[-1]], + devices=[device], ) - # Skip only an output produced from the requested patch model. - if not force and slide_embedding_is_complete( - group, - slide_model_name, - patch_model=selected_patch_model, - device=device, - output_name=output_name, - ): + except (RuntimeError, ValueError) as exc: + typer.echo(f"Error: {exc}", err=True) + raise typer.Exit(1) from exc + + for ( + slide_model_name, + selected_grid, + group, + selected_patch_model, + output_name, + complete, + ) in plans: + if complete: typer.echo( f"{slide_model_name} [{selected_grid}]: already complete (skipping)" ) continue - try: - validate_slide_encoder_runtime( - [slide_model_name], - devices=[device], - ) - except (RuntimeError, ValueError) as exc: - typer.echo(f"Error: {exc}", err=True) - raise typer.Exit(1) from exc - typer.echo( f"{slide_model_name} [{selected_grid}]: encoding from " f"'{selected_patch_model}' patch features …" @@ -154,7 +172,7 @@ def slide_embed( device, patch_model=selected_patch_model, ) - except (KeyError, ValueError) as exc: + except (KeyError, RuntimeError, ValueError) as exc: typer.echo(f"Error: {exc}", err=True) raise typer.Exit(1) from exc if encoding is None: diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index 036f2e8..db3d1cc 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -1454,29 +1454,14 @@ def run_slide( models_to_do = list(cfg.models) if cfg.slide_encoders: - from raw2features.slide_embedders.encoding import ( - slide_encoders_requiring_compute, - ) - from raw2features.slide_embedders.model_registry import ( - get_slide_spec, - validate_slide_encoder_runtime, - ) - - runtime_names = ( - list(cfg.slide_encoders) - if force - else slide_encoders_requiring_compute( - out_path, - cfg.slide_encoders, - device=cfg.device, - ) - ) - for name in cfg.slide_encoders: - if get_slide_spec(name).patch_encoder in models_to_do: - runtime_names.append(name) - validate_slide_encoder_runtime( - list(dict.fromkeys(runtime_names)), - devices=devices, + available = sorted(set(present_valid) | set(models_to_do)) + _preflight_slide_encoders_for_grid( + out_path, + grid_key_existing, + cfg.slide_encoders, + available, + models_to_do, + cfg.device, ) qc_to_do = list(cfg.qc) @@ -2228,23 +2213,6 @@ def embed_slide( "grids": grids, } - if cfg.slide_encoders: - from raw2features.slide_embedders.encoding import ( - slide_encoders_requiring_compute, - ) - from raw2features.slide_embedders.model_registry import ( - validate_slide_encoder_runtime, - ) - - validate_slide_encoder_runtime( - slide_encoders_requiring_compute( - out_path, - cfg.slide_encoders, - device=cfg.device, - ), - devices=devices, - ) - started = time.time() results = [] for i, gc in enumerate(group_cfgs): @@ -3823,6 +3791,75 @@ def _slide_encoders_for(names: list[str], available: list[str]) -> list[str]: return out +def _preflight_slide_encoders_for_grid( + out_path: str, + grid_key_existing: str | None, + names: list[str], + available_patch_models: list[str], + patch_models_to_write: list[str], + slide_device: str, +) -> list[str]: + """Validate runtimes needed by this exact grid before patch work starts.""" + + from raw2features.slide_embedders.model_registry import ( + get_slide_spec, + validate_slide_encoder_runtime, + ) + + candidates = [ + name + for name in _slide_encoders_for(names, available_patch_models) + if get_slide_spec(name).family == "prism2" + ] + if not candidates: + return [] + + required = [ + name + for name in candidates + if grid_key_existing is None + or get_slide_spec(name).patch_encoder in patch_models_to_write + ] + if grid_key_existing is not None and len(required) != len(candidates): + try: + import zarr + + from raw2features.core.store import open_grid + from raw2features.slide_embedders.encoding import ( + resolve_slide_patch_model, + slide_embedding_is_complete, + slide_output_key, + ) + + root = zarr.open_group(out_path, mode="r", use_consolidated=False) + group = open_grid(root, grid_key_existing) + for name in candidates: + if name in required: + continue + patch_model = resolve_slide_patch_model( + group, + name, + available_patch_models=available_patch_models, + ) + output_name = slide_output_key(group, name, patch_model) + if not slide_embedding_is_complete( + group, + name, + patch_model=patch_model, + device=slide_device, + output_name=output_name, + # _inspect_store already validated this exact patch array. + patch_is_complete=True, + ): + required.append(name) + except Exception: # noqa: BLE001 - unreadable state cannot prove completion + required = list(candidates) + + required = list(dict.fromkeys(required)) + validate_slide_encoder_runtime(required, devices=[slide_device]) + return required + + def _run_qc( qc_tools, reader, @@ -3911,6 +3948,7 @@ def _run_slide_encoders( patch_model=patch_model, device=device, output_name=output_name, + patch_is_complete=True, ): results[slide_model_name] = f"slide/{output_name}" continue @@ -3921,6 +3959,7 @@ def _run_slide_encoders( device, patch_model=patch_model, available_patch_models=available_patch_models, + patch_is_complete=True, ) if encoding is None: continue diff --git a/src/raw2features/slide_embedders/encoding.py b/src/raw2features/slide_embedders/encoding.py index da903d7..419e96d 100644 --- a/src/raw2features/slide_embedders/encoding.py +++ b/src/raw2features/slide_embedders/encoding.py @@ -7,7 +7,6 @@ from __future__ import annotations -import os import warnings from dataclasses import dataclass @@ -51,7 +50,12 @@ def slide_output_key(group, slide_model: str, patch_model: str) -> str: return f"{slide_model}__{patch_model}" -def _validated_patch_fingerprint(group, patch_model: str) -> dict: +def _validated_patch_fingerprint( + group, + patch_model: str, + *, + validate_data: bool = True, +) -> dict: """Return a committed, complete patch-output fingerprint. Slide outputs inherit the exact patch-output identity, so a self-consistent @@ -102,7 +106,7 @@ def _validated_patch_fingerprint(group, patch_model: str) -> dict: "The selected grid has no valid coordinate array; rerun raw2features " "embed before slide encoding." ) - if not validate_model( + if validate_data and not validate_model( group, patch_model, int(group["coords"].shape[0]), @@ -123,6 +127,7 @@ def slide_embedding_is_complete( patch_model: str | None = None, device: str = "cpu", output_name: str | None = None, + patch_is_complete: bool = False, ) -> bool: """Return whether ``slide/`` is valid for the requested patch model.""" stored_name = output_name or slide_model @@ -148,7 +153,11 @@ def slide_embedding_is_complete( patch_array = group["features"][stored_patch_model] try: - patch_fingerprint = _validated_patch_fingerprint(group, stored_patch_model) + patch_fingerprint = _validated_patch_fingerprint( + group, + stored_patch_model, + validate_data=not patch_is_complete, + ) spec = get_slide_spec(slide_model) expected_dim = slide_output_dim(spec, int(patch_array.shape[1])) expected_fingerprint = slide_output_fingerprint( @@ -196,46 +205,6 @@ def slide_embedding_is_complete( return bool(vector.size and np.isfinite(vector).all() and (vector != 0).any()) -def slide_encoders_requiring_compute( - store_path: str, - slide_models: list[str], - *, - device: str = "cpu", -) -> list[str]: - """Return requested slide encoders without a complete current store output.""" - - if not slide_models: - return [] - path = store_path.removeprefix("file://") - if not os.path.exists(path): - return list(slide_models) - - import zarr - - try: - root = zarr.open_group(path, mode="r", use_consolidated=False) - except Exception: # noqa: BLE001 - an unreadable store cannot prove completion - return list(slide_models) - - missing = [] - for slide_model in slide_models: - try: - _, group, patch_model = resolve_slide_grid(root, slide_model) - output_name = slide_output_key(group, slide_model, patch_model) - complete = slide_embedding_is_complete( - group, - slide_model, - patch_model=patch_model, - device=device, - output_name=output_name, - ) - except (KeyError, TypeError, ValueError): - complete = False - if not complete: - missing.append(slide_model) - return missing - - def resolve_slide_grid( root, slide_model: str, @@ -331,6 +300,7 @@ def encode_slide_embedding( *, patch_model: str | None = None, available_patch_models: list[str] | None = None, + patch_is_complete: bool = False, ) -> SlideEncoding | None: """Encode ``slide_model`` from patch features in one open grid group. @@ -367,7 +337,11 @@ def encode_slide_embedding( # Refuse to mint current-looking slide provenance over a legacy/unknown patch # array. The patch model must first be recomputed under the current contract. slide_embedder = build_slide_embedder(slide_model) - patch_fingerprint = _validated_patch_fingerprint(group, selected_patch_model) + patch_fingerprint = _validated_patch_fingerprint( + group, + selected_patch_model, + validate_data=not patch_is_complete, + ) output_fingerprint = slide_output_fingerprint( slide_embedder.spec, patch_model=selected_patch_model, diff --git a/src/raw2features/slide_embedders/prism2.py b/src/raw2features/slide_embedders/prism2.py index 19c8766..94a6d08 100644 --- a/src/raw2features/slide_embedders/prism2.py +++ b/src/raw2features/slide_embedders/prism2.py @@ -9,6 +9,7 @@ from __future__ import annotations +import importlib import sys import threading from importlib import metadata as importlib_metadata @@ -268,11 +269,24 @@ def _verified_snapshot( what: str, ) -> str: expected = tuple(artifacts) + local_dir = pinned_model_cache_dir(source, revision) + if Path(local_dir).is_dir() and all( + (Path(local_dir) / filename).is_file() for filename in expected + ): + _assert_snapshot_files( + local_dir, + expected, + allowed_files=allowed_files, + what=what, + ) + _verify_snapshot_artifacts_once(local_dir, artifacts, what=what) + return local_dir + snapshot = download_pinned_hf_snapshot( source, revision, allow_patterns=expected, - local_dir=pinned_model_cache_dir(source, revision), + local_dir=local_dir, ) _assert_snapshot_files( snapshot, @@ -349,9 +363,7 @@ def _require_prism2_runtime() -> None: installed = importlib_metadata.version(distribution) except importlib_metadata.PackageNotFoundError: installed = None - if installed is None or not SpecifierSet(f"=={expected}").contains( - installed, prereleases=True - ): + if installed is None or not SpecifierSet(f"=={expected}").contains(installed): found = installed or "missing" mismatches.append(f"{distribution}=={expected} (found {found})") if mismatches: @@ -360,6 +372,15 @@ def _require_prism2_runtime() -> None: "then install flash-attn==2.8.3 with --no-build-isolation. Mismatch: " + ", ".join(mismatches) ) + try: + importlib.import_module("flash_attn") + except Exception as exc: # noqa: BLE001 - surface binary/ABI import failures early + raise RuntimeError( + "PRISM2 found flash-attn==2.8.3, but it could not be imported. " + "Rebuild flash-attn against the active Torch/CUDA environment with " + "--no-build-isolation. Import error: " + f"{exc}" + ) from exc class _Prism2SlideEmbedder(SlideEmbedder): diff --git a/tests/test_embed_many.py b/tests/test_embed_many.py index 0c7ceef..67228a0 100644 --- a/tests/test_embed_many.py +++ b/tests/test_embed_many.py @@ -86,10 +86,12 @@ def test_embed_many_runtime_preflight_runs_before_warm_model_load( slides = tmp_path / "slides" slides.mkdir() (slides / "S.zarr").mkdir() + import raw2features.slide_embedders.prism2 as prism2 + monkeypatch.setattr( - em, - "_preflight_slide_encoder_runtime", - lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("runtime missing")), + prism2, + "_require_prism2_runtime", + lambda: (_ for _ in ()).throw(RuntimeError("runtime missing")), ) monkeypatch.setattr( em, @@ -117,36 +119,6 @@ def test_embed_many_runtime_preflight_runs_before_warm_model_load( assert "models must not be loaded" not in result.output -def test_embed_many_runtime_preflight_ignores_complete_slide_outputs( - tmp_path, monkeypatch -): - import raw2features.slide_embedders.encoding as encoding - import raw2features.slide_embedders.model_registry as slide_registry - from raw2features.cli.embed_many import _preflight_slide_encoder_runtime - - checked = [] - monkeypatch.setattr( - encoding, - "slide_encoders_requiring_compute", - lambda *args, **kwargs: [], - ) - monkeypatch.setattr( - slide_registry, - "validate_slide_encoder_runtime", - lambda names, *, devices: checked.append((names, devices)), - ) - - _preflight_slide_encoder_runtime( - [{"path": str(tmp_path / "S.zarr")}], - str(tmp_path / "out"), - RunConfig(models=["mock"], slide_encoders=["prism2"]), - ["cpu"], - force=False, - ) - - assert checked == [([], ["cpu"])] - - @pytest.mark.parametrize( "rows", [ diff --git a/tests/test_prism2.py b/tests/test_prism2.py index 351a4a6..2512044 100644 --- a/tests/test_prism2.py +++ b/tests/test_prism2.py @@ -17,10 +17,8 @@ slide_output_fingerprint, ) from raw2features.embedders.model_registry import get_spec -from raw2features.slide_embedders.encoding import ( - slide_encoders_requiring_compute, - write_slide_embedding, -) +from raw2features.pipeline.runner import _preflight_slide_encoders_for_grid +from raw2features.slide_embedders.encoding import write_slide_embedding from raw2features.slide_embedders.model_registry import ( get_slide_spec, validate_slide_encoder_runtime, @@ -49,6 +47,39 @@ ) +def _write_prism2_patch_grid(root, key: str): + group = root["grids"].create_group(key) + coords = group.create_array("coords", shape=(2, 2), dtype="int32") + coords[:] = np.asarray([[0, 0], [224, 0]], dtype=np.int32) + patch_spec = get_spec("virchow2") + patch_fingerprint = patch_output_fingerprint(patch_spec, "fp16") + features = group.create_group("features").create_array( + "virchow2", + shape=(2, patch_spec.embedding_dim), + chunks=(2, patch_spec.embedding_dim), + dtype="float32", + ) + features[:] = 1.0 + features.attrs.update( + { + "role": "features", + "model": "virchow2", + "output_fingerprint": patch_fingerprint, + } + ) + group.attrs["raw2features"] = { + "schema_version": "0.1", + "models": { + "virchow2": { + "embedding_dim": patch_spec.embedding_dim, + "output_fingerprint": patch_fingerprint, + } + }, + "patching": {"level0_patch": 224}, + } + return group, patch_spec, patch_fingerprint + + def test_prism2_registry_records_both_published_embeddings(): base = get_slide_spec("prism2") diagnostic = get_slide_spec("prism2_diagnostic") @@ -276,10 +307,54 @@ def test_prism2_runtime_pin_accepts_pep440_equivalent_versions(monkeypatch): "version", lambda distribution: versions[distribution], ) + monkeypatch.setattr(module.importlib, "import_module", lambda name: object()) _require_prism2_runtime() +def test_prism2_runtime_preflight_rejects_an_incompatible_flash_attn_abi( + monkeypatch, +): + import raw2features.slide_embedders.prism2 as module + + versions = {"transformers": "4.56.0", "flash-attn": "2.8.3"} + monkeypatch.setattr( + module.importlib_metadata, + "version", + lambda distribution: versions[distribution], + ) + + def import_module(name): + assert name == "flash_attn" + raise ImportError("undefined symbol: _ZN3c105ErrorC2E") + + monkeypatch.setattr(module.importlib, "import_module", import_module) + + with pytest.raises( + RuntimeError, + match=( + r"found flash-attn==2\.8\.3, but it could not be imported.*" + r"Rebuild flash-attn against the active Torch/CUDA environment.*" + r"undefined symbol" + ), + ): + _require_prism2_runtime() + + +def test_prism2_runtime_pin_rejects_unvalidated_prereleases(monkeypatch): + import raw2features.slide_embedders.prism2 as module + + versions = {"transformers": "4.56.0.dev0", "flash-attn": "2.8.3"} + monkeypatch.setattr( + module.importlib_metadata, + "version", + lambda distribution: versions[distribution], + ) + + with pytest.raises(RuntimeError, match=r"found 4\.56\.0\.dev0"): + _require_prism2_runtime() + + def test_prism2_runtime_preflight_checks_packages_and_device(monkeypatch): import raw2features.slide_embedders.prism2 as module @@ -337,6 +412,14 @@ def test_prism2_restores_released_phi3_mask_contract(): def test_prism2_binds_phi3_helper_only_to_verified_dynamic_module(monkeypatch): pytest.importorskip("transformers") + from transformers.models.phi3.modeling_phi3 import Phi3Model + + sentinel = object() + shared_before = getattr( + Phi3Model, + "_prepare_4d_causal_attention_mask_with_cache_position", + sentinel, + ) module_name = "transformers_modules.paige_prism2.modeling_prism2" original = type( "Phi3Model", @@ -352,10 +435,13 @@ def test_prism2_binds_phi3_helper_only_to_verified_dynamic_module(monkeypatch): helper = dynamic_module.Phi3Model assert helper is not original assert callable(helper._prepare_4d_causal_attention_mask_with_cache_position) - from transformers.models.phi3.modeling_phi3 import Phi3Model - - assert not hasattr( - Phi3Model, "_prepare_4d_causal_attention_mask_with_cache_position" + assert ( + getattr( + Phi3Model, + "_prepare_4d_causal_attention_mask_with_cache_position", + sentinel, + ) + is shared_before ) @@ -388,17 +474,27 @@ def test_prism2_snapshot_allowlist_fails_closed(tmp_path): def test_prism2_runtime_code_is_sha_bound_and_cached(monkeypatch, tmp_path): - artifact = tmp_path / "modeling_prism2.py" - artifact.write_bytes(b"verified code") - calls = [] + cache = tmp_path / "cache" + artifact = cache / "modeling_prism2.py" + downloads = [] + hashes = [] + + def download(*args, **kwargs): + downloads.append(kwargs["local_dir"]) + cache.mkdir() + artifact.write_bytes(b"verified code") + return str(cache) monkeypatch.setattr( - "raw2features.slide_embedders.prism2.download_pinned_hf_snapshot", - lambda *args, **kwargs: str(tmp_path), + "raw2features.slide_embedders.prism2.download_pinned_hf_snapshot", download + ) + monkeypatch.setattr( + "raw2features.slide_embedders.prism2.pinned_model_cache_dir", + lambda *args: str(cache), ) def verify(path, expected, *, what): - calls.append((Path(path).name, expected, what)) + hashes.append((Path(path).name, expected, what)) monkeypatch.setattr( "raw2features.slide_embedders.prism2.verify_sha256", @@ -409,11 +505,13 @@ def verify(path, expected, *, what): _verified_snapshot("hf-hub:test/model", "a" * 40, artifacts, what="test") _verified_snapshot("hf-hub:test/model", "a" * 40, artifacts, what="test") - assert len(calls) == 1 + assert len(downloads) == 1 + assert len(hashes) == 1 artifact.write_bytes(b"changed code") _verified_snapshot("hf-hub:test/model", "a" * 40, artifacts, what="test") - assert len(calls) == 2 + assert len(downloads) == 1 + assert len(hashes) == 2 def test_prism2_rejects_bad_runtime_digest_before_custom_code(monkeypatch, tmp_path): @@ -425,6 +523,10 @@ def test_prism2_rejects_bad_runtime_digest_before_custom_code(monkeypatch, tmp_p "raw2features.slide_embedders.prism2.download_pinned_hf_snapshot", lambda *args, **kwargs: str(tmp_path), ) + monkeypatch.setattr( + "raw2features.slide_embedders.prism2.pinned_model_cache_dir", + lambda *args: str(tmp_path), + ) monkeypatch.setattr( "raw2features.slide_embedders.prism2._require_prism2_runtime", lambda: None, @@ -521,40 +623,15 @@ def test_prism2_amp_contract_is_device_independent(): assert resolved_slide_amp(spec, "cuda:0") == "bf16" -def test_prism2_complete_output_can_be_recognized_without_gpu_runtime(tmp_path): +def test_prism2_grid_preflight_uses_exact_grid_and_slide_device( + tmp_path, monkeypatch +): path = str(tmp_path / "complete.embeddings.zarr") root = zarr.open_group(path, mode="w", zarr_format=2) - grids = root.create_group("grids") - group = grids.create_group("mpp0.5_px224") - coords = group.create_array("coords", shape=(2, 2), dtype="int32") - coords[:] = np.asarray([[0, 0], [224, 0]], dtype=np.int32) - - patch_spec = get_spec("virchow2") - patch_fingerprint = patch_output_fingerprint(patch_spec, "fp16") - features = group.create_group("features").create_array( - "virchow2", - shape=(2, patch_spec.embedding_dim), - chunks=(2, patch_spec.embedding_dim), - dtype="float32", - ) - features[:] = 1.0 - features.attrs.update( - { - "role": "features", - "model": "virchow2", - "output_fingerprint": patch_fingerprint, - } + root.create_group("grids") + group, patch_spec, patch_fingerprint = _write_prism2_patch_grid( + root, "mpp0.5_px224" ) - group.attrs["raw2features"] = { - "schema_version": "0.1", - "models": { - "virchow2": { - "embedding_dim": patch_spec.embedding_dim, - "output_fingerprint": patch_fingerprint, - } - }, - "patching": {"level0_patch": 224}, - } root.attrs["raw2features"] = { "schema_version": "0.1", "grids": {"mpp0.5_px224": {}}, @@ -579,11 +656,31 @@ def test_prism2_complete_output_can_be_recognized_without_gpu_runtime(tmp_path): }, ) - assert slide_encoders_requiring_compute(path, ["prism2"], device="cpu") == [] - del group["slide"]["prism2"] - assert slide_encoders_requiring_compute(path, ["prism2"], device="cpu") == [ - "prism2" - ] + assert ( + _preflight_slide_encoders_for_grid( + path, + "mpp0.5_px224", + ["prism2"], + ["virchow2"], + [], + "cpu", + ) + == [] + ) + + _write_prism2_patch_grid(root, "mpp1_px224") + import raw2features.slide_embedders.prism2 as module + + monkeypatch.setattr(module, "_require_prism2_runtime", lambda: None) + with pytest.raises(ValueError, match="requires a CUDA GPU"): + _preflight_slide_encoders_for_grid( + path, + "mpp1_px224", + ["prism2"], + ["virchow2"], + [], + "cpu", + ) def test_prism2_fields_do_not_change_existing_prism_fingerprint_shape(): diff --git a/tests/test_registry_contract.py b/tests/test_registry_contract.py index 4a04b4d..7c5f363 100644 --- a/tests/test_registry_contract.py +++ b/tests/test_registry_contract.py @@ -80,9 +80,9 @@ def test_every_slide_spec_is_complete(): if spec.family != "pool": assert spec.doi and spec.doi.startswith("10."), f"{name}: needs a DOI" assert spec.weights_filename, f"{name}: needs the exact weights_filename" - if spec.weights_filename.endswith(".safetensors.index.json"): + if spec.weights_filename.endswith(".index.json"): assert spec.weights_manifest is not None, ( - f"{name}: a sharded safetensors checkpoint needs a weights_manifest" + f"{name}: a sharded checkpoint needs a weights_manifest" ) if spec.weights_manifest is not None: assert spec.weights_filename in spec.weights_manifest diff --git a/tests/test_release_metadata.py b/tests/test_release_metadata.py index 754048c..c031533 100644 --- a/tests/test_release_metadata.py +++ b/tests/test_release_metadata.py @@ -61,3 +61,13 @@ def selected(python_version: str) -> bool: "python_full_version": "3.12.0", } assert not requirement.marker.evaluate(pypy_environment) + + +def test_prism2_extra_supplies_nonisolated_flash_build_prerequisites() -> None: + project = tomllib.loads((ROOT / "pyproject.toml").read_text()) + requirements = { + Requirement(item).name + for item in project["project"]["optional-dependencies"]["prism2"] + } + + assert {"torch", "packaging", "ninja", "setuptools", "wheel"} <= requirements diff --git a/tests/test_slide_embedders.py b/tests/test_slide_embedders.py index 51f00a7..425e767 100644 --- a/tests/test_slide_embedders.py +++ b/tests/test_slide_embedders.py @@ -1020,6 +1020,75 @@ def test_cli_slide_embed_validates_all_encoder_names_before_work( assert "slide" not in open_grid(path, mode="r+") +def test_cli_slide_embed_runtime_preflight_precedes_all_encoder_writes( + tmp_path, monkeypatch, recording_slide_embedder +): + path = _write_slide_store( + tmp_path, + { + "mpp0.5_px224": ( + 224, + { + "conch_v1_5": np.ones((2, 768), np.float32), + "virchow2": np.ones((2, 2560), np.float32), + }, + ) + }, + ) + import raw2features.slide_embedders.prism2 as prism2 + + monkeypatch.setattr( + prism2, + "_require_prism2_runtime", + lambda: (_ for _ in ()).throw(RuntimeError("PRISM2 runtime missing")), + ) + + result = CliRunner().invoke( + app, + [ + "slide-embed", + path, + "-s", + "titan", + "-s", + "prism2", + "--device", + "cpu", + ], + ) + + assert result.exit_code == 1 + assert "PRISM2 runtime missing" in result.output + assert recording_slide_embedder["loads"] == 0 + assert "slide" not in open_grid(path, mode="r+") + + +def test_embed_cli_reports_pipeline_runtime_errors_without_traceback( + tmp_path, monkeypatch +): + monkeypatch.setattr( + "raw2features.cli.embed.embed_slide", + lambda *args, **kwargs: (_ for _ in ()).throw( + RuntimeError("PRISM2 runtime missing") + ), + ) + + result = CliRunner().invoke( + app, + [ + "embed", + str(tmp_path / "slide.zarr"), + str(tmp_path / "out"), + "--device", + "cpu", + ], + ) + + assert result.exit_code == 1 + assert "Error: PRISM2 runtime missing" in result.output + assert "Traceback" not in result.output + + def test_cli_slide_embed_consolidates_once_after_all_writes( tmp_path, recording_slide_embedder, monkeypatch ): diff --git a/uv.lock b/uv.lock index e0e4739..607d437 100644 --- a/uv.lock +++ b/uv.lock @@ -3484,10 +3484,12 @@ prism2 = [ { name = "ninja" }, { name = "packaging" }, { name = "safetensors" }, + { name = "setuptools" }, { name = "timm" }, { name = "torch" }, { name = "torchvision" }, { name = "transformers" }, + { name = "wheel" }, ] s3 = [ { name = "s3fs" }, @@ -3611,6 +3613,7 @@ requires-dist = [ { name = "scanpy", marker = "extra == 'seal'" }, { name = "scikit-learn", marker = "extra == 'stain'", specifier = ">=1.3" }, { name = "segmentation-models-pytorch", marker = "extra == 'grandqc'", specifier = ">=0.3" }, + { name = "setuptools", marker = "extra == 'prism2'", specifier = ">=64" }, { name = "shapely", marker = "extra == 'image'", specifier = ">=2.0" }, { name = "spatialdata", marker = "extra == 'spatialdata'", specifier = ">=0.2" }, { name = "tifffile", marker = "extra == 'kronos'", specifier = ">=2024" }, @@ -3627,6 +3630,7 @@ requires-dist = [ { name = "transformers", marker = "extra == 'models'", specifier = ">=4.40" }, { name = "transformers", marker = "extra == 'prism2'", specifier = "==4.56.0" }, { name = "typer", specifier = ">=0.12" }, + { name = "wheel", marker = "extra == 'prism2'", specifier = ">=0.41" }, { name = "xformers", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux' and extra == 'kronos2'", specifier = "==0.0.29.post3" }, { name = "zarr", marker = "extra == 'zarr'", specifier = ">=3.1" }, ] @@ -4996,6 +5000,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fb/17/c68b6cbcfeadbf420b3c3edaf8fda51335bc9c38732adb2d3ba8984dc607/wcwidth-0.8.0-py3-none-any.whl", hash = "sha256:8c75e6099cefd197c4bcc67a486f70b5dbc68f997c05f34a811d853910450d64", size = 324935, upload-time = "2026-06-05T21:19:33.999Z" }, ] +[[package]] +name = "wheel" +version = "0.48.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d0/20/50ed6bdf27dec98b568a8ae25dc599f35baa3d9709f9e83fd1edb56b9a90/wheel-0.48.0.tar.gz", hash = "sha256:94800765601e9171bf5d58d066e640662842bcedcbab982b2c90787a2c987322", size = 66471, upload-time = "2026-08-11T22:02:27.327Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2e/29/69cfbb602cd91690c55d38ba9fe53e6a7e76a6fa647bf38f19c138d25449/wheel-0.48.0-py3-none-any.whl", hash = "sha256:3217dcc807155e45db462d7ef2431f5ddda0d7273b700d05a67b271ceb1287ab", size = 33320, upload-time = "2026-08-11T22:02:26.1Z" }, +] + [[package]] name = "wrapt" version = "2.2.1" From 352bf24965de03d16285b601b21590e58ce8e5ec Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Sun, 16 Aug 2026 22:29:59 +0100 Subject: [PATCH 5/7] Fix slide runtime preflight regressions --- src/raw2features/cli/embed_many.py | 13 +- src/raw2features/cli/slide_embed.py | 13 +- src/raw2features/pipeline/runner.py | 95 +++---------- .../slide_embedders/model_registry.py | 7 +- tests/test_embed_many.py | 19 ++- tests/test_prism2.py | 128 +++++------------- tests/test_slide_embedders.py | 37 ++++- 7 files changed, 131 insertions(+), 181 deletions(-) diff --git a/src/raw2features/cli/embed_many.py b/src/raw2features/cli/embed_many.py index 9500144..0f67bff 100644 --- a/src/raw2features/cli/embed_many.py +++ b/src/raw2features/cli/embed_many.py @@ -14,6 +14,7 @@ from __future__ import annotations +import dataclasses import glob as globmod import math import os @@ -390,6 +391,7 @@ def embed_many( ) cli = sanitize_argv(sys.argv) device_list = cfg.device_list() + cfg = _resolved_serial_config(cfg, device_list) try: from raw2features.slide_embedders.model_registry import ( @@ -483,6 +485,14 @@ def _classify(summary: dict) -> str: return "skipped" if summary["status"] == "skipped" else "done" +def _resolved_serial_config(cfg: RunConfig, devices: list[str]) -> RunConfig: + """Make a lone --devices value the actual serial worker device.""" + + if len(devices) != 1: + return cfg + return dataclasses.replace(cfg, device=devices[0], devices=None) + + def _resolve_manifest_sources(rows: list[dict], slide_dir: str) -> list[dict]: """Resolve only relative local manifest paths against ``slide_dir`` in place.""" @@ -560,8 +570,6 @@ def _validate_source_mpps(rows: list[dict]) -> None: def _with_source_mpp(cfg: RunConfig, row: dict) -> RunConfig: """Apply a manifest row's per-slide ``source_mpp`` override (else cfg unchanged).""" if "source_mpp" in row: - import dataclasses - return dataclasses.replace(cfg, source_mpp=row["source_mpp"]) return cfg @@ -637,7 +645,6 @@ def _embed_shard_parallel( rest of the shard). Distinct slides write distinct ``*.embeddings.zarr`` stores, so the concurrent writes never collide. """ - import dataclasses import queue import threading diff --git a/src/raw2features/cli/slide_embed.py b/src/raw2features/cli/slide_embed.py index 63c80cf..105a648 100644 --- a/src/raw2features/cli/slide_embed.py +++ b/src/raw2features/cli/slide_embed.py @@ -67,6 +67,7 @@ def slide_embed( import zarr from raw2features.core.device import resolve_device + from raw2features.core.store import open_grid from raw2features.slide_embedders.encoding import ( encode_slide_embedding, resolve_slide_grid, @@ -128,7 +129,6 @@ def slide_embed( ( slide_model_name, selected_grid, - group, selected_patch_model, output_name, complete, @@ -147,10 +147,10 @@ def slide_embed( typer.echo(f"Error: {exc}", err=True) raise typer.Exit(1) from exc + live_groups = {} for ( slide_model_name, selected_grid, - group, selected_patch_model, output_name, complete, @@ -161,6 +161,15 @@ def slide_embed( ) continue + # Zarr group handles cache their attrs metadata. Planning deliberately opens + # every target before any writes, so retaining those handles would let a later + # encoder rebuild the header from a stale snapshot and erase an earlier + # encoder's metadata. Reacquire once per grid for the execution phase, then + # reuse that live handle for every write to the same grid. + if selected_grid not in live_groups: + live_groups[selected_grid] = open_grid(root, selected_grid) + group = live_groups[selected_grid] + typer.echo( f"{slide_model_name} [{selected_grid}]: encoding from " f"'{selected_patch_model}' patch features …" diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index db3d1cc..bbae4ab 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -1454,14 +1454,18 @@ def run_slide( models_to_do = list(cfg.models) if cfg.slide_encoders: - available = sorted(set(present_valid) | set(models_to_do)) - _preflight_slide_encoders_for_grid( - out_path, - grid_key_existing, + from raw2features.slide_embedders.model_registry import ( + validate_slide_encoder_runtime, + ) + + # Validate every explicitly requested slide runtime before this grid does + # patch work. This intentionally does not infer relevance from cfg.models: + # embed_slide may satisfy an encoder later from a compatible grid already in + # the store, and that fallback must not discover a missing runtime only after + # an unrelated patch encoder has run. + validate_slide_encoder_runtime( cfg.slide_encoders, - available, - models_to_do, - cfg.device, + devices=[cfg.device], ) qc_to_do = list(cfg.qc) @@ -1909,6 +1913,7 @@ def embedder_factory(device: str) -> list: cfg.device, available, preferred_patch_models=list(cfg.models), + validated_patch_models=present_valid, ) final_dims = sink.feature_dims() @@ -3791,75 +3796,6 @@ def _slide_encoders_for(names: list[str], available: list[str]) -> list[str]: return out -def _preflight_slide_encoders_for_grid( - out_path: str, - grid_key_existing: str | None, - names: list[str], - available_patch_models: list[str], - patch_models_to_write: list[str], - slide_device: str, -) -> list[str]: - """Validate runtimes needed by this exact grid before patch work starts.""" - - from raw2features.slide_embedders.model_registry import ( - get_slide_spec, - validate_slide_encoder_runtime, - ) - - candidates = [ - name - for name in _slide_encoders_for(names, available_patch_models) - if get_slide_spec(name).family == "prism2" - ] - if not candidates: - return [] - - required = [ - name - for name in candidates - if grid_key_existing is None - or get_slide_spec(name).patch_encoder in patch_models_to_write - ] - if grid_key_existing is not None and len(required) != len(candidates): - try: - import zarr - - from raw2features.core.store import open_grid - from raw2features.slide_embedders.encoding import ( - resolve_slide_patch_model, - slide_embedding_is_complete, - slide_output_key, - ) - - root = zarr.open_group(out_path, mode="r", use_consolidated=False) - group = open_grid(root, grid_key_existing) - for name in candidates: - if name in required: - continue - patch_model = resolve_slide_patch_model( - group, - name, - available_patch_models=available_patch_models, - ) - output_name = slide_output_key(group, name, patch_model) - if not slide_embedding_is_complete( - group, - name, - patch_model=patch_model, - device=slide_device, - output_name=output_name, - # _inspect_store already validated this exact patch array. - patch_is_complete=True, - ): - required.append(name) - except Exception: # noqa: BLE001 - unreadable state cannot prove completion - required = list(candidates) - - required = list(dict.fromkeys(required)) - validate_slide_encoder_runtime(required, devices=[slide_device]) - return required - - def _run_qc( qc_tools, reader, @@ -3909,6 +3845,7 @@ def _run_slide_encoders( available_patch_models: list[str], *, preferred_patch_models: list[str] | None = None, + validated_patch_models: list[str] | None = None, ) -> dict[str, str]: """Run slide-level encoders on patch features already written to *sink*. @@ -3924,6 +3861,7 @@ def _run_slide_encoders( ) results: dict[str, str] = {} + trusted_patch_models = set(validated_patch_models or ()) for slide_model_name in slide_encoder_names: from raw2features.slide_embedders.model_registry import get_slide_spec @@ -3942,13 +3880,14 @@ def _run_slide_encoders( available_patch_models=available_patch_models, ) output_name = slide_output_key(sink._group, slide_model_name, patch_model) + patch_is_complete = patch_model in trusted_patch_models if slide_embedding_is_complete( sink._group, slide_model_name, patch_model=patch_model, device=device, output_name=output_name, - patch_is_complete=True, + patch_is_complete=patch_is_complete, ): results[slide_model_name] = f"slide/{output_name}" continue @@ -3959,7 +3898,7 @@ def _run_slide_encoders( device, patch_model=patch_model, available_patch_models=available_patch_models, - patch_is_complete=True, + patch_is_complete=patch_is_complete, ) if encoding is None: continue diff --git a/src/raw2features/slide_embedders/model_registry.py b/src/raw2features/slide_embedders/model_registry.py index c5d6b88..a1c68ce 100644 --- a/src/raw2features/slide_embedders/model_registry.py +++ b/src/raw2features/slide_embedders/model_registry.py @@ -64,9 +64,6 @@ def validate_slide_encoder_runtime( if not any(spec.family == "prism2" for spec in specs): return - from raw2features.slide_embedders.prism2 import _require_prism2_runtime - - _require_prism2_runtime() if devices is not None: incompatible = [device for device in devices if not device.startswith("cuda")] if incompatible: @@ -75,6 +72,10 @@ def validate_slide_encoder_runtime( f"{incompatible}" ) + from raw2features.slide_embedders.prism2 import _require_prism2_runtime + + _require_prism2_runtime() + def build_slide_embedder(name: str) -> SlideEmbedder: from raw2features.core.plugins import get diff --git a/tests/test_embed_many.py b/tests/test_embed_many.py index 67228a0..2d90389 100644 --- a/tests/test_embed_many.py +++ b/tests/test_embed_many.py @@ -9,7 +9,7 @@ from typer.testing import CliRunner from conftest import MockEmbedder, build_ngff_v04 -from raw2features.cli.embed_many import _shard +from raw2features.cli.embed_many import _resolved_serial_config, _shard from raw2features.cli.main import app from raw2features.pipeline.runner import RunConfig, run_slide @@ -48,6 +48,18 @@ def test_shard_partitions_disjoint_and_complete(): assert shards[0] == items[0::4] # strided +def test_single_devices_value_becomes_the_serial_worker_device(): + cfg = RunConfig(models=["mock"], device="cpu", devices="cuda:7") + + resolved = _resolved_serial_config(cfg, ["cuda:7"]) + + assert resolved.device == "cuda:7" + assert resolved.devices is None + assert cfg.device == "cpu" + assert cfg.devices == "cuda:7" + assert _resolved_serial_config(cfg, ["cuda:0", "cuda:1"]) is cfg + + def test_embed_many_validates_slide_encoders_before_discovery_or_model_load( tmp_path, monkeypatch ): @@ -93,6 +105,9 @@ def test_embed_many_runtime_preflight_runs_before_warm_model_load( "_require_prism2_runtime", lambda: (_ for _ in ()).throw(RuntimeError("runtime missing")), ) + monkeypatch.setattr( + "raw2features.core.device._accelerators", lambda: (True, False) + ) monkeypatch.setattr( em, "load_embedders", @@ -110,7 +125,7 @@ def test_embed_many_runtime_preflight_runs_before_warm_model_load( "-s", "prism2", "--device", - "cpu", + "cuda:0", ], ) diff --git a/tests/test_prism2.py b/tests/test_prism2.py index 2512044..d61e2fa 100644 --- a/tests/test_prism2.py +++ b/tests/test_prism2.py @@ -8,17 +8,15 @@ import numpy as np import pytest -import zarr +from conftest import MockEmbedder from raw2features.embedders.fingerprint import ( patch_output_fingerprint, resolved_slide_amp, - slide_output_dim, slide_output_fingerprint, ) from raw2features.embedders.model_registry import get_spec -from raw2features.pipeline.runner import _preflight_slide_encoders_for_grid -from raw2features.slide_embedders.encoding import write_slide_embedding +from raw2features.pipeline.runner import RunConfig, run_slide from raw2features.slide_embedders.model_registry import ( get_slide_spec, validate_slide_encoder_runtime, @@ -47,39 +45,6 @@ ) -def _write_prism2_patch_grid(root, key: str): - group = root["grids"].create_group(key) - coords = group.create_array("coords", shape=(2, 2), dtype="int32") - coords[:] = np.asarray([[0, 0], [224, 0]], dtype=np.int32) - patch_spec = get_spec("virchow2") - patch_fingerprint = patch_output_fingerprint(patch_spec, "fp16") - features = group.create_group("features").create_array( - "virchow2", - shape=(2, patch_spec.embedding_dim), - chunks=(2, patch_spec.embedding_dim), - dtype="float32", - ) - features[:] = 1.0 - features.attrs.update( - { - "role": "features", - "model": "virchow2", - "output_fingerprint": patch_fingerprint, - } - ) - group.attrs["raw2features"] = { - "schema_version": "0.1", - "models": { - "virchow2": { - "embedding_dim": patch_spec.embedding_dim, - "output_fingerprint": patch_fingerprint, - } - }, - "patching": {"level0_patch": 224}, - } - return group, patch_spec, patch_fingerprint - - def test_prism2_registry_records_both_published_embeddings(): base = get_slide_spec("prism2") diagnostic = get_slide_spec("prism2_diagnostic") @@ -363,13 +328,14 @@ def test_prism2_runtime_preflight_checks_packages_and_device(monkeypatch): validate_slide_encoder_runtime(["mean"], devices=["cpu"]) assert calls == [] - validate_slide_encoder_runtime(["prism2"], devices=["cuda:0"]) - assert calls == [True] with pytest.raises(ValueError, match="requires a CUDA GPU"): validate_slide_encoder_runtime( ["prism2_diagnostic"], devices=["cpu"], ) + assert calls == [] + validate_slide_encoder_runtime(["prism2"], devices=["cuda:0"]) + assert calls == [True] def test_prism2_restores_released_phi3_mask_contract(): @@ -623,65 +589,43 @@ def test_prism2_amp_contract_is_device_independent(): assert resolved_slide_amp(spec, "cuda:0") == "bf16" -def test_prism2_grid_preflight_uses_exact_grid_and_slide_device( - tmp_path, monkeypatch +def test_run_slide_preflights_prism2_before_unrelated_patch_work( + synthetic_ngff, tmp_path, monkeypatch ): - path = str(tmp_path / "complete.embeddings.zarr") - root = zarr.open_group(path, mode="w", zarr_format=2) - root.create_group("grids") - group, patch_spec, patch_fingerprint = _write_prism2_patch_grid( - root, "mpp0.5_px224" - ) - root.attrs["raw2features"] = { - "schema_version": "0.1", - "grids": {"mpp0.5_px224": {}}, - } + import raw2features.slide_embedders.model_registry as slide_registry - slide_spec = get_slide_spec("prism2") - slide_fingerprint = slide_output_fingerprint( - slide_spec, - patch_model="virchow2", - patch_output_fingerprint=patch_fingerprint, - patch_dim=patch_spec.embedding_dim, - resolved_amp="bf16", - ) - write_slide_embedding( - group, - "prism2", - np.ones(slide_output_dim(slide_spec, patch_spec.embedding_dim)), - { - "patch_encoder": "virchow2", - "embedding_dim": slide_spec.embedding_dim, - "output_fingerprint": slide_fingerprint, - }, - ) + calls = [] - assert ( - _preflight_slide_encoders_for_grid( - path, - "mpp0.5_px224", - ["prism2"], - ["virchow2"], - [], - "cpu", - ) - == [] - ) + def reject(names, *, devices): + calls.append((list(names), list(devices))) + raise RuntimeError("PRISM2 runtime missing") - _write_prism2_patch_grid(root, "mpp1_px224") - import raw2features.slide_embedders.prism2 as module - - monkeypatch.setattr(module, "_require_prism2_runtime", lambda: None) - with pytest.raises(ValueError, match="requires a CUDA GPU"): - _preflight_slide_encoders_for_grid( - path, - "mpp1_px224", - ["prism2"], - ["virchow2"], - [], - "cpu", + monkeypatch.setitem(sys.modules, "torch", SimpleNamespace()) + monkeypatch.setattr( + "raw2features.core.device._accelerators", lambda: (False, False) + ) + monkeypatch.setattr(slide_registry, "validate_slide_encoder_runtime", reject) + cfg = RunConfig( + models=["mock"], + slide_encoders=["prism2"], + no_seg=True, + target_mpp=0.5, + patch_px=64, + device="cpu", + amp="fp32", + ) + + with pytest.raises(RuntimeError, match="PRISM2 runtime missing"): + run_slide( + synthetic_ngff, + str(tmp_path / "out"), + cfg, + embedders=[MockEmbedder(dim=8, input_size=64, name="mock")], ) + assert calls == [(["prism2"], ["cpu"])] + assert not list((tmp_path / "out").glob("*.embeddings.zarr")) + def test_prism2_fields_do_not_change_existing_prism_fingerprint_shape(): patch = patch_output_fingerprint(get_spec("virchow"), "fp16") diff --git a/tests/test_slide_embedders.py b/tests/test_slide_embedders.py index 425e767..e92b931 100644 --- a/tests/test_slide_embedders.py +++ b/tests/test_slide_embedders.py @@ -1042,6 +1042,9 @@ def test_cli_slide_embed_runtime_preflight_precedes_all_encoder_writes( "_require_prism2_runtime", lambda: (_ for _ in ()).throw(RuntimeError("PRISM2 runtime missing")), ) + monkeypatch.setattr( + "raw2features.core.device._accelerators", lambda: (True, False) + ) result = CliRunner().invoke( app, @@ -1053,7 +1056,7 @@ def test_cli_slide_embed_runtime_preflight_precedes_all_encoder_writes( "-s", "prism2", "--device", - "cpu", + "cuda:0", ], ) @@ -1123,6 +1126,10 @@ def _record_consolidation(store): assert result.exit_code == 0, result.output assert recording_slide_embedder["loads"] == 2 assert len(calls) == 1 + group = open_grid(path, mode="r+") + header = dict(group.attrs["raw2features"]) + assert set(header["slide_embeddings"]) == {"mean", "max"} + assert set(group["slide"].keys()) == {"mean", "max"} def test_embed_slide_validates_all_encoder_names_before_work(tmp_path, monkeypatch): @@ -1161,6 +1168,34 @@ def test_inline_slide_encoder_rerun_is_idempotent( assert recording_slide_embedder["loads"] == 1 +def test_inline_slide_encoder_validates_new_patch_data_before_use( + tmp_path, recording_slide_embedder +): + from raw2features.pipeline.runner import _run_slide_encoders + from raw2features.sinks.zarr_sink import ZarrSink + + values = np.ones((2, 3), dtype=np.float32) + values[-1, 0] = np.nan + path = _write_slide_store( + tmp_path, + {"mpp0.5_px64": (128, {"mock": values})}, + ) + sink = ZarrSink() + sink._group = open_grid(path, mode="r+") + + with pytest.raises(ValueError, match="incomplete or contains invalid data"): + _run_slide_encoders( + sink, + ["mean"], + "cpu", + ["mock"], + validated_patch_models=[], + ) + + assert recording_slide_embedder["loads"] == 0 + assert "slide" not in sink._group + + @pytest.mark.skipif(not _TORCH, reason="torch not installed") def test_cli_slide_embed_on_existing_zarr(synthetic_ngff, tmp_path): # First produce patch embeddings. From b9d0a0bf6613071c12c7360ba89b8a77ec19386e Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Sun, 16 Aug 2026 23:16:29 +0100 Subject: [PATCH 6/7] Preserve idempotent slide preflight --- src/raw2features/cli/embed_many.py | 21 +++- src/raw2features/cli/slide_embed.py | 9 +- src/raw2features/pipeline/receipt.py | 15 ++- src/raw2features/pipeline/runner.py | 47 ++++++-- src/raw2features/slide_embedders/encoding.py | 57 +++++++++ tests/test_embed_many.py | 40 +++++++ tests/test_output_fingerprints.py | 10 ++ tests/test_prism2.py | 98 ++++++++++++++++ tests/test_slide_embedders.py | 115 ++++++++++++++++++- 9 files changed, 392 insertions(+), 20 deletions(-) diff --git a/src/raw2features/cli/embed_many.py b/src/raw2features/cli/embed_many.py index 0f67bff..7341d13 100644 --- a/src/raw2features/cli/embed_many.py +++ b/src/raw2features/cli/embed_many.py @@ -394,6 +394,9 @@ def embed_many( cfg = _resolved_serial_config(cfg, device_list) try: + from raw2features.slide_embedders.encoding import ( + slide_encoders_requiring_compute, + ) from raw2features.slide_embedders.model_registry import ( validate_slide_encoder_runtime, ) @@ -401,7 +404,23 @@ def embed_many( # embed-many loads patch models once before visiting individual stores. Its # slide-parallel workers run slide encoders on these resolved devices, so # validate that actual runtime now rather than after an expensive warm load. - validate_slide_encoder_runtime(cfg.slide_encoders, devices=device_list) + # A fully complete shard remains inspectable/resumable without reinstalling + # an optional slide runtime or scheduling a GPU solely for the skip check. + runtime_names = list(cfg.slide_encoders) if force else [] + if not force: + for row in shard: + slide_id = slide_id_from_path(row["path"]) + runtime_names.extend( + slide_encoders_requiring_compute( + os.path.join(out_dir, f"{slide_id}.embeddings.zarr"), + cfg.slide_encoders, + device=device_list[0], + ) + ) + validate_slide_encoder_runtime( + list(dict.fromkeys(runtime_names)), + devices=device_list, + ) except (RuntimeError, ValueError) as exc: typer.echo(f"Error: slide-encoder runtime preflight failed: {exc}", err=True) raise typer.Exit(1) from exc diff --git a/src/raw2features/cli/slide_embed.py b/src/raw2features/cli/slide_embed.py index 105a648..2977d8c 100644 --- a/src/raw2features/cli/slide_embed.py +++ b/src/raw2features/cli/slide_embed.py @@ -166,22 +166,21 @@ def slide_embed( # encoder rebuild the header from a stale snapshot and erase an earlier # encoder's metadata. Reacquire once per grid for the execution phase, then # reuse that live handle for every write to the same grid. - if selected_grid not in live_groups: - live_groups[selected_grid] = open_grid(root, selected_grid) - group = live_groups[selected_grid] - typer.echo( f"{slide_model_name} [{selected_grid}]: encoding from " f"'{selected_patch_model}' patch features …" ) try: + if selected_grid not in live_groups: + live_groups[selected_grid] = open_grid(root, selected_grid) + group = live_groups[selected_grid] encoding = encode_slide_embedding( group, slide_model_name, device, patch_model=selected_patch_model, ) - except (KeyError, RuntimeError, ValueError) as exc: + except (KeyError, OSError, RuntimeError, ValueError) as exc: typer.echo(f"Error: {exc}", err=True) raise typer.Exit(1) from exc if encoding is None: diff --git a/src/raw2features/pipeline/receipt.py b/src/raw2features/pipeline/receipt.py index 4adaa46..29412a4 100644 --- a/src/raw2features/pipeline/receipt.py +++ b/src/raw2features/pipeline/receipt.py @@ -202,11 +202,10 @@ def validate_model( """True iff ``features/`` in an open zarr group holds complete data. Streams the array in row-blocks to bound memory: every row must be finite. A zarr - shape is metadata only; an unwritten tail reads back as the finite fill value (0.0), - so a finiteness check alone would pass a truncated store as "complete". Writes go in - coord order, so the unwritten part is always a contiguous *suffix* -- an all-zero - **last** row signals truncation. Checking only the last row (not every row) lets a - model emit a legitimately all-zero feature row mid-array without making resume loop. + shape is metadata only; an unwritten tail reads back as the finite fill value (0.0). + Legacy arrays without an expected output fingerprint therefore retain the all-zero + final-row truncation heuristic. Under the current contract the matching array/header + fingerprint is a post-write commit marker, so a committed zero final row is valid. """ import numpy as np @@ -250,7 +249,11 @@ def validate_model( rows = np.asarray(arr[s : s + block]).astype(np.float32, copy=False) if not np.isfinite(rows).all(): return False - if n_patches > 0 and not np.asarray(arr[n_patches - 1]).any(): # all-zero tail + if ( + expected_fingerprint is None + and n_patches > 0 + and not np.asarray(arr[n_patches - 1]).any() + ): return False except Exception: # noqa: BLE001 - any failure means "not valid" return False diff --git a/src/raw2features/pipeline/runner.py b/src/raw2features/pipeline/runner.py index bbae4ab..7ec3273 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -1454,17 +1454,34 @@ def run_slide( models_to_do = list(cfg.models) if cfg.slide_encoders: + from raw2features.slide_embedders.encoding import ( + slide_encoders_requiring_compute, + ) from raw2features.slide_embedders.model_registry import ( + get_slide_spec, validate_slide_encoder_runtime, ) - # Validate every explicitly requested slide runtime before this grid does - # patch work. This intentionally does not infer relevance from cfg.models: - # embed_slide may satisfy an encoder later from a compatible grid already in - # the store, and that fallback must not discover a missing runtime only after - # an unrelated patch encoder has run. + # A complete current slide output must remain resumable on a host without its + # optional runtime. Incomplete outputs, including ones discovered on another + # compatible grid, still fail before this grid does unrelated patch work. + runtime_names = ( + list(cfg.slide_encoders) + if force + else slide_encoders_requiring_compute( + out_path, + cfg.slide_encoders, + device=cfg.device, + ) + ) + # A patch output being replaced can invalidate a slide output that is still + # self-consistent with the old patch fingerprint at preflight time. + for name in cfg.slide_encoders: + required = get_slide_spec(name).patch_encoder + if required in models_to_do or (required == "any" and models_to_do): + runtime_names.append(name) validate_slide_encoder_runtime( - cfg.slide_encoders, + list(dict.fromkeys(runtime_names)), devices=[cfg.device], ) @@ -1913,7 +1930,10 @@ def embedder_factory(device: str) -> list: cfg.device, available, preferred_patch_models=list(cfg.models), - validated_patch_models=present_valid, + # Existing arrays passed _inspect_store; fresh arrays passed + # producer-side shape/finiteness checks before their commit + # marker was stamped. Neither needs an immediate disk rescan. + validated_patch_models=available, ) final_dims = sink.feature_dims() @@ -3050,6 +3070,8 @@ def _run_batches( collector that stores into its shard buffer). Behaviour for the single-device case is identical -- ``device == cfg.device`` and the writes are 1:1 the same. """ + import numpy as np + prof = prof or null_profiler() n = int(coords.shape[0]) batch_size = _effective_batch_size(cfg.batch_size, embedders) @@ -3139,6 +3161,17 @@ def _producer() -> None: .numpy() .astype(cfg.features_dtype) ) + expected_shape = (len(patches), emb.embedding_dim) + if tuple(feats.shape) != expected_shape: + raise ValueError( + f"model {emb.name!r} returned feature shape " + f"{tuple(feats.shape)}; expected {expected_shape}" + ) + if not np.isfinite(feats).all(): + raise ValueError( + f"model {emb.name!r} returned non-finite patch " + "features" + ) with prof.stage("write"): write_block(emb.name, start, feats) finally: diff --git a/src/raw2features/slide_embedders/encoding.py b/src/raw2features/slide_embedders/encoding.py index 419e96d..edb65ec 100644 --- a/src/raw2features/slide_embedders/encoding.py +++ b/src/raw2features/slide_embedders/encoding.py @@ -7,6 +7,7 @@ from __future__ import annotations +import os import warnings from dataclasses import dataclass @@ -205,6 +206,62 @@ def slide_embedding_is_complete( return bool(vector.size and np.isfinite(vector).all() and (vector != 0).any()) +def slide_encoders_requiring_compute( + store_path: str, + slide_models: list[str], + *, + device: str = "cpu", +) -> list[str]: + """Return requested slide encoders without a complete current store output.""" + + if not slide_models: + return [] + path = store_path.removeprefix("file://") + if not os.path.exists(path): + return list(slide_models) + + import zarr + + try: + root = zarr.open_group(path, mode="r", use_consolidated=False) + except Exception: # noqa: BLE001 - an unreadable store cannot prove completion + return list(slide_models) + + missing = [] + valid_patches: set[tuple[str, str]] = set() + invalid_patches: set[tuple[str, str]] = set() + for slide_model in slide_models: + try: + selected_grid, group, patch_model = resolve_slide_grid(root, slide_model) + patch_key = (selected_grid, patch_model) + if patch_key in invalid_patches: + complete = False + else: + if patch_key not in valid_patches: + try: + _validated_patch_fingerprint(group, patch_model) + except ValueError: + invalid_patches.add(patch_key) + complete = False + else: + valid_patches.add(patch_key) + if patch_key in valid_patches: + output_name = slide_output_key(group, slide_model, patch_model) + complete = slide_embedding_is_complete( + group, + slide_model, + patch_model=patch_model, + device=device, + output_name=output_name, + patch_is_complete=True, + ) + except (KeyError, OSError, RuntimeError, TypeError, ValueError): + complete = False + if not complete: + missing.append(slide_model) + return missing + + def resolve_slide_grid( root, slide_model: str, diff --git a/tests/test_embed_many.py b/tests/test_embed_many.py index 2d90389..827d470 100644 --- a/tests/test_embed_many.py +++ b/tests/test_embed_many.py @@ -134,6 +134,46 @@ def test_embed_many_runtime_preflight_runs_before_warm_model_load( assert "models must not be loaded" not in result.output +def test_embed_many_complete_shard_skips_optional_runtime_preflight( + tmp_path, monkeypatch +): + import raw2features.cli.embed_many as em + import raw2features.slide_embedders.encoding as encoding + import raw2features.slide_embedders.prism2 as prism2 + + slides = tmp_path / "slides" + slides.mkdir() + (slides / "S.zarr").mkdir() + monkeypatch.setattr( + encoding, + "slide_encoders_requiring_compute", + lambda *args, **kwargs: [], + ) + monkeypatch.setattr( + prism2, + "_require_prism2_runtime", + lambda: pytest.fail("complete shard must not require the PRISM2 runtime"), + ) + monkeypatch.setattr(em, "_embed_shard_serial", lambda *args, **kwargs: (0, 1, 0)) + + result = CliRunner().invoke( + app, + [ + "embed-many", + str(slides), + str(tmp_path / "out"), + "-f", + "mock", + "-s", + "prism2", + "--device", + "cpu", + ], + ) + + assert result.exit_code == 0, result.output + + @pytest.mark.parametrize( "rows", [ diff --git a/tests/test_output_fingerprints.py b/tests/test_output_fingerprints.py index aae7d23..2e3197b 100644 --- a/tests/test_output_fingerprints.py +++ b/tests/test_output_fingerprints.py @@ -388,6 +388,16 @@ def test_validate_model_requires_current_fingerprint_and_expected_dimension(tmp_ expected_dim=3, expected_fingerprint=fingerprint, ) + array[-1] = 0 + assert validate_model( + group, + "m", + 4, + expected_dim=3, + expected_fingerprint=fingerprint, + ) + assert not validate_model(group, "m", 4) + array[-1] = 1 assert not validate_model( group, "m", diff --git a/tests/test_prism2.py b/tests/test_prism2.py index d61e2fa..446597e 100644 --- a/tests/test_prism2.py +++ b/tests/test_prism2.py @@ -8,6 +8,7 @@ import numpy as np import pytest +import zarr from conftest import MockEmbedder from raw2features.embedders.fingerprint import ( @@ -17,6 +18,10 @@ ) from raw2features.embedders.model_registry import get_spec from raw2features.pipeline.runner import RunConfig, run_slide +from raw2features.slide_embedders.encoding import ( + slide_encoders_requiring_compute, + write_slide_embedding, +) from raw2features.slide_embedders.model_registry import ( get_slide_spec, validate_slide_encoder_runtime, @@ -45,6 +50,39 @@ ) +def _write_prism2_patch_grid(root, key: str): + group = root["grids"].create_group(key) + coords = group.create_array("coords", shape=(2, 2), dtype="int32") + coords[:] = np.asarray([[0, 0], [224, 0]], dtype=np.int32) + patch_spec = get_spec("virchow2") + patch_fingerprint = patch_output_fingerprint(patch_spec, "fp16") + features = group.create_group("features").create_array( + "virchow2", + shape=(2, patch_spec.embedding_dim), + chunks=(2, patch_spec.embedding_dim), + dtype="float32", + ) + features[:] = 1.0 + features.attrs.update( + { + "role": "features", + "model": "virchow2", + "output_fingerprint": patch_fingerprint, + } + ) + group.attrs["raw2features"] = { + "schema_version": "0.1", + "models": { + "virchow2": { + "embedding_dim": patch_spec.embedding_dim, + "output_fingerprint": patch_fingerprint, + } + }, + "patching": {"level0_patch": 224}, + } + return group, patch_spec, patch_fingerprint + + def test_prism2_registry_records_both_published_embeddings(): base = get_slide_spec("prism2") diagnostic = get_slide_spec("prism2_diagnostic") @@ -627,6 +665,66 @@ def reject(names, *, devices): assert not list((tmp_path / "out").glob("*.embeddings.zarr")) +def test_prism2_runtime_planning_skips_complete_multigrid_outputs_once_per_patch( + tmp_path, monkeypatch +): + path = str(tmp_path / "complete.embeddings.zarr") + root = zarr.open_group(path, mode="w", zarr_format=2) + root.create_group("grids") + group, patch_spec, patch_fingerprint = _write_prism2_patch_grid( + root, "mpp0.5_px224" + ) + unrelated = root["grids"].create_group("mpp1_px64") + unrelated.create_array("coords", shape=(1, 2), dtype="int32")[:] = 0 + unrelated.create_group("features") + unrelated.attrs["raw2features"] = { + "schema_version": "0.1", + "models": {}, + "patching": {"level0_patch": 64}, + } + root.attrs["raw2features"] = { + "schema_version": "0.1", + "grids": {"mpp0.5_px224": {}, "mpp1_px64": {}}, + } + + for name in ("prism2", "prism2_diagnostic"): + slide_spec = get_slide_spec(name) + slide_fingerprint = slide_output_fingerprint( + slide_spec, + patch_model="virchow2", + patch_output_fingerprint=patch_fingerprint, + patch_dim=patch_spec.embedding_dim, + resolved_amp="bf16", + ) + write_slide_embedding( + group, + name, + np.ones(slide_spec.embedding_dim, dtype=np.float32), + { + "patch_encoder": "virchow2", + "embedding_dim": slide_spec.embedding_dim, + "output_fingerprint": slide_fingerprint, + }, + ) + + import raw2features.pipeline.receipt as receipt + + real_validate_model = receipt.validate_model + calls = 0 + + def count_validation(*args, **kwargs): + nonlocal calls + calls += 1 + return real_validate_model(*args, **kwargs) + + monkeypatch.setattr(receipt, "validate_model", count_validation) + names = ["prism2", "prism2_diagnostic"] + assert slide_encoders_requiring_compute(path, names, device="cpu") == [] + assert calls == 1 + del group["slide"]["prism2"] + assert slide_encoders_requiring_compute(path, names, device="cpu") == ["prism2"] + + def test_prism2_fields_do_not_change_existing_prism_fingerprint_shape(): patch = patch_output_fingerprint(get_spec("virchow"), "fp16") prism = get_slide_spec("prism") diff --git a/tests/test_slide_embedders.py b/tests/test_slide_embedders.py index e92b931..7f9e14f 100644 --- a/tests/test_slide_embedders.py +++ b/tests/test_slide_embedders.py @@ -677,7 +677,7 @@ def test_cli_slide_embed_rejects_incomplete_patch_output( {"mpp0.5_px64": (137, {"mock": np.ones((3, 4), dtype=np.float32)})}, ) group = open_grid(path, mode="r+") - group["features"]["mock"][-1] = 0 + group["features"]["mock"][-1, 0] = np.nan result = CliRunner().invoke( app, @@ -1132,6 +1132,37 @@ def _record_consolidation(store): assert set(group["slide"].keys()) == {"mean", "max"} +def test_cli_slide_embed_reports_execution_grid_reopen_error( + tmp_path, monkeypatch, recording_slide_embedder +): + path = _write_slide_store( + tmp_path, + {"mpp0.5_px64": (128, {"mock": np.ones((2, 3), np.float32)})}, + ) + import raw2features.core.store as store + + real_open_grid = store.open_grid + calls = 0 + + def fail_execution_reopen(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise OSError("grid disappeared") + return real_open_grid(*args, **kwargs) + + monkeypatch.setattr(store, "open_grid", fail_execution_reopen) + result = CliRunner().invoke( + app, + ["slide-embed", path, "-s", "mean", "--device", "cpu"], + ) + + assert result.exit_code == 1 + assert "Error: grid disappeared" in result.output + assert "Traceback" not in result.output + assert recording_slide_embedder["loads"] == 0 + + def test_embed_slide_validates_all_encoder_names_before_work(tmp_path, monkeypatch): from raw2features.pipeline.runner import embed_slide @@ -1196,6 +1227,88 @@ def test_inline_slide_encoder_validates_new_patch_data_before_use( assert "slide" not in sink._group +@pytest.mark.skipif(not _TORCH, reason="torch not installed") +def test_inline_slide_encoder_accepts_committed_zero_final_feature_row( + synthetic_ngff, tmp_path +): + state = {"forbid_forward": False} + + class ZeroTailEmbedder(MockEmbedder): + def embed_batch(self, batch): + if state["forbid_forward"]: + pytest.fail("committed zero-tail patch output must not be recomputed") + features = super().embed_batch(batch) + features[-1] = 0 + return features + + cfg = RunConfig( + models=["zero_tail"], + slide_encoders=["mean"], + no_seg=True, + target_mpp=0.5, + patch_px=64, + batch_size=4096, + features_dtype="float32", + device="cpu", + amp="fp32", + ) + receipts = str(tmp_path / "receipts") + summary = run_slide( + synthetic_ngff, + str(tmp_path / "out"), + cfg, + receipts_dir=receipts, + embedders=[ZeroTailEmbedder(name="zero_tail")], + ) + + group = open_grid(summary["output_uri"]) + assert not np.asarray(group["features"]["zero_tail"][-1]).any() + assert "mean" in group["slide"] + state["forbid_forward"] = True + rerun = run_slide( + synthetic_ngff, + str(tmp_path / "out"), + cfg, + receipts_dir=receipts, + embedders=[ZeroTailEmbedder(name="zero_tail")], + ) + assert rerun["models_added"] == [] + assert rerun["models_skipped"] == ["zero_tail"] + + +@pytest.mark.skipif(not _TORCH, reason="torch not installed") +def test_nonfinite_patch_features_fail_before_model_commit(synthetic_ngff, tmp_path): + class NonfiniteTailEmbedder(MockEmbedder): + def embed_batch(self, batch): + features = super().embed_batch(batch) + features[-1, 0] = float("nan") + return features + + out = tmp_path / "out" + with pytest.raises(ValueError, match="non-finite patch features"): + run_slide( + synthetic_ngff, + str(out), + RunConfig( + models=["nonfinite"], + slide_encoders=["mean"], + no_seg=True, + target_mpp=0.5, + patch_px=64, + batch_size=4096, + features_dtype="float32", + device="cpu", + amp="fp32", + ), + embedders=[NonfiniteTailEmbedder(name="nonfinite")], + ) + + stores = list(out.glob("*.embeddings.zarr")) + assert len(stores) == 1 + array = open_grid(str(stores[0]))["features"]["nonfinite"] + assert "output_fingerprint" not in array.attrs + + @pytest.mark.skipif(not _TORCH, reason="torch not installed") def test_cli_slide_embed_on_existing_zarr(synthetic_ngff, tmp_path): # First produce patch embeddings. From 96f7ebad40adf2310b1d83e0162720a59a20a53d Mon Sep 17 00:00:00 2001 From: Craig Myles Date: Sun, 16 Aug 2026 23:42:32 +0100 Subject: [PATCH 7/7] Finalize v0.2.1 release metadata --- CHANGELOG.md | 18 ++++++++++++++++-- CITATION.cff | 1 + docs/usage.md | 5 +++-- 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f4e872a..4538a98 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ Notable changes to raw2features, newest first. This project follows [Semantic Versioning](https://semver.org). -## [0.2.1] - Unreleased +## [0.2.1] - 2026-08-16 This focused model update adds KRONOS2 to the existing native multiplex path and PRISM2 to slide embedding. The store layout, geometry-only `grid_hash`, and brightfield @@ -48,6 +48,20 @@ extraction path are unchanged. the correct source-channel view. Additive runs can still write their outputs into the same store when their grid geometry matches. +### Reliability and compatibility + +- Keep standalone and inline slide embedding additive and idempotent. Multi-encoder + standalone runs retain every output's metadata, while a completed optional-runtime + output can be inspected or resumed without reinstalling that runtime or scheduling a + GPU. When computation is required, runtime validation still precedes unrelated patch + work. +- Validate patch-feature shape and finiteness before writing and committing each output. + A matching array/header output fingerprint acts as the post-write completion marker, + so a legitimate all-zero final feature row remains valid; legacy arrays without that + marker retain the conservative unwritten-tail check. +- Make a single resolved `embed-many --devices` value the serial worker's actual device, + and report device incompatibility before optional package/ABI errors. + ## [0.2.0] - 2026-07-23 This release extends named-channel multiplex support with a `channelwise` strategy, adds @@ -175,4 +189,4 @@ What 0.1.0 provides: [0.1.0]: https://github.com/CraigMyles/raw2features/releases/tag/v0.1.0 [0.2.0]: https://github.com/CraigMyles/raw2features/compare/v0.1.0...v0.2.0 -[0.2.1]: https://github.com/CraigMyles/raw2features/compare/v0.2.0...HEAD +[0.2.1]: https://github.com/CraigMyles/raw2features/compare/v0.2.0...v0.2.1 diff --git a/CITATION.cff b/CITATION.cff index d123f87..534aee7 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -21,3 +21,4 @@ keywords: - feature extraction license: MIT version: "0.2.1" +date-released: 2026-08-16 diff --git a/docs/usage.md b/docs/usage.md index cb26c41..bb2c007 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -482,8 +482,9 @@ the header. Model-specific settings such as resolved AMP live in each output fingerprint instead, so changing them replaces that model in the same grid. "Present" is strict: a model counts only if `features/` has the current expected dimension and a matching output fingerprint in both the array and -grid header, is fully finite, has no unwritten (all-zero) tail, and has -`len(coords) == n_patches`. The fingerprint covers the effective weights, +grid header, is fully finite, and has `len(coords) == n_patches`. For legacy arrays +without a post-write fingerprint commit marker, an all-zero final row is additionally +treated as an unwritten tail. The fingerprint covers the effective weights, preprocessing, pooling, resolved AMP, and loader construction. A truncated, damaged, legacy-unfingerprinted, or stale-contract array is re-embedded; unrelated model arrays and coordinates remain untouched. A derived `slide/` fingerprint includes its