From ff557007136861fc941ad7e9c5470fdc8046ee60 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 21 Sep 2026 11:38:27 +0100 Subject: [PATCH 01/14] Make the container image usable from Nextflow Nextflow requires /bin/bash to be the container entrypoint, so ENTRYPOINT ["adata"] made every Docker- or Podman-backed process fail with `No such command '/bin/bash'` -- Nextflow invokes `docker run IMG /bin/bash -ue .command.sh`. Apptainer users were unaffected, since `singularity exec` ignores the entrypoint, which is probably why this went unnoticed. Drop the entrypoint and spell the command out in CMD instead. Nextflow also needs bash, ps, awk, date, grep, sed, tail and tee in the task container to collect metrics. procps is not in bookworm-slim, so every task silently lost its trace row. Install it, and assert the whole set at build time so base-image drift fails the build rather than every task. Two further fixes for bind-mounted runtimes: - PYTHONNOUSERSITE, because Apptainer bind-mounts the host $HOME and a user's ~/.local site-packages would otherwise shadow the venv. - XDG_CACHE_HOME, because Nextflow is commonly configured with `-u $(id -u):$(id -g)`, leaving no writable $HOME. Add the missing .dockerignore. Without one, `COPY . .` pulled the host's .venv, .git, .pytest_cache and .claude/worktrees (two full repo copies) into every local build; CI never hit this because a fresh checkout has none of them. Also move the apt and duckdb layers ahead of `COPY . .`, so a source edit no longer re-downloads duckdb. Verified against Nextflow 26.04.6: the pipeline completes and the trace is populated (%cpu=296.5%, peak_rss=11.2 MB). Co-Authored-By: Claude Opus 5 (1M context) --- .dockerignore | 18 +++++++++++++++ Dockerfile | 64 ++++++++++++++++++++++++++++++++++++--------------- README.md | 2 +- docs/index.md | 2 +- 4 files changed, 66 insertions(+), 20 deletions(-) create mode 100644 .dockerignore diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..f841322 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,18 @@ +# Keep the build context to what `uv sync --locked` actually needs. In +# particular .venv/ would otherwise ship a host-platform virtualenv into the +# image, and .claude/worktrees/ holds full copies of this repo. +# +# Do NOT add README.md or LICENSE here: pyproject.toml references them via +# `readme` and `license-files`, so the build fails without them. +.venv/ +.git/ +.claude/ +.pytest_cache/ +__pycache__/ +*.py[cod] +.coverage +coverage.xml +htmlcov/ +build/ +dist/ +*.egg-info/ diff --git a/Dockerfile b/Dockerfile index 5955f68..a9b68e3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,24 +1,33 @@ # Base image: Python 3.12 + uv preinstalled (Debian slim) FROM ghcr.io/astral-sh/uv:python3.12-bookworm-slim -ENV UV_NO_DEV=1 +# PYTHONNOUSERSITE: Apptainer bind-mounts the host $HOME by default, so a user's +# ~/.local/lib/python3.12/site-packages would otherwise shadow this venv. +# XDG_CACHE_HOME: Nextflow is commonly configured with `-u $(id -u):$(id -g)`, +# which leaves the container with no writable $HOME. +# UV_COMPILE_BYTECODE: bake .pyc at build time, so nothing writes to a +# read-only rootfs on first import. +ENV UV_NO_DEV=1 \ + UV_COMPILE_BYTECODE=1 \ + PYTHONNOUSERSITE=1 \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + XDG_CACHE_HOME=/tmp/.cache -WORKDIR /cli - -# Copy the project files (from the GitHub Actions checkout context) -COPY . . - -# --locked asserts that uv.lock is in sync with pyproject.toml, so an image -# can never be built from a lockfile that drifted. -RUN uv sync --locked +# procps supplies `ps`, which Nextflow needs to collect per-task metrics. +# curl, unzip and ca-certificates fetch duckdb below, and are left in place +# rather than purged: pipeline scripts routinely reach for curl, and TLS +# roots are worth having in any container that may touch the network. +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + ca-certificates curl procps mawk unzip \ + && rm -rf /var/lib/apt/lists/* # duckdb, for the filtering workflows in the docs: export obs to CSV, query it, # feed the names back to `adata subset --obs`. A single static binary, so it # needs no venv and cannot conflict with the project's dependencies. ARG DUCKDB_VERSION=v1.1.3 -RUN apt-get update \ - && apt-get install -y --no-install-recommends ca-certificates curl unzip \ - && ARCH="$(dpkg --print-architecture)" \ +RUN ARCH="$(dpkg --print-architecture)" \ && case "$ARCH" in \ amd64) DUCKDB_ARCH=amd64 ;; \ arm64) DUCKDB_ARCH=aarch64 ;; \ @@ -28,13 +37,32 @@ RUN apt-get update \ "https://github.com/duckdb/duckdb/releases/download/${DUCKDB_VERSION}/duckdb_cli-linux-${DUCKDB_ARCH}.zip" \ && unzip -q /tmp/duckdb.zip -d /usr/local/bin \ && chmod +x /usr/local/bin/duckdb \ - && rm /tmp/duckdb.zip \ - && apt-get purge -y curl unzip \ - && apt-get autoremove -y \ - && rm -rf /var/lib/apt/lists/* + && rm /tmp/duckdb.zip + +# Fail the build, rather than every Nextflow task, if the base image ever drops +# one of the tools Nextflow requires in a task container. +RUN set -eu; for t in bash ps awk date grep sed tail tee; do \ + command -v "$t" >/dev/null || { echo "missing required tool: $t" >&2; exit 1; }; \ + done + +WORKDIR /cli + +# Copy the project files (from the GitHub Actions checkout context) +COPY . . + +# --locked asserts that uv.lock is in sync with pyproject.toml, so an image +# can never be built from a lockfile that drifted. +RUN uv sync --locked # Put the project venv on PATH so `adata` is directly runnable ENV PATH="/cli/.venv/bin:${PATH}" -ENTRYPOINT ["adata"] -CMD ["--help"] +# No ENTRYPOINT on purpose: Nextflow requires /bin/bash to be the container +# entrypoint, so the image must not set one of its own. This is why invocations +# spell out the command: `docker run IMAGE adata view file.h5ad`. +# +# Deliberately NOT set here: OMP_NUM_THREADS / OPENBLAS_NUM_THREADS. NumPy's +# BLAS sizes its thread pool to the whole host, which oversubscribes a shared +# LSF node. This workload is streaming I/O, so capping it would cost nothing -- +# but it belongs in the pipeline's `env` scope, not baked into the image. +CMD ["adata", "--help"] diff --git a/README.md b/README.md index fb6c2af..1b55ad9 100644 --- a/README.md +++ b/README.md @@ -91,5 +91,5 @@ adata concat per_sample/*.h5ad -o merged.h5ad --join outer --label sample A docker image is available on QUAY: `quay.io/cellgeni/adata-cli:latest`. Pull and run with: ```bash -docker run --rm -it -v /path/to/data:/data quay.io/cellgeni/adata-cli:latest view /data/your_file.h5ad +docker run --rm -it -v /path/to/data:/data quay.io/cellgeni/adata-cli:latest adata view /data/your_file.h5ad ``` \ No newline at end of file diff --git a/docs/index.md b/docs/index.md index b37ac86..ed04b98 100644 --- a/docs/index.md +++ b/docs/index.md @@ -22,7 +22,7 @@ Or run it without installing anything: ```bash docker run --rm -it -v /path/to/data:/data \ - quay.io/cellgeni/adata-cli:latest view /data/your_file.h5ad + quay.io/cellgeni/adata-cli:latest adata view /data/your_file.h5ad ``` ## Documentation From dcd4d2479a50a498a3803952a2bf98588aad8020 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 21 Sep 2026 11:38:57 +0100 Subject: [PATCH 02/14] Release 0.5.1 A container-only release: the Python package is unchanged. Tagging it is what republishes the image, since .github/workflows/quay-on-tag.yml only builds on a tag push and the 0.5.0 tag must not be moved -- dropping the entrypoint is a breaking change for anyone pinned to it. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 29 +++++++++++++++++++++++++++++ pyproject.toml | 2 +- uv.lock | 2 +- 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1bd7f7d..03a92e5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,35 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no `v` prefix. +## 0.5.1 + +A container-only release. No changes to the Python package; the PyPI +artifact is identical in behaviour to 0.5.0. + +### Fixed + +- **The image could not be used from a Nextflow process.** Nextflow requires + `/bin/bash` to be the container entrypoint, so `ENTRYPOINT ["adata"]` made + every Docker- and Podman-backed task fail with `No such command + '/bin/bash'`. Apptainer was unaffected, as `singularity exec` ignores the + entrypoint. +- **Task metrics were silently lost.** `procps` is absent from the base image, + so Nextflow could not run `ps` to collect them. The required tool set + (`bash`, `ps`, `awk`, `date`, `grep`, `sed`, `tail`, `tee`) is now installed + and asserted at build time. +- `PYTHONNOUSERSITE` is set, so a bind-mounted `$HOME` under Apptainer can no + longer shadow the image's virtualenv with the user's `~/.local` packages. +- `XDG_CACHE_HOME` points at `/tmp`, so the image tolerates being run under an + arbitrary UID with no writable `$HOME`. +- Added a `.dockerignore`. Local builds were copying the host's `.venv`, + `.git` and `.pytest_cache` into the image. + +### Changed + +- **The image no longer sets an entrypoint, so the command must be named + explicitly:** `docker run IMAGE adata view file.h5ad`, where `docker run + IMAGE view file.h5ad` previously worked. + ## 0.5.0 Renamed from `h5ad` to `adata-cli`, restored compatibility with current diff --git a/pyproject.toml b/pyproject.toml index 1d1ebaf..7767eca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ # so the distribution is published as `pyadata-cli`. The import package # and the command are both still `adata`. name = "pyadata-cli" -version = "0.5.0" +version = "0.5.1" description = "Streaming CLI for exploring and editing large AnnData .h5ad and .zarr stores" readme = "README.md" requires-python = ">=3.12" diff --git a/uv.lock b/uv.lock index adb6d8d..6c20152 100644 --- a/uv.lock +++ b/uv.lock @@ -486,7 +486,7 @@ wheels = [ [[package]] name = "pyadata-cli" -version = "0.5.0" +version = "0.5.1" source = { editable = "." } dependencies = [ { name = "h5py" }, From 4859445ef93fa87f12d4cfa5ac55da31daced3e3 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 21 Sep 2026 11:46:26 +0100 Subject: [PATCH 03/14] Size streaming reads for the filesystem, not the source chunk _chunk_step returned the source's chunk height verbatim, so a store chunked (1, n_cols) was copied one row per read. On a local disk that is merely wasteful; on Lustre or NFS every read is a round-trip costing milliseconds, so a million-row copy spent nearly all of its time waiting. Reads are now grown to a 32 MiB budget and rounded down to a whole number of source chunks, since a partial read still decompresses the whole chunk. A (1_000_000, 30_000) float32 store chunked (1, 30_000) goes from 1 row per read to 279. Sizing needs the dtype, which h5py misreports for variable-length strings: itemsize is 8 there because the value is a pointer, not the text. Assume VLEN_ELEMENT_BYTES instead, so the row count is not overestimated by an order of magnitude and the memory bound holds. The step is floored at one whole chunk, which makes the budget a target rather than a cap for a source whose own chunk already exceeds it. That matches the previous behaviour and is now commented as such. Co-Authored-By: Claude Opus 5 (1M context) --- CHANGELOG.md | 17 +++++++- src/adata/storage/__init__.py | 67 ++++++++++++++++++++++++++++--- tests/test_storage.py | 74 +++++++++++++++++++++++++++++++++++ 3 files changed, 151 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 03a92e5..2a129b5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,11 +5,24 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no ## 0.5.1 -A container-only release. No changes to the Python package; the PyPI -artifact is identical in behaviour to 0.5.0. +Makes the container image usable from Nextflow, and stops `copy_dataset` +reading one row at a time from row-chunked stores. ### Fixed +- **Copying a row-chunked store was dominated by read latency.** The read step + was the source's chunk height verbatim, so a store chunked `(1, n_cols)` was + copied one row per read. On a network filesystem (Lustre, NFS) each read is a + round-trip, so a million-row copy spent nearly all of its time waiting. Reads + are now sized to a 32 MiB budget, rounded down to a whole number of source + chunks. A `(1_000_000, 30_000)` float32 store chunked `(1, 30_000)` goes from + 1 row per read to 279. +- Read sizing no longer trusts `itemsize` for variable-length strings. h5py + reports 8 there because the value is a pointer, which overestimated the row + count by an order of magnitude and broke the memory bound. + +### Container + - **The image could not be used from a Nextflow process.** Nextflow requires `/bin/bash` to be the container entrypoint, so `ENTRYPOINT ["adata"]` made every Docker- and Podman-backed task fail with `No such command diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index 58e7bad..cfea148 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -440,12 +440,69 @@ def _is_string_src(src: Any) -> bool: return is_string_dtype(getattr(src, "dtype", None)) -def _chunk_step(shape: Sequence[int], chunks: Optional[Sequence[int]]) -> int: - if chunks is not None and len(chunks) > 0 and chunks[0]: - return int(chunks[0]) +#: Byte budget for a single read when streaming a dataset. +#: +#: Using a source's chunk height as the read size means inheriting whatever +#: the writer chose, and a store chunked `(1, n_cols)` is then copied one row +#: per read. On a local disk that is merely wasteful; on a network filesystem +#: (Lustre, NFS) every read is a round-trip costing milliseconds, so a +#: million-row copy spends nearly all of its time waiting. 32 MiB sits +#: comfortably above a typical 1 MiB Lustre stripe while bounding peak memory. +TARGET_READ_BYTES = 32 * 1024 * 1024 + +#: Assumed width of a variable-length string element, for read sizing only. +#: h5py reports `itemsize` 8 for vlen strings because the value is a pointer, +#: which would overestimate the row count by an order of magnitude and break +#: the memory bound. Real cell and gene names sit well under this. +VLEN_ELEMENT_BYTES = 64 + + +def _row_bytes(shape: Sequence[int], dtype: Any) -> int: + """In-memory size of one row along the first axis.""" + width = 1 + for dim in shape[1:]: + width *= max(1, int(dim)) + + itemsize = int(getattr(dtype, "itemsize", 0) or 0) + # 'O' is how h5py spells vlen str, 'T' is numpy StringDType as reported by + # zarr-python 3. Neither itemsize reflects the bytes actually stored. + if itemsize <= 0 or getattr(dtype, "kind", None) in ("O", "T"): + itemsize = VLEN_ELEMENT_BYTES + + return max(1, width * itemsize) + + +def _chunk_step( + shape: Sequence[int], chunks: Optional[Sequence[int]], dtype: Any +) -> int: + """Rows to copy per read, sized for the filesystem rather than the source. + + The step is grown to `TARGET_READ_BYTES`, then rounded down to a whole + number of source chunks: reading part of a chunk still costs decompressing + all of it, so a step that splits one wastes the remainder. + """ if not shape: return 1 - return max(1, min(1024, int(shape[0]))) + + n_rows = int(shape[0]) + if n_rows <= 0: + return 1 + + chunk_rows = 0 + if chunks is not None and len(chunks) > 0 and chunks[0]: + chunk_rows = max(1, int(chunks[0])) + + step = max(1, TARGET_READ_BYTES // _row_bytes(shape, dtype)) + if chunk_rows: + # Never go below a single chunk: a partial read still decompresses the + # whole thing, so a smaller step costs the same I/O for less data. That + # makes TARGET_READ_BYTES a target rather than a cap -- a source whose + # own chunk already exceeds the budget (say 1000 x 1e6 float32, chunked + # whole) reads that chunk regardless. This matches the previous + # behaviour, which used the chunk height verbatim. + step = max(chunk_rows, (step // chunk_rows) * chunk_rows) + + return min(step, n_rows) def copy_dataset(src: Any, dst_group: Any, name: str) -> Any: @@ -475,7 +532,7 @@ def copy_dataset(src: Any, dst_group: Any, name: str) -> Any: ds[()] = src[()] return ds - step = _chunk_step(shape, getattr(src, "chunks", None)) + step = _chunk_step(shape, getattr(src, "chunks", None), src.dtype) for start in range(0, shape[0], step): end = min(start + step, shape[0]) if len(shape) == 1: diff --git a/tests/test_storage.py b/tests/test_storage.py index 926627f..9c57d17 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -393,3 +393,77 @@ def test_clamping_handles_one_dimensional_chunks(): assert _clamp_chunks({"chunks": (100,)}, 5)["chunks"] == (5,) assert _clamp_chunks({}, 5) == {} + + +# --------------------------------------------------------------------------- +# read sizing + + +def test_read_step_grows_past_a_pathological_source_chunk(): + """A `(1, n_cols)` source must not be copied one row per read. + + Forwarding the source's chunk height meant inheriting whatever the writer + chose. On a network filesystem each read is a round-trip, so a row-chunked + million-cell store spent the whole copy waiting on latency. + """ + from adata.storage import TARGET_READ_BYTES, _chunk_step, _row_bytes + + shape, chunks, dtype = (1_000_000, 30_000), (1, 30_000), np.dtype("float32") + step = _chunk_step(shape, chunks, dtype) + + assert step > 1, "the source chunk height is no longer used verbatim" + assert step * _row_bytes(shape, dtype) >= TARGET_READ_BYTES // 2 + + +def test_read_step_stays_a_whole_number_of_source_chunks(): + """Reading part of a chunk still costs decompressing all of it.""" + from adata.storage import _chunk_step + + for chunk_rows in (1, 7, 100, 65_536): + step = _chunk_step((10_000_000,), (chunk_rows,), np.dtype("int64")) + assert step % chunk_rows == 0 + + +def test_read_step_never_exceeds_the_dataset(): + from adata.storage import _chunk_step + + assert _chunk_step((10,), (1,), np.dtype("int64")) == 10 + assert _chunk_step((), None, np.dtype("int64")) == 1 + assert _chunk_step((0,), None, np.dtype("int64")) == 1 + + +def test_read_step_bounds_peak_memory(): + """A wide unchunked source used to read 1024 rows however wide they were.""" + from adata.storage import TARGET_READ_BYTES, _chunk_step, _row_bytes + + cases = [ + ((1_000_000, 30_000), None, np.dtype("float32")), + ((1_000_000, 50), (1, 50), np.dtype("float32")), + ((200_000_000,), (65_536,), np.dtype("int64")), + ] + for shape, chunks, dtype in cases: + step = _chunk_step(shape, chunks, dtype) + assert step * _row_bytes(shape, dtype) <= 2 * TARGET_READ_BYTES + + +def test_row_bytes_does_not_trust_a_vlen_itemsize(): + """h5py reports itemsize 8 for vlen str -- that is the pointer, not the text.""" + import h5py + + from adata.storage import VLEN_ELEMENT_BYTES, _row_bytes + + assert _row_bytes((10,), h5py.string_dtype(encoding="utf-8")) == VLEN_ELEMENT_BYTES + assert _row_bytes((10,), np.dtype("O")) == VLEN_ELEMENT_BYTES + assert _row_bytes((10, 4), np.dtype("float32")) == 16 + + +def test_copy_dataset_of_a_row_chunked_source_round_trips(new_store): + """The larger read step must not change what lands on disk.""" + path, opener = new_store() + with opener("a") as root: + from adata.storage import create_dataset + + values = np.arange(400, dtype="float32").reshape(100, 4) + create_dataset(root["uns"], "src", data=values, chunks=(1, 4)) + copy_dataset(root["uns"]["src"], root["uns"], "dst") + assert np.array_equal(root["uns"]["dst"][...], values) From 9a774508198d7b016286f65daa8ba8dd123a19a0 Mon Sep 17 00:00:00 2001 From: Aljes Date: Mon, 21 Sep 2026 11:48:36 +0100 Subject: [PATCH 04/14] Make XDG_CACHE_HOME actually writable at runtime Setting XDG_CACHE_HOME was self-defeating as written: uv honours it, so `uv sync` created /tmp/.cache root-owned and mode 0755 during the build. A task running under `-u $(id -u):$(id -g)` then could not write to the very path the image advertises as its cache, which is worse than leaving the variable unset. Clear the directory and recreate it world-writable in the same layer as the sync. Verified: `mkdir $XDG_CACHE_HOME/probe` now succeeds as an arbitrary UID, where it failed with EACCES before. Reported by Codex review on #12. Co-Authored-By: Claude Opus 5 (1M context) --- Dockerfile | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/Dockerfile b/Dockerfile index a9b68e3..8d0e79b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -52,7 +52,18 @@ COPY . . # --locked asserts that uv.lock is in sync with pyproject.toml, so an image # can never be built from a lockfile that drifted. -RUN uv sync --locked +# +# uv honours XDG_CACHE_HOME, so the sync leaves a root-owned package cache at +# /tmp/.cache -- which made the variable self-defeating, as a task running +# under an arbitrary UID then could not write to the very path it advertises. +# Clear it and leave an empty world-writable directory behind. Done in this +# same layer because a later `rm` would mask the files without reclaiming +# them; that reclaims about 9 MB, the cache being mostly hardlinks into the +# venv rather than separate copies. +RUN uv sync --locked \ + && rm -rf /tmp/.cache /tmp/uv-*.lock \ + && mkdir -p /tmp/.cache \ + && chmod 1777 /tmp/.cache # Put the project venv on PATH so `adata` is directly runnable ENV PATH="/cli/.venv/bin:${PATH}" From 1db401f40af462d787b0f55e52839a1abfea96b3 Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 15:33:15 +0100 Subject: [PATCH 05/14] Make concat --merge finish: read each var column once, not once per variable `_write_var` aligned each candidate var column onto the target index with `tuple(read_str_all(group[name])[i] for i in where)`. Only the outermost iterable of a generator expression is evaluated eagerly, so `read_str_all` ran once per target variable -- a full read of the column from disk, per element of the same column. The cost is quadratic in the number of variables. Measured here on two synthetic inputs with two var columns: 0.48 s at 500 vars, 1.40 s at 1,000, 4.55 s at 2,000, 16.57 s at 4,000. Extrapolated to the 36,601 vars of REQ-71798 that is tens of minutes to hours of pure CPU with the output file never growing past the header it wrote first, which is exactly what was reported against 0.5.1: 12 of 13 pipeline tasks killed after 98 minutes, and byte counts identical between `--merge same` and `--merge first`. Reading the column once makes the same case 0.02 s, and 0.25 s at the full 36,601 x 8,766 of the ticket. `first` and `only` decide on presence alone and now read no column values at all, which is why their cost matched `same` before. Also accept `--merge drop` / `--uns-merge drop`. `drop` is the documented default behaviour but was rejected as a value, so a config could not state it. The three new tests count full-column reads rather than timing the merge: the defect is a complexity bug, invisible to every existing concat test because they all use two or three variables, and a wall-clock assertion would be flaky on shared CI. `test_concat_merge_same_reads_each_var_column_once_per_input` is parametrised over 4 and 64 variables so that a cost which grows with the var count fails the second case; against the old code it reports 24 and 384 reads where 6 are expected. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 19 +++++++ docs/COMMANDS.md | 2 +- src/adata/cli.py | 17 ++++-- src/adata/commands/__init__.py | 2 +- src/adata/commands/concat.py | 4 +- src/adata/core/concat.py | 26 ++++++++- tests/test_commands_phase2.py | 101 +++++++++++++++++++++++++++++++++ 7 files changed, 161 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1bd7f7d..ac2c44b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,25 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no `v` prefix. +## Unreleased + +### Fixed + +- **`concat --merge` never finished on a real store.** Aligning a var column + onto the target index re-read the whole column from disk once per target + variable, so the cost was quadratic: at 36,601 variables a merge that should + take a fraction of a second ran for hours at 100% CPU with the output file + never growing past its header. Reported against 0.5.1 (REQ-71798), where 12 + of 13 pipeline tasks had to be killed after 98 minutes. The column is now + read once per input, and `--merge first` / `--merge only`, which decide on + presence alone, read no column values at all. + +### Added + +- **`--merge drop` and `--uns-merge drop` are accepted.** `drop` was already + the documented default behaviour but was rejected as a value, so a config + could not state it explicitly. + ## 0.5.0 Renamed from `h5ad` to `adata-cli`, restored compatibility with current diff --git a/docs/COMMANDS.md b/docs/COMMANDS.md index a5d27c6..77e4540 100644 --- a/docs/COMMANDS.md +++ b/docs/COMMANDS.md @@ -118,7 +118,7 @@ adata concat a.h5ad b.h5ad -o m.h5ad --keys a,b --index-unique - --uns-merge sam | `--label` | Add an obs column recording each cell's source | | `--keys` | Names for the inputs; defaults to their filenames | | `--index-unique` | Delimiter for suffixing obs names with their key | -| `--merge` / `--uns-merge` | `same`, `unique`, `first`, `only`; default drops | +| `--merge` / `--uns-merge` | `drop` (default), `same`, `unique`, `first`, `only` | | `--fill-value` | Value for dense cells introduced by an outer join | obs columns keep their dtypes: categoricals union their category sets, nullable diff --git a/src/adata/cli.py b/src/adata/cli.py index ed7dc8d..45f1f30 100644 --- a/src/adata/cli.py +++ b/src/adata/cli.py @@ -7,7 +7,7 @@ import typer from adata.commands import ( - MERGE_STRATEGIES, + MERGE_CHOICES, concat_stores, create_store, split_store, @@ -413,12 +413,14 @@ def concat( merge: Optional[str] = typer.Option( None, "--merge", - help="How to reconcile var columns: same, unique, first, only (default: drop)", + help="How to reconcile var columns: drop, same, unique, first, only " + "(default: drop)", ), uns_merge: Optional[str] = typer.Option( None, "--uns-merge", - help="How to reconcile uns: same, unique, first, only (default: drop)", + help="How to reconcile uns: drop, same, unique, first, only " + "(default: drop)", ), fill_value: float = typer.Option( 0.0, "--fill-value", help="Value for dense cells introduced by an outer join" @@ -442,13 +444,18 @@ def concat( adata concat a.h5ad b.h5ad -o m.h5ad --keys a,b --index-unique - --uns-merge same """ for name, value in (("--merge", merge), ("--uns-merge", uns_merge)): - if value is not None and value not in MERGE_STRATEGIES: + if value is not None and value not in MERGE_CHOICES: console.print( f"[bold red]Error:[/] {name} must be one of: " - f"{', '.join(MERGE_STRATEGIES)}" + f"{', '.join(MERGE_CHOICES)}" ) raise typer.Exit(code=1) + # "drop" is the default, and naming it explicitly has to be allowed: a + # config that spells out the default should not be rejected. + merge = None if merge == "drop" else merge + uns_merge = None if uns_merge == "drop" else uns_merge + try: concat_stores( files, diff --git a/src/adata/commands/__init__.py b/src/adata/commands/__init__.py index 70e0690..3cc464b 100644 --- a/src/adata/commands/__init__.py +++ b/src/adata/commands/__init__.py @@ -5,4 +5,4 @@ from adata.commands.ls import list_store from adata.commands.create import create_store from adata.commands.split import split_store -from adata.commands.concat import MERGE_STRATEGIES, concat_stores +from adata.commands.concat import MERGE_CHOICES, MERGE_STRATEGIES, concat_stores diff --git a/src/adata/commands/concat.py b/src/adata/commands/concat.py index 650b854..befda91 100644 --- a/src/adata/commands/concat.py +++ b/src/adata/commands/concat.py @@ -7,9 +7,9 @@ from rich.console import Console -from adata.core.concat import MERGE_STRATEGIES, concat_on_disk +from adata.core.concat import MERGE_CHOICES, MERGE_STRATEGIES, concat_on_disk -__all__ = ["MERGE_STRATEGIES", "concat_stores"] +__all__ = ["MERGE_CHOICES", "MERGE_STRATEGIES", "concat_stores"] def concat_stores( diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index 4982916..e8b4c44 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -48,6 +48,8 @@ ) MERGE_STRATEGIES = ("same", "unique", "first", "only") +#: What the CLI accepts. "drop" is the default and maps to no merge at all. +MERGE_CHOICES = ("drop",) + MERGE_STRATEGIES def _index_union(per_input: Sequence[List[str]]) -> List[str]: @@ -174,6 +176,16 @@ def __repr__(self) -> str: # pragma: no cover - debugging aid _MISSING = _Missing() +class _Present: + """Stands for a column whose value was not read, only its presence.""" + + def __repr__(self) -> str: # pragma: no cover - debugging aid + return "" + + +_PRESENT = _Present() + + def _equal(a: Any, b: Any) -> bool: if isinstance(a, _Incomparable) or isinstance(b, _Incomparable): return False @@ -945,13 +957,25 @@ def _write_var( positions = [_column_map(target_var, names) for names in var_names] written: List[str] = [] + # "first" and "only" decide on presence alone, so the column values are + # never read for them. + compares = merge in ("same", "unique") + for name in candidates: aligned: List[Any] = [] for group, where in zip(groups, positions): if name not in group or (where < 0).any(): aligned.append(_MISSING) continue - aligned.append(tuple(read_str_all(group[name])[i] for i in where)) + if not compares: + aligned.append(_PRESENT) + continue + # Read the column once and index the result. Reading it inside the + # generator -- as an earlier version did -- re-read the whole + # column for every target variable, which is quadratic and turns a + # 36k-var merge into hours of pure CPU. + values = read_str_all(group[name]) + aligned.append(tuple(values[i] for i in where)) keep, _ = _merge_values(aligned, merge) if not keep: diff --git a/tests/test_commands_phase2.py b/tests/test_commands_phase2.py index cf7bd8b..7f1b367 100644 --- a/tests/test_commands_phase2.py +++ b/tests/test_commands_phase2.py @@ -758,3 +758,104 @@ def test_split_rejects_a_bad_zarr_format(tmp_path, sample): ) assert result.exit_code == 1 assert "must be 2 or 3" in _out(result) + + +# --------------------------------------------------------------------------- +# concat --merge cost +# +# These assert how often a var column is read, not how long the merge takes. +# The hang reported against 0.5.1 was invisible to every correctness test in +# this file because all of them use two or three variables, and the defect was +# quadratic in the number of variables: `tuple(read_str_all(col)[i] for i in +# where)` re-evaluated `read_str_all` once per target variable. A wall-clock +# assertion would be flaky on shared CI; a read count is exact. + + +def _count_var_column_reads(monkeypatch): + """Count full-column reads performed while merging var.""" + from adata.core import concat as concat_mod + + calls: list = [] + original = concat_mod.read_str_all + + def counting(obj, *args, **kwargs): + calls.append(getattr(obj, "name", "?")) + return original(obj, *args, **kwargs) + + monkeypatch.setattr(concat_mod, "read_str_all", counting) + return calls + + +def _two_stores_with_var_columns(tmp_path, n_var: int, n_col: int) -> list: + genes = [f"g{i}" for i in range(n_var)] + var = pd.DataFrame( + {f"col{c}": [f"v{c}-{i}" for i in range(n_var)] for c in range(n_col)}, + index=genes, + ) + paths = [] + for name, cells in (("a", ["c1", "c2"]), ("b", ["c3", "c4"])): + path = tmp_path / f"{name}.h5ad" + ad.AnnData( + X=np.ones((2, n_var), dtype="float32"), + obs=pd.DataFrame(index=cells), + var=var.copy(), + ).write_h5ad(path) + paths.append(path) + return paths + + +@pytest.mark.parametrize("n_var", [4, 64]) +def test_concat_merge_same_reads_each_var_column_once_per_input( + tmp_path, monkeypatch, n_var +): + """`--merge same` must not scale with the number of variables. + + Two inputs and three var columns is six reads whatever n_var is. The + parametrisation is the whole point: an implementation whose cost grows + with n_var fails the second case while passing the first. + """ + a, b = _two_stores_with_var_columns(tmp_path, n_var=n_var, n_col=3) + calls = _count_var_column_reads(monkeypatch) + + out = tmp_path / "m.h5ad" + result = runner.invoke( + app, ["concat", str(a), str(b), "-o", str(out), "--merge", "same"] + ) + assert result.exit_code == 0, _out(result) + assert len(calls) == 6, f"{len(calls)} reads for {n_var} vars: {calls[:10]}" + + var = ad.read_h5ad(out).var + assert list(var.columns) == ["col0", "col1", "col2"] + assert var["col0"].tolist() == [f"v0-{i}" for i in range(n_var)] + + +def test_concat_merge_first_does_not_read_var_column_values( + tmp_path, monkeypatch +): + """`first` decides on presence alone, so it reads no column values.""" + a, b = _two_stores_with_var_columns(tmp_path, n_var=64, n_col=3) + calls = _count_var_column_reads(monkeypatch) + + out = tmp_path / "m.h5ad" + result = runner.invoke( + app, ["concat", str(a), str(b), "-o", str(out), "--merge", "first"] + ) + assert result.exit_code == 0, _out(result) + assert calls == [], f"'first' read {len(calls)} columns it did not compare" + assert list(ad.read_h5ad(out).var.columns) == ["col0", "col1", "col2"] + + +def test_concat_merge_drop_is_accepted_and_keeps_no_var_columns(tmp_path): + """`drop` is the documented default, so it has to be sayable.""" + a, b = _two_stores_with_var_columns(tmp_path, n_var=4, n_col=2) + out = tmp_path / "m.h5ad" + result = runner.invoke( + app, + ["concat", str(a), str(b), "-o", str(out), + "--merge", "drop", "--uns-merge", "drop"], + ) + assert result.exit_code == 0, _out(result) + + got = ad.read_h5ad(out) + assert list(got.var.columns) == [] + assert dict(got.uns) == {} From ea81d6e8698f566647cf8853a15c75c00d9de08b Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 17:01:35 +0100 Subject: [PATCH 06/14] Add complexity guards, and fix the two cost defects they found The --merge hang was not a one-off bug but a defect class the suite could not see: output correct, cost wrong. Every fixture here is a few hundred elements and nothing measured cost, so no test could have failed. tests/perf_counters.py instruments four seams, all verified against the pinned h5py 3.15.1 and zarr 3.1.5: h5py.Dataset.__getitem__, zarr.Array.__getitem__, and LocalStore.get/set/delete (coroutines, wrapped as such). Patching the libraries rather than adding a seam inside src/adata is deliberate -- an in-repo helper would only see the call sites that remembered to use it, and the matrix paths in subset.py and concat.py slice the backend objects directly. Elements are counted, not just calls, so a vectorised-but-quadratic read is caught too. The known bypasses (read_direct, np.asarray(dataset), asstr) are documented, and two canary tests fail loudly if a hook stops firing, since otherwise every ratio below would pass on zeros. The invariant compares successive increments rather than raw counts: d1 = c(4n) - c(n); d2 = c(16n) - c(4n); assert d2 <= 6 * d1 The increment form cancels any fixed setup cost exactly, so there is no slack constant to tune and no floor for a small-coefficient quadratic to hide under. At 4x spacing the ratio is 4.0 for linear work, 4.4 for n log n, 8 for n**1.5 and 16 for quadratic, so 6 sits in the gap with room either side; a test asserts that calibration rather than leaving it as a comment. Each guard scales exactly one axis -- n_var, n_obs, n_inputs, n_columns, n_categories -- because scaling two at once makes legitimate work look quadratic. Reintroducing the REQ-71798 line makes the n_var guard fail at ratio 15.9 on all four merge strategies. Writing the guards turned up two more defects of the same class, both fixed here: _concat_categorical unioned categories with `if category not in categories` on a list -- O(k^2). Neither instrument above sees it: the category lists are read once either way, and `x not in lst` is a single bytecode, so the quadratic lives inside C-level list membership. Counting string comparisons via a str subclass is what makes it visible: 2,096,128 comparisons at k=1024, and around 5e9 for a 100k-category obs column. A dict takes it to 0. _concat_masked filled a Python list of length n_obs one element at a time and then walked it twice more. A typed numpy buffer filled by slice takes it from 1.21x the executed Python lines of the numeric path to 1.00x. The string path is 1.35x and stays there -- read_str_all materialises Python str objects, which is inherent to reading strings rather than a per-row loop -- so its guard is set at that measured level with the reason recorded. Also registers the perf marker (--strict-markers is on) and notes what these tests deliberately do not claim: obs columns are read whole, so peak allocation is O(n_obs) and not O(chunk). The streaming guarantee holds for X, not for obs annotation. benchmarks/ will report that curve. 853 tests pass, coverage 92.6%. Co-Authored-By: Claude Opus 5 --- pytest.ini | 1 + src/adata/core/concat.py | 55 +++-- tests/perf_counters.py | 311 +++++++++++++++++++++++ tests/test_performance.py | 504 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 847 insertions(+), 24 deletions(-) create mode 100644 tests/perf_counters.py create mode 100644 tests/test_performance.py diff --git a/pytest.ini b/pytest.ini index e7b2969..a116cbf 100644 --- a/pytest.ini +++ b/pytest.ini @@ -6,4 +6,5 @@ python_functions = test_* addopts = -v --strict-markers --tb=short --timeout=300 --timeout-method=thread markers = slow: marks tests as slow (deselect with '-m "not slow"') + perf: complexity guards that trace Python execution (10-50x slower) integration: builds environments with uv and needs the network on first run diff --git a/src/adata/core/concat.py b/src/adata/core/concat.py index e8b4c44..f1c4ecc 100644 --- a/src/adata/core/concat.py +++ b/src/adata/core/concat.py @@ -224,15 +224,20 @@ def _concat_categorical( concatenated directly. A row from an input lacking the column gets code -1, which anndata reads back as a missing value. """ - categories: List[str] = [] + # Hash lookup, not list membership. `if category not in categories` on a + # list is O(k) per probe and so O(k^2) over the union -- around 2 million + # string comparisons at 1024 categories, and 5 billion at 100k. It reads + # the same number of bytes either way, so only a comparison count sees it. + lookup: Dict[str, int] = {} for col in columns: if col is None: continue for category in read_categories(col): - if category not in categories: - categories.append(str(category)) + text = str(category) + if text not in lookup: + lookup[text] = len(lookup) - lookup = {c: i for i, c in enumerate(categories)} + categories = list(lookup) ordered = all(is_ordered(c) for c in columns if c is not None) codes = np.full(sum(lengths), -1, dtype=np.int64) @@ -289,33 +294,35 @@ def _concat_masked( total = sum(lengths) mask = np.ones(total, dtype=bool) - values: List[Any] = [None] * total + present = [c for c in columns if c is not None] + + # Fill a typed buffer by slice, rather than a Python list of length + # n_obs one element at a time. The list cost an object per row and three + # full passes over it, which is what made obs concatenation the most + # allocation-hungry part of a streamed concat. + if enc == spec.NULLABLE_STRING_ARRAY: + from adata.elements.read import decode_str_array + + filled = np.empty(total, dtype=object) + filled[:] = "" + else: + decode_str_array = None + dtype = np.result_type(*[c["values"].dtype for c in present]) + filled = np.zeros(total, dtype=dtype) offset = 0 for col, length in zip(columns, lengths): if col is not None: chunk = np.asarray(col["values"][...]) - chunk_mask = np.asarray(col["mask"][...], dtype=bool) - for i in range(length): - values[offset + i] = chunk[i] - mask[offset : offset + length] = chunk_mask + if decode_str_array is not None: + chunk = decode_str_array(chunk) + filled[offset : offset + length] = chunk + mask[offset : offset + length] = np.asarray( + col["mask"][...], dtype=bool + ) offset += length - if enc == spec.NULLABLE_STRING_ARRAY: - from adata.elements.read import decode_str_array - - filled = [ - "" if v is None else decode_str_array(np.asarray([v]))[0] for v in values - ] - write_masked(parent, name, filled, mask, enc) - return - - present = [c for c in columns if c is not None] - dtype = np.result_type(*[c["values"].dtype for c in present]) - filled_num = np.array( - [0 if v is None else v for v in values], dtype=dtype - ) - write_masked(parent, name, filled_num, mask, enc) + write_masked(parent, name, filled, mask, enc) def _concat_string( diff --git a/tests/perf_counters.py b/tests/perf_counters.py new file mode 100644 index 0000000..0a582d1 --- /dev/null +++ b/tests/perf_counters.py @@ -0,0 +1,311 @@ +"""Deterministic cost counters for the complexity guards. + +These measure *how much work* an operation does, never how long it takes. A +test that can fail because a CI runner was busy does not belong in a merge +gate, so nothing here touches the clock. + +Why patch the libraries rather than `adata` +------------------------------------------- +The hooks wrap h5py and zarr themselves. A seam inside `src/adata/` would only +measure the call sites that remembered to use it, which is precisely the wrong +property for a guard whose job is to catch the read you did not think of. The +matrix paths in particular slice the backend objects directly -- see +`core/subset.py` `_copy_dense_rows` and `core/concat.py` `_concat_csr` -- and +no in-repo helper sees them at all. + +What is counted +--------------- +Elements and bytes, not just calls. The `--merge` hang was n calls each reading +n elements; counting calls alone would also miss a vectorised-but-quadratic +variant that reads n elements n times in one call. + +Known bypasses +-------------- +These reach the file without passing through any hook below: + +* ``h5py.Dataset.read_direct`` +* ``np.asarray(dataset)`` and ``dataset[...]`` via ``__array__`` (verified: + ``np.asarray`` adds nothing to the element count) +* ``dataset.asstr()[...]``, which goes through ``AsStrWrapper.__getitem__`` +* ``dataset.fields(...)`` + +None are used in `src/adata/` today. If you add one, the counters will silently +report less work than actually happened -- which is why every guard also +asserts a *lower* bound, and why `test_performance.py` keeps a canary test with +a known absolute count. +""" + +from __future__ import annotations + +import sys +import tracemalloc +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Iterator, List + +import numpy as np + + +@dataclass +class IOCounts: + """Live tally of backend reads and writes. + + The object is yielded before the work runs and mutated in place, so read + the fields after the `with` block closes. + """ + + h5_calls: int = 0 + h5_elements: int = 0 + zarr_calls: int = 0 + zarr_elements: int = 0 + store_get: int = 0 + store_get_bytes: int = 0 + store_set: int = 0 + store_delete: int = 0 + + #: Per-dataset element counts, for diagnosing which column blew up. + by_name: Dict[str, int] = field(default_factory=dict) + + @property + def elements(self) -> int: + """Elements read through either backend's array interface.""" + return self.h5_elements + self.zarr_elements + + @property + def calls(self) -> int: + return self.h5_calls + self.zarr_calls + + @property + def reads(self) -> int: + """The headline number: array elements plus store-level fetches. + + Covers both the "read the same column n times" shape and the + "re-fetch the same chunk n times" shape. + """ + return self.elements + self.store_get + + def __str__(self) -> str: # pragma: no cover - diagnostic only + return ( + f"elements={self.elements} (h5={self.h5_elements} " + f"zarr={self.zarr_elements}) calls={self.calls} " + f"store: get={self.store_get} set={self.store_set} " + f"delete={self.store_delete}" + ) + + +def _size(value: Any) -> int: + try: + return int(np.asarray(value).size) + except Exception: # pragma: no cover - exotic dtypes + return 1 + + +@contextmanager +def count_io() -> Iterator[IOCounts]: + """Count backend reads and writes for the duration of the block. + + Patches four seams, all verified against h5py 3.15.1 and zarr 3.1.5: + + =================================== ==================================== + ``h5py.Dataset.__getitem__`` HDF5 reads + ``zarr.Array.__getitem__`` Zarr reads, v2 and v3 alike + ``zarr.storage.LocalStore.get`` chunk and metadata fetches + ``LocalStore.set`` / ``.delete`` write storms + =================================== ==================================== + + The store hooks are the ones that see a metadata write storm: attribute + writes never touch `Array.__getitem__`, so an O(k^2) attribute rewrite is + invisible at the array level and obvious at the store level. + + The `LocalStore` methods are coroutines and are wrapped as coroutines. + """ + import h5py + import zarr + from zarr.storage import LocalStore + + counts = IOCounts() + + h5_get = h5py.Dataset.__getitem__ + zarr_get = zarr.Array.__getitem__ + store_get = LocalStore.get + store_set = LocalStore.set + store_delete = LocalStore.delete + + def _record(name: str, n: int) -> None: + counts.by_name[name] = counts.by_name.get(name, 0) + n + + def h5_wrapper(self: Any, key: Any) -> Any: + result = h5_get(self, key) + n = _size(result) + counts.h5_calls += 1 + counts.h5_elements += n + _record(getattr(self, "name", "?"), n) + return result + + def zarr_wrapper(self: Any, key: Any) -> Any: + result = zarr_get(self, key) + n = _size(result) + counts.zarr_calls += 1 + counts.zarr_elements += n + _record(getattr(self, "path", "?") or "/", n) + return result + + async def get_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + result = await store_get(self, *args, **kwargs) + counts.store_get += 1 + if result is not None: + try: + counts.store_get_bytes += len(result) + except TypeError: # pragma: no cover - prototype buffers + pass + return result + + async def set_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + counts.store_set += 1 + return await store_set(self, *args, **kwargs) + + async def delete_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: + counts.store_delete += 1 + return await store_delete(self, *args, **kwargs) + + h5py.Dataset.__getitem__ = h5_wrapper + zarr.Array.__getitem__ = zarr_wrapper + LocalStore.get = get_wrapper + LocalStore.set = set_wrapper + LocalStore.delete = delete_wrapper + try: + yield counts + finally: + h5py.Dataset.__getitem__ = h5_get + zarr.Array.__getitem__ = zarr_get + LocalStore.get = store_get + LocalStore.set = store_set + LocalStore.delete = store_delete + + +@contextmanager +def count_allocations() -> Iterator[List[int]]: + """Peak Python allocation during the block, in bytes, as `[peak]`. + + `tracemalloc` sees Python-level allocation, which includes the numpy array + *objects* and every Python list and string built along the way, but not + memory a C extension takes outside the allocator. That is the right scope + here: the defects this catches are whole-column Python lists, and h5py's + own buffers are not what we are policing. + + Deterministic for deterministic input, which is why it can gate a merge. + """ + peak = [0] + started = not tracemalloc.is_tracing() + if started: + tracemalloc.start() + else: # pragma: no cover - nested use + tracemalloc.reset_peak() + try: + yield peak + finally: + peak[0] = tracemalloc.get_traced_memory()[1] + if started: + tracemalloc.stop() + + +@contextmanager +def count_lines(prefix: str) -> Iterator[List[int]]: + """Count executed Python line events in files under `prefix`, as `[n]`. + + A deterministic stand-in for CPU time. It is what catches a quadratic that + performs no extra I/O and allocates nothing extra -- `if x not in a_list` + inside a loop, say. Costs a 10-50x slowdown, so keep it to the few paths + with known Python-level inner loops and to small inputs. + + Blind to work done inside numpy's C code, so it complements the allocation + and I/O counters rather than replacing them. + """ + total = [0] + + def local_trace(frame: Any, event: str, arg: Any) -> Any: + if event == "line": + total[0] += 1 + return local_trace + + def trace(frame: Any, event: str, arg: Any) -> Any: + if event == "call" and frame.f_code.co_filename.startswith(prefix): + return local_trace + return None + + previous = sys.gettrace() + sys.settrace(trace) + try: + yield total + finally: + sys.settrace(previous) + + +# --------------------------------------------------------------------------- +# the growth assertion + + +#: Ratio of successive increments that counts as super-linear. +#: +#: Measured at 4x spacing, the increment ratio is 4.0 for linear work, about +#: 4.4 for n log n, 8 for n**1.5 and 16 for quadratic. 6 sits in the gap with +#: room on both sides, so it needs no per-test tuning. +GROWTH_LIMIT = 6.0 + +#: Sizes every growth guard runs at. 4x spacing is what makes GROWTH_LIMIT +#: mean what it says; changing one without the other invalidates the constant. +SIZES = (64, 256, 1024) + + +def assert_grows_linearly( + measure: Callable[[int], int], + *, + what: str, + axis: str, + sizes: tuple = SIZES, + limit: float = GROWTH_LIMIT, +) -> None: + """Assert `measure` grows no faster than linearly in one axis. + + Compares successive *increments* rather than raw counts:: + + d1 = c(4n) - c(n) + d2 = c(16n) - c(4n) + assert d2 <= limit * d1 + + The increment form is what removes the arbitrary additive constant. Any + fixed setup cost -- opening the store, writing the skeleton, reading the + index once -- appears in both differences and cancels exactly, so there is + no slack parameter to pick and no floor for a small-coefficient quadratic + to hide under. + + `d1 >= sizes[1]` is a vacuity guard. If a future refactor moves a read onto + a path the counters do not see (see the module docstring), every ratio + would pass trivially; this fails loudly instead. + + Scale exactly one axis per call and hold the others fixed. Scaling two at + once makes legitimate work look quadratic -- an outer-join concat really + does produce n_obs x n_var_union cells. + """ + small, mid, large = sizes + c_small, c_mid, c_large = (measure(n) for n in (small, mid, large)) + + d1 = c_mid - c_small + d2 = c_large - c_mid + series = ( + f"{what}, scaling {axis}: " + f"{small}->{c_small}, {mid}->{c_mid}, {large}->{c_large} " + f"(increments {d1} then {d2}" + + (f", ratio {d2 / d1:.1f}" if d1 > 0 else "") + + ")" + ) + + assert d1 >= mid, ( + f"Too little measured work for the growth ratio to mean anything -- " + f"the counters are probably not seeing this operation. {series}" + ) + assert d2 <= limit * d1, ( + f"Cost grows faster than linearly in {axis}. Expect an increment " + f"ratio near 4.0 for linear work; {limit} is the limit and quadratic " + f"would be about 16. {series}" + ) diff --git a/tests/test_performance.py b/tests/test_performance.py new file mode 100644 index 0000000..9f98dc5 --- /dev/null +++ b/tests/test_performance.py @@ -0,0 +1,504 @@ +"""Complexity guards: what an operation *costs*, never how long it takes. + +Why this file exists +-------------------- +`concat --merge` hung for hours on a 36,601-var input (REQ-71798). One line in +`_write_var` re-read a whole var column from disk once per target variable. +The output would have been correct; only the cost was wrong. 1,029 tests +missed it, and they were never going to catch it: every fixture in the suite +is at most a few hundred elements, and nothing measured cost at all. + +So this file closes a defect class, not a bug. The instruments are in +`tests/perf_counters.py`; the rule they enforce is that cost must grow no +faster than linearly in each axis, measured in operations rather than seconds. + +Why not wall clock +------------------ +A test that can fail because a CI runner was busy does not belong in a merge +gate. Every number here is deterministic: the same input produces the same +count on every machine. That is what lets these run in the ordinary test job +on both interpreters and block a merge when they fail. + +Reading a failure +----------------- +A failure here says cost grew super-linearly on the named axis. Usually the +code is wrong. Occasionally the expectation is -- an operation legitimately +gains work -- and then the new number needs a comment saying why, in the same +style as the exact counts in `test_commands_phase2.py`. + +One axis at a time. Scaling two at once makes legitimate work look quadratic: +an outer-join concat really does produce n_obs x n_var_union cells. +""" + +from __future__ import annotations + +import pathlib +from pathlib import Path +from typing import List, Optional + +import numpy as np +import pytest +from rich.console import Console + +from adata.core.concat import concat_on_disk +from adata.core.subset import subset_h5ad + +from tests.perf_counters import ( + GROWTH_LIMIT, + SIZES, + assert_grows_linearly, + count_allocations, + count_io, + count_lines, +) + +ad = pytest.importorskip("anndata", reason="anndata builds the fixtures") +pd = pytest.importorskip("pandas") +sparse = pytest.importorskip("scipy.sparse") + +QUIET = Console(quiet=True) + +#: Only Python executed inside the package counts towards the CPU proxy. +SOURCE_PREFIX = str(pathlib.Path(__file__).resolve().parent.parent / "src" / "adata") + + +# --------------------------------------------------------------------------- +# fixtures built to scale exactly one axis + + +def _store( + path: Path, + *, + name: str, + n_obs: int = 4, + n_var: int = 8, + n_var_columns: int = 2, + n_obs_columns: int = 0, + n_categories: int = 0, + obs_kind: Optional[str] = None, + shared_var: bool = True, +) -> Path: + """One input store. Every argument is an axis a guard can scale.""" + obs = pd.DataFrame(index=[f"{name}c{i}" for i in range(n_obs)]) + if n_categories: + obs["ct"] = pd.Categorical( + [f"{name}-t{i % n_categories}" for i in range(n_obs)] + ) + if obs_kind == "numeric": + obs["col"] = np.arange(n_obs, dtype="int32") + elif obs_kind == "masked": + obs["col"] = pd.array( + [i if i % 2 else None for i in range(n_obs)], dtype="Int32" + ) + elif obs_kind == "string": + obs["col"] = [f"{name}-{i}" for i in range(n_obs)] + if n_obs_columns: + extra = pd.DataFrame( + {f"m{c}": np.arange(n_obs, dtype="int32") for c in range(n_obs_columns)}, + index=obs.index, + ) + obs = pd.concat([obs, extra], axis=1) + + prefix = "" if shared_var else name + var = pd.DataFrame( + {f"v{c}": [f"{c}-{i}" for i in range(n_var)] for c in range(n_var_columns)}, + index=[f"{prefix}g{i}" for i in range(n_var)], + ) + + ad.AnnData( + X=sparse.csr_matrix(np.ones((n_obs, n_var), dtype="float32")), + obs=obs, + var=var, + ).write_h5ad(path) + return path + + +def _inputs(tmp_path: Path, tag: str, count: int = 2, **kwargs) -> List[Path]: + directory = tmp_path / tag + directory.mkdir(parents=True, exist_ok=True) + return [ + _store(directory / f"{chr(ord('a') + i)}.h5ad", name=chr(ord("a") + i), **kwargs) + for i in range(count) + ] + + +# --------------------------------------------------------------------------- +# canary +# +# Every growth guard below is a comparison between measurements. If a refactor +# moves a read onto a path the counters cannot see -- `read_direct`, +# `np.asarray(dataset)`, `asstr()[...]`; see perf_counters' docstring -- the +# comparisons would all pass on zeros. The vacuity check inside +# `assert_grows_linearly` catches most of that; this catches the rest, and +# says plainly what broke. + + +def test_the_counters_see_a_known_operation(tmp_path): + """A concat of two 4 x 64 stores must register substantial work. + + If this fails, the hooks have stopped matching the libraries and every + other test in this file has become meaningless. Fix the hooks first. + """ + files = _inputs(tmp_path, "canary", n_obs=4, n_var=64) + with count_io() as io: + concat_on_disk(files, tmp_path / "out.h5ad", QUIET, merge="same") + + assert io.h5_calls > 0, "h5py.Dataset.__getitem__ hook is not firing" + assert io.h5_elements >= 64, f"suspiciously few elements read: {io}" + + with count_allocations() as peak: + concat_on_disk(files, tmp_path / "out2.h5ad", QUIET, merge="same") + assert peak[0] > 0, "tracemalloc recorded no allocation" + + with count_lines(SOURCE_PREFIX) as lines: + concat_on_disk(files, tmp_path / "out3.h5ad", QUIET, merge="same") + assert lines[0] > 100, f"line tracer saw only {lines[0]} events in src/adata" + + +def test_the_zarr_store_hooks_fire(tmp_path): + """Metadata and chunk traffic must be visible, not just array reads. + + An attribute-write storm never touches `Array.__getitem__` -- it is all + `LocalStore.set`. That shape of bug has bitten this repo before, so the + store hooks get their own canary. + """ + files = _inputs(tmp_path, "zcanary", n_obs=4, n_var=32) + with count_io() as io: + concat_on_disk(files, tmp_path / "out.zarr", QUIET, merge="same") + + assert io.store_set > 0, "LocalStore.set hook is not firing" + assert io.store_get > 0, "LocalStore.get hook is not firing" + + +# --------------------------------------------------------------------------- +# concat: reads must not grow faster than linearly on any axis +# +# The n_var case is the direct regression guard for REQ-71798. The others +# exist because concat's cost surface is n_obs x n_var x n_columns x n_inputs +# and a defect can hide in any of them; scaling only the axis that broke last +# time would be fighting the last war. + + +@pytest.mark.parametrize("merge", [None, "same", "unique", "first", "only"]) +def test_concat_reads_grow_linearly_in_var_count(tmp_path, merge): + """The REQ-71798 axis. Every merge strategy, not just the one reported.""" + + def measure(n: int) -> int: + files = _inputs(tmp_path, f"var{merge}{n}", n_var=n) + with count_io() as io: + concat_on_disk(files, tmp_path / f"o{merge}{n}.h5ad", QUIET, merge=merge) + return io.reads + + assert_grows_linearly(measure, what=f"concat --merge {merge}", axis="n_var") + + +@pytest.mark.parametrize("join", ["inner", "outer"]) +def test_concat_reads_grow_linearly_in_obs_count(tmp_path, join): + def measure(n: int) -> int: + files = _inputs(tmp_path, f"obs{join}{n}", n_obs=n, n_var=4) + with count_io() as io: + concat_on_disk(files, tmp_path / f"o{join}{n}.h5ad", QUIET, join=join) + return io.reads + + assert_grows_linearly(measure, what=f"concat --join {join}", axis="n_obs") + + +def test_concat_reads_grow_linearly_in_input_count(tmp_path): + """Pairwise work across inputs would show up here and nowhere else.""" + + def measure(n: int) -> int: + files = _inputs(tmp_path, f"in{n}", count=n, n_obs=2, n_var=4) + with count_io() as io: + concat_on_disk(files, tmp_path / f"oin{n}.h5ad", QUIET, merge="same") + return io.reads + + # Sizes are scaled down: n inputs means n files on disk. + assert_grows_linearly( + measure, what="concat", axis="n_inputs", sizes=(4, 16, 64) + ) + + +def test_concat_reads_grow_linearly_in_obs_column_count(tmp_path): + def measure(n: int) -> int: + files = _inputs(tmp_path, f"oc{n}", n_obs=8, n_var=4, n_obs_columns=n) + with count_io() as io: + concat_on_disk(files, tmp_path / f"ooc{n}.h5ad", QUIET) + return io.reads + + assert_grows_linearly(measure, what="concat", axis="n_obs_columns") + + +def test_concat_reads_grow_linearly_in_var_column_count(tmp_path): + def measure(n: int) -> int: + files = _inputs(tmp_path, f"vc{n}", n_obs=4, n_var=8, n_var_columns=n) + with count_io() as io: + concat_on_disk(tmp_path / f"vc{n}" and files, tmp_path / f"ovc{n}.h5ad", + QUIET, merge="same") + return io.reads + + assert_grows_linearly(measure, what="concat --merge same", axis="n_var_columns") + + +# --------------------------------------------------------------------------- +# concat: the category union +# +# This one needs its own instrument. Merging categories reads each category +# list exactly once however the union is computed, so a read counter sees +# nothing; and `x not in some_list` is a single bytecode, so the line tracer +# sees nothing either -- the quadratic lives inside C-level list membership. +# Counting string comparisons is the only thing that makes it visible. + + +class _CountingStr(str): + """A string that tallies its own comparisons. + + `list.__contains__` rich-compares each element, and a subclass's `__eq__` + takes priority, so the tally reflects real membership-probe cost. + """ + + count = 0 + + def __eq__(self, other: object) -> bool: + _CountingStr.count += 1 + return str.__eq__(self, other) + + def __ne__(self, other: object) -> bool: + return not self.__eq__(other) + + def __hash__(self) -> int: + return str.__hash__(self) + + +def _count_category_comparisons(monkeypatch, tmp_path, k: int) -> int: + import adata.core.concat as concat_module + + original = concat_module.read_categories + monkeypatch.setattr( + concat_module, + "read_categories", + lambda col: [_CountingStr(c) for c in original(col)], + ) + + files = _inputs(tmp_path, f"cat{k}", n_obs=k, n_var=2, n_categories=k) + _CountingStr.count = 0 + concat_on_disk(files, tmp_path / f"ocat{k}.h5ad", QUIET) + return _CountingStr.count + + +def test_concat_unions_categories_without_quadratic_comparisons( + tmp_path, monkeypatch +): + """Category merging must use hash lookup, not list membership. + + `if category not in categories` on a list was O(k) per probe and so + O(k^2) over the union: 2,096,128 string comparisons at k=1024, and around + 5e9 for an obs column with 100k categories -- the same silent hang as + REQ-71798, in a different function. + + A dict makes it O(k) with comparisons only on hash collision, so the + budget is generous and still three orders of magnitude below the old cost. + """ + k = 1024 + comparisons = _count_category_comparisons(monkeypatch, tmp_path, k) + assert comparisons <= 10 * k, ( + f"{comparisons} string comparisons to union {k} categories across two " + f"inputs. Linear would be near zero (hash collisions only); the " + f"list-membership implementation cost 2,096,128." + ) + + +# --------------------------------------------------------------------------- +# concat: per-row cost of an obs column +# +# Growth ratios cannot see these: building a Python object per row is linear, +# just with a fat constant. Comparing a column kind against a plain numeric +# column in the same run calibrates that constant away, so the assertion holds +# across interpreters and platforms without a hand-tuned byte budget. + + +def _obs_column_cost(tmp_path, kind: str, n_obs: int) -> tuple: + files = _inputs(tmp_path, f"k{kind}{n_obs}", n_obs=n_obs, n_var=2, obs_kind=kind) + with count_lines(SOURCE_PREFIX) as lines, count_allocations() as peak: + concat_on_disk(files, tmp_path / f"ok{kind}{n_obs}.h5ad", QUIET) + return lines[0], peak[0] + + +@pytest.mark.perf +@pytest.mark.parametrize("kind", ["masked", "string"]) +def test_obs_columns_cost_no_more_python_work_per_row_than_a_numeric_one( + tmp_path, kind +): + """No Python-level loop over rows when concatenating an obs column. + + `_concat_masked` used to fill a `List[Any]` of length n_obs one element at + a time and then walk it twice more. That is invisible to a read counter + and to a growth ratio, but it shows up immediately as executed line events + per row relative to the numeric path, which has always been vectorised. + + Measured after the fix: masked is 1.00x numeric, string 1.35x (the string + path still materialises Python `str` objects via `read_str_all`, which is + inherent to how strings are read, not a per-row loop). Before the fix + masked was above 1.2x. + """ + n_obs = 4000 + numeric_lines, _ = _obs_column_cost(tmp_path, "numeric", n_obs) + kind_lines, _ = _obs_column_cost(tmp_path, kind, n_obs) + + limit = 1.10 if kind == "masked" else 1.45 + ratio = kind_lines / numeric_lines + assert ratio <= limit, ( + f"a {kind} obs column executes {ratio:.2f}x the Python lines of a " + f"numeric one ({kind_lines} vs {numeric_lines} at n_obs={n_obs}); " + f"limit {limit}. Look for a per-element loop over rows." + ) + + +@pytest.mark.parametrize("kind", ["numeric", "masked", "string"]) +def test_obs_column_allocation_grows_linearly(tmp_path, kind): + """Peak allocation must not grow faster than the data itself. + + Note what this does *not* claim: obs columns are read whole, so peak + allocation is O(n_obs) and not O(chunk). The streaming guarantee that the + README makes holds for X, not for obs annotation. `benchmarks/` reports + the real curve; this only stops it getting worse than linear. + """ + + def measure(n: int) -> int: + _, peak = _obs_column_cost(tmp_path, kind, n) + return peak + + assert_grows_linearly( + measure, + what=f"concat allocation, {kind} obs column", + axis="n_obs", + sizes=(256, 1024, 4096), + ) + + +# --------------------------------------------------------------------------- +# subset, split and conversion +# +# concat is where the defect was found, but nothing about the defect was +# specific to concat: every command that aligns one index onto another has the +# same opportunity to re-read a column per element. + + +@pytest.mark.parametrize("axis", ["obs", "var"]) +def test_subset_by_name_reads_grow_linearly(tmp_path, axis): + def measure(n: int) -> int: + source = _store( + tmp_path / f"s{axis}{n}.h5ad", + name="s", + n_obs=n if axis == "obs" else 4, + n_var=n if axis == "var" else 4, + ) + # A name file, not indices: resolving names onto the stored index is + # exactly the kind of alignment that went quadratic in concat. + names = [f"sc{i}" for i in range(0, n, 2)] if axis == "obs" else [ + f"g{i}" for i in range(0, n, 2) + ] + name_file = tmp_path / f"n{axis}{n}.txt" + name_file.write_text("\n".join(names) + "\n") + with count_io() as io: + subset_h5ad( + source, + tmp_path / f"os{axis}{n}.h5ad", + name_file if axis == "obs" else None, + name_file if axis == "var" else None, + console=QUIET, + ) + return io.reads + + assert_grows_linearly( + measure, what=f"subset --{axis}", axis=f"n_{axis}", sizes=(64, 256, 1024) + ) + + +def test_subset_by_query_reads_grow_linearly(tmp_path): + def measure(n: int) -> int: + source = _store( + tmp_path / f"q{n}.h5ad", name="q", n_obs=n, n_var=4, obs_kind="numeric" + ) + with count_io() as io: + subset_h5ad( + source, + tmp_path / f"oq{n}.h5ad", + None, + None, + console=QUIET, + obs_query="col >= 0", + ) + return io.reads + + assert_grows_linearly(measure, what="subset --obs-query", axis="n_obs") + + +def test_h5ad_to_zarr_conversion_reads_grow_linearly(tmp_path): + def measure(n: int) -> int: + source = _store(tmp_path / f"c{n}.h5ad", name="c", n_obs=n, n_var=4) + # subset is also the conversion path, but it insists on a selection; + # selecting every row is the identity and still rewrites the store. + keep = tmp_path / f"ck{n}.txt" + keep.write_text("\n".join(f"cc{i}" for i in range(n)) + "\n") + with count_io() as io: + subset_h5ad( + source, tmp_path / f"oc{n}.zarr", keep, None, console=QUIET + ) + return io.reads + + assert_grows_linearly(measure, what="h5ad to zarr", axis="n_obs") + + +def test_split_reads_grow_linearly_in_group_count(tmp_path): + """Splitting into k groups must stay linear in k at fixed store size. + + n_obs is held constant deliberately. Scaling rows alongside groups would + make the total legitimately quadratic -- k passes over 4k rows -- and the + guard would be measuring the fixture, not the implementation. + """ + from adata.commands.split import split_store + + n_obs = 512 + + def measure(n: int) -> int: + source = _store( + tmp_path / f"sp{n}.h5ad", + name="sp", + n_obs=n_obs, + n_var=4, + n_categories=n, + ) + out = tmp_path / f"osp{n}" + with count_io() as io: + split_store(source, "ct", out, QUIET, axis="obs", manifest=False) + return io.reads + + assert_grows_linearly( + measure, what="split --by", axis="n_groups", sizes=(8, 32, 128) + ) + + +# --------------------------------------------------------------------------- +# the constants themselves + + +def test_growth_limit_separates_linear_from_super_linear(): + """Documenting why K is 6, as an executable statement rather than a note. + + At the 4x spacing in SIZES the increment ratio is 4.0 for linear work, + about 4.4 for n log n, 8 for n**1.5 and 16 for quadratic. + """ + small, mid, large = SIZES + assert mid == small * 4 and large == mid * 4, ( + "GROWTH_LIMIT is calibrated for 4x spacing; changing SIZES without " + "recomputing it invalidates every guard in this file." + ) + + def ratio(exponent: float) -> float: + cost = [float(n) ** exponent for n in SIZES] + return (cost[2] - cost[1]) / (cost[1] - cost[0]) + + assert ratio(1.0) < GROWTH_LIMIT, "linear work must pass" + assert ratio(1.0) == pytest.approx(4.0, rel=0.01) + assert ratio(2.0) > GROWTH_LIMIT, "quadratic work must fail" + assert ratio(1.5) > GROWTH_LIMIT, "n**1.5 must fail" From 20a570f8bb5d9d586d711c79ad7679d1fb5a2b37 Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 17:14:00 +0100 Subject: [PATCH 07/14] Add the comparative benchmark, run and published on every tag Complements the complexity guards rather than duplicating them. The guards count operations, are deterministic, and gate merges; this measures wall time and peak RSS at realistic size, runs on tags, and never fails a build -- timing on a shared runner is too noisy to gate a release on, and a benchmark that can block a publish stops being run. Peak RSS is the headline. adata-cli exists so that memory is set by --chunk rather than input size, and for data that fits in RAM loading the whole thing is frequently faster; a table reporting only wall time would misrepresent the tool in the direction of flattery and then in the direction of failure. Measurement (benchmarks/_measure.py): os.wait4, not resource.getrusage. RUSAGE_CHILDREN is a running maximum over every child a process has reaped, so a 400 MB case followed by a 1 kB one reports 400 MB twice and every later row inherits the largest earlier peak. Output goes to temporary files rather than pipes, because communicate() reaps the child and there is then nothing for wait4 to report. ru_maxrss is normalised -- KiB on Linux, bytes on macOS. RLIMIT_AS at 12 GiB on every child. Cases the in-memory baseline cannot survive are the point of the comparison, but an uncontained OOM kills the runner agent and the job ends with no report at all; with a ceiling it is a row that says "out of memory" at a limit we can state. A timeout records the output size at the kill, which is how the 0.5.1 hang actually presented -- 1,489,960 bytes, never growing -- and distinguishes it from slow progress. tests/test_benchmark_harness.py covers exactly this and nothing else. If peak RSS were attributed to the wrong process, every published table would be wrong and would still look plausible. Fairness rules, written into cases.py and docs/TESTING.md because this is what decays first: use the best idiom the baseline has (read_elem, backed mode) and never a strawman full load as the primary row; pin compression on both sides, since adata-cli forwards the source's settings while write_h5ad defaults to none; print n/a with a reason where scanpy or concat_on_disk has no equivalent, because an omitted row reads as an oversight; include a startup floor, as the CLI costs 0.3-1 s to import and scanpy 3-8 s; say whether the page cache was dropped; and leave the rows where anndata wins exactly as measured. _concat_csr loops per row in Python and scipy's C vstack will often beat it on time at several times the memory -- that trade is the argument for this tool, and hiding it would make the table worthless. Baselines run from venvs built up front rather than `uv run --with`. The latter is right for reference_stores.py, where fixture cost is irrelevant, and wrong here: the first invocation would put hundreds of megabytes of wheel downloads into the measured wall time and uv's own memory into the measured peak. scanpy stays out of uv.lock and out of the image either way. Publishing goes to docs, not an orphan branch: docs/benchmarks/.json for the series and docs/BENCHMARKS.md for the rendered page, both already served by Pages from docs/, plus the step summary and a best-effort idempotent append to the release notes. Artifacts expire in 90 days and one absolute number with nothing to compare against says very little, so the durable series is the part that matters. Note for review: the docs commit is the one outward-facing write here -- a tag build is on a detached HEAD, so the job commits to main as the bot with [skip ci] and rebase-retries on a race. Its own workflow file rather than a job in publish.yml: at the ci tier this takes the better part of an hour, and hanging that off the release graph would either delay the PyPI publish or paint the release run red for a report. It benchmarks the checked-out source, not the published wheel, which would mean waiting on publish-pypi and then on index propagation for a measurement that comes out the same. 861 tests pass, coverage 92.6%. Co-Authored-By: Claude Opus 5 --- .github/workflows/benchmark.yml | 137 +++++++++++++ CHANGELOG.md | 18 ++ benchmarks/__init__.py | 10 + benchmarks/_measure.py | 223 +++++++++++++++++++++ benchmarks/cases.py | 342 ++++++++++++++++++++++++++++++++ benchmarks/datasets.py | 132 ++++++++++++ benchmarks/report.py | 257 ++++++++++++++++++++++++ benchmarks/run.py | 274 +++++++++++++++++++++++++ docs/BENCHMARKS.md | 25 +++ docs/TESTING.md | 174 ++++++++++++++++ tests/test_benchmark_harness.py | 104 ++++++++++ 11 files changed, 1696 insertions(+) create mode 100644 .github/workflows/benchmark.yml create mode 100644 benchmarks/__init__.py create mode 100644 benchmarks/_measure.py create mode 100644 benchmarks/cases.py create mode 100644 benchmarks/datasets.py create mode 100644 benchmarks/report.py create mode 100644 benchmarks/run.py create mode 100644 docs/BENCHMARKS.md create mode 100644 tests/test_benchmark_harness.py diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml new file mode 100644 index 0000000..938aba5 --- /dev/null +++ b/.github/workflows/benchmark.yml @@ -0,0 +1,137 @@ +# Comparative benchmark: adata-cli against anndata and scanpy. +# +# Report-only, by design. Wall time on a shared runner is too noisy to gate a +# release on, and a benchmark that can block a publish stops being run. What +# *does* gate merges is tests/test_performance.py, which counts operations +# rather than seconds and lives in the ordinary test job. +# +# Its own file rather than a job in publish.yml: at the `ci` tier this takes +# the better part of an hour, and hanging that off the release graph would +# either delay the PyPI publish or paint the release run red for a report. +name: Benchmark + +on: + push: + tags: ["*"] + workflow_dispatch: + inputs: + tier: + description: Input size + type: choice + options: [smoke, ci, large] + default: ci + publish: + description: Commit the results to docs/ on main + type: boolean + default: false + +concurrency: + group: benchmark-${{ github.ref }} + cancel-in-progress: true + +jobs: + benchmark: + name: adata-cli vs anndata/scanpy + runs-on: ubuntu-latest + timeout-minutes: 90 + # Job level, not step level: a report must never mark a release red. + continue-on-error: true + permissions: + contents: write + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - uses: astral-sh/setup-uv@v3 + with: + enable-cache: true + + # The checked-out source, not the published wheel. Benchmarking the + # wheel would mean waiting on publish-pypi and then on index + # propagation, for a measurement that would come out the same. + - name: Install adata-cli + run: uv sync --extra dev --frozen + + - name: Show free space + run: df -h / + + - name: Run + run: | + uv run python -m benchmarks.run \ + --tier "${{ inputs.tier || 'ci' }}" \ + --out results.json \ + --work "$RUNNER_TEMP/bench" + + - name: Render + run: | + uv run python -m benchmarks.report results.json --out summary.md + cat summary.md >> "$GITHUB_STEP_SUMMARY" + + - uses: actions/upload-artifact@v4 + if: always() + with: + name: benchmark-${{ github.ref_name }} + path: | + results.json + summary.md + retention-days: 90 + + # docs/ is the durable home: artifacts expire, and one absolute number + # with nothing to compare it against says very little. The page is + # already served by Pages from docs/, so this needs no extra machinery. + - name: Publish to docs + if: github.ref_type == 'tag' || inputs.publish + run: | + git config user.name "github-actions[bot]" + git config user.email "41898282+github-actions[bot]@users.noreply.github.com" + # A tag build is on a detached HEAD; the page belongs on main. + git fetch origin main + git checkout -B main origin/main + uv run python -m benchmarks.report results.json --docs docs --publish + git add docs/BENCHMARKS.md docs/benchmarks + if git diff --cached --quiet; then + echo "nothing to publish"; exit 0 + fi + git commit -m "Benchmark results for ${{ github.ref_name }} [skip ci]" + # Another job may have landed on main in the meantime. + for attempt in 1 2 3; do + if git push origin main; then exit 0; fi + git pull --rebase origin main + done + echo "could not push benchmark results"; exit 1 + + # Best effort, and last: releases here are cut by hand, so this job may + # well run before one exists. The docs page is authoritative either way. + - name: Append to the release notes + if: github.ref_type == 'tag' + continue-on-error: true + env: + GH_TOKEN: ${{ github.token }} + run: | + tag="${{ github.ref_name }}" + gh release view "$tag" >/dev/null 2>&1 || { + echo "no release for $tag yet; skipping"; exit 0 + } + gh release view "$tag" --json body -q .body > body.md + # Idempotent: replace any block this job wrote before. + python - <<'PY' + import pathlib, re + body = pathlib.Path("body.md").read_text() + summary = pathlib.Path("summary.md").read_text() + block = ( + "\n\n" + "
Benchmark vs anndata/scanpy\n\n" + f"{summary}\n
\n\n" + ) + pattern = re.compile( + r".*?", re.S + ) + body = ( + pattern.sub(block, body) + if pattern.search(body) + else body.rstrip() + "\n\n" + block + "\n" + ) + pathlib.Path("body.md").write_text(body) + PY + gh release edit "$tag" --notes-file body.md diff --git a/CHANGELOG.md b/CHANGELOG.md index ac2c44b..9867304 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,9 +15,27 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no of 13 pipeline tasks had to be killed after 98 minutes. The column is now read once per input, and `--merge first` / `--merge only`, which decide on presence alone, read no column values at all. +- **`concat` was quadratic in the number of categories** in an obs column. + Category merging probed a list rather than a dict: 2,096,128 string + comparisons to union 1,024 categories, and around 5e9 for a 100k-category + column. Found by the new guards. +- **`concat` built a Python object per row** for nullable and string obs + columns, then walked the list twice more. Filling a typed buffer by slice + removes three full passes over every such column. ### Added +- **Complexity guards in the test suite** (`tests/test_performance.py`). + Cost regressions now fail at merge time. They count operations rather than + seconds -- h5py and zarr reads, Zarr store traffic, Python allocation and + executed lines -- and assert that successive increments grow no faster than + linearly, so nothing here can fail because a CI runner was busy. See + [docs/TESTING.md](docs/TESTING.md#performance). +- **A comparative benchmark** (`benchmarks/`), run on every tag against + anndata and against scanpy where scanpy has a real equivalent. Reports peak + RSS, wall time and output size; publishes to + [docs/BENCHMARKS.md](docs/BENCHMARKS.md) and the release notes. Report-only + -- it never fails a build. - **`--merge drop` and `--uns-merge drop` are accepted.** `drop` was already the documented default behaviour but was rejected as a value, so a config could not state it explicitly. diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..f87cdbe --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1,10 @@ +"""Comparative benchmark for adata-cli, against anndata and scanpy. + +Report-only, and deliberately separate from `tests/`. The tests measure +operation counts and gate merges; this measures wall time and peak RSS, runs +on tags, and never fails a build -- timing on a shared runner is too noisy to +gate on. + + uv run python -m benchmarks.run --tier smoke + uv run python -m benchmarks.report results.json +""" diff --git a/benchmarks/_measure.py b/benchmarks/_measure.py new file mode 100644 index 0000000..c421e41 --- /dev/null +++ b/benchmarks/_measure.py @@ -0,0 +1,223 @@ +"""Run one command as a child process and report what it cost. + +Peak RSS is the number that matters here. adata-cli exists so that memory is +set by chunk size rather than by input size, and a comparison reporting only +wall time would misrepresent it -- for anything that fits in RAM, loading the +whole thing is usually faster. + +Measuring it correctly needs `os.wait4`, not `resource.getrusage`. +`RUSAGE_CHILDREN` is a running maximum over every child the process has ever +reaped: run a 400 MB case and then a 17 MB one and it still reports 400 MB. +`wait4` returns rusage for one specific child. (Checked on this machine: +wait4 gives 435 MB then 17 MB where RUSAGE_CHILDREN stays at 435 MB.) +""" + +from __future__ import annotations + +import json +import os +import resource +import subprocess +import sys +import tempfile +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Dict, List, Optional, Sequence + +#: Address-space ceiling for every child, in bytes. +#: +#: Cases where the in-memory baseline cannot cope are the whole point of the +#: comparison, but an uncontained OOM on a hosted runner kills the runner +#: agent and the job ends with no report at all. With a ceiling the baseline +#: dies with a MemoryError and a non-zero exit, which is a *result* -- recorded +#: as `oom` at a limit we can state. +#: +#: Linux only -- macOS refuses RLIMIT_AS, so a local run is unbounded and an +#: `oom` row cannot be reproduced there. +DEFAULT_MEMORY_LIMIT = 12 * 1024**3 + +#: `ru_maxrss` is KiB on Linux and bytes on macOS. Nothing in the man pages +#: says so; it is simply what each kernel does. +_MAXRSS_SCALE = 1024 if sys.platform.startswith("linux") else 1 + +#: How often to check whether the child has exited. Fine enough that it adds +#: no meaningful error to a run measured in seconds. +POLL_INTERVAL_S = 0.01 + + +@dataclass +class Measurement: + """What one run of one contender cost.""" + + wall_s: float + maxrss_bytes: int + exit_code: int + status: str # "ok" | "failed" | "oom" | "timeout" + output_bytes: int = 0 + output_files: int = 0 + stderr_tail: str = "" + + def as_dict(self) -> Dict: + return asdict(self) + + +def _tree_size(path: Path) -> tuple: + """Apparent bytes and file count for a file or a directory. + + Zarr stores are directories of many small chunks. Reporting only `du` + would make them look several times larger than the data they hold, since + every chunk rounds up to a filesystem block -- so the file count is + reported alongside, and both go in the table. + """ + if not path.exists(): + return 0, 0 + if path.is_file(): + return path.stat().st_size, 1 + total = 0 + count = 0 + for child in path.rglob("*"): + if child.is_file(): + total += child.stat().st_size + count += 1 + return total, count + + +def drop_page_cache() -> bool: + """Try to drop the page cache, so a read is not served from RAM. + + A fixture written seconds ago is entirely in cache, which systematically + understates how much streaming helps. GitHub-hosted runners have + passwordless sudo, so this usually works there and usually does not + locally; the report says which happened rather than pretending. + """ + try: + subprocess.run( + ["sudo", "-n", "sh", "-c", "sync; echo 3 > /proc/sys/vm/drop_caches"], + check=True, + capture_output=True, + timeout=30, + ) + return True + except Exception: + return False + + +def measure( + command: Sequence[str], + *, + output: Optional[Path] = None, + timeout_s: float = 900.0, + memory_limit: Optional[int] = DEFAULT_MEMORY_LIMIT, + env: Optional[Dict[str, str]] = None, + cwd: Optional[Path] = None, +) -> Measurement: + """Run `command`, returning its wall time, peak RSS and output size.""" + + def limit() -> None: # pragma: no cover - runs in the forked child + if memory_limit is None: + return + try: + resource.setrlimit(resource.RLIMIT_AS, (memory_limit, memory_limit)) + except (ValueError, OSError): + # macOS refuses RLIMIT_AS outright. The ceiling only has to hold + # on the runner, which is Linux; raising here would abort the + # spawn and lose the measurement entirely. + pass + + if output is not None and output.exists(): + raise FileExistsError(f"{output} exists; the harness never overwrites.") + + full_env = {**os.environ, **(env or {})} + + # Output goes to temporary files, not pipes. A pipe would have to be + # drained with communicate(), and communicate() reaps the child -- after + # which wait4 has nothing to report and the per-child rusage is lost. + with tempfile.TemporaryFile("w+") as out_f, tempfile.TemporaryFile("w+") as err_f: + started = time.perf_counter() + process = subprocess.Popen( + list(command), + stdout=out_f, + stderr=err_f, + preexec_fn=limit, + env=full_env, + cwd=str(cwd) if cwd else None, + text=True, + ) + + status = "ok" + deadline = started + timeout_s + while True: + pid, wait_status, usage = os.wait4(process.pid, os.WNOHANG) + if pid != 0: + process.returncode = ( + -os.WTERMSIG(wait_status) + if os.WIFSIGNALED(wait_status) + else os.WEXITSTATUS(wait_status) + ) + maxrss = usage.ru_maxrss + exit_code = process.returncode + break + if time.perf_counter() > deadline: + process.kill() + _, wait_status, usage = os.wait4(process.pid, 0) + process.returncode = -9 + maxrss = usage.ru_maxrss + exit_code = -9 + status = "timeout" + break + time.sleep(POLL_INTERVAL_S) + + err_f.seek(0) + stderr = err_f.read() + + wall_s = time.perf_counter() - started + stderr = stderr or "" + + if status == "ok" and exit_code != 0: + status = "oom" if _looks_like_oom(stderr) else "failed" + + size, files = _tree_size(output) if output is not None else (0, 0) + return Measurement( + wall_s=round(wall_s, 3), + maxrss_bytes=int(maxrss) * _MAXRSS_SCALE, + exit_code=exit_code, + status=status, + output_bytes=size, + output_files=files, + stderr_tail="\n".join(stderr.strip().splitlines()[-6:]), + ) + + +def _looks_like_oom(stderr: str) -> bool: + """Did the child die against the address-space ceiling? + + Under RLIMIT_AS an allocation failure surfaces as MemoryError, or as one + of numpy's or HDF5's own phrasings of the same thing. + """ + markers = ( + "MemoryError", + "Unable to allocate", + "bad_alloc", + "Cannot allocate memory", + "out of memory", + ) + return any(m in stderr for m in markers) + + +def main(argv: List[str]) -> int: # pragma: no cover - CLI entry + """Measure a command given after `--`, printing JSON.""" + if "--" not in argv: + print("usage: python -m benchmarks._measure [--output P] -- CMD...") + return 2 + split = argv.index("--") + head, command = argv[:split], argv[split + 1 :] + output = None + if "--output" in head: + output = Path(head[head.index("--output") + 1]) + print(json.dumps(measure(command, output=output).as_dict(), indent=2)) + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main(sys.argv[1:])) diff --git a/benchmarks/cases.py b/benchmarks/cases.py new file mode 100644 index 0000000..a9ee7a9 --- /dev/null +++ b/benchmarks/cases.py @@ -0,0 +1,342 @@ +"""What gets compared, and the rules that keep the comparison honest. + +The rules, because this is the part that decays fastest +------------------------------------------------------- + +**Use the best idiom the baseline has, not the naive one.** Comparing +`adata export dataframe obs` against a full `ad.read_h5ad()` is a strawman: +anndata reads just `obs` cheaply via `read_elem` on the h5py group. The good +idiom is the primary baseline. A naive full load may appear only as a clearly +labelled second row, where the point is the gap between the two. + +**Pin compression on both sides.** adata-cli forwards the source's settings to +its output; `write_h5ad` defaults to none. Left alone, adata-cli's output +looks smaller for reasons unrelated to the tool. + +**`n/a` is a result.** `concat_on_disk` is CSR/CSC-oriented and its outer-join +support has varied by version; scanpy has no streaming concat at all. Print +what refused and why. An omitted row reads as an oversight; a stated `n/a` +reads as a finding. + +**Startup is a floor, not noise.** The CLI takes 0.3-1 s to import typer, rich, +h5py and zarr, and `import scanpy` takes 3-8 s. On a small tier that is the +entire measurement, so every run includes a `--version` row to read the rest +against. + +**Do not hide the rows where the baseline wins.** `_concat_csr` loops per row +in Python; scipy's C `vstack` will very likely be several times faster in wall +time at many times the memory. That trade *is* the argument for this tool. +The table is "peak RSS against wall time", not a leaderboard. + +**scanpy's filter functions are not our subset.** `sc.pp.filter_cells` +*computes* `n_genes` by scanning X; `adata subset --obs-query` filters a +column that must already exist. Head to head, scanpy looks slow for doing +strictly more work. The like-for-like baseline is a plain boolean mask over a +precomputed column, and scanpy appears in a separately labelled row. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Dict, List, Optional + +#: Written into every generated fixture, and forced on every baseline write. +COMPRESSION = "lzf" + + +@dataclass(frozen=True) +class Contender: + """One way of doing the case: a command template, or a script.""" + + label: str + #: Python source run by the baseline interpreter. `None` for the CLI. + script: Optional[str] = None + #: argv for the CLI, templated with {input0}, {input1}, {output}. + argv: Optional[List[str]] = None + #: Set when the operation has no equivalent, explaining what is missing. + unsupported: Optional[str] = None + #: Which prebuilt environment runs `script`. + env: str = "anndata" + + +@dataclass(frozen=True) +class Case: + """One operation, and every way of performing it.""" + + name: str + question: str + contenders: List[Contender] + #: How many inputs the case needs. + inputs: int = 1 + #: Extension of the output the case writes ("" for read-only cases). + output_suffix: str = ".h5ad" + #: Use the var-heavy shape rather than the tier's main shape. + var_heavy: bool = False + tags: List[str] = field(default_factory=list) + + +_PRELUDE = f""" +import sys, warnings +warnings.filterwarnings("ignore") +import anndata as ad, numpy as np, h5py +IN = sys.argv[1:-1] +OUT = sys.argv[-1] +COMPRESSION = {COMPRESSION!r} +""" + + +def _py(body: str) -> str: + return _PRELUDE + body + + +CASES: List[Case] = [ + # -- startup floor ---------------------------------------------------- + Case( + name="startup", + question="What does each contender cost before it does any work?", + output_suffix="", + contenders=[ + Contender("adata-cli", argv=["adata", "--version"]), + Contender("anndata", script="import anndata"), + Contender("scanpy", script="import scanpy", env="scanpy"), + ], + tags=["floor"], + ), + # -- concat ----------------------------------------------------------- + Case( + name="concat-inner", + question="Concatenate two stores on the obs axis, inner join.", + inputs=2, + contenders=[ + Contender( + "adata-cli", + argv=["adata", "concat", "{input0}", "{input1}", "-o", "{output}"], + ), + Contender( + "anndata (concat_on_disk)", + script=_py( + "from anndata.experimental import concat_on_disk\n" + "concat_on_disk(IN, OUT, join='inner')\n" + ), + ), + Contender( + "anndata (in memory)", + script=_py( + "parts = [ad.read_h5ad(p) for p in IN]\n" + "ad.concat(parts, join='inner')" + ".write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + Contender( + "scanpy", + unsupported="scanpy has no concat of its own; it re-exports anndata's", + ), + ], + ), + Case( + name="concat-merge-same", + question=( + "Concatenate carrying var columns forward -- the REQ-71798 case." + ), + inputs=2, + var_heavy=True, + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "concat", "{input0}", "{input1}", + "-o", "{output}", "--merge", "same", + ], + ), + Contender( + "anndata (concat_on_disk)", + script=_py( + "from anndata.experimental import concat_on_disk\n" + "concat_on_disk(IN, OUT, join='inner', merge='same')\n" + ), + ), + Contender( + "anndata (in memory)", + script=_py( + "parts = [ad.read_h5ad(p) for p in IN]\n" + "ad.concat(parts, join='inner', merge='same')" + ".write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + ], + tags=["regression"], + ), + # -- subset ----------------------------------------------------------- + Case( + name="subset-query", + question="Keep the obs rows matching a predicate on an existing column.", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "subset", "{input0}", "-o", "{output}", + "-q", "quality < 50", + ], + ), + Contender( + "anndata (backed)", + script=_py( + "obj = ad.read_h5ad(IN[0], backed='r')\n" + "keep = obj.obs['quality'] < 50\n" + "obj[keep].to_memory().write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + Contender( + "anndata (in memory)", + script=_py( + "obj = ad.read_h5ad(IN[0])\n" + "obj[obj.obs['quality'] < 50]" + ".write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + Contender( + "scanpy (filter_cells)", + env="scanpy", + script=_py( + "import scanpy as sc\n" + "# NOT like-for-like: filter_cells computes n_genes by\n" + "# scanning X, which the rows above take as given. Kept\n" + "# for scale, labelled so nobody reads it as a race.\n" + "obj = ad.read_h5ad(IN[0])\n" + "sc.pp.filter_cells(obj, min_genes=1)\n" + "obj.write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + ], + ), + # -- conversion ------------------------------------------------------- + Case( + name="h5ad-to-zarr", + question="Convert a store to Zarr.", + output_suffix=".zarr", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "subset", "{input0}", "-o", "{output}", + "-q", "quality >= 0", + ], + ), + Contender( + "anndata (in memory)", + script=_py("ad.read_h5ad(IN[0]).write_zarr(OUT)\n"), + ), + Contender( + "scanpy", + unsupported="no streaming converter; scanpy defers to anndata", + ), + ], + ), + # -- metadata --------------------------------------------------------- + Case( + name="inspect", + question="Report what is in the store, without reading the matrix.", + output_suffix="", + contenders=[ + Contender("adata-cli", argv=["adata", "view", "{input0}"]), + Contender( + "anndata (read_elem on obs)", + script=_py( + "from anndata.io import read_elem\n" + "with h5py.File(IN[0], 'r') as f:\n" + " obs = read_elem(f['obs']); var = read_elem(f['var'])\n" + "print(obs.shape, var.shape)\n" + ), + ), + Contender( + "anndata (full load)", + script=_py("print(ad.read_h5ad(IN[0]))\n"), + ), + ], + tags=["headline"], + ), + Case( + name="export-obs", + question="Write the obs table out as CSV.", + output_suffix=".csv", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "export", "dataframe", "{input0}", "obs", + "-o", "{output}", + ], + ), + Contender( + "anndata (read_elem on obs)", + script=_py( + "from anndata.io import read_elem\n" + "with h5py.File(IN[0], 'r') as f:\n" + " read_elem(f['obs']).to_csv(OUT)\n" + ), + ), + Contender( + "anndata (full load)", + script=_py("ad.read_h5ad(IN[0]).obs.to_csv(OUT)\n"), + ), + ], + ), + # -- split ------------------------------------------------------------ + Case( + name="split-by-sample", + question="Write one store per distinct value of an obs column.", + output_suffix="", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "split", "{input0}", "--by", "sample", + "-o", "{outdir}", + ], + ), + Contender( + "anndata (hand-written loop)", + script=_py( + "import os\n" + "obj = ad.read_h5ad(IN[0])\n" + "os.makedirs(OUT, exist_ok=True)\n" + "for key, idx in obj.obs.groupby('sample', observed=True)" + ".groups.items():\n" + " obj[idx].write_h5ad(\n" + " os.path.join(OUT, f'{key}.h5ad'), compression=COMPRESSION\n" + " )\n" + ), + ), + ], + ), + # -- the claim itself ------------------------------------------------- + Case( + name="rss-vs-size", + question=( + "Does peak memory track input size, at a fixed --chunk? " + "This is the README's claim, stated as a measurement." + ), + inputs=2, + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "concat", "{input0}", "{input1}", + "-o", "{output}", "--chunk", "1024", + ], + ), + Contender( + "anndata (in memory)", + script=_py( + "parts = [ad.read_h5ad(p) for p in IN]\n" + "ad.concat(parts).write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + ], + tags=["headline", "sweep"], + ), +] + + +def by_name() -> Dict[str, Case]: + return {case.name: case for case in CASES} diff --git a/benchmarks/datasets.py b/benchmarks/datasets.py new file mode 100644 index 0000000..68c64e7 --- /dev/null +++ b/benchmarks/datasets.py @@ -0,0 +1,132 @@ +"""Synthetic inputs for the benchmark, generated deterministically. + +Two choices here are not arbitrary. + +**Poisson counts, not uniform floats.** Random float32 is incompressible, so +a uniform-random fixture is three times the size on disk for the same shape +and the `ci` tier stops fitting in a runner's 14 GB. Counts are also what the +tool actually sees. + +**lzf, not gzip.** gzip on 10^8 values is single-threaded and costs minutes +per fixture on four vCPUs. The compression setting is pinned on *both* sides +of every comparison -- see `cases.py` -- because adata-cli forwards the +source's compression to its output while `write_h5ad` defaults to none, which +would otherwise make adata-cli's output look smaller for reasons that have +nothing to do with the tool. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Dict, List + +import numpy as np + +#: Nonzeros per row block while generating, to keep the generator itself +#: from needing the memory the benchmark is about to measure. +_GEN_ROWS = 4096 + + +@dataclass(frozen=True) +class Tier: + """One size of input, and what it is for.""" + + name: str + n_obs: int + n_var: int + density: float + note: str + + +TIERS: Dict[str, Tier] = { + "smoke": Tier("smoke", 1_000, 2_000, 0.05, "fast enough to debug the harness"), + "ci": Tier("ci", 50_000, 20_000, 0.05, "fits a 16 GB runner with the baseline"), + "large": Tier( + "large", + 500_000, + 20_000, + 0.05, + "in-memory baseline is expected to hit the address-space ceiling", + ), +} + +#: The REQ-71798 shape: few rows, many variables, var columns worth merging. +#: Small enough to run in every tier, and the case the hang was reported on. +VAR_HEAVY = Tier("var-heavy", 2_000, 36_601, 0.05, "the shape that hung in 0.5.1") + + +def _write(path: Path, tier: Tier, seed: int, *, n_var_columns: int = 2) -> Path: + """Write one CSR h5ad of `tier`'s shape, built a row block at a time.""" + import anndata as ad + import h5py + import pandas as pd + import scipy.sparse as sp + + rng = np.random.default_rng(seed) + blocks = [] + for start in range(0, tier.n_obs, _GEN_ROWS): + rows = min(_GEN_ROWS, tier.n_obs - start) + block = sp.random( + rows, + tier.n_var, + density=tier.density, + format="csr", + dtype="float32", + random_state=rng, + data_rvs=lambda k: rng.poisson(3.0, k).astype("float32") + 1, + ) + blocks.append(block) + matrix = sp.vstack(blocks, format="csr") + + obs = pd.DataFrame( + { + "sample": pd.Categorical( + [f"s{i % 8}" for i in range(tier.n_obs)] + ), + "n_counts": np.asarray(matrix.sum(axis=1)).ravel(), + "n_genes": matrix.getnnz(axis=1).astype("int32"), + # A predicate on n_genes would select a different fraction at + # every tier, since genes-per-cell tracks n_var. This one keeps + # exactly half the rows whatever the shape, so the subset case + # measures the same work across tiers. + "quality": (np.arange(tier.n_obs) % 100).astype("int32"), + "barcode": [f"bc{seed}-{i}" for i in range(tier.n_obs)], + }, + index=[f"c{seed}-{i}" for i in range(tier.n_obs)], + ) + var = pd.DataFrame( + { + "gene_symbol": [f"SYM{i}" for i in range(tier.n_var)], + "feature_type": ["Gene Expression"] * tier.n_var, + **{ + f"extra{c}": [f"e{c}-{i}" for i in range(tier.n_var)] + for c in range(max(0, n_var_columns - 2)) + }, + }, + index=[f"ENSG{i:011d}" for i in range(tier.n_var)], + ) + + obj = ad.AnnData(X=matrix, obs=obs, var=var) + obj.write_h5ad(path, compression="lzf") + del obj, matrix, blocks + + with h5py.File(path, "r") as handle: # cheap sanity check + assert handle["X"].attrs["encoding-type"] == "csr_matrix" + return path + + +def build(directory: Path, tier: Tier, *, count: int = 2) -> List[Path]: + """Build (or reuse) `count` inputs of this tier's shape. + + Reused if already present: generation is the slowest part of the job and + the content is fully determined by (tier, seed). + """ + directory.mkdir(parents=True, exist_ok=True) + paths = [] + for seed in range(count): + path = directory / f"{tier.name}-{seed}.h5ad" + if not path.exists(): + _write(path, tier, seed) + paths.append(path) + return paths diff --git a/benchmarks/report.py b/benchmarks/report.py new file mode 100644 index 0000000..254af98 --- /dev/null +++ b/benchmarks/report.py @@ -0,0 +1,257 @@ +"""Turn results.json into the markdown that gets published. + +Two sinks, and the docs page is the authoritative one: artifacts expire, and +a single absolute number with nothing to compare it to is close to +uninterpretable. `docs/benchmarks/.json` keeps the series, and +`docs/BENCHMARKS.md` renders the latest run plus a history table. Both are in +git and both are already served by GitHub Pages, so no extra machinery. + + uv run python -m benchmarks.report results.json + uv run python -m benchmarks.report results.json --docs docs --publish +""" + +from __future__ import annotations + +import argparse +import json +import shutil +from pathlib import Path +from typing import Dict, List, Optional + +from benchmarks.cases import by_name + + +def _bytes(n: int) -> str: + if not n: + return "-" + for unit, scale in (("GB", 1e9), ("MB", 1e6), ("kB", 1e3)): + if n >= scale: + return f"{n / scale:.1f} {unit}" + return f"{n} B" + + +def _cell(record: Dict) -> tuple: + """(wall, rss, output) as they should read in the table.""" + status = record.get("status") + if status == "n/a": + return "n/a", "n/a", "n/a" + if status == "timeout": + # This is how the 0.5.1 hang presented: still running, output stuck at + # the size it reached in the first second. Say so, rather than + # reporting a large time as though the run had completed. + return ( + f"**timeout** (>{record['wall_s']:.0f} s)", + _bytes(record["maxrss_bytes"]), + f"{_bytes(record['output_bytes'])} and not growing", + ) + if status == "oom": + return "**out of memory**", "hit the ceiling", "-" + if status == "failed": + return "**failed**", _bytes(record["maxrss_bytes"]), "-" + + output = _bytes(record["output_bytes"]) + if record.get("output_files", 0) > 1: + # Zarr: a directory of small chunks looks far larger under `du` than + # the bytes it holds, so the file count goes alongside. + output += f" in {record['output_files']} files" + return f"{record['wall_s']:.2f} s", _bytes(record["maxrss_bytes"]), output + + +def render(payload: Dict, previous: Optional[Dict] = None) -> str: + cases = by_name() + before = ( + {(r["case"], r["contender"]): r for r in previous["results"]} + if previous + else {} + ) + + lines: List[str] = [] + lines.append(f"# Benchmark: `{payload['ref']}`") + lines.append("") + lines.append( + "adata-cli against anndata, and scanpy where it has a real " + "equivalent. **Peak RSS is the headline, not wall time.** This tool " + "exists so that memory is set by `--chunk` rather than by input size; " + "for anything that fits in RAM, loading the whole thing is often " + "faster, and the rows below say so where it is true." + ) + lines.append("") + + n_obs, n_var, density = payload["tier_shape"] + vh_obs, vh_var = payload["var_heavy_shape"] + cache = ( + "dropped between runs" + if payload.get("page_cache_dropped") + else "**warm** (could not be dropped; reads are served from RAM, " + "which understates the streaming advantage)" + ) + lines.append("| | |") + lines.append("|---|---|") + lines.append(f"| Tier | `{payload['tier']}`, {n_obs:,} obs x {n_var:,} var, " + f"{density:.0%} dense CSR |") + lines.append(f"| Var-heavy shape | {vh_obs:,} obs x {vh_var:,} var |") + lines.append(f"| Page cache | {cache} |") + lines.append(f"| Address-space ceiling | " + f"{_bytes(payload['memory_limit_bytes'])} per process |") + lines.append(f"| Repeats | {payload['repeats']} (fastest shown) |") + lines.append(f"| Platform | {payload['platform']}, Python " + f"{payload['python']} |") + versions = ", ".join(f"{k} {v}" for k, v in payload["versions"].items()) + lines.append(f"| Versions | {versions} |") + lines.append(f"| Generated | {payload['generated']} |") + lines.append("") + + grouped: Dict[str, List[Dict]] = {} + for record in payload["results"]: + grouped.setdefault(record["case"], []).append(record) + + for name, records in grouped.items(): + case = cases.get(name) + lines.append(f"## `{name}`") + lines.append("") + if case: + lines.append(f"{case.question}") + lines.append("") + lines.append("| Contender | Wall time | Peak RSS | Output | vs previous |") + lines.append("|---|---|---|---|---|") + for record in records: + wall, rss, output = _cell(record) + lines.append( + f"| {record['contender']} | {wall} | {rss} | {output} | " + f"{_delta(record, before.get((name, record['contender'])))} |" + ) + if record.get("note"): + lines.append(f"| | *{record['note']}* | | | |") + lines.append("") + + lines.append("---") + lines.append("") + lines.append( + "Rows marked `n/a` are operations the baseline does not offer; that " + "is a result, not a gap in the measurement. Where a baseline is " + "faster, the row stands as it is -- the trade this tool makes is " + "memory for time, and hiding the cost would make the table useless." + ) + lines.append("") + lines.append( + "**Output size** is the file as the filesystem reports it, which for " + "HDF5 includes allocation slack. On a small tier that slack can be " + "several times the stored bytes and the column says more about the " + "writer's allocation strategy than about the data; at `ci` and above " + "it is noise. Zarr stores report a file count alongside, because a " + "directory of small chunks measures much larger than it holds." + ) + return "\n".join(lines) + "\n" + + +def _delta(record: Dict, previous: Optional[Dict]) -> str: + """Wall time and peak RSS against the last published run.""" + if not previous or record.get("status") != "ok" or previous.get("status") != "ok": + return "-" + parts = [] + for key, label in (("wall_s", "time"), ("maxrss_bytes", "RSS")): + old, new = previous.get(key), record.get(key) + if not old: + continue + change = (new - old) / old + if abs(change) < 0.10: # within run-to-run noise on a shared runner + continue + parts.append(f"{label} {change:+.0%}") + return ", ".join(parts) if parts else "no change" + + +def history_table(docs: Path) -> str: + """One row per published run, newest first.""" + runs = [] + for path in sorted((docs / "benchmarks").glob("*.json")): + try: + payload = json.loads(path.read_text()) + except Exception: # pragma: no cover + continue + rows = {(r["case"], r["contender"]): r for r in payload["results"]} + key = ("concat-inner", "adata-cli") + headline = rows.get(key, {}) + runs.append( + ( + payload.get("generated", ""), + payload.get("ref", path.stem), + payload.get("tier", "?"), + headline.get("wall_s"), + headline.get("maxrss_bytes"), + path.name, + ) + ) + if not runs: + return "" + + lines = [ + "## History", + "", + "`concat-inner` on adata-cli, run by run. Full results for each are " + "in [`docs/benchmarks/`](benchmarks/).", + "", + "| Run | Ref | Tier | Wall time | Peak RSS | Raw |", + "|---|---|---|---|---|---|", + ] + for generated, ref, tier, wall, rss, filename in sorted(runs, reverse=True): + lines.append( + f"| {generated} | `{ref}` | {tier} | " + f"{f'{wall:.2f} s' if wall else '-'} | " + f"{_bytes(rss) if rss else '-'} | [json](benchmarks/{filename}) |" + ) + return "\n".join(lines) + "\n" + + +def publish(payload: Dict, results: Path, docs: Path) -> Path: + """Copy the raw results in and rewrite the docs page.""" + store = docs / "benchmarks" + store.mkdir(parents=True, exist_ok=True) + ref = payload["ref"].replace("/", "-") + shutil.copyfile(results, store / f"{ref}.json") + + previous = _previous(store, skip=f"{ref}.json") + page = docs / "BENCHMARKS.md" + page.write_text(render(payload, previous) + "\n" + history_table(docs)) + return page + + +def _previous(store: Path, *, skip: str) -> Optional[Dict]: + candidates = sorted( + (p for p in store.glob("*.json") if p.name != skip), + key=lambda p: p.stat().st_mtime, + ) + if not candidates: + return None + try: + return json.loads(candidates[-1].read_text()) + except Exception: # pragma: no cover + return None + + +def main() -> int: # pragma: no cover - CLI entry + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("results", type=Path) + parser.add_argument("--docs", type=Path, default=Path("docs")) + parser.add_argument( + "--publish", + action="store_true", + help="write docs/benchmarks/.json and rewrite docs/BENCHMARKS.md", + ) + parser.add_argument("--out", type=Path, help="also write the markdown here") + args = parser.parse_args() + + payload = json.loads(args.results.read_text()) + if args.publish: + page = publish(payload, args.results, args.docs) + print(f"wrote {page}") + text = page.read_text() + else: + text = render(payload, _previous(args.docs / "benchmarks", skip="")) + print(text) + if args.out: + args.out.write_text(text) + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/benchmarks/run.py b/benchmarks/run.py new file mode 100644 index 0000000..6d7f501 --- /dev/null +++ b/benchmarks/run.py @@ -0,0 +1,274 @@ +"""Run the benchmark and write results.json. + +Not a pytest module. Each contender has to run in its own process for peak +RSS to mean anything, and the whole thing takes tens of minutes at the `ci` +tier -- neither of which belongs in a suite that gates merges. The guards +that *do* gate merges are in `tests/test_performance.py`, and they measure +operation counts rather than time. + + uv run python -m benchmarks.run --tier smoke + uv run python -m benchmarks.run --tier ci --out results.json + +Baselines run from environments built once, up front, rather than through +`uv run --with`. `reference_stores.py` uses the latter to build fixtures, +where cost is irrelevant; here the first invocation would resolve and download +several hundred megabytes of wheels straight into the measured wall time, and +uv's own memory into the measured peak. +""" + +from __future__ import annotations + +import argparse +import json +import platform +import shutil +import subprocess +import sys +import time +from pathlib import Path +from typing import Dict, List, Optional + +from benchmarks import datasets +from benchmarks._measure import ( + DEFAULT_MEMORY_LIMIT, + Measurement, + drop_page_cache, + measure, +) +from benchmarks.cases import CASES, Case, Contender + +#: Packages each baseline environment needs. +ENVIRONMENTS: Dict[str, List[str]] = { + "anndata": ["anndata", "scipy", "pandas", "h5py", "zarr"], + "scanpy": ["scanpy", "anndata", "scipy", "pandas", "h5py", "zarr"], +} + + +def build_environments(root: Path, wanted: List[str]) -> Dict[str, Path]: + """Create one venv per baseline and return its interpreter.""" + interpreters: Dict[str, Path] = {} + for name in wanted: + venv = root / name + python = venv / "bin" / "python" + if not python.exists(): + print(f"[env] building {name}", flush=True) + subprocess.run(["uv", "venv", str(venv)], check=True, capture_output=True) + subprocess.run( + ["uv", "pip", "install", "--python", str(python), *ENVIRONMENTS[name]], + check=True, + ) + interpreters[name] = python + return interpreters + + +def _versions(interpreters: Dict[str, Path]) -> Dict[str, str]: + """Record exactly what was compared, so a table is interpretable later.""" + found: Dict[str, str] = {} + try: + found["adata-cli"] = subprocess.run( + ["adata", "--version"], capture_output=True, text=True, timeout=120 + ).stdout.strip() + except Exception: # pragma: no cover + found["adata-cli"] = "unknown" + for name, python in interpreters.items(): + code = ( + "import importlib.metadata as m;" + f"print(m.version({name!r}))" + ) + try: + found[name] = subprocess.run( + [str(python), "-c", code], capture_output=True, text=True, timeout=120 + ).stdout.strip() + except Exception: # pragma: no cover + found[name] = "unknown" + return found + + +def _run_contender( + case: Case, + contender: Contender, + inputs: List[Path], + workdir: Path, + interpreters: Dict[str, Path], + *, + timeout_s: float, + scripts_dir: Path, +) -> Measurement: + workdir.mkdir(parents=True, exist_ok=True) + output = workdir / f"{case.name}{case.output_suffix}" + outdir = workdir / f"{case.name}-out" + + # The harness never overwrites, so clear any previous run first. + for stale in (output, outdir): + if stale.is_dir(): + shutil.rmtree(stale) + elif stale.exists(): + stale.unlink() + + substitutions = { + "output": str(output), + "outdir": str(outdir), + **{f"input{i}": str(p) for i, p in enumerate(inputs)}, + } + + if contender.argv is not None: + command = [part.format(**substitutions) for part in contender.argv] + target = outdir if "{outdir}" in " ".join(contender.argv) else output + else: + safe = f"{case.name}-{contender.label}".replace(" ", "_").replace("(", "").replace(")", "") + script = scripts_dir / f"{safe}.py" + script.write_text(contender.script or "") + # A script that makes a directory is handed the directory; everything + # else is handed the single output path. + target = outdir if "makedirs" in (contender.script or "") else output + command = [ + str(interpreters[contender.env]), str(script), *map(str, inputs), str(target) + ] + + # Read-only cases produce nothing to size. + watched = None if (case.output_suffix == "" and target == output) else target + return measure(command, output=watched, timeout_s=timeout_s) + + +def run( + tier_name: str, + *, + out: Path, + work: Path, + only: Optional[List[str]] = None, + repeats: int = 3, + timeout_s: float = 900.0, + drop_cache: bool = True, +) -> Dict: + tier = datasets.TIERS[tier_name] + work.mkdir(parents=True, exist_ok=True) + fixtures = work / "fixtures" + scripts = work / "scripts" + scripts.mkdir(parents=True, exist_ok=True) + + selected = [c for c in CASES if only is None or c.name in only] + needed = sorted( + {c.env for case in selected for c in case.contenders if c.script} + ) + interpreters = build_environments(work / "envs", needed) + + print(f"[data] building {tier.name} fixtures ({tier.n_obs} x {tier.n_var})", + flush=True) + main_inputs = datasets.build(fixtures, tier, count=2) + var_heavy_inputs = datasets.build(fixtures, datasets.VAR_HEAVY, count=2) + + cache_dropped = drop_page_cache() if drop_cache else False + + results: List[Dict] = [] + for case in selected: + pool = var_heavy_inputs if case.var_heavy else main_inputs + inputs = pool[: case.inputs] + print(f"[case] {case.name}", flush=True) + for contender in case.contenders: + if contender.unsupported: + results.append( + { + "case": case.name, + "contender": contender.label, + "status": "n/a", + "note": contender.unsupported, + } + ) + print(f" {contender.label}: n/a ({contender.unsupported})") + continue + + runs: List[Measurement] = [] + for _ in range(repeats): + if drop_cache: + drop_page_cache() + runs.append( + _run_contender( + case, + contender, + inputs, + work / "out", + interpreters, + timeout_s=timeout_s, + scripts_dir=scripts, + ) + ) + if runs[-1].status != "ok": + break # a failure repeats identically; do not pay for it twice + + best = min(runs, key=lambda m: m.wall_s) + record = { + "case": case.name, + "contender": contender.label, + # The fastest run, because a slower one differs only by what + # else the machine was doing. + **best.as_dict(), + "runs": len(runs), + "wall_s_all": [m.wall_s for m in runs], + } + results.append(record) + print( + f" {contender.label}: {best.status} " + f"{best.wall_s:.2f}s {best.maxrss_bytes / 1e6:.0f} MB", + flush=True, + ) + + payload = { + "tier": tier.name, + "tier_shape": [tier.n_obs, tier.n_var, tier.density], + "var_heavy_shape": [datasets.VAR_HEAVY.n_obs, datasets.VAR_HEAVY.n_var], + "generated": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "ref": _git_ref(), + "platform": f"{platform.system()} {platform.machine()}", + "python": sys.version.split()[0], + "memory_limit_bytes": DEFAULT_MEMORY_LIMIT, + "page_cache_dropped": cache_dropped, + "repeats": repeats, + "versions": _versions(interpreters), + "results": results, + } + out.parent.mkdir(parents=True, exist_ok=True) + out.write_text(json.dumps(payload, indent=2) + "\n") + print(f"\nwrote {out}") + return payload + + +def _git_ref() -> str: + for args in (["git", "describe", "--tags", "--exact-match"], + ["git", "rev-parse", "--short", "HEAD"]): + try: + done = subprocess.run(args, capture_output=True, text=True, timeout=30) + if done.returncode == 0 and done.stdout.strip(): + return done.stdout.strip() + except Exception: # pragma: no cover + pass + return "unknown" + + +def main() -> int: # pragma: no cover - CLI entry + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tier", default="smoke", choices=sorted(datasets.TIERS)) + parser.add_argument("--out", type=Path, default=Path("results.json")) + parser.add_argument("--work", type=Path, default=Path(".bench")) + parser.add_argument("--case", action="append", dest="only") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--timeout", type=float, default=900.0) + parser.add_argument( + "--keep-cache", + action="store_true", + help="do not try to drop the page cache between runs", + ) + args = parser.parse_args() + run( + args.tier, + out=args.out, + work=args.work, + only=args.only, + repeats=args.repeats, + timeout_s=args.timeout, + drop_cache=not args.keep_cache, + ) + return 0 + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md new file mode 100644 index 0000000..34dd33b --- /dev/null +++ b/docs/BENCHMARKS.md @@ -0,0 +1,25 @@ +# Benchmarks + +This page is written by the [`Benchmark`](../.github/workflows/benchmark.yml) +workflow on every tag. Until the next release it stands empty. + +What will appear here: adata-cli against anndata, and against scanpy wherever +scanpy has a real equivalent, measured on a GitHub-hosted runner at 50,000 +obs x 20,000 var plus the 2,000 x 36,601 shape that hung in 0.5.1. + +**Peak RSS is the headline, not wall time.** This tool exists so that memory +is set by `--chunk` rather than by input size. For data that fits in RAM, +loading the whole thing is often faster, and the tables say so where it is +true -- the trade is memory for time, and a table that hid the cost would be +worth nothing. + +To produce one locally: + +```bash +uv run python -m benchmarks.run --tier smoke --out results.json +uv run python -m benchmarks.report results.json +``` + +See [TESTING.md](TESTING.md#performance) for what is measured, how the +baselines are kept honest, and how this differs from the complexity guards +that run on every pull request. diff --git a/docs/TESTING.md b/docs/TESTING.md index 7159a14..5bcf1dd 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -4,6 +4,7 @@ uv sync --extra dev uv run pytest # everything uv run pytest -m "not integration" # fast: no environment building, ~40s +uv run pytest -m perf # just the tracing complexity guards uv run pytest -m integration # compatibility across anndata releases ``` @@ -22,6 +23,8 @@ and the compatibility suite as a separate job. | `test_anndata_roundtrip.py` | anndata writes the fixtures, reads back our output | | `test_anndata_versions.py` | Compatibility with six real anndata releases (see below) | | `test_commands_phase2.py`, `test_commands_coverage.py`, `test_cli.py` | Command surfaces and error paths | +| `test_performance.py` | Complexity guards -- what an operation costs, not how long it takes (see below) | +| `test_benchmark_harness.py` | That the benchmark's measurement is trustworthy | | `test_subset.py`, `test_export.py`, `test_import.py`, `test_info_read.py`, `test_zarr.py`, `test_query.py` | Per-feature unit tests | ## Writing a test @@ -44,6 +47,177 @@ For CLI surfaces, use the module-level `CliRunner` and assert on ANSI before matching message text — Rich also wraps long lines, so collapse whitespace. +## Performance + +Two mechanisms, and confusing them is the main way this gets misused. + +| | `tests/test_performance.py` | `benchmarks/` | +|---|---|---| +| Measures | operation **counts** | wall time, peak RSS, output size | +| Runs | every CI job, both interpreters | tags, and `workflow_dispatch` | +| Gates | **yes** -- it fails the build | never | +| Data | 64-4096 elements | 50,000 x 20,000 | + +### Why the guards do not use a clock + +A test that can fail because a runner was busy does not belong in a merge +gate. Every number in `test_performance.py` is deterministic: the same input +gives the same count on every machine. That is what lets it block a merge. + +This is not new -- `test_commands_phase2.py` already counts `read_str_all` +calls for exactly this reason. The performance file generalises the idea; it +does not replace those tests, and where an exact count is derivable an exact +count is still better than a ratio, because it says what the number *should* +be rather than only that it did not grow. + +### What is counted + +`tests/perf_counters.py` patches four seams, all verified against the pinned +h5py 3.15.1 and zarr 3.1.5: + +| Hook | Sees | +|---|---| +| `h5py.Dataset.__getitem__` | HDF5 reads | +| `zarr.Array.__getitem__` | Zarr reads, v2 and v3 alike | +| `zarr.storage.LocalStore.get` | chunk and metadata fetches | +| `LocalStore.set` / `.delete` | write storms | + +The libraries are patched rather than anything in `src/adata/`. An in-repo +seam would only see the call sites that remembered to use it, which is the +wrong property for a guard meant to catch the read nobody thought of -- and +the matrix paths slice the backend objects directly anyway. + +Elements are counted, not just calls: the `--merge` bug was n calls each +reading n elements, and counting calls alone would miss a vectorised variant +that reads n elements n times in one call. + +**Known bypasses.** `Dataset.read_direct`, `np.asarray(dataset)` (it goes via +`__array__`), `dataset.asstr()[...]` and `.fields()` reach the file without +passing any hook. None are used today. If you add one, the counters will +quietly report less work than happened -- which is why every guard also +asserts a lower bound, and why there are two canary tests with known absolute +counts. If a canary fails, fix the hooks before trusting anything else here. + +### The invariant + +Three sizes at 4x spacing, comparing successive **increments**: + +```python +d1 = c(4n) - c(n) +d2 = c(16n) - c(4n) +assert d2 <= 6 * d1 # grows no faster than linearly +assert d1 >= 4n # and the counters actually saw something +``` + +The increment form is the point. Comparing raw counts needs an additive +constant to absorb fixed setup cost, and there is no principled value for +one: too small and it is flaky, too large and a quadratic with a small +coefficient hides underneath at n=64. Any constant appears in both +differences and cancels exactly. + +Why 6: at 4x spacing the increment ratio is 4.0 for linear work, about 4.4 +for n log n, 8 for n^1.5 and 16 for quadratic. 6 sits in the gap with room on +both sides. `test_growth_limit_separates_linear_from_super_linear` asserts +that calibration rather than leaving it as a comment, and it fails if anyone +changes `SIZES` without recomputing the limit. + +**Scale one axis per test, and name it in the test id.** Scaling two at once +makes legitimate work look quadratic -- an outer-join concat really does +produce n_obs x n_var_union cells. The axes worth separate coverage are +`n_var`, `n_obs`, `n_inputs`, `n_obs_columns`, `n_var_columns`, +`n_categories` and `n_groups`. + +### Where counting is not enough + +Two real defects in `concat.py` were invisible to all of the above, and both +needed their own instrument: + +- **Category merging** used `if category not in categories` on a list. The + category lists are read once either way, so no read counter sees it, and + `x not in lst` is a single bytecode, so the line tracer does not either -- + the quadratic is inside C-level list membership. Counting string + comparisons through a `str` subclass is what made it visible: 2,096,128 + comparisons at k=1024. +- **A per-element loop over obs rows.** Linear, just with a fat constant, so + no ratio catches it. The guard compares executed Python lines per row + against the numeric column path measured in the same run -- self-calibrating, + so it needs no hand-tuned budget and holds across interpreters. + +`count_allocations` (tracemalloc) and `count_lines` (`sys.settrace`) are the +tools for these. The line tracer costs a 10-50x slowdown, so its tests carry +the `perf` marker and stay small. + +### What the guards deliberately do not claim + +Obs columns are read whole, so peak allocation for a concat is O(n_obs), not +O(`--chunk`). The streaming guarantee holds for X, not for obs annotation. +The guards only stop that getting worse than linear; `benchmarks/` reports +the actual curve. + +### Reading a failure + +It says cost grew super-linearly on the named axis. The assertion message +carries the whole measured series and the computed ratio. Usually the code is +wrong. Occasionally the expectation is -- an operation legitimately gained +work -- and then the new number needs a comment saying why, in the same style +as the exact counts in `test_commands_phase2.py`. + +## The comparative benchmark + +```bash +uv run python -m benchmarks.run --tier smoke --out results.json +uv run python -m benchmarks.report results.json +``` + +Tiers are `smoke` (1,000 x 2,000, seconds), `ci` (50,000 x 20,000, 35-50 +minutes for the whole suite) and `large` (500,000 x 20,000, dispatch only, +where the in-memory baseline is expected to hit the ceiling). Every tier also +builds the 2,000 x 36,601 var-heavy shape, which is what hung in 0.5.1. + +Results are published to `docs/BENCHMARKS.md` and `docs/benchmarks/.json` +on every tag, and appended to the GitHub release notes if a release exists. +Artifacts expire; the docs page is the durable series. + +### How it measures + +- **`os.wait4`, not `resource.getrusage`.** `RUSAGE_CHILDREN` is a running + maximum over every child ever reaped, so one large case would poison every + later row. Output goes to temporary files rather than pipes, because + `communicate()` reaps the child and loses its rusage. +- **`RLIMIT_AS` at 12 GiB on every child.** Cases where the in-memory + baseline cannot cope are the whole point, but an uncontained OOM kills the + runner agent and the job ends with no report. With a ceiling it is a row + that says `out of memory` at a limit we can state. macOS refuses + `RLIMIT_AS`, so a local run is unbounded. +- **Baselines run from venvs built up front**, not `uv run --with`. The + latter is right for building fixtures, as `reference_stores.py` does, and + wrong here: the first invocation would put several hundred megabytes of + wheel downloads into the measured wall time and uv's own memory into the + measured peak. scanpy therefore never enters `uv.lock` or the image. + +### Rules that keep the comparison honest + +These live in `benchmarks/cases.py` as well, because this is the part that +decays fastest. + +- **Use the best idiom the baseline has.** Comparing `export dataframe` + against a full `read_h5ad()` is a strawman -- anndata reads just `obs` via + `read_elem`. The good idiom is the primary row; the naive load may appear + only as a clearly labelled second row. +- **Pin compression on both sides.** adata-cli forwards the source's + settings; `write_h5ad` defaults to none. +- **`n/a` is a result.** Where scanpy has no equivalent, or + `concat_on_disk` refuses, say which and why. An omitted row reads as an + oversight. +- **Include the startup floor.** The CLI costs 0.3-1 s to import, and + `import scanpy` 3-8 s; on a small tier that is the entire measurement. +- **Do not hide the rows where the baseline wins.** `_concat_csr` loops per + row in Python and scipy's C `vstack` will often beat it on time at many + times the memory. That trade is the argument for the tool. +- **Say whether the page cache was dropped.** A fixture written seconds ago + is entirely in RAM, which understates streaming. The runner can drop it; a + laptop usually cannot, and the report states which happened. + ## Compatibility testing against real anndata releases `test_anndata_versions.py` does not trust this repo's idea of the format. For diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py new file mode 100644 index 0000000..092c3b2 --- /dev/null +++ b/tests/test_benchmark_harness.py @@ -0,0 +1,104 @@ +"""The benchmark harness has to be trustworthy before its numbers are. + +Only the measurement is tested here, not the cases. If peak RSS were +attributed to the wrong process every table the project publishes would be +wrong, and the failure would be invisible -- the numbers would still look +plausible. Everything else in `benchmarks/` is reporting, and a mistake there +shows up the moment anyone reads the page. +""" + +from __future__ import annotations + +import sys + +import pytest + +from benchmarks._measure import Measurement, _looks_like_oom, measure + + +def test_peak_rss_is_attributed_to_the_right_child(): + """The reason the harness uses `os.wait4` rather than `getrusage`. + + `RUSAGE_CHILDREN` is a running maximum over every child a process has + ever reaped, so a 400 MB case followed by a 1 kB one would report 400 MB + twice and every later row would inherit the largest earlier peak. + """ + big = measure([sys.executable, "-c", "x = bytearray(400 * 1024 * 1024)"]) + small = measure([sys.executable, "-c", "x = bytearray(1024)"]) + + assert big.status == "ok" and small.status == "ok" + assert big.maxrss_bytes > 300 * 1024**2, big + assert small.maxrss_bytes < big.maxrss_bytes / 4, ( + f"peak RSS leaked between children: {small.maxrss_bytes} after " + f"{big.maxrss_bytes}" + ) + + +def test_a_hang_is_reported_as_a_timeout_with_its_output_size(tmp_path): + """How the 0.5.1 hang presented: still running, output never growing. + + A harness that recorded this as a slow success would have made the + reported bug look like ordinary slowness. + """ + output = tmp_path / "partial.bin" + command = [ + sys.executable, + "-c", + f"open({str(output)!r}, 'wb').write(b'x' * 1024); import time; time.sleep(60)", + ] + result = measure(command, output=output, timeout_s=2.0) + + assert result.status == "timeout" + assert result.wall_s >= 2.0 + assert result.output_bytes == 1024, "output size at the kill must be recorded" + + +def test_a_failure_against_the_memory_ceiling_is_recorded_as_oom(): + """A baseline that cannot cope is a result, not a crashed job. + + Skipped on macOS, which refuses RLIMIT_AS; the ceiling only has to hold + on the Linux runner where the benchmark actually runs. + """ + if not sys.platform.startswith("linux"): + pytest.skip("RLIMIT_AS is not enforceable on this platform") + + result = measure( + [sys.executable, "-c", "x = bytearray(8 * 1024**3)"], + memory_limit=512 * 1024**2, + ) + assert result.status == "oom", result + + +@pytest.mark.parametrize( + "text", + ["MemoryError", "numpy: Unable to allocate 4.0 GiB", "std::bad_alloc"], +) +def test_out_of_memory_is_recognised_however_it_is_phrased(text): + assert _looks_like_oom(text) + + +def test_an_ordinary_failure_is_not_mistaken_for_an_oom(): + result = measure([sys.executable, "-c", "raise ValueError('nope')"]) + assert result.status == "failed" + assert "ValueError" in result.stderr_tail + + +def test_the_harness_refuses_to_overwrite_an_existing_output(tmp_path): + """A stale output would be sized and reported as this run's work.""" + existing = tmp_path / "already.h5ad" + existing.write_bytes(b"old") + with pytest.raises(FileExistsError): + measure([sys.executable, "-c", "pass"], output=existing) + + +def test_maxrss_is_normalised_to_bytes(): + """`ru_maxrss` is KiB on Linux and bytes on macOS; the field is bytes. + + A missing conversion would be a silent 1024x error in one direction on + one platform, which is exactly the kind of thing nobody notices in a + table of plausible-looking numbers. + """ + result = measure([sys.executable, "-c", "pass"]) + assert isinstance(result, Measurement) + # Any CPython start-up is well over 1 MB and well under 4 GB. + assert 1024**2 < result.maxrss_bytes < 4 * 1024**3, result.maxrss_bytes From 3d597a62b6253454d82133d189b2a9874fc9c6c8 Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 17:15:58 +0100 Subject: [PATCH 08/14] Replace the estimated benchmark runtimes with measured ones The ci tier is much cheaper than the plan assumed: 580 MB of fixtures and three cases ran end to end in 43 seconds, so the full set is minutes rather than the hour the workflow allows. The 90-minute timeout stays as headroom for the large tier; calling it an estimate would have been wrong. Records the numbers that run produced, including that adata-cli's 202 MB peak is not flat in input size -- obs columns are read whole and a dense block is --chunk x n_var. Better to state that next to the figures than to let the page imply a guarantee the code does not yet meet. Co-Authored-By: Claude Opus 5 --- docs/BENCHMARKS.md | 5 +++++ docs/TESTING.md | 25 +++++++++++++++++++++---- 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 34dd33b..8092b87 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -13,6 +13,11 @@ loading the whole thing is often faster, and the tables say so where it is true -- the trade is memory for time, and a table that hid the cost would be worth nothing. +A `ci` run taken during development, as an indication: `adata concat` of two +50,000 x 20,000 stores peaked at 202 MB against 1,778 MB for `ad.concat` in +memory and 439 MB for `anndata.experimental.concat_on_disk`, and was faster +than both. The published tables will say where it is slower, too. + To produce one locally: ```bash diff --git a/docs/TESTING.md b/docs/TESTING.md index 5bcf1dd..2fe6ae4 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -169,10 +169,27 @@ uv run python -m benchmarks.run --tier smoke --out results.json uv run python -m benchmarks.report results.json ``` -Tiers are `smoke` (1,000 x 2,000, seconds), `ci` (50,000 x 20,000, 35-50 -minutes for the whole suite) and `large` (500,000 x 20,000, dispatch only, -where the in-memory baseline is expected to hit the ceiling). Every tier also -builds the 2,000 x 36,601 var-heavy shape, which is what hung in 0.5.1. +Tiers are `smoke` (1,000 x 2,000, seconds), `ci` (50,000 x 20,000) and +`large` (500,000 x 20,000, dispatch only, where the in-memory baseline is +expected to hit the ceiling). Every tier also builds the 2,000 x 36,601 +var-heavy shape, which is what hung in 0.5.1. + +`ci` is cheaper than it looks: 580 MB of fixtures and three cases took 43 +seconds end to end on a laptop, so the full set with three repeats is minutes +rather than the hour the workflow allows. The 90-minute timeout is headroom +for `large`, not an estimate. + +A `ci` run measured while writing this, for a sense of what the tables say: + +| | adata-cli | anndata (`concat_on_disk`) | anndata (in memory) | +|---|---|---|---| +| `concat-inner`, peak RSS | **202 MB** | 439 MB | 1,778 MB | +| `concat-inner`, wall time | **2.35 s** | 16.44 s | 3.69 s | +| `inspect`, peak RSS | **63 MB** | 138 MB (`read_elem`) | 556 MB (full load) | + +Note that 202 MB is not flat in input size -- obs columns are read whole, and +a dense block is `--chunk` x n_var. The benchmark exists to keep that curve +visible rather than to assert a claim the code does not yet meet. Results are published to `docs/BENCHMARKS.md` and `docs/benchmarks/.json` on every tag, and appended to the GitHub release notes if a release exists. From 75e056b54e95045d28b970b2744200ecff4e3b7f Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 20:54:28 +0100 Subject: [PATCH 09/14] Guard every remaining command, and fix the quadratic split found doing it The previous round pointed both mechanisms almost entirely at concat. ls, view, create, the five exports and the five imports had no cost coverage at all. Extending to them needed three additions to the instruments, and turned up a third defect of the same class as the first two. split --by was O(n_rows x n_groups). core/select.py group_indices grouped rows with `np.nonzero(values == label)` inside a loop over distinct labels, so each chunk was rescanned once per label: measured at 4,096 rows, 16,384 elements scanned for 4 groups and 1,048,576 for 256 -- exactly n_rows per group. A million cells split by a thousand samples is 10^9 comparisons, and split would have looked like a hang for the same reason concat --merge did. np.unique with return_index and return_inverse does it in one pass per chunk: 16,388 elements at 4 groups and 16,640 at 256, flat to within 1.5%. Order of first appearance is preserved through argsort on the first-occurrence indices, because it names the output files. The existing split guard passed throughout. The chunk is already in memory, so no read counter moves -- the same blind spot as _concat_categorical, and the same lesson: one instrument is never enough. Three additions to tests/perf_counters.py: count_scanned_elements counts what is handed to numpy's scanning primitives. It is a floor, not a measurement -- an operator like `values == label` dispatches to the ufunc in C and never passes the patched np.equal -- so the guard using it asserts a lower bound, and the docstring says exactly what is and is not visible. HDF5 writes were not counted at all, which made every import and create guard silently vacuous at zero. Both Dataset.__setitem__ and Group.create_dataset are now hooked; the latter matters because create_dataset(name, data=...) writes its payload at creation and never touches __setitem__. assert_independent_of asserts cost does not grow at all, and assert_grows_slower_than_input asserts it grows by at least some factor less than the input. assert_grows_linearly could only catch super-linear growth, and would have accepted a 64x increase in a command that is supposed to read nothing. The vacuity floor in assert_grows_linearly moved from `d1 >= mid` to half the increment: work that is exactly one operation per element -- export dict reads each key once -- gives 0.75 * mid and was being rejected as unmeasurable. Two claims are now enforced rather than asserted in prose: view and ls read zero data elements, at 64 rows and at 4,096. Not "grows slowly" -- zero, an exact count needing no tolerance. A control test exports the same fixture to prove there was data there to read, so the zero cannot pass by accident, and injecting a single column read into show_info fails four of the ten cases with the offending dataset named. Streaming is bounded well below the input, which is weaker than the README implies and is what the measurements support. Over a 256x span at fixed chunk: export array 1.6x, export dataframe 2.2x, export sparse 6.5x, subset 46x. Only export array is close to flat, so only it is asserted as such, and subset's guard is deliberately the loosest -- obs columns are materialised per column, a known gap that benchmarks/ reports rather than this hiding. Streamed export sparse is also asserted to stay under a quarter of what --in-memory costs, measured in the same run so the factor holds anywhere. Coverage added for view, view --types, ls, ls --long, ls --plain, create (generated names and name file), export dataframe by rows and by columns, export array, export sparse both paths, export dict, export image, import dataframe/array/sparse/dict/image, concat --label, concat --index-unique and split --axis var. 895 tests pass, coverage 92.71%. The perf file is ~55 s, most of it building 65,536-row fixtures, so those three carry the previously unused slow marker. Co-Authored-By: Claude Opus 5 --- src/adata/core/select.py | 24 +- tests/perf_counters.py | 242 ++++++++++++++++- tests/test_performance.py | 547 +++++++++++++++++++++++++++++++++++++- 3 files changed, 797 insertions(+), 16 deletions(-) diff --git a/src/adata/core/select.py b/src/adata/core/select.py index f1049bc..49edb43 100644 --- a/src/adata/core/select.py +++ b/src/adata/core/select.py @@ -100,11 +100,31 @@ def group_indices( values = np.asarray( col_chunk_as_strings(group, column, start, end, cache), dtype=str ) - for label in dict.fromkeys(values.tolist()): + + # One pass over the chunk, whatever the number of distinct labels. + # Taking `np.nonzero(values == label)` per label -- as an earlier + # version did -- rescanned the whole chunk once per label, so the + # cost was O(n_rows * n_groups): a million cells split by a thousand + # samples came to 10^9 comparisons, and `split` looked like a hang. + uniques, first_seen, codes = np.unique( + values, return_index=True, return_inverse=True + ) + codes = codes.ravel() + + # Sorting the codes puts each label's positions in one contiguous + # run, so every group is a slice rather than a search. + by_code = np.argsort(codes, kind="stable") + run_starts = np.searchsorted(codes[by_code], np.arange(len(uniques)), "left") + run_ends = np.searchsorted(codes[by_code], np.arange(len(uniques)), "right") + + # `np.unique` sorts; `order` has to stay in order of first appearance, + # because it names the output files. + for code in np.argsort(first_seen, kind="stable"): + label = str(uniques[code]) if label not in buckets: buckets[label] = [] order.append(label) - buckets[label].append(np.nonzero(values == label)[0] + start) + buckets[label].append(by_code[run_starts[code] : run_ends[code]] + start) return ( {k: np.concatenate(v).astype(np.int64) for k, v in buckets.items()}, diff --git a/tests/perf_counters.py b/tests/perf_counters.py index 0a582d1..cd7e257 100644 --- a/tests/perf_counters.py +++ b/tests/perf_counters.py @@ -56,6 +56,8 @@ class IOCounts: h5_calls: int = 0 h5_elements: int = 0 + h5_write_calls: int = 0 + h5_written_elements: int = 0 zarr_calls: int = 0 zarr_elements: int = 0 store_get: int = 0 @@ -75,6 +77,26 @@ def elements(self) -> int: def calls(self) -> int: return self.h5_calls + self.zarr_calls + @property + def writes(self) -> int: + """Elements written through either backend. + + Zarr is counted in store keys rather than elements -- a chunk write + is one `set` -- so the two are not the same unit. For a guard that + only compares a command against itself at two sizes, that is fine. + """ + return self.h5_written_elements + self.store_set + + @property + def work(self) -> int: + """Reads plus writes, for commands whose job is mostly writing. + + `import` and `create` read almost nothing; measuring only reads made + their guards vacuous, which the lower bound in `assert_grows_linearly` + caught. + """ + return self.reads + self.writes + @property def reads(self) -> int: """The headline number: array elements plus store-level fetches. @@ -86,10 +108,10 @@ def reads(self) -> int: def __str__(self) -> str: # pragma: no cover - diagnostic only return ( - f"elements={self.elements} (h5={self.h5_elements} " - f"zarr={self.zarr_elements}) calls={self.calls} " - f"store: get={self.store_get} set={self.store_set} " - f"delete={self.store_delete}" + f"read={self.elements} (h5={self.h5_elements} " + f"zarr={self.zarr_elements}) wrote={self.h5_written_elements} " + f"calls={self.calls} store: get={self.store_get} " + f"set={self.store_set} delete={self.store_delete}" ) @@ -106,12 +128,18 @@ def count_io() -> Iterator[IOCounts]: Patches four seams, all verified against h5py 3.15.1 and zarr 3.1.5: - =================================== ==================================== - ``h5py.Dataset.__getitem__`` HDF5 reads - ``zarr.Array.__getitem__`` Zarr reads, v2 and v3 alike - ``zarr.storage.LocalStore.get`` chunk and metadata fetches - ``LocalStore.set`` / ``.delete`` write storms - =================================== ==================================== + ====================================== ================================= + ``h5py.Dataset.__getitem__`` HDF5 reads + ``h5py.Dataset.__setitem__`` HDF5 writes + ``h5py.Group.create_dataset`` HDF5 writes made at creation + ``zarr.Array.__getitem__`` Zarr reads, v2 and v3 alike + ``zarr.storage.LocalStore.get`` chunk and metadata fetches + ``LocalStore.set`` / ``.delete`` write storms + ====================================== ================================= + + Both write seams are needed: `create_dataset(name, data=...)` writes its + payload during creation and never touches `__setitem__`, so hooking only + the latter leaves `import` and `create` measuring zero. The store hooks are the ones that see a metadata write storm: attribute writes never touch `Array.__getitem__`, so an O(k^2) attribute rewrite is @@ -126,6 +154,8 @@ def count_io() -> Iterator[IOCounts]: counts = IOCounts() h5_get = h5py.Dataset.__getitem__ + h5_set = h5py.Dataset.__setitem__ + h5_create = h5py.Group.create_dataset zarr_get = zarr.Array.__getitem__ store_get = LocalStore.get store_set = LocalStore.set @@ -150,6 +180,21 @@ def zarr_wrapper(self: Any, key: Any) -> Any: _record(getattr(self, "path", "?") or "/", n) return result + def h5_set_wrapper(self: Any, key: Any, value: Any) -> Any: + counts.h5_write_calls += 1 + counts.h5_written_elements += _size(value) + return h5_set(self, key, value) + + def h5_create_wrapper(self: Any, name: Any, *args: Any, **kwargs: Any) -> Any: + counts.h5_write_calls += 1 + data = kwargs.get("data") + if data is None and args: + # positional signature is (name, shape, dtype, data, ...) + data = args[2] if len(args) > 2 else None + if data is not None: + counts.h5_written_elements += _size(data) + return h5_create(self, name, *args, **kwargs) + async def get_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: result = await store_get(self, *args, **kwargs) counts.store_get += 1 @@ -169,6 +214,8 @@ async def delete_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: return await store_delete(self, *args, **kwargs) h5py.Dataset.__getitem__ = h5_wrapper + h5py.Dataset.__setitem__ = h5_set_wrapper + h5py.Group.create_dataset = h5_create_wrapper zarr.Array.__getitem__ = zarr_wrapper LocalStore.get = get_wrapper LocalStore.set = set_wrapper @@ -177,6 +224,8 @@ async def delete_wrapper(self: Any, *args: Any, **kwargs: Any) -> Any: yield counts finally: h5py.Dataset.__getitem__ = h5_get + h5py.Dataset.__setitem__ = h5_set + h5py.Group.create_dataset = h5_create zarr.Array.__getitem__ = zarr_get LocalStore.get = store_get LocalStore.set = store_set @@ -241,6 +290,78 @@ def trace(frame: Any, event: str, arg: Any) -> Any: sys.settrace(previous) +#: numpy entry points that scan their operand, and the argument that is +#: scanned. Patched by name, so only calls written as `np.f(...)` are seen. +_SCANNING = { + "nonzero": 0, + "flatnonzero": 0, + "unique": 0, + "argsort": 0, + "sort": 0, + "searchsorted": 0, + "bincount": 0, + "equal": 0, + "isin": 0, +} + + +@contextmanager +def count_scanned_elements() -> Iterator[List[int]]: + """Elements handed to numpy's scanning primitives, as `[n]`. + + The third instrument, and the one the other two cannot replace. A loop of + the shape:: + + for label in distinct_labels: + np.nonzero(values == label) + + reads nothing extra (the chunk is already in memory), allocates nothing + that lasts, and executes a constant number of Python lines per label. It + is invisible to `count_io`, to `count_allocations` and to `count_lines` + alike, and it is O(len(values) * len(distinct_labels)). + + Counting what is passed into numpy makes it visible, in the same spirit + as the `_CountingStr` guard in `test_performance.py`: when the cost lives + inside C, count what is handed to C. + + What this does and does not see + ------------------------------- + Only calls that go through the `numpy` module namespace by name. An + operator -- `values == label` -- dispatches to `ndarray.__eq__` and then + to the ufunc in C, so it never passes the patched `np.equal`; a scan + written purely as `(a == b).sum()` is invisible. Method form, + `arr.argsort()`, is invisible for the same reason. + + So this is a floor on the real work, not a measurement of it, which is + exactly what a guard needs: it can only under-report, and the guards + using it assert a lower bound so that under-reporting to nothing fails + rather than passes. + """ + total = [0] + originals = {name: getattr(np, name) for name in _SCANNING} + + def wrap(name: str, position: int): + original = originals[name] + + def wrapper(*args: Any, **kwargs: Any) -> Any: + if len(args) > position: + try: + total[0] += int(np.size(args[position])) + except Exception: # pragma: no cover - exotic operands + pass + return original(*args, **kwargs) + + return wrapper + + for name, position in _SCANNING.items(): + setattr(np, name, wrap(name, position)) + try: + yield total + finally: + for name, original in originals.items(): + setattr(np, name, original) + + # --------------------------------------------------------------------------- # the growth assertion @@ -300,12 +421,109 @@ def assert_grows_linearly( + ")" ) - assert d1 >= mid, ( + # Half an operation per element added. Work that is exactly one + # operation per element -- `export dict` reads each key once -- gives + # d1 = mid - small, which is 0.75 * mid at 4x spacing, so a bound of + # `mid` would reject correct code. Anything the counters cannot see at + # all gives 0 and still fails. + floor = (mid - small) / 2 + assert d1 >= floor, ( f"Too little measured work for the growth ratio to mean anything -- " - f"the counters are probably not seeing this operation. {series}" + f"the counters are probably not seeing this operation. Expected at " + f"least {floor:.0f} more operations between {small} and {mid}. " + f"{series}" ) assert d2 <= limit * d1, ( f"Cost grows faster than linearly in {axis}. Expect an increment " f"ratio near 4.0 for linear work; {limit} is the limit and quadratic " f"would be about 16. {series}" ) + + +#: How much a "flat" cost may drift across the whole size span. +#: +#: Not 1.0: opening a store, resolving an index and printing a tree all cost +#: a little more when there is more to describe -- a longer index name, a +#: wider shape to format. 1.5 across a 64x span leaves room for that and none +#: at all for reading the data, which would be 64x. +FLAT_TOLERANCE = 1.5 + +#: Span for the flat guards. Wider than SIZES, because the claim is stronger +#: and a wide span is what makes it convincing. +FLAT_SIZES = (64, 4096) + + +def assert_independent_of( + measure: Callable[[int], int], + *, + what: str, + axis: str, + sizes: tuple = FLAT_SIZES, + tolerance: float = FLAT_TOLERANCE, +) -> None: + """Assert `measure` does not grow with `axis` at all. + + A stronger claim than `assert_grows_linearly`, and the right one wherever + the tool promises work proportional to something other than input size: + + * **Inspection.** `view` and `ls` read shapes, dtypes and attributes and + never the values behind them, which is the whole reason they return + instantly on a store too large to open. Linear growth here would mean + the promise had quietly stopped holding. + * **Grouping.** Work per row must not depend on how many groups there are. + * **Streaming.** At a fixed chunk size, peak memory must not track the + size of the input. + + Each of those reads as obvious prose and none of them is checked by a + linear-growth guard, which would happily accept a 64x increase. + """ + small, large = sizes[0], sizes[-1] + c_small, c_large = measure(small), measure(large) + + span = large / small + series = ( + f"{what}, scaling {axis}: {small}->{c_small}, {large}->{c_large} " + f"over a {span:.0f}x span" + ) + ratio = (c_large / c_small) if c_small else float("inf") if c_large else 1.0 + + assert ratio <= tolerance, ( + f"Cost tracks {axis}, and it is supposed to be independent of it. " + f"Growing in step with the input would be about {span:.0f}x; " + f"{tolerance}x is the limit. {series} (ratio {ratio:.2f})" + ) + + +def assert_grows_slower_than_input( + measure: Callable[[int], int], + *, + what: str, + axis: str, + sizes: tuple, + at_least: float, +) -> None: + """Assert cost grows at least `at_least` times slower than the input. + + The middle ground between `assert_grows_linearly` and + `assert_independent_of`, and the honest shape of most of this tool's + streaming: peak memory is not flat -- an index or an indptr is read whole + -- but it is far below the input curve, and that margin is the feature. + + Stating the margin as a factor rather than an absolute ceiling keeps the + guard meaningful across platforms and interpreters, and makes a drift + back towards linear fail long before it becomes a bug report. + """ + small, large = sizes[0], sizes[-1] + c_small, c_large = measure(small), measure(large) + + span = large / small + growth = (c_large / c_small) if c_small else float("inf") + budget = span / at_least + + assert growth <= budget, ( + f"{what} cost is tracking {axis} too closely. Input grew {span:.0f}x " + f"and cost grew {growth:.1f}x; the requirement is at least " + f"{at_least:.0f}x better than the input, i.e. no more than " + f"{budget:.0f}x. Series: {small}->{c_small}, {large}->{c_large}" + ) + diff --git a/tests/test_performance.py b/tests/test_performance.py index 9f98dc5..8dd7b07 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -40,6 +40,16 @@ import pytest from rich.console import Console +from adata.commands.create import create_store +from adata.commands.export import ( + export_json, + export_mtx, + export_npy, + export_table, +) +from adata.commands.import_data import import_object +from adata.commands.info import show_info +from adata.commands.ls import list_store from adata.core.concat import concat_on_disk from adata.core.subset import subset_h5ad @@ -47,7 +57,10 @@ GROWTH_LIMIT, SIZES, assert_grows_linearly, + assert_grows_slower_than_input, + assert_independent_of, count_allocations, + count_scanned_elements, count_io, count_lines, ) @@ -105,11 +118,13 @@ def _store( index=[f"{prefix}g{i}" for i in range(n_var)], ) - ad.AnnData( + obj = ad.AnnData( X=sparse.csr_matrix(np.ones((n_obs, n_var), dtype="float32")), obs=obs, var=var, - ).write_h5ad(path) + ) + obj.obsm["X_pca"] = np.zeros((n_obs, 3), dtype="float32") + obj.write_h5ad(path) return path @@ -502,3 +517,531 @@ def ratio(exponent: float) -> float: assert ratio(1.0) == pytest.approx(4.0, rel=0.01) assert ratio(2.0) > GROWTH_LIMIT, "quadratic work must fail" assert ratio(1.5) > GROWTH_LIMIT, "n**1.5 must fail" + + +# --------------------------------------------------------------------------- +# inspection is free +# +# `view` and `ls` reach only `.shape`, `.dtype` and `.attrs`: `axis_len` goes +# through `element_len`, which reads a shape, and `_array_details` and +# `_infer_untagged` never touch a value. That is the whole reason they return +# instantly on a store far too large to open, and it is the one claim in the +# README that can be stated as an exact number rather than a ratio. + + +def _reads_for_inspection(path: Path, run) -> "object": + with count_io() as io: + run(path) + return io + + +@pytest.mark.parametrize("n_obs", [64, 4096]) +@pytest.mark.parametrize( + "label,run", + [ + ("view", lambda p: show_info(p, QUIET, out_console=QUIET)), + ("view --types", lambda p: show_info(p, QUIET, show_types=True, + out_console=QUIET)), + ("ls", lambda p: list_store(p, QUIET)), + ("ls --long", lambda p: list_store(p, QUIET, long=True)), + ("ls --plain", lambda p: list_store(p, QUIET, plain=True)), + ], + ids=lambda v: v if isinstance(v, str) else "", +) +def test_inspection_reads_no_data_at_all(tmp_path, n_obs, label, run): + """Not "grows slowly" -- zero. An exact count, so it needs no tolerance. + + A ratio would be the wrong instrument here: 0 to 0 is not a meaningful + ratio, and the moment inspection reads *one* column the answer stops + being zero regardless of how the store scales. + """ + source = _store(tmp_path / f"i{label}{n_obs}.h5ad", name="i", n_obs=n_obs, + n_var=16, obs_kind="numeric", n_categories=4) + io = _reads_for_inspection(source, run) + + assert io.elements == 0, ( + f"`{label}` read {io.elements} data elements from a {n_obs}-row store; " + f"inspection is supposed to touch only shapes and attributes. " + f"Largest reader: {max(io.by_name.items(), key=lambda kv: kv[1], default=('-', 0))}" + ) + + +def test_the_inspection_fixture_really_does_hold_readable_data(tmp_path): + """Keeps the test above from passing because there was nothing to read. + + Same store, read by a command that is supposed to read: if this registers + nothing either, the fixture or the hooks are broken, not the claim. + """ + source = _store(tmp_path / "control.h5ad", name="i", n_obs=4096, n_var=16, + obs_kind="numeric", n_categories=4) + with count_io() as io: + export_table(source, "obs", None, tmp_path / "control.csv", 10_000, None, QUIET) + + assert io.elements >= 4096, ( + f"the control read only {io.elements} elements from a 4096-row store, " + "so the zero above proves nothing" + ) + + +# --------------------------------------------------------------------------- +# grouping work does not depend on the number of groups +# +# `split --by` groups rows through `core.select.group_indices`. Its cost is +# invisible to every other instrument in this file: the chunk is already in +# memory so no read counter moves, nothing lasting is allocated, and the +# Python line count per label is constant. Only counting what is handed to +# numpy shows it. + + +def test_grouping_does_not_rescan_the_column_once_per_group(tmp_path): + """One pass over the column, however many distinct values it holds. + + Scanning per label is O(n_rows * n_groups): measured at n_obs=4096 it was + 16,384 elements for 4 groups and 1,048,576 for 256 -- exactly n_rows per + group. On a million cells split by a thousand samples that is 10^9 + comparisons, and `split` would appear to hang for the same reason + `concat --merge` did. + + n_obs is fixed. The claim is about work per row, not total work. + """ + from adata.core.select import group_indices + from adata.storage import open_store + + n_obs = 4096 + + def measure(n_groups: int) -> int: + source = _store( + tmp_path / f"g{n_groups}.h5ad", + name="g", + n_obs=n_obs, + n_var=2, + n_categories=n_groups, + ) + with open_store(source, "r") as store, count_scanned_elements() as scanned: + groups, order = group_indices(store.root, "obs", "ct") + assert len(order) == n_groups, "fixture did not produce the groups asked for" + # A floor, not a measurement -- see count_scanned_elements. If the + # implementation stops calling numpy by name the count collapses to + # zero and the ratio below would pass for the wrong reason. + assert scanned[0] >= n_obs, ( + f"only {scanned[0]} elements scanned for {n_obs} rows; the counter " + "is no longer seeing this code path" + ) + return scanned[0] + + assert_independent_of( + measure, what="group_indices", axis="n_groups", sizes=(4, 256) + ) + + +# --------------------------------------------------------------------------- +# export +# +# Each export reads the thing it is asked for and nothing else. The axis +# differs per subcommand -- rows, elements, nonzeros, keys -- so each says +# which one it scales and holds the rest fixed. + + +def test_export_dataframe_reads_grow_linearly_in_rows(tmp_path): + def measure(n: int) -> int: + source = _store(tmp_path / f"ed{n}.h5ad", name="e", n_obs=n, n_var=4, + obs_kind="numeric") + with count_io() as io: + export_table(source, "obs", None, tmp_path / f"ed{n}.csv", + 10_000, None, QUIET) + return io.reads + + assert_grows_linearly(measure, what="export dataframe", axis="n_obs") + + +def test_export_dataframe_reads_grow_linearly_in_column_count(tmp_path): + def measure(n: int) -> int: + source = _store(tmp_path / f"ec{n}.h5ad", name="e", n_obs=16, n_var=4, + n_obs_columns=n) + with count_io() as io: + export_table(source, "obs", None, tmp_path / f"ec{n}.csv", + 10_000, None, QUIET) + return io.reads + + assert_grows_linearly(measure, what="export dataframe", axis="n_obs_columns") + + +def test_export_array_reads_grow_linearly_in_element_count(tmp_path): + def measure(n: int) -> int: + source = _store(tmp_path / f"ea{n}.h5ad", name="e", n_obs=n, n_var=4) + with count_io() as io: + export_npy(source, "obsm/X_pca", tmp_path / f"ea{n}.npy", + 100_000, QUIET) + return io.reads + + assert_grows_linearly(measure, what="export array", axis="n_obs") + + +@pytest.mark.parametrize("in_memory", [False, True]) +def test_export_sparse_reads_grow_linearly_in_nonzeros(tmp_path, in_memory): + """Both paths: the streamed one and the one that loads the matrix.""" + + def measure(n: int) -> int: + source = _store(tmp_path / f"es{in_memory}{n}.h5ad", name="e", + n_obs=n, n_var=4) + with count_io() as io: + export_mtx(source, "X", tmp_path / f"es{in_memory}{n}.mtx", + None, 1_000, in_memory, QUIET) + return io.reads + + assert_grows_linearly( + measure, what=f"export sparse (in_memory={in_memory})", axis="nnz" + ) + + +def test_export_dict_reads_grow_linearly_in_key_count(tmp_path): + def measure(n: int) -> int: + source = _store_with_uns_keys(tmp_path / f"ej{n}.h5ad", n) + with count_io() as io: + export_json(source, "uns", tmp_path / f"ej{n}.json", + 100_000, False, QUIET) + return io.reads + + assert_grows_linearly( + measure, what="export dict", axis="n_keys", sizes=(64, 256, 1024) + ) + + +def test_export_image_reads_grow_linearly_in_pixels(tmp_path): + from adata.commands.export import export_image + + def measure(n: int) -> int: + source = _store_with_image(tmp_path / f"ei{n}.h5ad", n) + with count_io() as io: + export_image(source, "uns/picture", tmp_path / f"ei{n}.png", QUIET) + return io.reads + + # n is the side of a square image, so pixels grow as n**2 -- the axis + # being scaled is the pixel count, and the sizes below keep it at 4x. + assert_grows_linearly( + measure, what="export image", axis="pixels", sizes=(16, 32, 64) + ) + + +# --------------------------------------------------------------------------- +# import + + +def test_import_dataframe_reads_grow_linearly_in_rows(tmp_path): + def measure(n: int) -> int: + source = _store(tmp_path / f"id{n}.h5ad", name="m", n_obs=n, n_var=4) + csv = tmp_path / f"id{n}.csv" + pd.DataFrame( + {"score": np.arange(n, dtype="int32")}, + index=[f"mc{i}" for i in range(n)], + ).to_csv(csv) + with count_io() as io: + import_object(source, "obs", csv, tmp_path / f"od{n}.h5ad", + False, None, QUIET) + # `work`, not `reads`: import is a write path, and measuring only + # reads made this vacuous. + return io.work + + assert_grows_linearly(measure, what="import dataframe", axis="n_obs") + + +def test_import_array_reads_grow_linearly_in_element_count(tmp_path): + def measure(n: int) -> int: + source = _store(tmp_path / f"ia{n}.h5ad", name="m", n_obs=n, n_var=4) + npy = tmp_path / f"ia{n}.npy" + np.save(npy, np.zeros((n, 3), dtype="float32")) + with count_io() as io: + import_object(source, "obsm/imported", npy, + tmp_path / f"oa{n}.h5ad", False, None, QUIET) + return io.work + + assert_grows_linearly(measure, what="import array", axis="n_obs") + + +def test_import_sparse_reads_grow_linearly_in_nonzeros(tmp_path): + def measure(n: int) -> int: + source = _store(tmp_path / f"is{n}.h5ad", name="m", n_obs=n, n_var=4) + mtx = tmp_path / f"is{n}.mtx" + sparse_io = pytest.importorskip("scipy.io") + sparse_io.mmwrite( + str(mtx), sparse.csr_matrix(np.ones((n, 4), dtype="float32")) + ) + with count_io() as io: + import_object(source, "layers/imported", mtx, + tmp_path / f"os{n}.h5ad", False, None, QUIET) + return io.work + + assert_grows_linearly(measure, what="import sparse", axis="nnz") + + +def test_import_dict_reads_grow_linearly_in_key_count(tmp_path): + import json + + def measure(n: int) -> int: + source = _store(tmp_path / f"ij{n}.h5ad", name="m", n_obs=8, n_var=4) + blob = tmp_path / f"ij{n}.json" + blob.write_text(json.dumps({f"k{i}": i for i in range(n)})) + with count_io() as io: + import_object(source, "uns/imported", blob, + tmp_path / f"oj{n}.h5ad", False, None, QUIET) + return io.work + + assert_grows_linearly( + measure, what="import dict", axis="n_keys", sizes=(64, 256, 1024) + ) + + +def test_import_image_reads_grow_linearly_in_pixels(tmp_path): + """Images take their own entry point. + + `import_object` dispatches on extension and `.png` is deliberately not in + `EXTENSION_FORMAT`; the CLI's `import image` calls `_import_image` + directly, and it edits in place rather than writing a copy. + """ + Image = pytest.importorskip("PIL.Image") + from adata.commands.import_data import _import_image + + def measure(n: int) -> int: + source = _store(tmp_path / f"ii{n}.h5ad", name="m", n_obs=8, n_var=4) + png = tmp_path / f"ii{n}.png" + Image.fromarray(np.zeros((n, n, 3), dtype="uint8")).save(png) + with count_io() as io: + _import_image(source, "uns/picture", png, QUIET) + return io.work + + assert_grows_linearly( + measure, what="import image", axis="pixels", sizes=(16, 32, 64) + ) + + +# --------------------------------------------------------------------------- +# create, and the concat options nothing else covers + + +@pytest.mark.parametrize("from_file", [False, True]) +def test_create_writes_grow_linearly_in_obs_count(tmp_path, from_file): + """Both the generated-names path and the name-file path.""" + + def measure(n: int) -> int: + names = None + if from_file: + names = tmp_path / f"cn{n}.txt" + names.write_text("\n".join(f"c{i}" for i in range(n)) + "\n") + output = tmp_path / f"cr{from_file}{n}.h5ad" + with count_io() as io: + create_store( + output, + QUIET, + n_obs=None if from_file else n, + n_var=4, + obs_names=names, + ) + return io.work + + assert_grows_linearly( + measure, what=f"create (from_file={from_file})", axis="n_obs" + ) + + +@pytest.mark.parametrize( + "option", ["label", "index_unique"], ids=["--label", "--index-unique"] +) +def test_concat_option_reads_grow_linearly_in_obs_count(tmp_path, option): + """Both build a per-row value, so both are worth a guard of their own.""" + + def measure(n: int) -> int: + files = _inputs(tmp_path, f"co{option}{n}", n_obs=n, n_var=4) + kwargs = {"label": "batch"} if option == "label" else {"index_unique": "-"} + with count_io() as io: + concat_on_disk(files, tmp_path / f"oco{option}{n}.h5ad", QUIET, **kwargs) + return io.reads + + assert_grows_linearly(measure, what=f"concat --{option}", axis="n_obs") + + +def test_split_by_var_reads_grow_linearly_in_group_count(tmp_path): + """The var axis takes a different path through `split` than obs does.""" + from adata.commands.split import split_store + + n_var = 512 + + def measure(n: int) -> int: + source = _store_with_var_groups(tmp_path / f"sv{n}.h5ad", n_var, n) + with count_io() as io: + split_store(source, "kind", tmp_path / f"osv{n}", QUIET, + axis="var", manifest=False) + return io.reads + + assert_grows_linearly( + measure, what="split --axis var", axis="n_groups", sizes=(8, 32, 128) + ) + + +# --------------------------------------------------------------------------- +# fixture variants the guards above need + + +def _store_with_uns_keys(path: Path, n_keys: int) -> Path: + obj = ad.AnnData( + X=np.ones((4, 2), dtype="float32"), + obs=pd.DataFrame(index=["a", "b", "c", "d"]), + var=pd.DataFrame(index=["g1", "g2"]), + ) + obj.uns.update({f"k{i}": i for i in range(n_keys)}) + obj.write_h5ad(path) + return path + + +def _store_with_image(path: Path, side: int) -> Path: + obj = ad.AnnData( + X=np.ones((4, 2), dtype="float32"), + obs=pd.DataFrame(index=["a", "b", "c", "d"]), + var=pd.DataFrame(index=["g1", "g2"]), + ) + obj.uns["picture"] = np.zeros((side, side, 3), dtype="uint8") + obj.write_h5ad(path) + return path + + +def _store_with_var_groups(path: Path, n_var: int, n_groups: int) -> Path: + ad.AnnData( + X=np.ones((4, n_var), dtype="float32"), + obs=pd.DataFrame(index=["a", "b", "c", "d"]), + var=pd.DataFrame( + {"kind": pd.Categorical([f"k{i % n_groups}" for i in range(n_var)])}, + index=[f"g{i}" for i in range(n_var)], + ), + ).write_h5ad(path) + return path + + +# --------------------------------------------------------------------------- +# streaming is bounded by the chunk, not by the input +# +# The README's actual claim, and the reason this tool exists. It is only +# wholly true in one place today, so these guards say what is true rather +# than what would be nice, and record the measured numbers so that drift is +# visible. Everything below holds the chunk size fixed and grows the input by +# 256x; a command that loaded everything would grow 256x with it. + +#: 256x span. Wide enough that "flat" and "linear" cannot be confused, and +#: the reason these carry the `slow` marker: building a 65,536-row store is +#: most of the ~45 s this section costs. They still gate merges; the marker +#: is so a local run can say `-m "not slow"`. +STREAM_SIZES = (256, 65536) + + +def _peak_for(tmp_path: Path, tag: str, n: int, run) -> int: + source = _store(tmp_path / f"{tag}{n}.h5ad", name="s", n_obs=n, n_var=8, + obs_kind="numeric") + with count_allocations() as peak: + run(source, n) + return peak[0] + + +@pytest.mark.slow +def test_export_array_peak_memory_is_set_by_the_chunk_not_the_input(tmp_path): + """The closest the tool comes to the claim outright. + + Measured at a fixed 1,000-element chunk: 21.5 KiB at 256 rows and + 34.3 KiB at 65,536 -- 1.6x for a 256x input. Not flat, so this does not + assert flat; the residue is fixed-size bookkeeping that grows with the + length of a formatted shape rather than with the data. Stating it as "at + least 64x better than the input" is both true and strong: loading + everything would be 256x. + """ + + def measure(n: int) -> int: + return _peak_for( + tmp_path, "sa", n, + lambda p, k: export_npy(p, "obsm/X_pca", tmp_path / f"sa{k}.npy", + 1_000, QUIET), + ) + + assert_grows_slower_than_input( + measure, + what="export array peak allocation at --chunk 1000", + axis="n_obs", + sizes=STREAM_SIZES, + at_least=64.0, + ) + + +@pytest.mark.parametrize( + "label,factor,run", + [ + # Measured over the 256x span: export sparse 6.5x, export dataframe + # 2.2x, subset 46x. The factors below sit roughly midway between the + # measurement and linear, so ordinary variation passes and a real + # drift towards loading everything fails. + ( + "export sparse", + 8.0, + lambda p, k, t: export_mtx(p, "X", t / f"ss{k}.mtx", None, 1_000, + False, QUIET), + ), + ( + "export dataframe", + 8.0, + lambda p, k, t: export_table(p, "obs", None, t / f"st{k}.csv", + 1_000, None, QUIET), + ), + ( + "subset", + 2.0, + lambda p, k, t: subset_h5ad(p, t / f"su{k}.h5ad", None, None, + console=QUIET, obs_query="col >= 0", + chunk_rows=256), + ), + ], + ids=["export-sparse", "export-dataframe", "subset"], +) +@pytest.mark.slow +def test_streamed_peak_memory_grows_far_slower_than_the_input( + tmp_path, label, factor, run +): + """Not flat, and the guards should not pretend otherwise. + + These paths read an index or an indptr whole, so peak allocation does + track n_obs -- just far below it. `subset` is the weakest of the three + because obs columns are materialised per column; that is a known gap, + reported by `benchmarks/` rather than hidden here. + """ + + def measure(n: int) -> int: + return _peak_for(tmp_path, f"st{label[-4:]}", n, + lambda p, k: run(p, k, tmp_path)) + + assert_grows_slower_than_input( + measure, + what=f"{label} peak allocation at a fixed chunk", + axis="n_obs", + sizes=STREAM_SIZES, + at_least=factor, + ) + + +@pytest.mark.slow +def test_streaming_export_sparse_costs_far_less_than_loading_the_matrix( + tmp_path, +): + """`--in-memory` exists as the fast path; the default must earn its place. + + Self-calibrating: both are measured in the same run, so the assertion + holds regardless of platform. Measured at 65,536 rows: 777 KiB streamed + against 31,763 KiB in memory, a factor of 41. + """ + n = 65_536 + source = _store(tmp_path / "cmp.h5ad", name="s", n_obs=n, n_var=8) + + with count_allocations() as streamed: + export_mtx(source, "X", tmp_path / "streamed.mtx", None, 1_000, False, QUIET) + with count_allocations() as loaded: + export_mtx(source, "X", tmp_path / "loaded.mtx", None, 1_000, True, QUIET) + + assert streamed[0] <= loaded[0] / 4, ( + f"streaming export used {streamed[0] // 1024} KiB against " + f"{loaded[0] // 1024} KiB for --in-memory, a factor of only " + f"{loaded[0] / max(streamed[0], 1):.1f}. The default path is supposed " + "to be the one you reach for when the matrix does not fit." + ) From 797bf60a1e432e39c7a5a89c8def6e6e80901fb9 Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 21:04:54 +0100 Subject: [PATCH 10/14] Benchmark the remaining commands, and document the three invariants Six cases added, chosen by whether a real baseline exists: ls, create, export-array, export-sparse, import-dataframe and concat-outer. export image, export dict, import image and import dict are deliberately left out -- no library offers them, so the rows would only ever read n/a while adding runtime to every tag. The complexity guards cover them instead. Two of the new rows are the reason the report is framed as "peak RSS against wall time" rather than as a leaderboard. At the ci tier, export sparse streams a 50,000 x 20,000 matrix in 68 MB and takes 10.8 s where loading it whole takes 759 MB and 1.9 s; and h5ls -r lists the file in 0.03 s and 7 MB against our 0.36 s and 63 MB, being C rather than a Python process that has to import typer, rich, h5py and zarr first. Both are published. A benchmark that showed only the rows we win would not be worth the runtime. The rest of the ci run: create 78 MB / 0.28 s against 2,225 MB / 2.85 s, import-dataframe 107 MB / 0.41 s against 571 MB / 1.84 s, concat-outer 202 MB / 2.60 s against 1,362 MB / 4.02 s in memory and 436 MB / 2.24 s for concat_on_disk, which is slightly the faster of the two. Harness changes the new cases needed: Baseline environments now install dask. concat_on_disk imports it to concatenate a dense element and raises ModuleNotFoundError without it, which surfaced as soon as the fixtures gained an obsm. Giving the baseline its best idiom is the standing rule; measuring a library crippled by a missing optional dependency would be measuring our own setup. A case can declare a sidecar input, built once from the real store, so import-dataframe reads a CSV the file could plausibly have held rather than an invented one. Contenders whose binary is absent -- h5ls is often not installed -- are recorded as n/a with the reason rather than crashing the run. create takes its shape from the tier. Hardcoding 50,000 x 20,000 made the smoke tier allocate a 4 GB dense array, so "smoke" was not smoke: its peak went from 2,324 MB to 64 MB once the shape followed the tier. Fixtures gained a 50-column obsm, without which export array and import array had nothing of realistic width to move. docs/TESTING.md gets the three invariants and when each applies, the two write seams and why both are needed, count_scanned_elements and the precise statement of what it cannot see, and a per-command coverage table so the next command's author knows what is expected. The measured streaming growth figures are recorded there too -- 1.6x, 2.2x, 6.5x, 46x over a 256x span -- because the guards are set from them. 895 tests pass, coverage 92.71%. All 15 benchmark cases run clean at smoke. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 18 ++++- benchmarks/cases.py | 176 +++++++++++++++++++++++++++++++++++++++++ benchmarks/datasets.py | 5 ++ benchmarks/report.py | 7 +- benchmarks/run.py | 52 +++++++++++- docs/BENCHMARKS.md | 15 +++- docs/TESTING.md | 127 ++++++++++++++++++++++++++--- 7 files changed, 376 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9867304..cf21eec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,13 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no Category merging probed a list rather than a dict: 2,096,128 string comparisons to union 1,024 categories, and around 5e9 for a 100k-category column. Found by the new guards. +- **`split --by` was quadratic**, O(n_rows x n_groups). `group_indices` grouped + rows with `np.nonzero(values == label)` inside a loop over distinct labels, + rescanning each chunk once per label: at 4,096 rows, 16,384 elements scanned + for 4 groups and 1,048,576 for 256. A million cells split by a thousand + samples is ~10^9 comparisons. One `np.unique` pass per chunk makes it flat + in the group count. Found by the new guards; order of first appearance, + which names the output files, is unchanged. - **`concat` built a Python object per row** for nullable and string obs columns, then walked the list twice more. Filling a typed buffer by slice removes three full passes over every such column. @@ -30,12 +37,19 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no seconds -- h5py and zarr reads, Zarr store traffic, Python allocation and executed lines -- and assert that successive increments grow no faster than linearly, so nothing here can fail because a CI runner was busy. See - [docs/TESTING.md](docs/TESTING.md#performance). + [docs/TESTING.md](docs/TESTING.md#performance). Every subcommand is covered: + `ls`, `view`, `create`, all five `export` and all five `import` variants, + `split` on both axes, and the `concat` options nothing else reached. + Two claims are now enforced rather than described -- `view` and `ls` read + **zero** data elements at any store size, and streaming stays far below the + input curve at a fixed `--chunk`. - **A comparative benchmark** (`benchmarks/`), run on every tag against anndata and against scanpy where scanpy has a real equivalent. Reports peak RSS, wall time and output size; publishes to [docs/BENCHMARKS.md](docs/BENCHMARKS.md) and the release notes. Report-only - -- it never fails a build. + -- it never fails a build. Fifteen cases, covering every command with a real + baseline, including `h5ls -r` for `ls` and the rows where adata-cli is the + slower of the two. - **`--merge drop` and `--uns-merge drop` are accepted.** `drop` was already the documented default behaviour but was rejected as a value, so a config could not state it explicitly. diff --git a/benchmarks/cases.py b/benchmarks/cases.py index a9ee7a9..7269b51 100644 --- a/benchmarks/cases.py +++ b/benchmarks/cases.py @@ -13,6 +13,11 @@ its output; `write_h5ad` defaults to none. Left alone, adata-cli's output looks smaller for reasons unrelated to the tool. +**Give the baseline everything it needs.** `concat_on_disk` imports `dask` +to handle a dense element and fails outright without it, so the baseline +environments install it. Measuring a library crippled by a missing optional +dependency would be measuring our own setup. + **`n/a` is a result.** `concat_on_disk` is CSR/CSC-oriented and its outer-join support has varied by version; scanpy has no streaming concat at all. Print what refused and why. An omitted row reads as an oversight; a stated `n/a` @@ -28,6 +33,12 @@ time at many times the memory. That trade *is* the argument for this tool. The table is "peak RSS against wall time", not a leaderboard. +**Not every command has a baseline, and four are left out on purpose.** +`export image`, `export dict`, `import image` and `import dict` have no +library equivalent, so their rows would only ever read `n/a` and they would +add runtime to every tag for nothing. They are covered by the complexity +guards in `tests/test_performance.py` instead. + **scanpy's filter functions are not our subset.** `sc.pp.filter_cells` *computes* `n_genes` by scanning X; `adata subset --obs-query` filters a column that must already exist. Head to head, scanpy looks slow for doing @@ -72,6 +83,10 @@ class Case: output_suffix: str = ".h5ad" #: Use the var-heavy shape rather than the tier's main shape. var_heavy: bool = False + #: An auxiliary input the case needs, e.g. the CSV an import reads. + #: Built once from the first input and offered to CLI contenders as + #: `{sidecar_csv}` and to scripts as $BENCH_SIDECAR_CSV. + sidecar: Optional[str] = None tags: List[str] = field(default_factory=list) @@ -309,6 +324,167 @@ def _py(body: str) -> str: ), ], ), + # -- structure and metadata ------------------------------------------- + Case( + name="ls", + question="List everything in the store.", + output_suffix="", + contenders=[ + Contender("adata-cli", argv=["adata", "ls", "{input0}", "--plain"]), + Contender( + "h5py (visit)", + script=_py( + "names = []\n" + "with h5py.File(IN[0], 'r') as f:\n" + " f.visit(names.append)\n" + "print(len(names))\n" + ), + ), + # The obvious tool someone already has. If the CLI cannot beat a + # 20-year-old C program at walking an HDF5 file, that is worth + # knowing and worth printing. + Contender("h5ls -r", argv=["h5ls", "-r", "{input0}"]), + ], + tags=["headline"], + ), + Case( + name="create", + question="Write an empty store with a given obs/var shape.", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "create", "{output}", + "--n-obs", "{n_obs}", "--n-var", "{n_var}", + ], + ), + Contender( + "anndata", + script=_py( + "import os, pandas as pd\n" + "n_obs = int(os.environ['BENCH_N_OBS'])\n" + "n_var = int(os.environ['BENCH_N_VAR'])\n" + "obs = pd.DataFrame(index=[f'cell_{i}' for i in range(n_obs)])\n" + "var = pd.DataFrame(index=[f'gene_{i}' for i in range(n_var)])\n" + "ad.AnnData(\n" + " X=np.zeros((n_obs, n_var), dtype='float32'), obs=obs, var=var\n" + ").write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + ], + ), + # -- export ------------------------------------------------------------- + Case( + name="export-array", + question="Write a dense element out as .npy.", + output_suffix=".npy", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "export", "array", "{input0}", "obsm/X_pca", + "-o", "{output}", + ], + ), + Contender( + "anndata (read_elem)", + script=_py( + "from anndata.io import read_elem\n" + "with h5py.File(IN[0], 'r') as f:\n" + " np.save(OUT, read_elem(f['obsm/X_pca']))\n" + ), + ), + Contender( + "anndata (full load)", + script=_py("np.save(OUT, ad.read_h5ad(IN[0]).obsm['X_pca'])\n"), + ), + ], + ), + Case( + name="export-sparse", + question="Write X out as Matrix Market text.", + output_suffix=".mtx", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "export", "sparse", "{input0}", "X", "-o", "{output}", + ], + ), + Contender( + "anndata (sparse_dataset)", + script=_py( + "import scipy.io as sio\n" + "from anndata.abc import CSRDataset\n" + "from anndata.io import sparse_dataset\n" + "with h5py.File(IN[0], 'r') as f:\n" + " sio.mmwrite(OUT, sparse_dataset(f['X'])[...])\n" + ), + ), + Contender( + "anndata (full load)", + script=_py( + "import scipy.io as sio\n" + "sio.mmwrite(OUT, ad.read_h5ad(IN[0]).X)\n" + ), + ), + ], + ), + # -- import ------------------------------------------------------------- + Case( + name="import-dataframe", + question="Replace obs from a CSV.", + sidecar="csv", + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "import", "dataframe", "{input0}", "obs", + "{sidecar_csv}", "-o", "{output}", + ], + ), + Contender( + "anndata", + script=_py( + "import pandas as pd, os\n" + "csv = os.environ['BENCH_SIDECAR_CSV']\n" + "obj = ad.read_h5ad(IN[0])\n" + "obj.obs = pd.read_csv(csv, index_col=0)\n" + "obj.write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + ], + ), + # -- concat, outer ------------------------------------------------------ + Case( + name="concat-outer", + question="Concatenate two stores keeping the union of variables.", + inputs=2, + contenders=[ + Contender( + "adata-cli", + argv=[ + "adata", "concat", "{input0}", "{input1}", + "-o", "{output}", "--join", "outer", + ], + ), + Contender( + "anndata (concat_on_disk)", + script=_py( + "from anndata.experimental import concat_on_disk\n" + "concat_on_disk(IN, OUT, join='outer')\n" + ), + ), + Contender( + "anndata (in memory)", + script=_py( + "parts = [ad.read_h5ad(p) for p in IN]\n" + "ad.concat(parts, join='outer')" + ".write_h5ad(OUT, compression=COMPRESSION)\n" + ), + ), + ], + ), # -- the claim itself ------------------------------------------------- Case( name="rss-vs-size", diff --git a/benchmarks/datasets.py b/benchmarks/datasets.py index 68c64e7..f3d9281 100644 --- a/benchmarks/datasets.py +++ b/benchmarks/datasets.py @@ -108,6 +108,11 @@ def _write(path: Path, tier: Tier, seed: int, *, n_var_columns: int = 2) -> Path ) obj = ad.AnnData(X=matrix, obs=obs, var=var) + # A dense element, so the `export array` / `import array` cases have + # something of realistic width to move. 50 columns is a typical PCA. + obj.obsm["X_pca"] = rng.standard_normal( + (tier.n_obs, 50), dtype="float32" + ) obj.write_h5ad(path, compression="lzf") del obj, matrix, blocks diff --git a/benchmarks/report.py b/benchmarks/report.py index 254af98..eebb5bb 100644 --- a/benchmarks/report.py +++ b/benchmarks/report.py @@ -120,8 +120,11 @@ def render(payload: Dict, previous: Optional[Dict] = None) -> str: f"| {record['contender']} | {wall} | {rss} | {output} | " f"{_delta(record, before.get((name, record['contender'])))} |" ) - if record.get("note"): - lines.append(f"| | *{record['note']}* | | | |") + note = record.get("note") or ( + record.get("stderr_tail") if record.get("status") == "n/a" else None + ) + if note: + lines.append(f"| | *{note}* | | | |") lines.append("") lines.append("---") diff --git a/benchmarks/run.py b/benchmarks/run.py index 6d7f501..102e566 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -38,9 +38,13 @@ from benchmarks.cases import CASES, Case, Contender #: Packages each baseline environment needs. +#: `dask` is there for `concat_on_disk`, which imports it to concatenate a +#: dense element such as obsm and raises ModuleNotFoundError without it. +#: Installing it is the fair thing to do -- the rule is to give the baseline +#: the best idiom it has -- and the dependency is itself worth knowing about. ENVIRONMENTS: Dict[str, List[str]] = { - "anndata": ["anndata", "scipy", "pandas", "h5py", "zarr"], - "scanpy": ["scanpy", "anndata", "scipy", "pandas", "h5py", "zarr"], + "anndata": ["anndata", "scipy", "pandas", "h5py", "zarr", "dask"], + "scanpy": ["scanpy", "anndata", "scipy", "pandas", "h5py", "zarr", "dask"], } @@ -93,6 +97,7 @@ def _run_contender( *, timeout_s: float, scripts_dir: Path, + shape: tuple, ) -> Measurement: workdir.mkdir(parents=True, exist_ok=True) output = workdir / f"{case.name}{case.output_suffix}" @@ -105,11 +110,21 @@ def _run_contender( elif stale.exists(): stale.unlink() + sidecar = _build_sidecar(case, inputs, workdir) + # The tier's shape, so a case that builds rather than reads -- `create` + # has no input to take its size from -- scales with the tier instead of + # making the smoke run as heavy as the ci one. substitutions = { "output": str(output), "outdir": str(outdir), + "sidecar_csv": str(sidecar) if sidecar else "", + "n_obs": str(shape[0]), + "n_var": str(shape[1]), **{f"input{i}": str(p) for i, p in enumerate(inputs)}, } + env = {"BENCH_N_OBS": str(shape[0]), "BENCH_N_VAR": str(shape[1])} + if sidecar: + env["BENCH_SIDECAR_CSV"] = str(sidecar) if contender.argv is not None: command = [part.format(**substitutions) for part in contender.argv] @@ -127,7 +142,33 @@ def _run_contender( # Read-only cases produce nothing to size. watched = None if (case.output_suffix == "" and target == output) else target - return measure(command, output=watched, timeout_s=timeout_s) + if shutil.which(command[0]) is None and not Path(command[0]).exists(): + # A contender whose binary is not installed -- `h5ls` comes with the + # HDF5 tools and is often absent. Missing is a result, not a crash. + return Measurement( + wall_s=0.0, maxrss_bytes=0, exit_code=127, status="n/a", + stderr_tail=f"{command[0]} is not installed", + ) + return measure(command, output=watched, timeout_s=timeout_s, env=env) + + +def _build_sidecar(case: Case, inputs: List[Path], workdir: Path) -> Optional[Path]: + """Build the auxiliary input a case declares, once per case. + + Derived from the real store rather than invented, so an import writes + back something the file could plausibly have held. + """ + if case.sidecar != "csv": + return None + path = workdir / f"{case.name}-sidecar.csv" + if path.exists(): + return path + import h5py + from anndata.io import read_elem + + with h5py.File(inputs[0], "r") as handle: + read_elem(handle["obs"]).to_csv(path) + return path def run( @@ -190,6 +231,11 @@ def run( interpreters, timeout_s=timeout_s, scripts_dir=scripts, + shape=( + (datasets.VAR_HEAVY.n_obs, datasets.VAR_HEAVY.n_var) + if case.var_heavy + else (tier.n_obs, tier.n_var) + ), ) ) if runs[-1].status != "ok": diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 8092b87..370d790 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -13,10 +13,17 @@ loading the whole thing is often faster, and the tables say so where it is true -- the trade is memory for time, and a table that hid the cost would be worth nothing. -A `ci` run taken during development, as an indication: `adata concat` of two -50,000 x 20,000 stores peaked at 202 MB against 1,778 MB for `ad.concat` in -memory and 439 MB for `anndata.experimental.concat_on_disk`, and was faster -than both. The published tables will say where it is slower, too. +A `ci` run taken during development, as an indication. `adata concat` of two +50,000 x 20,000 stores peaked at 202 MB against 1,778 MB for `ad.concat` in memory +and 439 MB for `anndata.experimental.concat_on_disk`, and was faster than both. +`adata create` used 78 MB against 2,225 MB. + +The tables also carry the rows where adata-cli loses, because those are the same +measurement. `export sparse` streams a 50,000 x 20,000 matrix in 68 MB and takes +10.8 s, where loading it whole takes 759 MB and 1.9 s; and `h5ls -r` lists the file +in 0.01 s and 7 MB against our 0.25 s and 63 MB, being C rather than a Python +process. The trade is memory for time, and a table that hid the cost would not be +worth publishing. To produce one locally: diff --git a/docs/TESTING.md b/docs/TESTING.md index 2fe6ae4..4dba3d1 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -78,10 +78,20 @@ h5py 3.15.1 and zarr 3.1.5: | Hook | Sees | |---|---| | `h5py.Dataset.__getitem__` | HDF5 reads | +| `h5py.Dataset.__setitem__` | HDF5 writes | +| `h5py.Group.create_dataset` | HDF5 writes made at creation | | `zarr.Array.__getitem__` | Zarr reads, v2 and v3 alike | | `zarr.storage.LocalStore.get` | chunk and metadata fetches | | `LocalStore.set` / `.delete` | write storms | +Both write seams are needed. `create_dataset(name, data=...)` writes its payload +during creation and never touches `__setitem__`, so hooking only the latter leaves +every `import` and `create` guard measuring zero — which is how they were first +written, and what the lower bound below caught. + +Use `io.reads` for a read path, `io.work` (reads plus writes) for `import` and +`create`. + The libraries are patched rather than anything in `src/adata/`. An in-repo seam would only see the call sites that remembered to use it, which is the wrong property for a guard meant to catch the read nobody thought of -- and @@ -127,6 +137,43 @@ produce n_obs x n_var_union cells. The axes worth separate coverage are `n_var`, `n_obs`, `n_inputs`, `n_obs_columns`, `n_var_columns`, `n_categories` and `n_groups`. +### Three invariants, not one + +`assert_grows_linearly` only catches super-linear growth. Most commands need exactly +that, but two claims this tool makes are stronger, and a linear guard would happily +accept a 64x increase in a command that is supposed to read nothing. + +| Helper | Claim | Used for | +|---|---|---| +| `assert_grows_linearly` | no worse than linear in one axis | most commands | +| `assert_grows_slower_than_input` | grows at least N times slower than the input | streaming at a fixed chunk | +| `assert_independent_of` | does not grow at all | inspection, and grouping work per row | + +**Inspection is free, and that is an exact number.** `view` and `ls` reach only +`.shape`, `.dtype` and `.attrs`: `axis_len` goes through `element_len`, which reads a +shape, and `_array_details` and `_infer_untagged` never touch a value. So the guard +asserts **zero** data elements read rather than a ratio — 0 against 0 is not a +meaningful ratio, and the moment inspection reads one column the answer stops being +zero however the store scales. A companion test exports the same fixture to prove +there was data there to read, so the zero cannot pass because the fixture was empty. + +**Streaming is bounded well below the input, which is weaker than it sounds and is +what the measurements support.** Over a 256x span at a fixed chunk: + +| | peak allocation growth | +|---|---| +| `export array` | 1.6x | +| `export dataframe` | 2.2x | +| `export sparse` | 6.5x | +| `subset` | 46x | + +Only `export array` is close to flat, so only it is held to a near-flat bound. None is +asserted as flat outright. `subset` is the weakest because obs columns are +materialised one at a time — a known gap that `benchmarks/` reports rather than these +guards conceal. Streamed `export sparse` is separately asserted to stay under a +quarter of what `--in-memory` costs, both measured in the same run so the factor holds +on any machine. + ### Where counting is not enough Two real defects in `concat.py` were invisible to all of the above, and both @@ -143,9 +190,23 @@ needed their own instrument: against the numeric column path measured in the same run -- self-calibrating, so it needs no hand-tuned budget and holds across interpreters. -`count_allocations` (tracemalloc) and `count_lines` (`sys.settrace`) are the -tools for these. The line tracer costs a 10-50x slowdown, so its tests carry -the `perf` marker and stay small. +A third case needed a third instrument. `split --by` grouped rows with +`np.nonzero(values == label)` inside a loop over distinct labels, rescanning each +chunk once per label: O(n_rows x n_groups), which at a million cells and a thousand +samples is 10^9 comparisons. No read counter moves, because the chunk is already in +memory; nothing lasting is allocated; and the Python line count per label is constant. +`count_scanned_elements` counts what is handed to numpy's scanning primitives and +makes it visible — 16,384 elements at 4 groups against 1,048,576 at 256. + +It is a **floor, not a measurement**: an operator such as `values == label` dispatches +to the ufunc in C and never passes the patched `np.equal`, and `arr.argsort()` is +invisible for the same reason. That is the right property for a guard — it can only +under-report — but it means a guard using it must also assert a lower bound, so that +under-reporting to nothing fails instead of passing. + +`count_allocations` (tracemalloc), `count_lines` (`sys.settrace`) and +`count_scanned_elements` are the tools for all of this. The line tracer costs a 10-50x +slowdown, so its tests carry the `perf` marker and stay small. ### What the guards deliberately do not claim @@ -154,6 +215,28 @@ O(`--chunk`). The streaming guarantee holds for X, not for obs annotation. The guards only stop that getting worse than linear; `benchmarks/` reports the actual curve. +### What is covered + +Every subcommand has a guard. When you add a command, add one: pick the axis its cost +should scale with, scale only that, and hold everything else fixed. + +| Command | Axis scaled | Invariant | +|---|---|---| +| `view`, `view --types`, `ls`, `ls --long`, `ls --plain` | n_obs, n_var | **zero** data reads | +| `create` | n_obs | linear in writes, both generated names and a name file | +| `concat` | n_var, n_obs, n_inputs, n_obs_columns, n_var_columns | linear; every `--merge` strategy and both joins | +| `concat --label`, `--index-unique` | n_obs | linear | +| `concat` category union | n_categories | bounded string comparisons | +| `subset` by name, by query | n_obs, n_var | linear | +| `split --by`, `--axis var` | n_groups | linear reads, **flat** scan work | +| `export dataframe` | n_obs, n_obs_columns | linear | +| `export array`, `sparse` (both paths), `dict`, `image` | elements, nnz, keys, pixels | linear | +| `import dataframe`, `array`, `sparse`, `dict`, `image` | rows, elements, nnz, keys, pixels | linear in writes | +| h5ad to zarr | n_obs | linear | + +The `slow` marker is on the three guards that build a 65,536-row store; they still gate +merges, and `-m "not slow"` skips them locally. `perf` is on the tracing guards. + ### Reading a failure It says cost grew super-linearly on the named axis. The assertion message @@ -179,17 +262,28 @@ seconds end to end on a laptop, so the full set with three repeats is minutes rather than the hour the workflow allows. The 90-minute timeout is headroom for `large`, not an estimate. -A `ci` run measured while writing this, for a sense of what the tables say: - -| | adata-cli | anndata (`concat_on_disk`) | anndata (in memory) | -|---|---|---|---| -| `concat-inner`, peak RSS | **202 MB** | 439 MB | 1,778 MB | -| `concat-inner`, wall time | **2.35 s** | 16.44 s | 3.69 s | -| `inspect`, peak RSS | **63 MB** | 138 MB (`read_elem`) | 556 MB (full load) | +A `ci` run measured while writing this, for a sense of what the tables say — and of +what they are for. Peak RSS first, wall time second: -Note that 202 MB is not flat in input size -- obs columns are read whole, and -a dense block is `--chunk` x n_var. The benchmark exists to keep that curve -visible rather than to assert a claim the code does not yet meet. +| Case | adata-cli | best baseline | +|---|---|---| +| `concat-inner` | **202 MB**, 2.35 s | 439 MB, 16.44 s (`concat_on_disk`) | +| `concat-outer` | **202 MB**, 2.60 s | 436 MB, **2.24 s** (`concat_on_disk`) | +| `inspect` | **63 MB**, 0.31 s | 138 MB, 0.79 s (`read_elem`) | +| `create` | **78 MB**, 0.28 s | 2,225 MB, 2.85 s | +| `import-dataframe` | **107 MB**, 0.41 s | 571 MB, 1.84 s | +| `export-sparse` | **68 MB**, 10.79 s | 759 MB, **1.90 s** (full load) | +| `ls` | 63 MB, 0.25 s | **7 MB, 0.01 s** (`h5ls -r`) | + +The last two rows are the reason the report is not a leaderboard. `export sparse` +streams in a tenth of the memory and takes five times as long; `h5ls` walks the file +in a hundredth of our time and a ninth of our memory, because it is C and does not +start a Python interpreter. Both belong in the table. A benchmark that only published +the rows we win would not be worth running. + +Note also that 202 MB is not flat in input size — obs columns are read whole, and a +dense block is `--chunk` x n_var. The benchmark exists to keep that curve visible +rather than to assert a claim the code does not yet meet. Results are published to `docs/BENCHMARKS.md` and `docs/benchmarks/.json` on every tag, and appended to the GitHub release notes if a release exists. @@ -231,6 +325,13 @@ decays fastest. - **Do not hide the rows where the baseline wins.** `_concat_csr` loops per row in Python and scipy's C `vstack` will often beat it on time at many times the memory. That trade is the argument for the tool. +- **Give the baseline everything it needs.** `concat_on_disk` imports `dask` to + concatenate a dense element and fails outright without it, so the baseline + environments install it. Measuring a library crippled by a missing optional + dependency would be measuring our own setup. +- **Four commands are deliberately not benchmarked.** `export image`, `export dict`, + `import image` and `import dict` have no library equivalent, so their rows would only + ever read `n/a` while adding runtime to every tag. The complexity guards cover them. - **Say whether the page cache was dropped.** A fixture written seconds ago is entirely in RAM, which understates streaming. The runner can drop it; a laptop usually cannot, and the report states which happened. From ab7de2c261b2e8e792d6bad3fb893fb0daa6f20c Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 21:12:24 +0100 Subject: [PATCH 11/14] Make the benchmarks page reachable, and survive being republished Three problems with docs/BENCHMARKS.md, all of which would have shown up first on the next tag. Nothing linked to it. Neither README.md nor docs/index.md mentioned the page at all, so the one document that says what the streaming claim actually costs was reachable only by guessing the URL. Both now link it, and the README's "streaming access to very large stores" bullet points straight at it -- including at the rows where loading the file outright is faster. `publish()` overwrote the whole page with bare tables. Every word explaining what the numbers mean would have been deleted the first time the workflow ran. The prose now lives in benchmarks/page_template.md with a `` marker, `build_page` fills it, and docs/BENCHMARKS.md is generated from that same template so the words exist once. A test asserts the marker is still there and that the framing and the trailing sections survive a republish. The rendered results opened with their own H1 and repeated the "peak RSS is the headline" paragraph the template already carries. Results are now an H2 under the page's own title, per-case tables are H3, and the duplicated framing is gone -- the published page has exactly one H1, which the test checks. The page itself now says what is measured, which four commands are deliberately absent and why, and carries the caveat that peak memory is not flat in input size: 1.6x to 46x over a 256x span depending on the command. Better for that to be on the page than discovered by a user. Co-Authored-By: Claude Opus 5 --- README.md | 5 ++- benchmarks/page_template.md | 50 +++++++++++++++++++++ benchmarks/report.py | 49 ++++++++++++++------ docs/BENCHMARKS.md | 80 +++++++++++++++++++++++---------- docs/index.md | 7 ++- tests/test_benchmark_harness.py | 24 ++++++++++ 6 files changed, 174 insertions(+), 41 deletions(-) create mode 100644 benchmarks/page_template.md diff --git a/README.md b/README.md index fb6c2af..60245c0 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ A command-line tool for exploring huge AnnData stores (`.h5ad` and `.zarr`) with ## Features -- Streaming access to very large `.h5ad` and `.zarr` stores +- Streaming access to very large `.h5ad` and `.zarr` stores — see [the benchmarks](docs/BENCHMARKS.md) for what that costs in practice, including where loading the file outright is faster - Auto-detects `.h5ad` files vs `.zarr` directories - Chunked processing for dense and sparse matrices (CSR/CSC) - Reads every AnnData on-disk layout, from 0.7.x through the current spec, and always writes the current one @@ -83,7 +83,8 @@ adata concat per_sample/*.h5ad -o merged.h5ad --join outer --label sample - [Get started](docs/GET_STARTED.md) — a short tutorial - [Command reference](docs/COMMANDS.md) — every command and flag - [Element spec: HDF5](docs/ELEMENTS_h5ad.md) / [Zarr](docs/ELEMENTS_zarr.md) — the on-disk format, and what this tool does with it -- [Testing](docs/TESTING.md) — how the suite is organised, and how compatibility is verified against six anndata releases +- [Testing](docs/TESTING.md) — how the suite is organised, how compatibility is verified against six anndata releases, and the complexity guards that keep cost regressions out +- [Benchmarks](docs/BENCHMARKS.md) — peak memory and wall time against anndata and scanpy, remeasured on every release - [Changelog](CHANGELOG.md) ## Docker diff --git a/benchmarks/page_template.md b/benchmarks/page_template.md new file mode 100644 index 0000000..20a07b8 --- /dev/null +++ b/benchmarks/page_template.md @@ -0,0 +1,50 @@ +# Benchmarks + +adata-cli measured against anndata, and against scanpy wherever scanpy has a real +equivalent. The [`Benchmark`](../.github/workflows/benchmark.yml) workflow rewrites +this page on every tag and keeps each run's raw numbers in `docs/benchmarks/`. + +**Peak RSS is the headline, not wall time.** This tool exists so that memory is set +by `--chunk` rather than by the size of the input. For data that fits in RAM, loading +the whole thing is often faster — so the tables below carry the rows where adata-cli +is the slower of the two, because those are the same measurement. The trade is memory +for time, and a table that hid the cost would not be worth publishing. + +## What is measured + +Fifteen cases on a GitHub-hosted runner: concat (inner, outer, `--merge same`), +subset by query, h5ad→zarr, `view`, `ls`, `create`, `export` dataframe / array / +sparse, `import dataframe`, `split`, and a peak-RSS-against-input-size sweep. Inputs +are 50,000 × 20,000 at 5% density, plus the 2,000 × 36,601 shape that hung in 0.5.1. + +Every run also reports a startup floor (`--version` for each contender), because the +CLI costs 0.3–1 s to import and `import scanpy` costs 3–8 s; on small inputs that is +the entire measurement. + +`export image`, `export dict`, `import image` and `import dict` are not here. No +library offers an equivalent, so the rows would only ever read `n/a`. They are covered +by the complexity guards instead. + + + +## A caveat on the streaming claim + +Peak memory is not flat in input size. Over a 256× span at a fixed chunk, peak +allocation grows 1.6× for `export array`, 2.2× for `export dataframe`, 6.5× for +`export sparse` and 46× for `subset` — far below the input curve, but not constant. +obs columns are read whole and a dense block is `--chunk` × n_var. The benchmark +exists to keep that curve visible rather than to assert a guarantee the code does +not yet meet. + +## Running it yourself + +```bash +uv run python -m benchmarks.run --tier smoke --out results.json +uv run python -m benchmarks.report results.json +``` + +Tiers are `smoke` (seconds), `ci` (50,000 × 20,000) and `large` (dispatch only, where +the in-memory baseline is expected to hit the address-space ceiling). + +See [TESTING.md](TESTING.md#performance) for how the baselines are kept honest, and +how this differs from the complexity guards that gate every pull request. diff --git a/benchmarks/report.py b/benchmarks/report.py index eebb5bb..9596dae 100644 --- a/benchmarks/report.py +++ b/benchmarks/report.py @@ -66,15 +66,10 @@ def render(payload: Dict, previous: Optional[Dict] = None) -> str: ) lines: List[str] = [] - lines.append(f"# Benchmark: `{payload['ref']}`") - lines.append("") - lines.append( - "adata-cli against anndata, and scanpy where it has a real " - "equivalent. **Peak RSS is the headline, not wall time.** This tool " - "exists so that memory is set by `--chunk` rather than by input size; " - "for anything that fits in RAM, loading the whole thing is often " - "faster, and the rows below say so where it is true." - ) + # H2, because this is embedded under the page's own H1 by `build_page`. + # The framing lives in `page_template.md`; repeating it here would print + # it twice on the published page. + lines.append(f"## Results — `{payload['ref']}`") lines.append("") n_obs, n_var, density = payload["tier_shape"] @@ -107,7 +102,7 @@ def render(payload: Dict, previous: Optional[Dict] = None) -> str: for name, records in grouped.items(): case = cases.get(name) - lines.append(f"## `{name}`") + lines.append(f"### `{name}`") lines.append("") if case: lines.append(f"{case.question}") @@ -188,7 +183,7 @@ def history_table(docs: Path) -> str: return "" lines = [ - "## History", + "### History", "", "`concat-inner` on adata-cli, run by run. Full results for each are " "in [`docs/benchmarks/`](benchmarks/).", @@ -205,8 +200,34 @@ def history_table(docs: Path) -> str: return "\n".join(lines) + "\n" +#: The standing prose of the docs page, with `` marking where +#: the tables go. Kept as a file rather than inline so the words that explain +#: the numbers live in one place and survive every republish -- an earlier +#: version overwrote the whole page with bare tables, which would have thrown +#: away the framing on the first tag. +PAGE_TEMPLATE = Path(__file__).with_name("page_template.md") + +#: What the template shows before any run has happened. +NOT_YET_RUN = ( + "## Results\n\nNo run has been published yet. The next tag fills this in; " + "until then, produce one locally with the commands below.\n" +) + + +def build_page(body: str) -> str: + """Wrap rendered results in the page's standing explanation.""" + template = PAGE_TEMPLATE.read_text() + if "" not in template: # pragma: no cover - template edited + return template.rstrip() + "\n\n" + body + return template.replace("", body.strip()) + + def publish(payload: Dict, results: Path, docs: Path) -> Path: - """Copy the raw results in and rewrite the docs page.""" + """Copy the raw results in and rewrite the docs page. + + Only the results section is replaced; everything explaining what the + numbers mean comes from `page_template.md` and is reinstated every time. + """ store = docs / "benchmarks" store.mkdir(parents=True, exist_ok=True) ref = payload["ref"].replace("/", "-") @@ -214,7 +235,9 @@ def publish(payload: Dict, results: Path, docs: Path) -> Path: previous = _previous(store, skip=f"{ref}.json") page = docs / "BENCHMARKS.md" - page.write_text(render(payload, previous) + "\n" + history_table(docs)) + page.write_text( + build_page(render(payload, previous) + "\n" + history_table(docs)) + ) return page diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 370d790..713da31 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -1,37 +1,69 @@ # Benchmarks -This page is written by the [`Benchmark`](../.github/workflows/benchmark.yml) -workflow on every tag. Until the next release it stands empty. +adata-cli measured against anndata, and against scanpy wherever scanpy has a real +equivalent. The [`Benchmark`](../.github/workflows/benchmark.yml) workflow rewrites +this page on every tag and keeps each run's raw numbers in `docs/benchmarks/`. -What will appear here: adata-cli against anndata, and against scanpy wherever -scanpy has a real equivalent, measured on a GitHub-hosted runner at 50,000 -obs x 20,000 var plus the 2,000 x 36,601 shape that hung in 0.5.1. +**Peak RSS is the headline, not wall time.** This tool exists so that memory is set +by `--chunk` rather than by the size of the input. For data that fits in RAM, loading +the whole thing is often faster — so the tables below carry the rows where adata-cli +is the slower of the two, because those are the same measurement. The trade is memory +for time, and a table that hid the cost would not be worth publishing. -**Peak RSS is the headline, not wall time.** This tool exists so that memory -is set by `--chunk` rather than by input size. For data that fits in RAM, -loading the whole thing is often faster, and the tables say so where it is -true -- the trade is memory for time, and a table that hid the cost would be -worth nothing. +## What is measured -A `ci` run taken during development, as an indication. `adata concat` of two -50,000 x 20,000 stores peaked at 202 MB against 1,778 MB for `ad.concat` in memory -and 439 MB for `anndata.experimental.concat_on_disk`, and was faster than both. -`adata create` used 78 MB against 2,225 MB. +Fifteen cases on a GitHub-hosted runner: concat (inner, outer, `--merge same`), +subset by query, h5ad→zarr, `view`, `ls`, `create`, `export` dataframe / array / +sparse, `import dataframe`, `split`, and a peak-RSS-against-input-size sweep. Inputs +are 50,000 × 20,000 at 5% density, plus the 2,000 × 36,601 shape that hung in 0.5.1. -The tables also carry the rows where adata-cli loses, because those are the same -measurement. `export sparse` streams a 50,000 x 20,000 matrix in 68 MB and takes -10.8 s, where loading it whole takes 759 MB and 1.9 s; and `h5ls -r` lists the file -in 0.01 s and 7 MB against our 0.25 s and 63 MB, being C rather than a Python -process. The trade is memory for time, and a table that hid the cost would not be -worth publishing. +Every run also reports a startup floor (`--version` for each contender), because the +CLI costs 0.3–1 s to import and `import scanpy` costs 3–8 s; on small inputs that is +the entire measurement. -To produce one locally: +`export image`, `export dict`, `import image` and `import dict` are not here. No +library offers an equivalent, so the rows would only ever read `n/a`. They are covered +by the complexity guards instead. + +## Results + +No tagged run has been published yet — the next tag fills this in. These +numbers are from a `ci` run during development, as an indication of what the +tables will say. + +| Case | adata-cli | best baseline | +|---|---|---| +| `concat-inner` | **202 MB**, 2.35 s | 439 MB, 16.44 s (`concat_on_disk`) | +| `concat-outer` | **202 MB**, 2.60 s | 436 MB, **2.24 s** (`concat_on_disk`) | +| `create` | **78 MB**, 0.28 s | 2,225 MB, 2.85 s | +| `inspect` | **63 MB**, 0.31 s | 138 MB, 0.79 s (`read_elem`) | +| `import-dataframe` | **107 MB**, 0.41 s | 571 MB, 1.84 s | +| `export-sparse` | **68 MB**, 10.79 s | 759 MB, **1.90 s** (full load) | +| `ls` | 63 MB, 0.36 s | **7 MB, 0.03 s** (`h5ls -r`) | + +The last two rows are why this is not a leaderboard. `export sparse` streams in a +tenth of the memory and takes five times as long. `h5ls -r` walks the file in a +hundredth of our time and a ninth of our memory, being a C program rather than a +Python process that must import typer, rich, h5py and zarr before it starts. + +## A caveat on the streaming claim + +Peak memory is not flat in input size. Over a 256× span at a fixed chunk, peak +allocation grows 1.6× for `export array`, 2.2× for `export dataframe`, 6.5× for +`export sparse` and 46× for `subset` — far below the input curve, but not constant. +obs columns are read whole and a dense block is `--chunk` × n_var. The benchmark +exists to keep that curve visible rather than to assert a guarantee the code does +not yet meet. + +## Running it yourself ```bash uv run python -m benchmarks.run --tier smoke --out results.json uv run python -m benchmarks.report results.json ``` -See [TESTING.md](TESTING.md#performance) for what is measured, how the -baselines are kept honest, and how this differs from the complexity guards -that run on every pull request. +Tiers are `smoke` (seconds), `ci` (50,000 × 20,000) and `large` (dispatch only, where +the in-memory baseline is expected to hit the address-space ceiling). + +See [TESTING.md](TESTING.md#performance) for how the baselines are kept honest, and +how this differs from the complexity guards that gate every pull request. diff --git a/docs/index.md b/docs/index.md index b37ac86..92289f3 100644 --- a/docs/index.md +++ b/docs/index.md @@ -34,8 +34,11 @@ docker run --rm -it -v /path/to/data:/data \ out in `.h5ad`, and what this tool does with it. - **[Element spec: Zarr](ELEMENTS_zarr.md)** — the same for `.zarr`, including the v2/v3 differences. -- **[Testing](TESTING.md)** — how the suite is organised, and how compatibility - is verified against six real anndata releases. +- **[Testing](TESTING.md)** — how the suite is organised, how compatibility is + verified against six real anndata releases, and the complexity guards that + keep cost regressions out. +- **[Benchmarks](BENCHMARKS.md)** — peak memory and wall time against anndata + and scanpy, remeasured and republished on every release. ## At a glance diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py index 092c3b2..1d406bc 100644 --- a/tests/test_benchmark_harness.py +++ b/tests/test_benchmark_harness.py @@ -102,3 +102,27 @@ def test_maxrss_is_normalised_to_bytes(): assert isinstance(result, Measurement) # Any CPython start-up is well over 1 MB and well under 4 GB. assert 1024**2 < result.maxrss_bytes < 4 * 1024**3, result.maxrss_bytes + + +def test_publishing_keeps_the_page_prose(): + """A republish must not reduce the docs page to bare tables. + + `publish` rewrites `docs/BENCHMARKS.md` in full on every tag. An earlier + version wrote only the rendered results, which would have thrown away + everything explaining what the numbers mean the first time it ran. + """ + from benchmarks.report import PAGE_TEMPLATE, build_page + + template = PAGE_TEMPLATE.read_text() + assert "" in template, ( + "the results marker is gone from page_template.md, so results would " + "be appended rather than placed" + ) + + page = build_page("## Results\n\nsome tables\n") + assert "some tables" in page + assert "Peak RSS is the headline" in page, "the framing was dropped" + assert "Running it yourself" in page, "the trailing sections were dropped" + assert page.count("\n# ") + page.startswith("# ") == 1, ( + "the published page must have exactly one H1" + ) From 7c38c41187a96399ff8979df2522c10a32033126 Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 21:18:54 +0100 Subject: [PATCH 12/14] Restore the blank line before the 0.5.1 heading Fallout from resolving the CHANGELOG merge by dropping conflict markers. Co-Authored-By: Claude Opus 5 --- .gitignore | 6 ++++++ CHANGELOG.md | 1 + 2 files changed, 7 insertions(+) diff --git a/.gitignore b/.gitignore index fb49bae..8626410 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,9 @@ htmlcov/ .pytest_cache/ pytest-results*.xml compat-results.xml + +# Benchmark working directory and results -- the defaults of +# `python -m benchmarks.run`, both written into the repo root. +.bench/ +results.json +summary.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d6f2ed..80d1d78 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -53,6 +53,7 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no - **`--merge drop` and `--uns-merge drop` are accepted.** `drop` was already the documented default behaviour but was rejected as a value, so a config could not state it explicitly. + ## 0.5.1 Makes the container image usable from Nextflow, and stops `copy_dataset` From 4db7d3383e019e4822eb45be166ef26342777369 Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 21:51:40 +0100 Subject: [PATCH 13/14] Measure peak RSS from a shim, so it is the child's own and not the runner's CI failed one test on both interpreters, and it was not a flaky threshold. test_peak_rss_is_attributed_to_the_right_child reported 146 MB for a child that allocated 1 KB, after a 410 MB child. Its docstring says why it exists: if peak RSS were attributed to the wrong process, every table this project publishes would be wrong and the numbers would still look plausible. A uniform interpreter floor was ruled out by the data -- at 140 MB the big child would have measured 540 MB, not 410. Reproduced under python:3.12-slim: parent 14.7 MB -> no-op child 11.8 MB parent 329.6 MB -> no-op child 326.4 MB parent 329.6 MB -> via shim 8.1 MB On Linux a forked child inherits its parent's resident pages, and execve folds that pre-exec high-water mark into the accumulated maxrss that wait4 reports. A child of a fat parent cannot appear small. macOS resets it at exec, which is why this passed locally and failed on CI -- the one platform the benchmark actually runs on. This was never only a test problem. run.py imports anndata, pandas and numpy to build fixtures in the same process that calls measure(), so on the runner every contender would have been floored at roughly 200 MB. The headline result -- 202 MB against 1,847 MB for ad.concat -- would have collapsed to "everything costs about the same", which is precisely the claim the benchmark exists to test, failing silently in the flattering direction. measure() now re-invokes benchmarks/_measure.py as a subprocess and that freshly-exec'd interpreter, about 8 MB, forks the command being measured. The in-process logic is unchanged, renamed _measure_here; main() grew the --timeout, --memory-limit, --cwd and --env options it needs to carry the call. RLIMIT_AS still applies via the shim's preexec_fn, so OOM containment for the large tier is intact. Dropping preexec_fn to get posix_spawn was the obvious alternative and does not work: the middle row above measures 329.5 MB that way too. It would also have given up the address-space ceiling. The test now asserts the property rather than a ratio. It holds 300 MB of ballast for its duration, measures a no-op baseline child, and requires both that the baseline is small in absolute terms -- the absence of an inherited floor -- and that the small child resembles the baseline rather than the big one before it. Against the old code on Linux it fails at 342 MB for a no-op child; the previous form only failed when the floor happened to exceed a quarter of the largest child. Two things found on the way: build_environments only checked that the interpreter existed, so a reused --work directory kept whatever was installed first. Adding dask to ENVIRONMENTS had no effect on an existing tree and concat_on_disk went on raising ModuleNotFoundError as though that were a finding about anndata. The package list is now recorded beside the venv and triggers a rebuild when it changes, with --clear so the rebuild does not die on the existing tree. A command that does not exist is classified n/a by the shim with the reason, rather than surfacing as a traceback; the duplicate check in run.py is gone. The refuse-to-overwrite guard stays in the parent so it still raises. Re-measured at the ci tier afterwards: peak RSS is unchanged to within a megabyte across every documented case, as expected since those figures came from macOS. They stand. 905 tests pass, coverage 92.78%. All 15 benchmark cases clean at smoke; the harness tests pass under python:3.12-slim, where they now also cover the RLIMIT_AS path that macOS skips. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 9 +++ benchmarks/_measure.py | 134 ++++++++++++++++++++++++++++---- benchmarks/page_template.md | 9 +++ benchmarks/run.py | 44 ++++++++--- docs/BENCHMARKS.md | 9 +++ docs/TESTING.md | 6 ++ tests/test_benchmark_harness.py | 72 ++++++++++++++--- 7 files changed, 246 insertions(+), 37 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 80d1d78..981c131 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,15 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no Two claims are now enforced rather than described -- `view` and `ls` read **zero** data elements at any store size, and streaming stays far below the input curve at a fixed `--chunk`. +- **Peak RSS in the benchmark was floored by the runner's own memory on Linux.** + A forked child inherits its parent's resident pages and `execve` folds that + into the `maxrss` the kernel reports, so every contender would have measured + at least what `benchmarks/run.py` used to build the fixtures — around + 200 MB — and the tables would have read "everything costs about the same". + Commands are now forked from a small shim: with a 330 MB parent, a no-op + child goes from 326 MB to 8 MB. Caught by `test_benchmark_harness.py`, which + exists for exactly this. The published figures were measured on macOS, which + resets the high-water mark at exec, and are unchanged. - **A comparative benchmark** (`benchmarks/`), run on every tag against anndata and against scanpy where scanpy has a real equivalent. Reports peak RSS, wall time and output size; publishes to diff --git a/benchmarks/_measure.py b/benchmarks/_measure.py index c421e41..072f196 100644 --- a/benchmarks/_measure.py +++ b/benchmarks/_measure.py @@ -5,18 +5,42 @@ wall time would misrepresent it -- for anything that fits in RAM, loading the whole thing is usually faster. -Measuring it correctly needs `os.wait4`, not `resource.getrusage`. -`RUSAGE_CHILDREN` is a running maximum over every child the process has ever -reaped: run a 400 MB case and then a 17 MB one and it still reports 400 MB. -`wait4` returns rusage for one specific child. (Checked on this machine: -wait4 gives 435 MB then 17 MB where RUSAGE_CHILDREN stays at 435 MB.) +Two things have to be right for the number to mean anything, and both were +wrong at some point. + +**Use `os.wait4`, not `resource.getrusage`.** `RUSAGE_CHILDREN` is a running +maximum over every child the process has ever reaped: run a 400 MB case and +then a 17 MB one and it still reports 400 MB. `wait4` returns rusage for one +specific child. + +**Fork the child from a small process.** On Linux a forked child inherits its +parent's resident pages, and `execve` folds that pre-exec high-water mark into +the accumulated `maxrss` that `wait4` reports. So a child of a fat parent can +never appear small. Measured under python:3.12-slim: + + parent 14.7 MB -> no-op child 11.8 MB + parent 329.6 MB -> no-op child 326.4 MB <- the parent's RSS, not the child's + parent 329.6 MB -> via shim 8.1 MB + +macOS resets the high-water mark at exec and shows none of this, which is why +it went unnoticed locally and failed on CI. It matters because `run.py` +imports anndata, pandas and numpy to build fixtures in the same process that +measures, so every contender would have been floored at ~200 MB and the +tables would have read "everything costs about the same". + +`posix_spawn` does not help -- the middle row above is 329.5 MB that way too. +The fix is the shim: `measure()` re-invokes this file as a subprocess, and +that freshly-exec'd interpreter, about 8 MB, is what forks the command being +measured. """ from __future__ import annotations +import argparse import json import os import resource +import shutil import subprocess import sys import tempfile @@ -112,7 +136,57 @@ def measure( env: Optional[Dict[str, str]] = None, cwd: Optional[Path] = None, ) -> Measurement: - """Run `command`, returning its wall time, peak RSS and output size.""" + """Run `command` from a small shim, and report what it cost. + + The shim is this same file, re-invoked. It exists so that the process + which forks `command` is a bare interpreter rather than whatever imported + anndata -- see the module docstring for the measurement it fixes. + """ + # Checked here rather than in the shim: a stale output would be sized and + # reported as this run's work, which is a bug in the caller and should be + # loud rather than turned into a "failed" row. + if output is not None and output.exists(): + raise FileExistsError(f"{output} exists; the harness never overwrites.") + + argv = [sys.executable, str(Path(__file__).resolve()), "--timeout", str(timeout_s)] + if output is not None: + argv += ["--output", str(output)] + if memory_limit is None: + argv += ["--no-memory-limit"] + else: + argv += ["--memory-limit", str(memory_limit)] + if cwd is not None: + argv += ["--cwd", str(cwd)] + for key, value in (env or {}).items(): + argv += ["--env", f"{key}={value}"] + argv += ["--", *command] + + done = subprocess.run(argv, capture_output=True, text=True) + try: + return Measurement(**json.loads(done.stdout)) + except (json.JSONDecodeError, TypeError, ValueError): + # The shim itself failed. Report it rather than crashing the run, and + # keep enough of its output to diagnose. + return Measurement( + wall_s=0.0, + maxrss_bytes=0, + exit_code=done.returncode, + status="failed", + stderr_tail="measurement shim failed:\n" + + "\n".join((done.stderr or done.stdout).strip().splitlines()[-6:]), + ) + + +def _measure_here( + command: Sequence[str], + *, + output: Optional[Path] = None, + timeout_s: float = 900.0, + memory_limit: Optional[int] = DEFAULT_MEMORY_LIMIT, + env: Optional[Dict[str, str]] = None, + cwd: Optional[Path] = None, +) -> Measurement: + """Run `command` in this process's own child. Only the shim calls this.""" def limit() -> None: # pragma: no cover - runs in the forked child if memory_limit is None: @@ -130,6 +204,12 @@ def limit() -> None: # pragma: no cover - runs in the forked child full_env = {**os.environ, **(env or {})} + if shutil.which(command[0]) is None and not Path(command[0]).exists(): + return Measurement( + wall_s=0.0, maxrss_bytes=0, exit_code=127, status="n/a", + stderr_tail=f"{command[0]} is not installed", + ) + # Output goes to temporary files, not pipes. A pipe would have to be # drained with communicate(), and communicate() reaps the child -- after # which wait4 has nothing to report and the per-child rusage is lost. @@ -205,17 +285,37 @@ def _looks_like_oom(stderr: str) -> bool: return any(m in stderr for m in markers) -def main(argv: List[str]) -> int: # pragma: no cover - CLI entry - """Measure a command given after `--`, printing JSON.""" - if "--" not in argv: - print("usage: python -m benchmarks._measure [--output P] -- CMD...") - return 2 - split = argv.index("--") - head, command = argv[:split], argv[split + 1 :] - output = None - if "--output" in head: - output = Path(head[head.index("--output") + 1]) - print(json.dumps(measure(command, output=output).as_dict(), indent=2)) +def main(argv: List[str]) -> int: # pragma: no cover - runs as the shim + """Measure the command after `--` and print one JSON object. + + This is the shim `measure()` re-invokes; it is also usable by hand: + + python benchmarks/_measure.py --timeout 60 -- adata view f.h5ad + """ + parser = argparse.ArgumentParser(description="Measure one command.") + parser.add_argument("--output", type=Path) + parser.add_argument("--timeout", type=float, default=900.0) + parser.add_argument("--memory-limit", type=int, default=DEFAULT_MEMORY_LIMIT) + parser.add_argument("--no-memory-limit", action="store_true") + parser.add_argument("--cwd", type=Path) + parser.add_argument("--env", action="append", default=[]) + parser.add_argument("command", nargs=argparse.REMAINDER) + args = parser.parse_args(argv) + + command = args.command[1:] if args.command[:1] == ["--"] else args.command + if not command: + parser.error("no command given after --") + + extra = dict(pair.split("=", 1) for pair in args.env) + result = _measure_here( + command, + output=args.output, + timeout_s=args.timeout, + memory_limit=None if args.no_memory_limit else args.memory_limit, + env=extra or None, + cwd=args.cwd, + ) + print(json.dumps(result.as_dict())) return 0 diff --git a/benchmarks/page_template.md b/benchmarks/page_template.md index 20a07b8..4964e0e 100644 --- a/benchmarks/page_template.md +++ b/benchmarks/page_template.md @@ -27,6 +27,15 @@ by the complexity guards instead. +## How peak memory is measured + +Each measured command is forked from a small shim process, not from the benchmark +runner itself. On Linux a forked child inherits its parent's resident pages and +`execve` folds that into the `maxrss` the kernel reports, so a child of a fat parent +cannot appear small: with the runner holding 330 MB, a process allocating nothing +measured 326 MB. The shim brings that floor down to about 8 MB, uniform across every +contender and visible in the `startup` row. + ## A caveat on the streaming claim Peak memory is not flat in input size. Over a 256× span at a fixed chunk, peak diff --git a/benchmarks/run.py b/benchmarks/run.py index 102e566..290667b 100644 --- a/benchmarks/run.py +++ b/benchmarks/run.py @@ -49,18 +49,45 @@ def build_environments(root: Path, wanted: List[str]) -> Dict[str, Path]: - """Create one venv per baseline and return its interpreter.""" + """Create one venv per baseline and return its interpreter. + + The package list is recorded beside the venv and rebuilt when it changes. + Checking only that the interpreter exists means a reused `--work` + directory keeps whatever was installed the first time: adding `dask` to + ENVIRONMENTS had no effect on an existing tree, and `concat_on_disk` + went on failing with ModuleNotFoundError as though that were a finding + about anndata. + """ interpreters: Dict[str, Path] = {} for name in wanted: venv = root / name python = venv / "bin" / "python" - if not python.exists(): - print(f"[env] building {name}", flush=True) - subprocess.run(["uv", "venv", str(venv)], check=True, capture_output=True) + stamp = venv / ".packages" + wanted_packages = "\n".join(sorted(ENVIRONMENTS[name])) + current = stamp.read_text() if stamp.exists() else None + + if not python.exists() or current != wanted_packages: + if current is not None and current != wanted_packages: + print(f"[env] {name}: package list changed, rebuilding", flush=True) + else: + print(f"[env] building {name}", flush=True) + # --clear because a rebuild runs over an existing tree; without + # it `uv venv` refuses and the run dies before any measurement. + done = subprocess.run( + ["uv", "venv", "--clear", str(venv)], + capture_output=True, + text=True, + ) + if done.returncode != 0: + raise RuntimeError( + f"could not create the {name} environment at {venv}:\n" + + (done.stderr or done.stdout) + ) subprocess.run( ["uv", "pip", "install", "--python", str(python), *ENVIRONMENTS[name]], check=True, ) + stamp.write_text(wanted_packages) interpreters[name] = python return interpreters @@ -142,13 +169,8 @@ def _run_contender( # Read-only cases produce nothing to size. watched = None if (case.output_suffix == "" and target == output) else target - if shutil.which(command[0]) is None and not Path(command[0]).exists(): - # A contender whose binary is not installed -- `h5ls` comes with the - # HDF5 tools and is often absent. Missing is a result, not a crash. - return Measurement( - wall_s=0.0, maxrss_bytes=0, exit_code=127, status="n/a", - stderr_tail=f"{command[0]} is not installed", - ) + # A contender whose binary is not installed -- `h5ls` ships with the HDF5 + # tools and is often absent -- comes back as `n/a` from `measure`. return measure(command, output=watched, timeout_s=timeout_s, env=env) diff --git a/docs/BENCHMARKS.md b/docs/BENCHMARKS.md index 713da31..de52293 100644 --- a/docs/BENCHMARKS.md +++ b/docs/BENCHMARKS.md @@ -46,6 +46,15 @@ tenth of the memory and takes five times as long. `h5ls -r` walks the file in a hundredth of our time and a ninth of our memory, being a C program rather than a Python process that must import typer, rich, h5py and zarr before it starts. +## How peak memory is measured + +Each measured command is forked from a small shim process, not from the benchmark +runner itself. On Linux a forked child inherits its parent's resident pages and +`execve` folds that into the `maxrss` the kernel reports, so a child of a fat parent +cannot appear small: with the runner holding 330 MB, a process allocating nothing +measured 326 MB. The shim brings that floor down to about 8 MB, uniform across every +contender and visible in the `startup` row. + ## A caveat on the streaming claim Peak memory is not flat in input size. Over a 256× span at a fixed chunk, peak diff --git a/docs/TESTING.md b/docs/TESTING.md index 4dba3d1..d363cb1 100644 --- a/docs/TESTING.md +++ b/docs/TESTING.md @@ -291,6 +291,12 @@ Artifacts expire; the docs page is the durable series. ### How it measures +- **Every command is forked from a small shim, not from the runner.** On Linux a + forked child inherits its parent's resident pages and `execve` folds that into the + reported `maxrss`, so a child of a process that has imported anndata cannot measure + below roughly 200 MB — which would have flattened every row in the table. Measured + with a 330 MB parent: 326 MB for a no-op child directly, 8 MB through the shim. + `posix_spawn` does not help (329 MB), so dropping `preexec_fn` is not a fix. - **`os.wait4`, not `resource.getrusage`.** `RUSAGE_CHILDREN` is a running maximum over every child ever reaped, so one large case would poison every later row. Output goes to temporary files rather than pipes, because diff --git a/tests/test_benchmark_harness.py b/tests/test_benchmark_harness.py index 1d406bc..7bc384b 100644 --- a/tests/test_benchmark_harness.py +++ b/tests/test_benchmark_harness.py @@ -16,23 +16,66 @@ from benchmarks._measure import Measurement, _looks_like_oom, measure -def test_peak_rss_is_attributed_to_the_right_child(): - """The reason the harness uses `os.wait4` rather than `getrusage`. +#: A child that allocates nothing should cost about what a bare interpreter +#: costs. Generous, because that depends on the build; far below the ballast +#: below, which is the point. +FLOOR_CEILING = 60 * 1024**2 + +#: Allocated by the parent before measuring, to make an inherited floor +#: impossible to miss. +BALLAST = 300 * 1024**2 + + +def test_peak_rss_is_the_child_s_own_and_carries_no_floor(): + """Two ways this number can lie, both of which it has. `RUSAGE_CHILDREN` is a running maximum over every child a process has - ever reaped, so a 400 MB case followed by a 1 kB one would report 400 MB - twice and every later row would inherit the largest earlier peak. + reaped, so a 400 MB case makes every later one look like 400 MB. `wait4` + fixes that. + + Worse, and what CI caught: on Linux a forked child inherits its parent's + resident pages, and `execve` folds that into the `maxrss` `wait4` reports. + A child of a fat parent could not appear small however little it used. + Measured under python:3.12-slim before the fix -- parent at 329.6 MB, a + no-op child reported 326.4 MB. `posix_spawn` reported 329.5 MB, so + dropping `preexec_fn` would not have helped; measuring from a shim + reported 8.1 MB. + + macOS resets the high-water mark at exec and shows none of this, so this + is mostly a Linux guard -- which is the platform the benchmark runs on. + + The test holds ballast in the parent for its whole duration so that an + inherited floor cannot hide. """ + ballast = bytearray(BALLAST) + ballast[::4096] = b"\x01" * len(ballast[::4096]) # make it resident + + baseline = measure([sys.executable, "-c", "pass"]) big = measure([sys.executable, "-c", "x = bytearray(400 * 1024 * 1024)"]) small = measure([sys.executable, "-c", "x = bytearray(1024)"]) - assert big.status == "ok" and small.status == "ok" + assert {baseline.status, big.status, small.status} == {"ok"} + + # No floor: a child that allocates nothing looks like nothing, even though + # this process is holding 300 MB. + assert baseline.maxrss_bytes < FLOOR_CEILING, ( + f"a no-op child reported {baseline.maxrss_bytes / 1e6:.0f} MB while " + f"this process held {BALLAST / 1e6:.0f} MB of ballast. The measured " + "peak is inheriting the parent's resident pages, so every benchmark " + "row would be floored at roughly the runner's own footprint." + ) + + # Isolation: the small child resembles the baseline, not the big one that + # ran before it. assert big.maxrss_bytes > 300 * 1024**2, big - assert small.maxrss_bytes < big.maxrss_bytes / 4, ( - f"peak RSS leaked between children: {small.maxrss_bytes} after " - f"{big.maxrss_bytes}" + assert small.maxrss_bytes < baseline.maxrss_bytes * 2, ( + f"peak RSS leaked between children: {small.maxrss_bytes / 1e6:.0f} MB " + f"for a 1 KB allocation, against a {baseline.maxrss_bytes / 1e6:.0f} MB " + f"baseline, after a {big.maxrss_bytes / 1e6:.0f} MB child" ) + del ballast + def test_a_hang_is_reported_as_a_timeout_with_its_output_size(tmp_path): """How the 0.5.1 hang presented: still running, output never growing. @@ -84,13 +127,24 @@ def test_an_ordinary_failure_is_not_mistaken_for_an_oom(): def test_the_harness_refuses_to_overwrite_an_existing_output(tmp_path): - """A stale output would be sized and reported as this run's work.""" + """A stale output would be sized and reported as this run's work. + + Checked in the parent, not the shim, so it surfaces as an exception the + caller must fix rather than as a quiet "failed" row. + """ existing = tmp_path / "already.h5ad" existing.write_bytes(b"old") with pytest.raises(FileExistsError): measure([sys.executable, "-c", "pass"], output=existing) +def test_a_missing_binary_is_a_result_not_a_crash(): + """`h5ls` ships with the HDF5 tools and is often simply absent.""" + result = measure(["/nonexistent-binary-for-tests"]) + assert result.status == "n/a" + assert "not installed" in result.stderr_tail + + def test_maxrss_is_normalised_to_bytes(): """`ru_maxrss` is KiB on Linux and bytes on macOS; the field is bytes. From e8c24fbcd40a8adfc34f96fb1ae35724dd0b23f0 Mon Sep 17 00:00:00 2001 From: Aljes Date: Wed, 23 Sep 2026 22:15:04 +0100 Subject: [PATCH 14/14] Size string reads from the real element width, not an assumed 64 bytes Flagged by an automated review on PR #14, on code that came in with the main merge -- 4859445, released in 0.5.1. Confirmed by measurement, and worse than the review described. _row_bytes assumed 64 bytes for a variable-length element, because h5py reports the itemsize of a pointer. An assumption is not a bound. The step was therefore the same 524,288 elements whatever the data actually held: 16 B strings, 200k rows -> 11.5 MB peak 256 B -> 59.4 MB 1 KiB -> 213.0 MB 4 KiB -> 827.4 MB against a stated 32 MiB budget and because the computed step exceeded the row count in every one of those cases, the whole array was read in a single go -- the opposite of what the budget is for. Real cell and gene names do sit under 64 bytes, but copy_tree carries arbitrary `uns` content, so this is not a width the storage layer can assume. The width is now sampled from the first 256 elements, one small read against a copy about to stream the whole array. VLEN_ELEMENT_BYTES stays as the fallback when sampling is not possible, and as a floor so a column of empty strings cannot produce an unbounded step. _row_bytes keeps its old two- argument form via a default, so the existing 0.5.1 tests still describe it. Two guards, because they catch different things. The parametrised one checks the arithmetic directly -- step x width must stay inside the budget at 16, 256 and 4096 bytes -- which is exact and costs nothing. The slow one checks it in practice on a 164 MB array, and is the one that fails against the old code, at 165 MB where the bound is 100 MB. This is the path the rest of test_performance.py did not reach, and every `copy:` task in subset and every uns entry goes through it. The coverage claimed in the previous commit was not as complete as it read. 909 tests pass, coverage 92.79%. Co-Authored-By: Claude Opus 5 --- CHANGELOG.md | 9 ++++ src/adata/storage/__init__.py | 86 +++++++++++++++++++++++++++----- tests/test_performance.py | 93 +++++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 981c131..437917c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,15 @@ Notable changes to `adata-cli`. Versions are `MAJOR.MINOR.PATCH`; tags carry no Two claims are now enforced rather than described -- `view` and `ls` read **zero** data elements at any store size, and streaming stays far below the input curve at a fixed `--chunk`. +- **Copying variable-length strings ignored its own read budget.** The width + of a vlen element was assumed to be 64 bytes, because h5py reports the + itemsize of a pointer, so the step was the same 524,288 elements whatever + the data held: 2 GiB per read at 4 KiB elements against a stated 32 MiB + budget, and for any array shorter than that step, the whole array in one + go. Copying 200,000 strings of 4 KiB peaked at 827 MB. The width is now + sampled from the first 256 elements. Reported by an automated review on + PR #14 and confirmed by measurement; `uns` can hold arbitrary text, so this + was not a width the layer could assume. - **Peak RSS in the benchmark was floored by the runner's own memory on Linux.** A forked child inherits its parent's resident pages and `execve` folds that into the `maxrss` the kernel reports, so every contender would have measured diff --git a/src/adata/storage/__init__.py b/src/adata/storage/__init__.py index cfea148..025bce5 100644 --- a/src/adata/storage/__init__.py +++ b/src/adata/storage/__init__.py @@ -450,30 +450,86 @@ def _is_string_src(src: Any) -> bool: #: comfortably above a typical 1 MiB Lustre stripe while bounding peak memory. TARGET_READ_BYTES = 32 * 1024 * 1024 -#: Assumed width of a variable-length string element, for read sizing only. -#: h5py reports `itemsize` 8 for vlen strings because the value is a pointer, -#: which would overestimate the row count by an order of magnitude and break -#: the memory bound. Real cell and gene names sit well under this. +#: Fallback width of a variable-length string element, used only when the +#: real width cannot be sampled. h5py reports `itemsize` 8 for vlen strings +#: because the value is a pointer, which would overestimate the row count by +#: an order of magnitude. Real cell and gene names sit well under this. VLEN_ELEMENT_BYTES = 64 +#: Elements read when sampling a variable-length array's real element width. +#: One small read, against a copy that is about to stream the whole array. +VLEN_SAMPLE_ROWS = 256 -def _row_bytes(shape: Sequence[int], dtype: Any) -> int: - """In-memory size of one row along the first axis.""" + +def _is_vlen(dtype: Any) -> bool: + """Does this dtype hide its real element width behind a pointer? + + 'O' is how h5py spells vlen str, 'T' is numpy StringDType as reported by + zarr-python 3. Neither `itemsize` reflects the bytes actually stored. + """ + itemsize = int(getattr(dtype, "itemsize", 0) or 0) + return itemsize <= 0 or getattr(dtype, "kind", None) in ("O", "T") + + +def _sample_element_bytes(src: Any, n_rows: int) -> int: + """Mean stored width of a variable-length element, by reading a few. + + An assumed width is not a bound. Estimating 64 bytes and reading + `TARGET_READ_BYTES // 64` elements means the read is 32 MiB only if the + guess holds: at 4 KiB elements it is 2 GiB, and for an array shorter than + the computed step the whole thing is read at once. Measured before this: + copying 200,000 strings of 4 KiB peaked at 827 MB against a stated 32 MiB + budget. + + One small read fixes that, and it is negligible beside the copy it is + about to size -- `uns` can hold arbitrary text, so the width is not + something this layer can assume. + """ + try: + sample = np.asarray(src[: min(VLEN_SAMPLE_ROWS, max(1, n_rows))]) + except Exception: # pragma: no cover - unreadable source + return VLEN_ELEMENT_BYTES + + total = 0 + count = 0 + for value in sample.reshape(-1)[:VLEN_SAMPLE_ROWS]: + try: + total += len(value) + except TypeError: # pragma: no cover - non-sized element + total += VLEN_ELEMENT_BYTES + count += 1 + + if not count: + return VLEN_ELEMENT_BYTES + # The floor keeps a column of empty strings from producing an unbounded + # step; the object header dominates at that size anyway. + return max(VLEN_ELEMENT_BYTES, total // count) + + +def _row_bytes( + shape: Sequence[int], dtype: Any, element_bytes: int = VLEN_ELEMENT_BYTES +) -> int: + """In-memory size of one row along the first axis. + + `element_bytes` is the measured width for a variable-length dtype; the + default is the fallback for callers that have no sample to offer. + """ width = 1 for dim in shape[1:]: width *= max(1, int(dim)) itemsize = int(getattr(dtype, "itemsize", 0) or 0) - # 'O' is how h5py spells vlen str, 'T' is numpy StringDType as reported by - # zarr-python 3. Neither itemsize reflects the bytes actually stored. - if itemsize <= 0 or getattr(dtype, "kind", None) in ("O", "T"): - itemsize = VLEN_ELEMENT_BYTES + if _is_vlen(dtype): + itemsize = element_bytes return max(1, width * itemsize) def _chunk_step( - shape: Sequence[int], chunks: Optional[Sequence[int]], dtype: Any + shape: Sequence[int], + chunks: Optional[Sequence[int]], + dtype: Any, + element_bytes: int = VLEN_ELEMENT_BYTES, ) -> int: """Rows to copy per read, sized for the filesystem rather than the source. @@ -492,7 +548,7 @@ def _chunk_step( if chunks is not None and len(chunks) > 0 and chunks[0]: chunk_rows = max(1, int(chunks[0])) - step = max(1, TARGET_READ_BYTES // _row_bytes(shape, dtype)) + step = max(1, TARGET_READ_BYTES // _row_bytes(shape, dtype, element_bytes)) if chunk_rows: # Never go below a single chunk: a partial read still decompresses the # whole thing, so a smaller step costs the same I/O for less data. That @@ -532,7 +588,11 @@ def copy_dataset(src: Any, dst_group: Any, name: str) -> Any: ds[()] = src[()] return ds - step = _chunk_step(shape, getattr(src, "chunks", None), src.dtype) + element_bytes = ( + _sample_element_bytes(src, shape[0]) if _is_vlen(src.dtype) + else VLEN_ELEMENT_BYTES + ) + step = _chunk_step(shape, getattr(src, "chunks", None), src.dtype, element_bytes) for start in range(0, shape[0], step): end = min(start + step, shape[0]) if len(shape) == 1: diff --git a/tests/test_performance.py b/tests/test_performance.py index 8dd7b07..87c0023 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -1045,3 +1045,96 @@ def test_streaming_export_sparse_costs_far_less_than_loading_the_matrix( f"{loaded[0] / max(streamed[0], 1):.1f}. The default path is supposed " "to be the one you reach for when the matrix does not fit." ) + + +# --------------------------------------------------------------------------- +# copying a store must stay inside its read budget +# +# `copy_dataset` sizes each read to TARGET_READ_BYTES. For variable-length +# strings it cannot ask the dtype how wide an element is -- h5py reports the +# itemsize of a pointer -- so it estimated. An estimate is not a bound: at +# 4 KiB elements the 32 MiB budget became an 827 MB peak, and at 200,000 rows +# the computed step exceeded the dataset, so the whole array was read at once. +# +# This is the one path the rest of this file did not reach, and it is the +# path every `copy:` task in `subset` and every uns entry goes through. + + +def _vlen_source(path: Path, width: int, n_rows: int): + """An HDF5 variable-length string dataset of `n_rows` x `width` bytes.""" + import h5py + + with h5py.File(path, "w") as handle: + handle.create_dataset( + "t", + data=np.array([("x" * width).encode()] * n_rows, dtype=object), + dtype=h5py.string_dtype(), + chunks=(1024,), + ) + return path + + +@pytest.mark.parametrize("width", [16, 256, 4096]) +def test_a_read_of_variable_length_strings_stays_inside_the_byte_budget( + tmp_path, width +): + """The budget is bytes per read, so a wider element means fewer elements. + + This is the exact invariant, and it costs nothing to check: the step is + computed, not measured. Before the fix the element width was assumed to + be 64 bytes, so the step was the same 524,288 elements whatever the data + actually held -- 2 GiB per read at 4 KiB elements, against a stated + 32 MiB budget. + """ + import h5py + + from adata.storage import ( + TARGET_READ_BYTES, + _chunk_step, + _sample_element_bytes, + ) + + n_rows = 50_000 + source = _vlen_source(tmp_path / f"s{width}.h5", width, n_rows) + with h5py.File(source, "r") as handle: + dataset = handle["t"] + sampled = _sample_element_bytes(dataset, n_rows) + step = _chunk_step(dataset.shape, dataset.chunks, dataset.dtype, sampled) + + assert sampled >= min(width, 64), ( + f"sampled {sampled} bytes for {width}-byte elements; the width is " + "being guessed rather than measured" + ) + # A read may round up to a whole source chunk, hence the slack. + assert step * width <= TARGET_READ_BYTES * 2, ( + f"one read would take {step * width / 1e6:.0f} MB of {width}-byte " + f"strings ({step} elements), against a {TARGET_READ_BYTES / 1e6:.0f} " + "MB budget" + ) + + +@pytest.mark.slow +def test_copying_wide_strings_does_not_read_the_whole_array(tmp_path): + """And the budget holds in practice, not just in the arithmetic. + + Sized so the array is several times the budget: before the fix the + computed step exceeded the row count, so the whole thing was read at once + -- 164 MB here, and 827 MB in the 200,000-row case that prompted this. + """ + import h5py + + from adata.storage import TARGET_READ_BYTES, copy_dataset + + width, n_rows = 4096, 40_000 # ~164 MB, about 5x the budget + source = _vlen_source(tmp_path / "wide.h5", width, n_rows) + + with h5py.File(source, "r") as src, h5py.File(tmp_path / "out.h5", "w") as dst: + with count_allocations() as peak: + copy_dataset(src["t"], dst, "t") + + assert peak[0] <= TARGET_READ_BYTES * 3, ( + f"copying a {width * n_rows / 1e6:.0f} MB array of {width}-byte " + f"strings peaked at {peak[0] / 1e6:.0f} MB, against a " + f"{TARGET_READ_BYTES / 1e6:.0f} MB read budget. A step larger than " + "the array means the whole array is read in one go." + )