Fix/all neighbors heuristics

docs/release-notes/0.18.0.md

@@ -11,9 +11,11 @@
 * Add GSVA via {func}`~rapids_singlecell.dcg.gsva`.{pr}`810` {smaller}`S Dicks`
 * Add Over Representation Analysis with {func}`~rapids_singlecell.dcg.ora` for matrices and {func}`~rapids_singlecell.dcg.query_set` for gene lists. {pr}`806` {smaller}`S Dicks`
 * Bundle a version-matched agent analysis skill with the package, installed with `rapids-singlecell-install-skills` and checked with `rapids-singlecell-check-kernel` (see {ref}`Agent skill <agent-skill>`). {pr}`739` {smaller}`S Dicks & L Heumos`
+* Speed up {func}`~rapids_singlecell.tl.umap` on large graphs; ``rng=None`` uses cuML's faster unseeded optimizer {pr}`822` {smaller}`S Dicks`
 
 ```{rubric} Bug fixes
 ```
+* Size the ``algorithm="all_neighbors"`` batching of {func}`~rapids_singlecell.pp.neighbors` to the data and free memory instead of the GPU count, avoiding crashes on massive datasets {pr}`822` {smaller}`S Dicks`
 * Match the LDA scaling of {class}`~rapids_singlecell.ptg.Mixscape` to scikit-learn 1.9.1 {pr}`799` {smaller}`S Dicks`
 * Preserve existing `.obs` metadata, including categorical dtypes, when running {func}`~rapids_singlecell.pp.scrublet` with `batch_key`.
 {pr}`817` {smaller}`S Dicks`

src/rapids_singlecell/preprocessing/_neighbors/__init__.py

@@ -145,13 +145,13 @@ def neighbors(
 
         * 'algo': The algorithm to use. Valid options are: 'ivf_pq' and 'nn_descent'. Default is 'nn_descent'. `ivf_pq` is restricted to the `euclidean` and `sqeuclidean` metrics; use `nn_descent` for `cosine` and `inner_product`.
 
-        * 'n_clusters': Number of clusters/batches to partition the dataset into (> overlap_factor). Default is 1 on a single GPU and the smallest multiple of the device count greater than `overlap_factor` otherwise.
+        * 'n_clusters': Number of clusters/batches to partition the dataset into (> overlap_factor). Default is 1 if the data fits on a single GPU, otherwise sized to the data and free memory (e.g. 24 for 100M cells on 8 GPUs).
 
-        * 'overlap_factor': Number of clusters each point is assigned to. Must be < n_clusters when the build is batched (`n_clusters > 1`). Default is 1 for an unbatched build and `min(max(2, ceil(log2(n_clusters))), n_clusters - 1)` otherwise. Lower values are faster but lose neighbors at cluster boundaries.
+        * 'overlap_factor': Number of clusters each point is assigned to. Must be < n_clusters when the build is batched (`n_clusters > 1`). Default is 1 unbatched, 2 with automatic `n_clusters` and `min(max(2, ceil(log2(n_clusters))), n_clusters - 1)` otherwise. Lower values are faster but lose neighbors at cluster boundaries.
 
         * 'n_lists': Number of inverted lists for IVF indexing. Default is 2 * next_power_of_2(sqrt(n_samples)). Only available for `ivf_pq` algorithm.
 
-        * 'graph_degree': The degree of the graph nn-descent builds before selecting the final `n_neighbors`. Default is 64, raised to `n_neighbors` if larger. Only available for `nn_descent` algorithm.
+        * 'graph_degree': The degree of the graph nn-descent builds before selecting the final `n_neighbors`. Default is 64 unbatched and `n_neighbors` batched, raised to `n_neighbors` if smaller. Only available for `nn_descent` algorithm.
 
         * 'intermediate_graph_degree': The degree of the intermediate graph. Default is `max(128, int(1.5 * graph_degree))`, following the recommended `>= 1.5 * graph_degree`. A smaller user-supplied value is raised to `graph_degree`. Only available for `nn_descent` algorithm.
 
@@ -212,7 +212,7 @@ def neighbors(
         )
 
     X = _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs)
-    X_contiguous = _check_neighbors_X(X, algorithm, algorithm_kwds)
+    X_contiguous = _check_neighbors_X(X, algorithm)
     _check_metrics(algorithm, metric)
 
     knn_indices, knn_dist = KNN_ALGORITHMS[algorithm](
@@ -410,7 +410,7 @@ def bbknn(
         adata._init_as_actual(adata.copy())
 
     X = _choose_representation(adata, use_rep=use_rep, n_pcs=n_pcs)
-    X_contiguous = _check_neighbors_X(X, algorithm, algorithm_kwds)
+    X_contiguous = _check_neighbors_X(X, algorithm)
     _check_metrics(algorithm, metric)
 
     n_obs = adata.shape[0]

src/rapids_singlecell/preprocessing/_neighbors/_algorithms/_all_neighbors.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import math
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import cupy as cp
@@ -22,6 +23,14 @@
 
 _CUVS_HOST_OUTPUT_MIN_VERSION = parse_version("26.08")
 
+# Automatic batching, following the cuVS out-of-core UMAP guidance
+_DEFAULT_OVERLAP_FACTOR = 2
+_MAX_ROWS_PER_CLUSTER = 10_000_000
+_CLUSTER_IMBALANCE = 2  # largest / average cluster size
+_MEMORY_FRACTION = 0.5  # share of free memory a build may use
+# cuVS nn-descent overflows int32 indices (rows x 32 samples) beyond this many rows
+_NN_DESCENT_MAX_ROWS = 2**31 // 32
+
 
 def _default_overlap_factor(n_clusters: int) -> int:
     """Overlap needed to hold recall as the dataset is split into more clusters."""
@@ -30,19 +39,80 @@ def _default_overlap_factor(n_clusters: int) -> int:
     return max(2, math.ceil(math.log2(n_clusters)))
 
 
-def _all_neighbors_batching(algorithm_kwds: Mapping) -> tuple[int, int]:
-    """Resolve ``(n_clusters, overlap_factor)`` for the cuVS all-neighbors build."""
+def _available_host_memory() -> float:
+    """MemAvailable, capped by cgroup v2 limits."""
+    meminfo = Path("/proc/meminfo").read_text()
+    available = float(meminfo.split("MemAvailable:")[1].split()[0]) * 1024
+    try:
+        cgroup = Path("/proc/self/cgroup").read_text().split("::", 1)[1].strip()
+    except (OSError, IndexError):
+        return available
+    path = Path("/sys/fs/cgroup", cgroup.lstrip("/"))
+    for group in (path, *path.parents):
+        try:
+            limit = (group / "memory.max").read_text().strip()
+            if limit != "max":
+                used = int((group / "memory.current").read_text())
+                available = min(available, int(limit) - used)
+        except (OSError, ValueError):
+            continue
+    return available
+
+
+def _batched_bytes_per_row(
+    n_features: int, k: int, graph_degree: int
+) -> tuple[int, int]:
+    """Host and device bytes per cluster row of cuVS's batched nn-descent."""
+    # padded node / intermediate degrees as in cuVS's nn-descent
+    node = 32 * math.ceil(
+        (graph_degree * 1.3 if graph_degree > 32 else graph_degree) / 32
+    )
+    internal = 32 * math.ceil(max(128, 1.5 * graph_degree) * 1.3 / 32)
+    # pinned buffers, graph, bloom filter, gathered rows, outputs
+    host = 912 + 8 * node + 2 * internal + 4 * n_features + 16 * k
+    # fp16 rows, graph buffers, outputs
+    device = 2 * n_features + 288 + 12 * k
+    return host, device
+
+
+def _all_neighbors_batching(
+    algorithm_kwds: Mapping, shape: tuple[int, int] = (0, 0), k: int = 0
+) -> tuple[int, int]:
+    """Resolve ``(n_clusters, overlap_factor)``, sized to the data and free memory."""
     n_devices = cp.cuda.runtime.getDeviceCount()
     n_clusters = algorithm_kwds.get("n_clusters")
     overlap_factor = algorithm_kwds.get("overlap_factor")
     if n_clusters is None:
-        n_clusters = 1 if n_devices == 1 else n_devices
-        while n_clusters > 1 and n_clusters <= (
-            _default_overlap_factor(n_clusters)
-            if overlap_factor is None
-            else overlap_factor
+        n_obs, n_features = shape
+        # data, nn-descent buffers and outputs on one GPU
+        unbatched_bytes = n_obs * (4 * n_features + 280 + 20 * k)
+        if (
+            n_devices == 1
+            and n_obs < _NN_DESCENT_MAX_ROWS
+            and unbatched_bytes <= _MEMORY_FRACTION * cp.cuda.runtime.memGetInfo()[0]
         ):
-            n_clusters += n_devices
+            n_clusters = 1
+        else:
+            if overlap_factor is None:
+                overlap_factor = _DEFAULT_OVERLAP_FACTOR
+            graph_degree = max(algorithm_kwds.get("graph_degree", k), k)
+            host_row, device_row = _batched_bytes_per_row(n_features, k, graph_degree)
+            device_memory = min(
+                cp.cuda.runtime.getDeviceProperties(i)["totalGlobalMem"]
+                for i in range(n_devices)
+            )
+            rows_per_cluster = min(
+                _MAX_ROWS_PER_CLUSTER,
+                _MEMORY_FRACTION
+                * _available_host_memory()
+                / (n_devices * _CLUSTER_IMBALANCE * host_row),
+                _MEMORY_FRACTION * device_memory / (_CLUSTER_IMBALANCE * device_row),
+            )
+            n_clusters = n_devices * max(
+                1, math.ceil(n_obs * overlap_factor / (n_devices * rows_per_cluster))
+            )
+            while n_clusters <= overlap_factor:
+                n_clusters += n_devices
     if overlap_factor is None:
         overlap_factor = _default_overlap_factor(n_clusters)
         if n_clusters > 1:
@@ -70,7 +140,9 @@ def _all_neighbors_knn(
             "Please update your cuvs installation."
         )
     algo = algorithm_kwds.get("algo", "nn_descent")
-    n_clusters, overlap_factor = _all_neighbors_batching(algorithm_kwds)
+    n_clusters, overlap_factor = _all_neighbors_batching(algorithm_kwds, X.shape, k)
+    # batched builds read from host, unbatched ones from device
+    X = cp.asnumpy(X) if n_clusters > 1 else cp.asarray(X)
     use_host_output = (
         n_clusters > 1
         and parse_version(cuvs.__version__) >= _CUVS_HOST_OUTPUT_MIN_VERSION
@@ -99,7 +171,9 @@ def _all_neighbors_knn(
     elif algo == "nn_descent":
         from cuvs.neighbors import nn_descent
 
-        graph_degree = max(algorithm_kwds.get("graph_degree", 64), k)
+        # the cluster overlap recovers the recall of degree 64
+        default_degree = 64 if n_clusters == 1 else k
+        graph_degree = max(algorithm_kwds.get("graph_degree", default_degree), k)
         intermediate_graph_degree = algorithm_kwds.get(
             "intermediate_graph_degree", max(128, int(1.5 * graph_degree))
         )

src/rapids_singlecell/preprocessing/_neighbors/_helper/__init__.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import math
-from types import MappingProxyType
 from typing import TYPE_CHECKING
 
 import cupy as cp
@@ -12,8 +11,6 @@
 from scipy import sparse as sc_sparse
 
 if TYPE_CHECKING:
-    from collections.abc import Mapping
-
     from rapids_singlecell.preprocessing._neighbors import _Algorithms, _Metrics
 
 
@@ -30,7 +27,6 @@ def _cuvs_switch():
 def _check_neighbors_X(
     X: cp_sparse.spmatrix | sc_sparse.spmatrix | np.ndarray | cp.ndarray,
     algorithm: _Algorithms,
-    algorithm_kwds: Mapping = MappingProxyType({}),
 ) -> cp_sparse.spmatrix | cp.ndarray | np.ndarray:
     """Check and convert input X to the expected format based on algorithm.
 
@@ -45,18 +41,15 @@ def _check_neighbors_X(
     X_contiguous (cupy.ndarray or sparse.csr_matrix): Contiguous array or CSR matrix.
 
     """
-    from rapids_singlecell.preprocessing._neighbors._algorithms._all_neighbors import (
-        _all_neighbors_batching,
-    )
-
     if cp_sparse.issparse(X) or sc_sparse.issparse(X):
         if algorithm != "brute":
             raise ValueError(
                 f"Sparse input is not supported for {algorithm} algorithm. Use 'brute' instead."
             )
         X_contiguous = X.tocsr()
+    # all_neighbors moves X itself once batching is resolved
     elif algorithm in ["mg_ivfflat", "mg_ivfpq"] or (
-        algorithm == "all_neighbors" and _all_neighbors_batching(algorithm_kwds)[0] > 1
+        algorithm == "all_neighbors" and isinstance(X, np.ndarray)
     ):
         if isinstance(X, np.ndarray):
             X_contiguous = np.asarray(X, order="C", dtype=np.float32)

src/rapids_singlecell/tools/_umap.py

@@ -4,11 +4,13 @@
 
 import cuml.internals.logger as logger
 import cupy as cp
+import cupyx
 import numpy as np
 from cuml.manifold.umap import find_ab_params, simplicial_set_embedding
 from cupyx.scipy import sparse
 from scanpy._utils import NeighborsView
 from scanpy.tools._utils import get_init_pos_from_paga
+from scipy import sparse as sc_sparse
 
 from rapids_singlecell._compat import _rng_kwargs
 from rapids_singlecell._keys import _embedding_keys
@@ -29,6 +31,31 @@
 
 _InitPos = Literal["auto", "spectral", "random", "paga"]
 
+_H2D_CHUNK = 1 << 28  # indices per host-to-device copy
+
+
+def _device_coo(graph) -> sparse.coo_matrix:
+    """Canonical device COO, expanded on the GPU for a host CSR."""
+    if not (
+        sc_sparse.issparse(graph)
+        and graph.format == "csr"
+        and graph.has_canonical_format
+    ):
+        return sparse.coo_matrix(graph)
+    row = cp.zeros(graph.nnz, dtype=cp.int32)
+    starts = cp.asarray(graph.indptr[1:-1])
+    cupyx.scatter_add(row, starts[starts < graph.nnz], 1)
+    cp.cumsum(row, out=row)
+    col = cp.empty(graph.nnz, dtype=cp.int32)
+    for start in range(0, graph.nnz, _H2D_CHUNK):
+        col[start : start + _H2D_CHUNK] = cp.asarray(
+            graph.indices[start : start + _H2D_CHUNK]
+        )
+    data = cp.asarray(graph.data, dtype=cp.float32)
+    coo = sparse.coo_matrix((data, (row, col)), shape=graph.shape)
+    coo.has_canonical_format = True
+    return coo
+
 
 @_accepts_legacy_random_state(0)
 def umap(
@@ -101,6 +128,7 @@ def umap(
         Random seed or :class:`~numpy.random.Generator` used by the random
         number generator.
         The superseded `random_state` argument is still accepted.
+        `rng=None` runs unseeded, which is faster but not reproducible.
     a
         More specific parameters controlling the embedding. If `None` these
         values are set automatically as determined by `min_dist` and
@@ -135,6 +163,8 @@ def umap(
             UMAP parameters `a`, `b`, and `random_state` (if specified).
     """
 
+    # a seed forces cuML's slower deterministic optimizer
+    unseeded = rng is None or (isinstance(rng, _LegacyRng) and rng.arg is None)
     rng = np.random.default_rng(rng)
 
     adata = adata.copy() if copy else adata
@@ -189,15 +219,15 @@ def umap(
         # from `graph` alone, so we pass a placeholder instead of materializing
         # the representation on the GPU.
         data=cp.zeros((n_obs, 1), dtype=cp.float32),
-        graph=sparse.coo_matrix(neighbors["connectivities"]),
+        graph=_device_coo(neighbors["connectivities"]),
         n_components=n_components,
         initial_alpha=alpha,
         a=a,
         b=b,
         negative_sample_rate=negative_sample_rate,
         n_epochs=n_epochs,
         init=init_coords,
-        random_state=_legacy_random_state(rng, always_state=True),
+        random_state=None if unseeded else _legacy_random_state(rng, always_state=True),
     )
     logger.set_level(logger_level)
     X_umap = cp.asarray(X_umap).get()

tests/test_embeddings.py

@@ -5,18 +5,38 @@
 import pytest
 import scanpy as sc
 from scanpy.datasets import pbmc68k_reduced
+from scipy import sparse
 
 from rapids_singlecell.tools import draw_graph, tsne, umap
+from rapids_singlecell.tools._umap import _device_coo
 from testing.rapids_singlecell._pytest import needs
 
 
-def test_umap():
+@pytest.mark.parametrize("kwargs", [{}, {"rng": None}])
+def test_umap(kwargs):
     pbmc = pbmc68k_reduced()
     del pbmc.obsm["X_umap"]
-    umap(pbmc)
+    umap(pbmc, **kwargs)
     assert pbmc.obsm["X_umap"].shape == (700, 2)
 
 
+@pytest.mark.parametrize("index_dtype", [np.int32, np.int64])
+def test_umap_device_coo(index_dtype):
+    graph = sparse.random(50, 50, density=0.2, format="csr", rng=0, dtype=np.float32)
+    # empty rows at the start, middle and end
+    graph = sparse.csr_matrix(
+        graph.multiply(np.isin(np.arange(50), [0, 1, 25, 49], invert=True)[:, None])
+    )
+    graph.indices = graph.indices.astype(index_dtype)
+    graph.indptr = graph.indptr.astype(index_dtype)
+    ref = graph.tocoo()
+    coo = _device_coo(graph)
+    assert coo.has_canonical_format
+    np.testing.assert_array_equal(coo.row.get(), ref.row)
+    np.testing.assert_array_equal(coo.col.get(), ref.col)
+    np.testing.assert_array_equal(coo.data.get(), ref.data)
+
+
 def test_tsne():
     pbmc = pbmc68k_reduced()
     tsne(pbmc)

tests/test_mg_neighbors.py

@@ -84,18 +84,34 @@ def test_all_neighbors(algo):
 
 
 @pytest.mark.parametrize(
-    ("n_devices", "expected"),
-    [(1, (1, 1)), (2, (4, 2)), (3, (3, 2)), (4, (4, 2)), (8, (8, 3)), (16, (16, 4))],
+    ("n_devices", "shape", "host_memory", "expected"),
+    [
+        (1, (0, 0), 2e12, (1, 1)),
+        (2, (0, 0), 2e12, (4, 2)),
+        (3, (0, 0), 2e12, (3, 2)),
+        (4, (0, 0), 2e12, (4, 2)),
+        (8, (0, 0), 2e12, (8, 2)),
+        (16, (0, 0), 2e12, (16, 2)),
+        pytest.param(8, (100_000_000, 100), 2e12, (24, 2), id="100M-8gpu"),
+        pytest.param(8, (100_000_000, 100), 2e11, (72, 2), id="100M-8gpu-low-ram"),
+        pytest.param(1, (100_000_000, 100), 2e12, (20, 2), id="100M-1gpu"),
+    ],
 )
-def test_all_neighbors_batching_defaults(monkeypatch, n_devices, expected):
+def test_all_neighbors_batching_defaults(
+    monkeypatch, n_devices, shape, host_memory, expected
+):
     import cupy as cp
 
-    from rapids_singlecell.preprocessing._neighbors._algorithms._all_neighbors import (
-        _all_neighbors_batching,
-    )
+    from rapids_singlecell.preprocessing._neighbors._algorithms import _all_neighbors
 
     monkeypatch.setattr(cp.cuda.runtime, "getDeviceCount", lambda: n_devices)
-    n_clusters, overlap_factor = _all_neighbors_batching({})
+    # 80 GB GPUs, independent of the test machine
+    monkeypatch.setattr(
+        cp.cuda.runtime, "getDeviceProperties", lambda i: {"totalGlobalMem": 80e9}
+    )
+    monkeypatch.setattr(cp.cuda.runtime, "memGetInfo", lambda: (78e9, 80e9))
+    monkeypatch.setattr(_all_neighbors, "_available_host_memory", lambda: host_memory)
+    n_clusters, overlap_factor = _all_neighbors._all_neighbors_batching({}, shape, 15)
     assert (n_clusters, overlap_factor) == expected
     assert n_clusters == 1 or overlap_factor < n_clusters