make p2p datatransfers safe

src/rapids_singlecell/_utils/__init__.py

@@ -8,19 +8,29 @@
 from dask.array import Array as DaskArray
 
 from ._multi_gpu import (
+    MultiGPUFallbackWarning,
     _calculate_blocks_per_pair,
+    _copy_to_device_via_host,
     _create_category_index_mapping,
     _get_device_attrs,
     _split_pairs,
     parse_device_ids,
+    peer_copy_verified,
+    peer_copy_works,
+    validate_multi_gpu,
 )
 
 __all__ = [
+    "MultiGPUFallbackWarning",
     "_calculate_blocks_per_pair",
+    "_copy_to_device_via_host",
     "_create_category_index_mapping",
     "_get_device_attrs",
     "_split_pairs",
     "parse_device_ids",
+    "peer_copy_verified",
+    "peer_copy_works",
+    "validate_multi_gpu",
 ]
 
 

src/rapids_singlecell/_utils/_multi_gpu.py

@@ -12,12 +12,221 @@
 
 from __future__ import annotations
 
+import warnings
+from functools import cache
+
 import cupy as cp
+import numpy as np
 
 # Cache for device attributes per device (lazy initialization)
 _DEVICE_ATTRS_CACHE: dict[int, dict] = {}
 
 
+CUDA_ERROR_PEER_ACCESS_UNSUPPORTED = 217
+CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED = 704
+CUDA_ERROR_PEER_ACCESS_NOT_ENABLED = 705
+CUDA_ERROR_TOO_MANY_PEERS = 711
+
+_SAFE_PEER_FALLBACK_ERRORS = {
+    CUDA_ERROR_PEER_ACCESS_UNSUPPORTED,
+    CUDA_ERROR_PEER_ACCESS_NOT_ENABLED,
+    CUDA_ERROR_TOO_MANY_PEERS,
+}
+_P2P_CANARY_EXPECTED = np.arange(1, 9, dtype=np.float64)
+_P2P_CANARY_POISON = np.full(8, -7.0, dtype=np.float64)
+_WARNED_P2P_FAILURES: set[tuple[int, tuple[tuple[int, int], ...]]] = set()
+
+
+class MultiGPUFallbackWarning(RuntimeWarning):
+    """Warning emitted when unsafe P2P makes an operation use one GPU."""
+
+
+def _copy_to_device_via_host(array: cp.ndarray, device_id: int) -> cp.ndarray:
+    """Copy a device array without touching a peer link.
+
+    This is intentionally reserved for small control arrays whose owner may
+    differ from the main input's device. Capturing the source before changing
+    devices is essential: calling ``cp.asarray`` first would already attempt
+    the peer transfer this helper exists to avoid.
+    """
+    source_device = array.device.id
+    if source_device == device_id:
+        return array
+
+    with cp.cuda.Device(source_device):
+        host = array.get()
+    with cp.cuda.Device(device_id):
+        copied = cp.asarray(host)
+        cp.cuda.get_current_stream().synchronize()
+        return copied
+
+
+def _run_peer_copy_canary(destination: int, source: int) -> bool:
+    """Return whether a small peer copy arrives intact.
+
+    The explicit synchronizations make the canary independent of whichever
+    CuPy stream was current when the multi-GPU operation was entered.
+    """
+    with cp.cuda.Device(source):
+        expected = cp.asarray(_P2P_CANARY_EXPECTED)
+        cp.cuda.get_current_stream().synchronize()
+
+    with cp.cuda.Device(destination):
+        probe = cp.asarray(_P2P_CANARY_POISON)
+        cp.cuda.get_current_stream().synchronize()
+        try:
+            cp.cuda.runtime.memcpyPeer(
+                probe.data.ptr,
+                destination,
+                expected.data.ptr,
+                source,
+                expected.nbytes,
+            )
+            cp.cuda.runtime.deviceSynchronize()
+            arrived = cp.asnumpy(probe)
+        except cp.cuda.runtime.CUDARuntimeError as error:
+            if error.status in _SAFE_PEER_FALLBACK_ERRORS:
+                return False
+            # memcpy/readback can surface an earlier asynchronous failure. In
+            # particular, an illegal address poisons the CUDA context and
+            # cannot be repaired by switching to one GPU in this process.
+            raise
+
+    return bool(np.array_equal(arrived, _P2P_CANARY_EXPECTED))
+
+
+@cache
+def peer_copy_works(destination: int, source: int) -> bool:
+    """Return whether ``source -> destination`` P2P transfers are usable.
+
+    ``deviceCanAccessPeer`` is only a capability query. Some affected systems
+    report support and return success from ``memcpyPeer`` while leaving the
+    destination unchanged, so this function verifies the transferred bytes.
+    The result is cached per ordered device pair for the lifetime of the
+    process.
+    """
+    if destination == source:
+        return True
+
+    with cp.cuda.Device(destination):
+        if not cp.cuda.runtime.deviceCanAccessPeer(destination, source):
+            return False
+        try:
+            cp.cuda.runtime.deviceEnablePeerAccess(source)
+        except cp.cuda.runtime.CUDARuntimeError as error:
+            if error.status == CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED:
+                pass
+            elif error.status in _SAFE_PEER_FALLBACK_ERRORS:
+                return False
+            else:
+                raise
+
+    return _run_peer_copy_canary(destination, source)
+
+
+# Compatibility name used by the container-level P2P diagnostics.
+peer_copy_verified = peer_copy_works
+
+
+def validate_multi_gpu(
+    device_ids: list[int],
+    *,
+    source_device: int | None = None,
+    gather_device: int | None = None,
+) -> list[int]:
+    """Return usable devices, falling back to one GPU when P2P is unsafe.
+
+    Multi-GPU implementations in this package fan input buffers out from a
+    source GPU and gather results onto a selected GPU. This validates exactly
+    those ordered transfers before sharding or worker threads start.
+    If a required link is unavailable or silently corrupts the canary, the
+    operation falls back to the source GPU and emits one warning per failed
+    link set. Unexpected CUDA errors propagate because they may indicate an
+    already-poisoned context.
+
+    Parameters
+    ----------
+    device_ids
+        Requested execution devices.
+    source_device
+        Device holding shared input buffers. Defaults to the current device.
+    gather_device
+        Device receiving worker results. Defaults to the first requested
+        device.
+
+    Returns
+    -------
+    list[int]
+        The requested devices when all required P2P directions work, otherwise
+        a one-element list containing ``source_device``.
+    """
+    if not device_ids:
+        raise ValueError("device_ids must contain at least one device")
+
+    device_ids = list(dict.fromkeys(device_ids))
+    if source_device is None:
+        source_device = cp.cuda.Device().id
+
+    n_available = cp.cuda.runtime.getDeviceCount()
+    invalid_ids = [
+        device_id
+        for device_id in device_ids
+        if device_id < 0 or device_id >= n_available
+    ]
+    if invalid_ids:
+        raise ValueError(
+            f"Invalid GPU device ID(s): {invalid_ids}. "
+            f"Available devices: {list(range(n_available))}"
+        )
+    if source_device < 0 or source_device >= n_available:
+        raise ValueError(
+            f"Invalid source GPU device ID {source_device}. "
+            f"Available devices: {list(range(n_available))}"
+        )
+
+    if gather_device is None:
+        gather_device = device_ids[0]
+    if gather_device < 0 or gather_device >= n_available:
+        raise ValueError(
+            f"Invalid gather GPU device ID {gather_device}. "
+            f"Available devices: {list(range(n_available))}"
+        )
+
+    required_pairs: set[tuple[int, int]] = set()
+    for device_id in device_ids:
+        if device_id != source_device:
+            required_pairs.add((device_id, source_device))
+        if device_id != gather_device:
+            required_pairs.add((gather_device, device_id))
+
+    failed_pairs = ()
+    for pair in sorted(required_pairs):
+        if not peer_copy_works(*pair):
+            # Once fallback is required, do not exercise any more suspect
+            # links in the caller's CUDA context.
+            failed_pairs = (pair,)
+            break
+    if not failed_pairs:
+        return device_ids
+
+    warning_key = (source_device, failed_pairs)
+    if warning_key not in _WARNED_P2P_FAILURES:
+        transfers = ", ".join(
+            f"GPU {source} -> GPU {destination}" for destination, source in failed_pairs
+        )
+        warnings.warn(
+            "Multi-GPU execution was disabled because the required P2P "
+            f"transfer(s) failed validation: {transfers}. Falling back to "
+            f"GPU {source_device} for this operation. This avoids silent "
+            "result corruption but may be slower.",
+            MultiGPUFallbackWarning,
+            stacklevel=2,
+        )
+        _WARNED_P2P_FAILURES.add(warning_key)
+
+    return [source_device]
+
+
 def parse_device_ids(*, multi_gpu: bool | list[int] | str | None) -> list[int]:
     """Parse multi_gpu parameter into a list of device IDs.
 

src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py

@@ -7,7 +7,11 @@
 import numpy as np
 
 from rapids_singlecell._keys import _preset_obsm_names, _resolve_obsm_key
-from rapids_singlecell._utils import _create_category_index_mapping, parse_device_ids
+from rapids_singlecell._utils import (
+    _copy_to_device_via_host,
+    _create_category_index_mapping,
+    parse_device_ids,
+)
 from rapids_singlecell.squidpy_gpu._utils import _assert_categorical_obs
 
 if TYPE_CHECKING:
@@ -119,13 +123,21 @@ def _subset_to_groups(
             Ordered group names matching the category indices.
         """
         embedding_raw = self._get_embedding(adata)
+        source_device = (
+            embedding_raw.device.id
+            if isinstance(embedding_raw, cp.ndarray)
+            else cp.cuda.Device().id
+        )
         mask, cat_offsets, cell_indices, groups_list = self._subset_indices(
             adata, groupby, needed_groups
         )
-        if mask is not None:
-            embedding = cp.asarray(embedding_raw[mask])
-        else:
-            embedding = cp.asarray(embedding_raw)
+        with cp.cuda.Device(source_device):
+            if mask is not None:
+                embedding = cp.asarray(embedding_raw[mask])
+            else:
+                embedding = cp.asarray(embedding_raw)
+        cat_offsets = _copy_to_device_via_host(cat_offsets, source_device)
+        cell_indices = _copy_to_device_via_host(cell_indices, source_device)
         return embedding, cat_offsets, cell_indices, groups_list
 
     def _subset_indices(