diff --git a/CHANGELOG.md b/CHANGELOG.md index bcb6e80..4538a98 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,13 +3,22 @@ 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. 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 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 @@ -39,6 +48,20 @@ layout, geometry-only `grid_hash`, and brightfield execution 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 @@ -166,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/MODELS.md b/docs/MODELS.md index ae8f519..39bd99a 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,32 @@ 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]" +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 +``` + +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.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 +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. 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. + 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..a139b5f 100644 --- a/docs/SPEC.md +++ b/docs/SPEC.md @@ -216,8 +216,13 @@ 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. +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 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..5414562 100644 --- a/docs/adding_a_model.md +++ b/docs/adding_a_model.md @@ -63,3 +63,8 @@ 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. +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/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 diff --git a/pyproject.toml b/pyproject.toml index 1c7eff1..c686667 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,22 @@ 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.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", + "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): # pip install git+https://github.com/mahmoodlab/MADELEINE.git@419287dc60a57296d959840b893481019c4f0d21 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 34d6075..7341d13 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,39 @@ 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.encoding import ( + slide_encoders_requiring_compute, + ) + 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. + # 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 # 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 @@ -470,6 +504,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.""" @@ -547,8 +589,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 @@ -624,7 +664,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 dab837c..2977d8c 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, @@ -76,6 +77,7 @@ def slide_embed( ) from raw2features.slide_embedders.model_registry import ( validate_slide_encoder_names, + validate_slide_encoder_runtime, ) try: @@ -101,50 +103,84 @@ 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, 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, + 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 + + live_groups = {} + for ( + slide_model_name, + selected_grid, + selected_patch_model, + output_name, + complete, + ) in plans: + if complete: typer.echo( f"{slide_model_name} [{selected_grid}]: already complete (skipping)" ) 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. 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, 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/embedders/fingerprint.py b/src/raw2features/embedders/fingerprint.py index 2e84bf8..de5d236 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,7 +713,54 @@ def _effective_slide_checkpoint(spec) -> 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, + ) + + contract = { + "repo": _hf_repo(PRISM2_PHI3_SOURCE), + "revision": PRISM2_PHI3_REVISION, + "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" + if diagnostic + else "pinned_sha256_verified_local_config" + ), + "weights_source": "included_in_PRISM2_checkpoint", + } + 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, @@ -733,6 +787,34 @@ def _slide_constructor(spec) -> dict[str, Any]: "features_dtype": "float32", "batched": True, }, + "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", + "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 +902,12 @@ 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( + diagnostic=spec.name == "prism2_diagnostic" + ) + return contract def slide_output_dim(spec, patch_dim: int) -> int: @@ -836,6 +923,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 spec.family == "prism2": + return "bf16" if str(device).startswith("cuda") and spec.family in { "gigapath_slide", "prism", @@ -862,6 +951,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 +967,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/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 481dc97..7ec3273 100644 --- a/src/raw2features/pipeline/runner.py +++ b/src/raw2features/pipeline/runner.py @@ -1453,6 +1453,38 @@ 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, + ) + + # 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( + list(dict.fromkeys(runtime_names)), + devices=[cfg.device], + ) + qc_to_do = list(cfg.qc) thumbnail_to_do = bool(cfg.emit_thumbnail) thumb_meta = None @@ -1898,6 +1930,10 @@ def embedder_factory(device: str) -> list: cfg.device, available, preferred_patch_models=list(cfg.models), + # 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() @@ -3034,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) @@ -3123,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: @@ -3829,6 +3878,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*. @@ -3844,6 +3894,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 @@ -3862,12 +3913,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=patch_is_complete, ): results[slide_model_name] = f"slide/{output_name}" continue @@ -3878,6 +3931,7 @@ def _run_slide_encoders( device, patch_model=patch_model, available_patch_models=available_patch_models, + patch_is_complete=patch_is_complete, ) if encoding is None: continue 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..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 @@ -50,7 +51,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 @@ -101,7 +107,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]), @@ -122,6 +128,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 @@ -147,7 +154,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( @@ -195,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, @@ -290,6 +357,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. @@ -326,7 +394,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, @@ -367,6 +439,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/model_registry.py b/src/raw2features/slide_embedders/model_registry.py index b62e20b..a1c68ce 100644 --- a/src/raw2features/slide_embedders/model_registry.py +++ b/src/raw2features/slide_embedders/model_registry.py @@ -53,6 +53,30 @@ 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 + + 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}" + ) + + 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/src/raw2features/slide_embedders/prism2.py b/src/raw2features/slide_embedders/prism2.py new file mode 100644 index 0000000..94a6d08 --- /dev/null +++ b/src/raw2features/slide_embedders/prism2.py @@ -0,0 +1,513 @@ +"""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 + +import importlib +import sys +import threading +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.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" + ), + "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_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" +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" + ), +} +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( + 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 + + +class _Prism2Phi3MaskCompatibility: + """Private namespace for the Phi-3 helper used by pinned PRISM2 code.""" + + _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", + 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, ...], + *, + allowed_files: tuple[str, ...] | None = None, + what: str, +) -> None: + """Reject a stale managed directory before its custom code can execute.""" + + root = Path(snapshot) + allowed = {*(allowed_files or 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], + *, + allowed_files: tuple[str, ...] | None = None, + 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=local_dir, + ) + _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, + } + mismatches = [] + for distribution, expected in required.items(): + try: + installed = importlib_metadata.version(distribution) + except importlib_metadata.PackageNotFoundError: + installed = None + if installed is None or not SpecifierSet(f"=={expected}").contains(installed): + found = installed or "missing" + mismatches.append(f"{distribution}=={expected} (found {found})") + if mismatches: + raise RuntimeError( + "PRISM2 requires its validated runtime. Install raw2features[prism2], " + "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): + 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), + 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, + phi3_artifacts, + allowed_files=tuple(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", + ) + if self.spec.name == "prism2_diagnostic": + _bind_phi3_mask_compatibility(model) + 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: + 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_embed_many.py b/tests/test_embed_many.py index 4198f64..827d470 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 ): @@ -78,6 +90,90 @@ 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() + import raw2features.slide_embedders.prism2 as prism2 + + monkeypatch.setattr( + prism2, + "_require_prism2_runtime", + lambda: (_ for _ in ()).throw(RuntimeError("runtime missing")), + ) + monkeypatch.setattr( + "raw2features.core.device._accelerators", lambda: (True, False) + ) + 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", + "cuda:0", + ], + ) + + 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_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 new file mode 100644 index 0000000..446597e --- /dev/null +++ b/tests/test_prism2.py @@ -0,0 +1,741 @@ +"""PRISM2 registry, loader, and forward-contract tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +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_fingerprint, +) +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, +) +from raw2features.slide_embedders.prism2 import ( + _VERIFIED_SNAPSHOT_CACHE, + PRISM2_CHECKPOINT_SHA256, + 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, + _bind_phi3_mask_compatibility, + _prepare_phi3_4d_causal_mask, + _require_prism2_runtime, + _verified_snapshot, +) + + +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") + + 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, *, 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") + + 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_RUNTIME_SHA256, + (), + "prism2", + ) + 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): + 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"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\)" + ), + ): + _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], + ) + 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 + + calls = [] + monkeypatch.setattr(module, "_require_prism2_runtime", lambda: calls.append(True)) + + validate_slide_encoder_runtime(["mean"], devices=["cpu"]) + assert calls == [] + 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(): + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + + mask = _prepare_phi3_4d_causal_mask( + 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_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", + (), + {"__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) + assert ( + getattr( + Phi3Model, + "_prepare_4d_causal_attention_mask_with_cache_position", + sentinel, + ) + is shared_before + ) + + +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") + _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_runtime_code_is_sha_bound_and_cached(monkeypatch, tmp_path): + 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", download + ) + monkeypatch.setattr( + "raw2features.slide_embedders.prism2.pinned_model_cache_dir", + lambda *args: str(cache), + ) + + def verify(path, expected, *, what): + hashes.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(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(downloads) == 1 + assert len(hashes) == 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.pinned_model_cache_dir", + lambda *args: 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") + 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_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" + 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_run_slide_preflights_prism2_before_unrelated_patch_work( + synthetic_ngff, tmp_path, monkeypatch +): + import raw2features.slide_embedders.model_registry as slide_registry + + calls = [] + + def reject(names, *, devices): + calls.append((list(names), list(devices))) + raise RuntimeError("PRISM2 runtime missing") + + 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_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") + 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..7c5f363 100644 --- a/tests/test_registry_contract.py +++ b/tests/test_registry_contract.py @@ -80,6 +80,21 @@ 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(".index.json"): + assert spec.weights_manifest is not None, ( + f"{name}: a sharded 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] == ( + 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..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 @@ -28,6 +29,14 @@ ) 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}") _SHA256 = re.compile(r"[0-9a-f]{64}") @@ -238,6 +247,71 @@ 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: + 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, + 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 + + 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(): + 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.""" 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..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, @@ -1020,6 +1020,78 @@ 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")), + ) + monkeypatch.setattr( + "raw2features.core.device._accelerators", lambda: (True, False) + ) + + result = CliRunner().invoke( + app, + [ + "slide-embed", + path, + "-s", + "titan", + "-s", + "prism2", + "--device", + "cuda:0", + ], + ) + + 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 ): @@ -1054,6 +1126,41 @@ 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_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): @@ -1092,6 +1199,116 @@ 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_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. diff --git a/uv.lock b/uv.lock index af8a22a..607d437 100644 --- a/uv.lock +++ b/uv.lock @@ -1916,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" @@ -3451,6 +3477,20 @@ prism = [ { name = "torchvision" }, { name = "transformers" }, ] +prism2 = [ + { name = "einops" }, + { name = "einops-exts" }, + { name = "huggingface-hub" }, + { name = "ninja" }, + { name = "packaging" }, + { name = "safetensors" }, + { name = "setuptools" }, + { name = "timm" }, + { name = "torch" }, + { name = "torchvision" }, + { name = "transformers" }, + { name = "wheel" }, +] s3 = [ { name = "s3fs" }, ] @@ -3535,6 +3575,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" }, @@ -3542,6 +3583,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'" }, @@ -3560,6 +3602,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'" }, @@ -3570,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" }, @@ -3577,17 +3621,20 @@ 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 = "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" }, ] -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 = [ @@ -4953,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"