diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..4ff402c6 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,224 @@ +# AGENTS.md + +Guidance for AI coding agents working in this repository. + +## What this is + +STAMP (Solid Tumor Associative Modeling in Pathology) is an end-to-end, weakly +supervised deep-learning pipeline for computational pathology: gigapixel +whole-slide images (WSIs) in, biomarker predictions out. It ships as a single +CLI (`stamp`) driven by one YAML config file, plus an optional MCP server. + +The pipeline stages, in order, map 1:1 onto CLI subcommands and onto top-level +keys in the config file: + +``` +WSIs → preprocess → [encode_slides | encode_patients] → train / crossval + → deploy → statistics + → heatmaps +``` + +Feature files are HDF5 (`.h5`) with a `feats` dataset plus metadata attributes +(`extractor`, `encoder`, `stamp_version`, `code_hash`, `feat_type`). Downstream +stages read those attributes to validate that features came from a compatible +extractor. + +## Branches and PRs + +**Never commit directly to `main`.** STAMP is developed branch-first: `main` is +the integration branch and only receives changes through pull requests. + +- Branch off an up-to-date `main` (`git fetch origin && git switch -c + origin/main`). +- Name branches with a topic prefix, as in the existing history: + `fix/…`, `feature/…`, `dev/…` (e.g. `fix/build-6`, `dev/multitask`, + `feature/random-seed`). +- Open a PR against `main`. CI is wired to `push` and `pull_request` on `main` + only, so a branch gets its first full check when the PR opens — expect the + dependency-config, ruff and Linux test jobs to gate the merge. +- Recent history is squash-merged, giving one commit per PR with a `(#NNN)` + suffix. Write the commit message accordingly: imperative mood, one line + describing the change ("Pin every dependency to an exact version", "Run the + test suite on Linux only"). + +If you are an agent asked to "commit" work, check what branch you are on first +and create one off `main` if the answer is `main`. + +## Repository layout + +``` +src/stamp/ + __main__.py CLI entry point (`stamp`); argparse + dispatch on config + config.yaml Factory-settings config, copied by `stamp init` + types.py Shared type aliases (Microns, SlidePixels, Task, Bags, …) + utils/ config.py (StampConfig root model), seed.py, cache.py, path.py + preprocessing/ Tiling and tile-level feature extraction + __init__.py `extract_()` — dispatches ExtractorName → extractor + config.py `ExtractorName` enum + `PreprocessingConfig` + tiling.py WSI → tiles, background rejection + extractor/ One module per foundation model (uni2, virchow2, …) + encoding/ Tile features → slide-/patient-level embeddings + config.py `EncoderName` enum + encoding configs + encoder/__init__.py `Encoder` ABC; subclasses in encoder/{titan,cobra,…}.py + modeling/ Training, cross-validation, deployment + data.py Dataloaders, clini/slide table parsing, feature typing + train.py, crossval.py, deploy.py + registry.py `ModelName` enum + (feature_type, task) → Lightning class + models/ VisionTransformer, MLP, TransMIL, Barspoon, Cox heads + statistics/ AUROC/AUPRC/regression/survival metrics and plots + heatmaps/ Attention heatmaps and top/bottom-scoring tile exports +tests/ pytest suite; `random_data.py` builds synthetic inputs +scripts/ check_dependency_config.py, verify_installed_stack.py, + remote_gpu_check.sh (manual GPU-host harness, not CI) +mcp/ FastMCP server exposing the CLI stages as MCP tools +.github/workflows/build.yml CI +``` + +`README.md` covers installation, `getting-started.md` is the user-facing tutorial. + +## Environment and commands + +The project uses **uv** (pinned to `>=0.12,<0.13`) and Python 3.14 (3.13 also +supported). uv is not optional — the CUDA wheel indexes and build exclusions +live in `[tool.uv.*]` and are invisible to pip. See "Dependency rules" below. + +```bash +uv sync --extra cpu --dev # what you normally want locally +uv sync --locked --extra cpu --dev # what CI runs; never rewrites uv.lock +uv sync --extra gpu # Linux + CUDA 13.0 +uv sync --extra gpu_all # + conchv1_5, gigapath, musk +``` + +`cpu`, `gpu` and `gpu_all` are declared as mutually conflicting extras; pick one. + +```bash +# Tests. The full suite is slow (model downloads, real training runs). +uv run pytest tests/ --ignore=tests/test_feature_extractors.py # what CI runs +uv run pytest -m "not slow" tests/ # quick pass +uv run pytest tests/test_data.py -k some_case --verbose # one case +uv run pytest tests/test_feature_extractors.py -k uni2 # one extractor + +# Lint, format, types +uv run ruff check --target-version=py313 +uv run ruff format --target-version=py313 # add --diff to only check +uv run pyright + +# Dependency-configuration checks (stdlib only, no install needed) +python scripts/check_dependency_config.py +uv run python scripts/verify_installed_stack.py --expect-no-cuda +``` + +Extractor tests download gated model weights from Hugging Face and need +`HF_TOKEN` in the environment; without it they are expected to fail, not to be +"fixed". `tests/conftest.py` adds a `--extractor` option to select which ones +run. + +CI (`.github/workflows/build.yml`) runs: dependency-config check, ruff +lint+format, the test suite on Linux (3.13 and 3.14), a macOS install/import +smoke test, GPU-extra installs on a machine with no CUDA toolkit (x86_64 and +aarch64), and one job per extractor. macOS runs installation and imports only — +it is a supported development platform, not a tested one. + +## Dependency rules + +These are enforced by `scripts/check_dependency_config.py` (and mirrored in +`tests/test_dependency_config.py`). Breaking one turns CI red: + +- **Every requirement is pinned with `==`.** No floors, no ranges. A loose + specifier lets an unreviewed release into a fresh resolution. +- **Git dependencies must be HTTPS + a full 40-char commit SHA.** No tags, no + branches. +- **`flash-attn`, `mamba-ssm` and `causal-conv1d` come from the Astral index** + (`https://wheels.astral.sh/simple/cu130/`) as pre-built wheels, never from a + pinned wheel URL and never from source. `[tool.uv.no-build-package]` forbids + building them; their local versions (`+cu.13.0.torch.2.11`) must agree with + the PyTorch minor and CUDA channel the project selects. +- **CPU installs must contain zero CUDA packages.** That is what + `[tool.uv.exclude-dependencies]` is for — upstream forks (UNI, GigaPath, + COBRA) declare CUDA extras unconditionally, so STAMP drops them and + re-declares them with platform markers in its own extras. +- Several pins carry a comment explaining *why* they are held back (numpy 2.4 + vs beartype, transformers 4.57 vs `trust_remote_code` models). Read the + comment before bumping; if you bump anyway, update or remove the comment. +- Changing a dependency means regenerating `uv.lock` (`uv lock`) in the same + commit. + +## Code conventions + +- **Runtime type checking is on.** `src/stamp/__init__.py` calls + `beartype_this_package()`, so annotations are checked at runtime (violations + surface as `UserWarning`). Annotate accurately; a sloppy annotation is a bug, + not cosmetic. +- **Array shapes are annotated with jaxtyping**, e.g. + `Float[Tensor, "batch tile feature"]`. Ruff's `F722` is globally ignored + because of this — don't re-enable it. +- **Domain units are `NewType`s** in `stamp/types.py`: `Microns`, + `SlidePixels`, `TilePixels`, `SlideMPP`. Use them rather than bare `float`/ + `int` so slide-space and tile-space pixels can't be mixed up. +- **A trailing underscore means the function has side effects** (writes files, + mutates its arguments): `extract_()`, `train_categorical_model_()`, + `encode_slides_()`, `filter_complete_patient_data_()`. Keep the convention + when adding functions. +- **Configs are Pydantic models with `extra="forbid"`**, composed under + `StampConfig` in `stamp/utils/config.py`. A new CLI option means: add the + field to the right config model, document it in `src/stamp/config.yaml`, and + pass it through in `__main__.py`. +- **Heavy imports are deferred.** `__main__.py` imports torch/pydantic-heavy + modules inside the `match` arms so `stamp init` and `stamp --help` stay fast; + `extract_()` imports each extractor inside its `case`. Keep new imports local + in the same way. +- **Optional backends fail loudly at import.** Every `extractor/*.py` and + `encoder/*.py` wraps its third-party imports in `try: … except + ModuleNotFoundError` and re-raises with the extra to install + (`pip install 'stamp[uni2]'`). Follow that pattern; the macOS CI job relies on + it, treating only `cobra` and `gigapath` as legitimately absent. +- **Logging, not printing.** Use `_logger = logging.getLogger("stamp")`. Each + command also attaches a file handler writing `logfile.log` into its + `output_dir`. +- Ruff is the formatter and linter (line length and style are its defaults). + Import sorting is on in the VS Code config; run `ruff format` before + committing. + +## Adding things + +**A tile-level feature extractor**: add a module under +`preprocessing/extractor/` returning an `Extractor(model=…, transform=…, +identifier=…)`; add a variant to `ExtractorName`; add a `case` to `extract_()`; +declare its dependencies as an optional extra in `pyproject.toml` and add it to +the `cpu`/`gpu`/`gpu_all` blanket extras as appropriate; list it in +`src/stamp/config.yaml`'s comment and in the `test_extractors` CI matrix. + +**A slide/patient encoder**: subclass `Encoder` in `encoding/encoder/`, +implementing `_generate_slide_embedding` and `_generate_patient_embedding` +(override `encode_slides_`/`encode_patients_` only if you need tile +coordinates); add a variant to `EncoderName` and wire it into +`init_slide_encoder_`/`init_patient_encoder_`. `required_extractors` gates which +tile features the encoder accepts. + +**A model**: add it under `modeling/models/`, add a `ModelName` variant, and +wire it into `load_model_class()` in `modeling/registry.py`. The Lightning +wrapper is chosen by `(feature_type, task)` from `MODEL_REGISTRY`; feature type +is detected from the `.h5` metadata, task is one of `classification`, +`regression`, `survival`. Only `barspoon` supports multi-target classification. + +## Gotchas + +- `identifier` on an `Extractor` must uniquely identify model *and* weights — + it is written into every feature file and checked downstream. +- `generate_hash: True` appends a hash of the preprocessing source (see + `get_processing_code_hash`) to output directory names. Editing any file in + `preprocessing/` therefore changes where features land; that is deliberate. +- Reproducibility goes through `Seed.set(seed)`, called once from `__main__.py`. + `torch.use_deterministic_algorithms()` is intentionally *not* enabled (large + performance cost), so same-seed runs follow the same trajectory but are not + bit-identical. `Seed.get_loader_worker_init()` raises if the seed was never + set. +- Tests force the `spawn` multiprocessing start method (`tests/conftest.py`). + Don't switch to `fork`; it warns and misbehaves with the threaded loaders. +- `deploy` accepts multiple checkpoints and majority-votes across them — that's + how cross-validation folds are combined. +- Slides that raise during feature extraction are skipped, not fatal; + `extract_()` is deliberately fail-safe. Check `logfile.log` when output looks + short. +- Don't commit `config.yaml` at the repo root — it's gitignored on purpose + (it's a user's local, path-laden copy of `src/stamp/config.yaml`). diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..22940e1c --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,4 @@ +# CLAUDE.md + +See [AGENTS.md](AGENTS.md) for all repository guidance — project layout, build +and test commands, dependency rules, code conventions, and gotchas. diff --git a/README.md b/README.md index bf074a52..a0d4e5c3 100644 --- a/README.md +++ b/README.md @@ -48,25 +48,13 @@ To setup STAMP you need [uv](https://docs.astral.sh/uv/) 0.12 or newer. | CUDA (`--extra gpu`, `--extra gpu_all`) | 3.13, 3.14 | Linux x86_64, Linux aarch64 | Python **3.14 is recommended** and is what `.python-version` selects; 3.13 remains -supported. The CUDA builds are pinned to one ABI stack: **CUDA 13.0 with PyTorch -2.11.0 and TorchVision 0.26.0**. - -CI runs the test suite on Linux. macOS is checked for installation and imports -only, so it stays usable for development but is not a tested target. - -There is no CUDA build for macOS or Windows. `flash-attn`, `mamba-ssm` and -`causal-conv1d` are only published as pre-built wheels for Linux, and STAMP -refuses to compile them (see [below](#why-uv-is-required)), so a GPU extra on -those platforms fails with a clear resolution error rather than a compiler error. +supported. The CUDA builds are pinned to **CUDA 13.0 with PyTorch 2.11.0 and +TorchVision 0.26.0**. There is no CUDA build for macOS or Windows — use +`--extra cpu` there. > [!IMPORTANT] -> uv is required for the GPU workflow. The PyTorch and Astral wheel indexes are -> configured through `[tool.uv.sources]` in `pyproject.toml`, which is -> uv-specific and invisible to other installers. Installing with plain `pip` -> would mean pointing it at -> `https://download.pytorch.org/whl/cu130` and -> `https://wheels.astral.sh/simple/cu130/` yourself, and re-deriving the exact -> pins by hand. See [Why uv is required](#why-uv-is-required). +> Install with uv, not pip. The wheel indexes STAMP needs are configured in +> `[tool.uv.sources]` in `pyproject.toml`, which pip does not read. ### Install or Update uv: @@ -107,17 +95,13 @@ uv sync --extra gpu_all source .venv/bin/activate ``` -Both GPU extras install **pre-built wheels by default**. `flash-attn`, -`mamba-ssm` and `causal-conv1d` come from the -[Astral wheel index](https://wheels.astral.sh/simple/cu130/) already compiled -against CUDA 13.0 and PyTorch 2.11, so nothing is compiled locally and no CUDA -toolkit needs to be installed to *install* STAMP. You still need an NVIDIA -driver to *run* on a GPU. +Both GPU extras install pre-built wheels, so nothing is compiled locally and no +CUDA toolkit is needed to install STAMP. You still need an NVIDIA driver to run +on a GPU. > [!NOTE] -> `--extra gpu_prebuilt` is deprecated and now just an alias for -> `--extra gpu_all`, which is pre-built anyway. It will be removed in the next -> breaking release. +> `--extra gpu_prebuilt` is a deprecated alias for `--extra gpu_all` and will be +> removed in the next breaking release. If you encounter errors during installation please read Installation Troubleshooting [below](#installation-troubleshooting). @@ -137,31 +121,6 @@ If you encounter errors during installation please read Installation Troubleshoo > apt update && apt install -y libgl1 libglx-mesa0 libglib2.0-0 > ``` -### Why uv is required - -The GPU workflow depends on configuration that only uv reads: - -* `[tool.uv.sources]` routes `torch` and `torchvision` to - `https://download.pytorch.org/whl/cu130` (or `.../cpu` for the `cpu` extra) - and the three compiled extensions to `https://wheels.astral.sh/simple/cu130/`. -* `[tool.uv.exclude-dependencies]` drops the unconditional CUDA requirements - that the UNI, GigaPath and COBRA forks declare, so a CPU install stays free of - CUDA-only packages. -* `[tool.uv.no-build-package]` forbids source builds of `flash-attn`, - `mamba-ssm` and `causal-conv1d`. - -None of this is visible to `pip`, which reads only `[project]`. Installing with -pip would resolve `flash-attn` from PyPI and try to compile it — which is -exactly what this configuration exists to prevent. If you must use pip, you have -to add both indexes yourself and pin the extensions to the same -`+cu.13.0.torch.2.11` local versions listed in `pyproject.toml`. - -Source builds of the three extensions are refused on purpose: they are -ABI-locked to one PyTorch build, take a long time, need a matching CUDA toolkit, -and were the most common cause of broken installs. An unsupported platform now -fails during resolution with a clear message instead of part-way through a -compile. - ## Basic Usage If the installation was successful, running `stamp` in your terminal should yield the following output: @@ -228,8 +187,7 @@ please consider citing our [Nature Protocols publication](https://www.nature.com These only exist as pre-built wheels for Linux x86_64 and Linux aarch64 on Python 3.13/3.14. On any other platform uv reports that no compatible version -was found. This is deliberate — STAMP does not fall back to compiling them. -Use `--extra cpu` on macOS and Windows. +was found. Use `--extra cpu` on macOS and Windows. #### Triton Errors @@ -258,10 +216,8 @@ was compiled against: E ImportError: [...]/.venv/lib/python3.14/site-packages/flash_attn_2_cuda.cpython-314-x86_64-linux-gnu.so: undefined symbol: _ZN3c105ErrorC2ENS_14SourceLocationENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE ``` -This should no longer happen, because the extension versions carry the PyTorch -build in their version string (`+cu.13.0.torch.2.11`) and the lockfile pins them -together with torch. If you do hit it, you are almost certainly in an -environment left over from an older STAMP release. Check what is installed: +This usually means the environment is left over from an older STAMP release. +Check what is installed: ```bash uv run python scripts/verify_installed_stack.py --expect-cuda