Central scanpy↔rapids-singlecell GPU routing; route marker DE (#38)

skills/cell-type-annotation/pixi.toml

@@ -100,6 +100,22 @@ extra-index-urls = ["https://pypi.nvidia.com"]
 # — no override, no source build, no build-host `-march=native` fragility.
 [feature.gpu.target.linux-64.dependencies]
 hnswlib = ">=0.8"
+# rapids-singlecell's compiled DE kernels (rapids_singlecell._cuda._aggr_cuda, used by the GPU
+# `rank_genes_groups`) are built against CUDA 13 and link `libcudart.so.13`, but the rest of the
+# GPU stack is CUDA 12 (cuml-cu12 / cupy-cuda12x / torch-cu12, which load libcudart.so.12). Without
+# a CUDA-13 runtime present the kernel .so fails to load, rapids-singlecell sets it to None, and the
+# GPU DE path silently falls back to scanpy (verified on Euler). Ship the CUDA-13 cudart runtime from
+# conda-forge so libcudart.so.13 lands in $ENV/lib alongside .so.12 — the two majors coexist (the
+# driver supports both) and the kernel loads. This is the actual enabler of issue #38's GPU path.
+cuda-cudart = "13.*"
+
+# Put the env's own lib dir on the dynamic-loader path for the gpu env so the CUDA-13 DE kernel
+# above finds `libcudart.so.13` in $ENV/lib. pixi/`pixi run` does NOT add `$CONDA_PREFIX/lib` to
+# LD_LIBRARY_PATH by default, so without this the kernel .so can't resolve libcudart.so.13 and the
+# GPU DE path silently falls back to scanpy. Scoped to the gpu environment only; the CPU `default`
+# env never sets this.
+[feature.gpu.activation.env]
+LD_LIBRARY_PATH = "$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
 
 [feature.gpu.target.linux-64.pypi-dependencies]
 # Reference-model option (scANVI/scArches + scvi-hub). torch+OpenMP coexist fine with the
@@ -108,6 +124,17 @@ scvi-tools = "*"
 # CellMapper's `rapids` GPU k-NN backend calls cuML directly (validated on an RTX 4090).
 cuml-cu12 = "*"
 cudf-cu12 = "*"
+# GPU marker DE: `cta markers compute` routes `rank_genes_groups` to rapids-singlecell on a GPU
+# profile (atlas-scale ranking in seconds vs ~minutes on scanpy/CPU). Drop-in — same wilcoxon
+# method, same `uns['rank_genes_groups']` output the CPU path writes; scanpy stays the CPU default.
+# The `rapids-cu12` extra pulls the FULL RAPIDS stack it imports at load (cudf/cuml/cugraph/cuvs);
+# the bare cuml-cu12/cudf-cu12 above (for CellMapper's kNN) are not enough — verified on Euler that
+# `import rapids_singlecell` needs cuvs too.
+rapids-singlecell = { version = "*", extras = ["rapids-cu12"] }
+# rapids-singlecell 0.15.2 also imports `dask.array` at module load (rapids_singlecell/_compat.py)
+# but only declares `dask` under its `doc`/`test` extras, not core — so the import ALSO needs dask
+# present. Pull it in explicitly so the GPU DE path loads instead of silently falling back to scanpy.
+dask = "*"
 # Optional large-memory zero-shot annotator (additional cross-check only). Its hnswlib dep is
 # pinned to the portable conda-forge build above (see comment) to avoid the AVX-512 wheel SIGILL.
 scimilarity = "*"

skills/cell-type-annotation/reference/compute_environments.md

@@ -114,8 +114,19 @@ CPU host that merely has the client tools installed. `cta detect-env` flags it;
 
 The per-tool step / GPU-requirement / profile mapping has **one home**:
 `reference/tool_registry.md`. Read it to decide what to run where. In short: most of the
-pipeline is CPU (`none`); only the reference-mapping cross-check and (optionally)
-large-query scoring touch a GPU. `cta crosscheck scimilarity` is CPU too and needs only moderate
+pipeline is CPU (`none`); the reference-mapping cross-check, (optionally) large-query scoring,
+and the **full-object marker DE** touch a GPU. The scanpy↔rapids-singlecell split is **profile-
+routed with no knob** through one central mechanism (`cta/gpu.py`): scanpy on CPU,
+`rapids-singlecell` on a GPU profile, with automatic CPU fallback on any GPU failure. Adding a
+future GPU-accelerable step is a two-callable `gpu.run(adata, cpu=…, gpu=…, label=…, equivalent=…)`
+call, not a re-implementation of the check/transfer/fallback. Today only **marker DE**
+(`cta markers compute`, plus the `markers score` overlap-DE fallback) is routed — atlas-scale,
+full-object, and bit-identical CPU/GPU (ranking in seconds vs ~minutes). Steps that run on a
+small subset (e.g. `cta cluster subcluster`) deliberately stay on CPU: no real speedup at that
+size, and GPU Leiden would diverge from scanpy's igraph for no gain. The mechanism's
+`equivalent=False` flag is reserved for any future routed op whose GPU result differs (clustering/
+UMAP), so the divergence is logged. On a `scheduler` profile the atlas marker job is submitted to
+a GPU node just like the reference-model steps. `cta crosscheck scimilarity` is CPU too and needs only moderate
 RAM (~16–27 GB), so it runs on any reasonably-sized machine — it is HUMAN-only, not GPU- or
 large-memory-gated. The profile decides *how* the `preferred`/`required` rows execute — inline on
 `local_gpu`, handed off on `scheduler` (next section).
@@ -160,12 +171,15 @@ the scripts encode the matching best practices so a SLURM run stays clean and fa
 - GPU results are **not bit-identical** to CPU, and deep models (scVI/scANVI) are
   **stochastic** — seed them (`scvi.settings.seed = 0`; the scripts do) but still report
   predictions as a **hypothesis**, never a verdict.
-- GPU k-NN (`cta crosscheck reference-map --knn-method rapids`) and any optional GPU clustering/UMAP
-  (e.g. rapids-singlecell, not a default extra) are not bit-identical to their CPU
-  equivalents. This skill annotates **pre-existing** clusters and does not re-cluster the
-  main object, so this mostly doesn't bite — but if a GPU-accelerated step ever feeds
-  clustering, treat the clusters as the same kind of input the skill validates, not as
-  ground truth.
+- The GPU-routed steps (`cta/gpu.py`) are not bit-identical to their CPU equivalents, but the
+  ones routed today are **result-equivalent for annotation**: marker DE (`cta markers compute`)
+  and GPU k-NN (`cta crosscheck reference-map --knn-method rapids`) differ only by minor float/tie
+  noise that doesn't change rankings. Marker DE writes the same `uns` schema either way and
+  degrades to scanpy on any GPU failure (`equivalent=True`). Clustering/UMAP are NOT routed (they
+  use different GPU backends and would diverge); the mechanism's `equivalent=False` flag exists so
+  that if such an op is ever routed, the divergence is logged and the clusters are treated as
+  input to validate, not as truth. This skill annotates **pre-existing** clusters and does not
+  re-cluster the main object, so this mostly doesn't bite.
 
 ## The invariant that does not change with the hardware
 

skills/cell-type-annotation/reference/tool_registry.md

@@ -29,7 +29,7 @@ A separate **RAM** note flags steps whose real cost is memory, not GPU.
 | Canonicalize (1e) | `cta data canonicalize` | none | all | **Standard Stage-1 step** (not optional): write the one canonical sidecar every step reads — symbol var_names (+ kept Ensembl map), counts in `layers['counts']` (from a layer or `.raw`), log-norm `X`, slim embeddings/obs. Also applies the **upfront absolute gene filter** (`sc.pp.filter_genes(min_cells=GENE_FILTER_MIN_CELLS=5)`) so every downstream step sees one clean, consistent gene set. Replaces the old slim-cache step. |
 | QC metrics (1e) | `cta data qc` | none | all | compute (when absent) standard QC into the sidecar obs: `n_genes_by_counts`/`total_counts`/`pct_counts_mt` + cell-cycle `S_score`/`G2M_score`/`phase` (Tirosh lists, case-insensitive). Detection only — no doublet detection, normalization, or filtering. |
 | Seed markers (3.1, opt.) | `cta markers assemble` | none | all | seed `markers_<level>.json` from a CellTypist model / reference DE. A `--reference` is aligned to query symbols first (`--ref-symbol-col`/`--ref-symbol-map`) so its DE marker names panel-gate. |
-| Score — bottom-up (3.3) | `cta markers compute` | none | all | `rank_genes_groups` (wilcoxon, scanpy defaults) → `markers_<level>.csv`; `filter_rank_genes_groups` specificity → `markers_filtered_<level>.csv` (CPU). GPU acceleration would need rapids-singlecell (full RAPIDS+dask, not in the default extras) — niche, since the main object isn't re-preprocessed. |
+| Score — bottom-up (3.3) | `cta markers compute` | optional | all | `rank_genes_groups` (wilcoxon, scanpy defaults) → `markers_<level>.csv`; `filter_rank_genes_groups` specificity → `markers_filtered_<level>.csv`. **Compute-profile-routed, no knob:** CPU profile → scanpy; GPU profile → `rapids_singlecell.tl.rank_genes_groups` (same wilcoxon + same `uns` output, atlas-scale ranking in seconds vs ~minutes). Auto-detected via `torch.cuda`; degrades to scanpy on any GPU failure. |
 | Score — top-down (3.3) | `cta markers profile` | none | all | score the curated panel across clusters → pct/mean + dotplot. |
 | Score — signatures (3.3) | `cta markers score` | none | all | decoupler `aucell` + pyUCell + DE-overlap → `crosscheck`. |
 | Cross-check — in-object (3.4) | `cta crosscheck obs-pred` | none | all | summarize an existing `obs` prediction column. Preferred when one fits. |
@@ -41,7 +41,7 @@ A separate **RAM** note flags steps whose real cost is memory, not GPU.
 | Provision SCimilarity model (before scimilarity) | `cta reference fetch-scimilarity` | none | all (HUMAN-only) | download + extract the model bundle for `cta crosscheck scimilarity`. Annotation-only by default (~9 GB on disk; ~30 GB resumable download, skips the 32 GB cellsearch). Needs network. |
 | Fetch reference (3.4 fallback) | `cta reference fetch-census` | none | all (larger slices on large-mem) | capped, tissue-matched CELLxGENE Census slice → feeds `cta crosscheck reference-map`. Needs network. |
 | Investigate before Unknown (3.7) | `cta cluster profile` | none | all | QC medians + categorical crosstabs + optional spatial. |
-| Sub-cluster (3.7) | `cta cluster subcluster` | none | all | Leiden on the anchor embedding + per-sub summary. |
+| Sub-cluster (3.7) | `cta cluster subcluster` | none | all | neighbors + Leiden on the anchor embedding + per-sub summary. CPU (scanpy igraph) — runs on a small subset, so not GPU-routed. |
 | Confidence calibration (3.7) | `cta cluster confidence` | none | all | marker-logFC-margin + agreement proxy (advisory; from `markers_filtered_<level>.csv`). |
 | Evidence cards (3.6) | `cta report scaffold-cards` | none | all | render cards from the labels CSV; `--evidence` embeds the numbers. |
 | Report + gate (3.9) | `cta report build` | none | all | per-level HTML; `--strict` = level-completeness gate. |

skills/cell-type-annotation/src/cta/compute_markers.py

@@ -4,12 +4,18 @@
 + a `filter_rank_genes_groups` specificity pass, written as tidy marker tables to
 <workdir>/tables/. The descriptive marker evidence that feeds annotation. NOT condition DE
 (Job B = pseudobulk + PyDESeq2; see reference/differential_expression.md).
+
+The Wilcoxon ranking is compute-profile-routed (no user knob): on a CUDA GPU it runs via
+rapids-singlecell (atlas-scale: seconds vs ~minutes), otherwise via scanpy on CPU — same
+method and same `uns['rank_genes_groups']` output, so the specificity pass and every
+downstream reader are identical on both paths.
 """
 
 from __future__ import annotations
 
 import typer
 
+from cta import gpu
 from cta.io import (
     ADATA_OPT,
     CLUSTER_KEY_OPT,
@@ -52,20 +58,26 @@ def cli(
     adata_obj.obs[cluster_key] = adata_obj.obs[cluster_key].astype("category")
     grps = split_csv(groups) or "all"
 
-    eprint(">> rank_genes_groups (wilcoxon, pts=True)")
-    # scanpy assembles its results frame one column per group, so with many clusters pandas
-    # raises 'DataFrame is highly fragmented' once per cluster. It is internal to scanpy
-    # (scanpy/tools/_rank_genes_groups.py) and has no effect on the output — silence exactly
-    # that message for this one call, leaving every other warning visible.
-    with suppressed_warnings(FRAGMENTED_FRAME_WARNING):
-        sc.tl.rank_genes_groups(
-            adata_obj,
-            groupby=cluster_key,
-            groups=grps,
-            method="wilcoxon",
-            pts=True,
-            use_raw=False,
+    # Compute-profile routing via the central cta.gpu mechanism (no user knob): rapids-singlecell
+    # on a CUDA GPU (atlas-scale: seconds vs ~minutes), else scanpy on CPU, with automatic CPU
+    # fallback. Same wilcoxon method + `uns['rank_genes_groups']` either way, so everything
+    # downstream (extraction, #36's filter_rank_genes_groups, reports) is identical.
+    def _rgg_cpu(a):
+        # scanpy assembles its results frame one column per group, so with many clusters pandas
+        # raises 'DataFrame is highly fragmented' once per cluster. It is internal to scanpy
+        # (scanpy/tools/_rank_genes_groups.py) and has no effect on the output — silence exactly
+        # that message for this one call, leaving every other warning visible.
+        with suppressed_warnings(FRAGMENTED_FRAME_WARNING):
+            sc.tl.rank_genes_groups(
+                a, groupby=cluster_key, groups=grps, method="wilcoxon", pts=True, use_raw=False
+            )
+
+    def _rgg_gpu(a, rsc):
+        rsc.tl.rank_genes_groups(
+            a, groupby=cluster_key, groups=grps, method="wilcoxon", pts=True, use_raw=False
         )
+
+    gpu.run(adata_obj, cpu=_rgg_cpu, gpu=_rgg_gpu, label="rank_genes_groups (wilcoxon, pts=True)")
     frames = []
     grp_list = adata_obj.obs[cluster_key].cat.categories if grps == "all" else grps
     for g in grp_list:

skills/cell-type-annotation/src/cta/gpu.py

@@ -0,0 +1,88 @@
+"""Central scanpy <-> rapids-singlecell routing for the compute-profile split.
+
+Every scVerse step that has a rapids-singlecell GPU equivalent (marker DE, sub-cluster Leiden,
+…) routes through `run()` here instead of re-implementing the env check + host<->device transfer
++ CPU fallback. One place to reason about "CPU or GPU", so adding a future GPU-accelerable step
+is a two-callable call, not a copy of the plumbing.
+
+Routing rule (no user knob): use the GPU iff a CUDA GPU is reachable in this process AND
+rapids-singlecell imports; otherwise scanpy on CPU. The GPU path moves the object's matrices to
+the device for the op and ALWAYS restores them to host afterwards (so downstream CPU code and any
+fallback see numpy/scipy), and falls back to CPU on ANY failure — the GPU is an accelerator, never
+a correctness dependency.
+
+Result equivalence is per-op and the caller declares it. `rank_genes_groups` is bit-identical
+CPU/GPU (only float/tie noise); Leiden/UMAP use different backends and genuinely DIVERGE — pass
+`equivalent=False` for those so the chosen engine is logged (the run records which one produced
+the divergent output). See reference/compute_environments.md.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from cta.io import eprint, resolve_accelerator
+
+_GPU_CACHE: bool | None = None
+
+
+def gpu_available() -> bool:
+    """True iff a CUDA GPU is reachable in this process AND rapids-singlecell imports.
+
+    The single source of truth for the CPU/GPU split. Cached per process — the answer can't
+    change mid-run. Importing rapids-singlecell is the real gate: `resolve_accelerator` only
+    sees torch, but rapids-singlecell may be absent (CPU env) or fail to load (e.g. a
+    CUDA-mismatched compiled kernel), and either way we must use scanpy.
+    """
+    global _GPU_CACHE
+    if _GPU_CACHE is None:
+        ok = False
+        if resolve_accelerator("auto") == "cuda":
+            try:
+                import rapids_singlecell  # noqa: F401
+
+                ok = True
+            except Exception:  # noqa: BLE001 — any import/load failure -> use CPU
+                ok = False
+        _GPU_CACHE = ok
+    return _GPU_CACHE
+
+
+def run(
+    adata,
+    *,
+    cpu: Callable,
+    gpu: Callable,
+    label: str,
+    equivalent: bool = True,
+) -> str:
+    """Run `gpu(adata, rsc)` on a CUDA GPU, else `cpu(adata)`; return the engine used ('gpu'/'cpu').
+
+    On a GPU profile, moves `adata`'s matrices to the device (rapids-singlecell `anndata_to_GPU`),
+    runs `gpu`, then ALWAYS restores them to host (`anndata_to_CPU`) — even on failure — so the rest
+    of the step and any fallback see CPU arrays. Any exception in the GPU path logs and falls back
+    to `cpu`. `equivalent=False` marks ops whose GPU result differs from CPU (clustering/UMAP); it
+    only annotates the log line. Both callables mutate `adata` in place; `gpu` additionally receives
+    the imported `rapids_singlecell` module so it needn't import it itself.
+    """
+    if gpu_available():
+        try:
+            import rapids_singlecell as rsc
+
+            note = "" if equivalent else "  (result differs from CPU — distinct GPU backend)"
+            eprint(f">> {label} [GPU / rapids-singlecell]{note}")
+            rsc.get.anndata_to_GPU(adata)
+            try:
+                gpu(adata, rsc)
+                return "gpu"
+            finally:
+                # Restore X to host even on failure, so the CPU path / downstream see numpy/scipy.
+                try:
+                    rsc.get.anndata_to_CPU(adata)
+                except Exception:  # noqa: BLE001
+                    pass
+        except Exception as e:  # noqa: BLE001 — any GPU failure falls back to scanpy CPU
+            eprint(f"   GPU {label} failed ({type(e).__name__}: {e}) — falling back to scanpy CPU")
+    eprint(f">> {label} [CPU / scanpy]")
+    cpu(adata)
+    return "cpu"

skills/cell-type-annotation/src/cta/score_signatures.py

@@ -38,6 +38,7 @@
 import pandas as pd
 import typer
 
+from cta import gpu
 from cta.constants import MIN_ON_PANEL_MARKERS
 from cta.io import (
     ADATA_OPT,
@@ -150,9 +151,17 @@ def cli(
         else:
             eprint("   overlap: markers table not found — computing rank_genes_groups")
             # Default-param Wilcoxon, matching `cta markers compute` (no tie_correct) so this
-            # fallback can't silently disagree with the canonical markers table.
-            sc.tl.rank_genes_groups(
-                adata_obj, groupby=cluster_key, method="wilcoxon", use_raw=False
+            # fallback can't silently disagree with the canonical markers table. Routed through the
+            # central cta.gpu mechanism (rapids-singlecell on GPU, scanpy on CPU, auto-fallback).
+            gpu.run(
+                adata_obj,
+                cpu=lambda a: sc.tl.rank_genes_groups(
+                    a, groupby=cluster_key, method="wilcoxon", use_raw=False
+                ),
+                gpu=lambda a, rsc: rsc.tl.rank_genes_groups(
+                    a, groupby=cluster_key, method="wilcoxon", use_raw=False
+                ),
+                label="rank_genes_groups (overlap fallback)",
             )
             for g in clusters:
                 ranked[str(g)] = (

skills/cell-type-annotation/src/cta/subcluster.py

@@ -94,6 +94,9 @@ def cli(
         f"on '{embedding_key}' @ resolution {resolution}"
     )
 
+    # Sub-clustering runs on a SMALL subset (one/few clusters), so it stays on CPU — GPU offers no
+    # real speedup at this size and cuGraph Leiden would diverge from scanpy's igraph for no gain.
+    # (GPU routing via cta.gpu is reserved for the full-object heavy steps, e.g. marker DE.)
     sc.pp.neighbors(sub, use_rep=embedding_key, n_neighbors=n_neighbors)
     # igraph flavor = CPU-friendly, deterministic; never requires GPU/leidenalg quirks
     sc.tl.leiden(