From 598070d1dc15a43fbd9e2c60477858ef6fd350d7 Mon Sep 17 00:00:00 2001
From: Jim Crist-Harif <jcristharif@gmail.com>
Date: Fri, 27 Feb 2026 10:38:47 -0600
Subject: [PATCH 1/4] Add builtin support for patches to accelerator

Previously the accelerator only supported "overrides" (though we weren't
consistent in our terminology). An override:

- Doesn't mutate the original (non-accelerated) module
- Is only visible to consumers not in the accelerator's exclude list.
  For example, the override for `sklearn.linear_models.LinearRegression`
  is visible to external consumers, but any usage within sklearn itself
  will see the original (non-accelerated) version.

Sometimes though we also need to patch sklearn itself, mutating the
original module. Previously we accomplished this in a one-off hack, but
now I have need to make this something more supported. Overrides should
be preferred to patches when possible, but are sometimes necessary to
achieve robust acceleration in the simplest way.

To do this, we:

- Standardize our internal terminology. An `override` is a non-mutating,
  overlay layer in an accelerated module, only visible to external
  consumers. A `patch` is a mutation applied to the original module, and
  is visible to both internal and external consumers.
- Add builtin support for patches to the accelerator, including a test.
- Rename the `_wrappers` directory to `_overrides`.
- Add a new `_patches` directory to contain any patches.
- Move the `sklearn.utils.all_estimators` patch to the new `_patches`
  directory.
---
 .../{_wrappers => _overrides}/__init__.py     |   0
 .../{_wrappers => _overrides}/hdbscan.py      |   0
 .../sklearn/__init__.py                       |   0
 .../sklearn/cluster.py                        |   0
 .../sklearn/covariance.py                     |   0
 .../sklearn/decomposition.py                  |   0
 .../sklearn/ensemble.py                       |   0
 .../sklearn/kernel_ridge.py                   |   0
 .../sklearn/linear_model.py                   |   0
 .../sklearn/manifold.py                       |   0
 .../sklearn/neighbors.py                      |   0
 .../sklearn/preprocessing.py                  |   0
 .../{_wrappers => _overrides}/sklearn/svm.py  |   0
 .../accel/{_wrappers => _overrides}/umap.py   |   0
 python/cuml/cuml/accel/_patches/__init__.py   |   0
 .../cuml/accel/_patches/sklearn/__init__.py   |   0
 .../sklearn/utils.py}                         |  45 +++---
 python/cuml/cuml/accel/accelerator.py         | 138 ++++++++++--------
 python/cuml/cuml/accel/core.py                |  20 ++-
 python/cuml/cuml/accel/estimator_proxy.py     |   2 +-
 python/cuml/cuml/accel/pytest_plugin.py       |  14 +-
 .../cuml/cuml_accel_tests/test_accelerator.py |  26 +++-
 22 files changed, 138 insertions(+), 107 deletions(-)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/__init__.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/hdbscan.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/__init__.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/cluster.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/covariance.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/decomposition.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/ensemble.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/kernel_ridge.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/linear_model.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/manifold.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/neighbors.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/preprocessing.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/sklearn/svm.py (100%)
 rename python/cuml/cuml/accel/{_wrappers => _overrides}/umap.py (100%)
 create mode 100644 python/cuml/cuml/accel/_patches/__init__.py
 create mode 100644 python/cuml/cuml/accel/_patches/sklearn/__init__.py
 rename python/cuml/cuml/accel/{_sklearn_patch.py => _patches/sklearn/utils.py} (56%)

diff --git a/python/cuml/cuml/accel/_wrappers/__init__.py b/python/cuml/cuml/accel/_overrides/__init__.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/__init__.py
rename to python/cuml/cuml/accel/_overrides/__init__.py
diff --git a/python/cuml/cuml/accel/_wrappers/hdbscan.py b/python/cuml/cuml/accel/_overrides/hdbscan.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/hdbscan.py
rename to python/cuml/cuml/accel/_overrides/hdbscan.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/__init__.py b/python/cuml/cuml/accel/_overrides/sklearn/__init__.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/__init__.py
rename to python/cuml/cuml/accel/_overrides/sklearn/__init__.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/cluster.py b/python/cuml/cuml/accel/_overrides/sklearn/cluster.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/cluster.py
rename to python/cuml/cuml/accel/_overrides/sklearn/cluster.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/covariance.py b/python/cuml/cuml/accel/_overrides/sklearn/covariance.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/covariance.py
rename to python/cuml/cuml/accel/_overrides/sklearn/covariance.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/decomposition.py b/python/cuml/cuml/accel/_overrides/sklearn/decomposition.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/decomposition.py
rename to python/cuml/cuml/accel/_overrides/sklearn/decomposition.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/ensemble.py b/python/cuml/cuml/accel/_overrides/sklearn/ensemble.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/ensemble.py
rename to python/cuml/cuml/accel/_overrides/sklearn/ensemble.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/kernel_ridge.py b/python/cuml/cuml/accel/_overrides/sklearn/kernel_ridge.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/kernel_ridge.py
rename to python/cuml/cuml/accel/_overrides/sklearn/kernel_ridge.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/linear_model.py b/python/cuml/cuml/accel/_overrides/sklearn/linear_model.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/linear_model.py
rename to python/cuml/cuml/accel/_overrides/sklearn/linear_model.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/manifold.py b/python/cuml/cuml/accel/_overrides/sklearn/manifold.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/manifold.py
rename to python/cuml/cuml/accel/_overrides/sklearn/manifold.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/neighbors.py b/python/cuml/cuml/accel/_overrides/sklearn/neighbors.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/neighbors.py
rename to python/cuml/cuml/accel/_overrides/sklearn/neighbors.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/preprocessing.py b/python/cuml/cuml/accel/_overrides/sklearn/preprocessing.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/preprocessing.py
rename to python/cuml/cuml/accel/_overrides/sklearn/preprocessing.py
diff --git a/python/cuml/cuml/accel/_wrappers/sklearn/svm.py b/python/cuml/cuml/accel/_overrides/sklearn/svm.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/sklearn/svm.py
rename to python/cuml/cuml/accel/_overrides/sklearn/svm.py
diff --git a/python/cuml/cuml/accel/_wrappers/umap.py b/python/cuml/cuml/accel/_overrides/umap.py
similarity index 100%
rename from python/cuml/cuml/accel/_wrappers/umap.py
rename to python/cuml/cuml/accel/_overrides/umap.py
diff --git a/python/cuml/cuml/accel/_patches/__init__.py b/python/cuml/cuml/accel/_patches/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/python/cuml/cuml/accel/_patches/sklearn/__init__.py b/python/cuml/cuml/accel/_patches/sklearn/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/python/cuml/cuml/accel/_sklearn_patch.py b/python/cuml/cuml/accel/_patches/sklearn/utils.py
similarity index 56%
rename from python/cuml/cuml/accel/_sklearn_patch.py
rename to python/cuml/cuml/accel/_patches/sklearn/utils.py
index d9acd2fe15..16439b5735 100644
--- a/python/cuml/cuml/accel/_sklearn_patch.py
+++ b/python/cuml/cuml/accel/_patches/sklearn/utils.py
@@ -1,24 +1,32 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
+import functools
 from collections import defaultdict
 from operator import itemgetter
 
-from sklearn.utils.discovery import all_estimators as sklearn_all_estimators
+from sklearn.utils.discovery import all_estimators as _all_estimators
 
 from cuml.accel.estimator_proxy import is_proxy
 
+__all__ = ("all_estimators",)
 
-def _patched_all_estimators(*args, **kwargs):
-    """Monkeypatch sklearn's all_estimators to prioritize proxy estimators.
 
-    This function replaces sklearn's all_estimators function with a version
-    that filters out duplicate estimator names, keeping only proxy estimators
-    when both proxy and non-proxy versions exist.
-    """
+@functools.wraps(_all_estimators)
+def all_estimators(*args, **kwargs):
+    # This function replaces sklearn's all_estimators function with a version
+    # that filters out duplicate estimator names, keeping only proxy estimators
+    # when both proxy and non-proxy versions exist.
+    #
+    # When the accelerator is installed, sklearn's all_estimators() returns
+    # duplicate entries for the same estimator name (e.g., LinearSVC appears
+    # both with sklearn.svm.LinearSVC and sklearn._classes.LinearSVC). This causes
+    # a TypeError during test collection when sklearn tries to sort the
+    # estimators as it attempts to sort based on the type.
+
     # Obtain the list of all estimators from sklearn
-    ret = sklearn_all_estimators(*args, **kwargs)
+    ret = _all_estimators(*args, **kwargs)
 
     # Group estimators by name
     estimator_groups = defaultdict(list)
@@ -37,22 +45,3 @@ def _patched_all_estimators(*args, **kwargs):
 
     # Return the sorted list of estimators like the original
     return sorted(set(estimators), key=itemgetter(0))
-
-
-def apply_sklearn_patches():
-    """Apply all sklearn patches necessary for the accelerator testing."""
-
-    # Monkeypatch sklearn's all_estimators to prioritize proxy estimators
-    #
-    # When the accelerator is installed, sklearn's all_estimators() returns
-    # duplicate entries for the same estimator name (e.g., LinearSVC appears
-    # both with sklearn.svm.LinearSVC and sklearn._classes.LinearSVC). This causes
-    # a TypeError during test collection when sklearn tries to sort the
-    # estimators as it attempts to sort based on the type.
-    #
-    # The patch filters out duplicates by keeping only proxy estimators when
-    # multiple classes with the same name exist, ensuring test collection
-    # succeeds.
-    import sklearn.utils
-
-    sklearn.utils.all_estimators = _patched_all_estimators
diff --git a/python/cuml/cuml/accel/accelerator.py b/python/cuml/cuml/accel/accelerator.py
index 2343b8d9a8..7f9fd44017 100644
--- a/python/cuml/cuml/accel/accelerator.py
+++ b/python/cuml/cuml/accel/accelerator.py
@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -18,7 +18,7 @@
 __all__ = ("Accelerator",)
 
 
-PatchType = Union[
+LazyNamespace = Union[
     Callable[[types.ModuleType], dict[str, Any]],
     dict[str, Any],
     str,
@@ -110,38 +110,43 @@ def wrap_module(
     return out
 
 
-def patch_module(module: AccelModule, patch: PatchType) -> None:
-    """Patch an accelerated module.
+class ModuleTransform:
+    def __init__(
+        self,
+        override: LazyNamespace | None = None,
+        patch: LazyNamespace | None = None,
+    ):
+        self.override = override
+        self.patch = patch
 
-    Parameters
-    ----------
-    module : AccelModule
-        The accelerated module.
-    patch : mapping, callable, or str
-        One of the following:
-
-        - A mapping of attributes to override in the accelerated module.
-        - A callable taking the original module and returning a mapping of attributes
-          to override in the accelerated module.
-        - A string name of a module to import containing overrides. All names
-          specified in `__all__` in the module will be used to override attributes in
-          the accelerated module.
-    """
-    assert isinstance(module, AccelModule)
-    if callable(patch):
-        overrides = patch(module)
-    elif isinstance(patch, str):
-        new_module = importlib.import_module(patch)
-        if (names := getattr(new_module, "__all__", None)) is None:
-            raise ValueError(
-                f"Module `{patch}` must define `__all__` to specify the names of "
-                "the attributes to override in the accelerated module"
-            )
-        overrides = {name: getattr(new_module, name) for name in names}
-    else:
-        overrides = patch
-
-    module._accel_overrides.update(overrides)
+    @staticmethod
+    def _load_namespace(
+        module: AccelModule, namespace: LazyNamespace
+    ) -> dict[str, Any]:
+        assert isinstance(module, AccelModule)
+        if callable(namespace):
+            return namespace(module)
+        elif isinstance(namespace, str):
+            new_module = importlib.import_module(namespace)
+            if (names := getattr(new_module, "__all__", None)) is None:
+                raise ValueError(
+                    f"Module `{namespace}` must define `__all__` to specify the names of "
+                    "the attributes to override/patch in the accelerated module"
+                )
+            return {name: getattr(new_module, name) for name in names}
+        else:
+            return namespace
+
+    def apply(self, module: AccelModule) -> None:
+        """Load and apply patches/overrides to an accelerated module."""
+        if self.patch is not None:
+            ns = self._load_namespace(module, self.patch)
+            for k, v in ns.items():
+                setattr(module._accel_module, k, v)
+
+        if self.override is not None:
+            ns = self._load_namespace(module, self.override)
+            module._accel_overrides.update(ns)
 
 
 class AccelLoader(importlib.abc.Loader):
@@ -150,11 +155,11 @@ class AccelLoader(importlib.abc.Loader):
     def __init__(
         self,
         spec: importlib.machinery.ModuleSpec,
-        patch: PatchType,
+        transform: ModuleTransform,
         exclude: Callable[[str], bool] | None = None,
     ) -> None:
         self._spec = spec
-        self._patch = patch
+        self._transform = transform
         self._exclude = exclude
 
     def create_module(
@@ -168,7 +173,7 @@ def exec_module(self, module: types.ModuleType) -> None:
         assert isinstance(module, AccelModule)
         assert self._spec.loader is not None
         self._spec.loader.exec_module(module._accel_module)
-        patch_module(module, self._patch)
+        self._transform.apply(module)
 
 
 class AccelFinder(importlib.abc.MetaPathFinder):
@@ -192,7 +197,7 @@ def find_spec(
         if fullname in self._importing:
             return None
 
-        if (patch := self.accelerator.patches.get(fullname)) is None:
+        if (transform := self.accelerator.transforms.get(fullname)) is None:
             return None
 
         try:
@@ -206,7 +211,7 @@ def find_spec(
 
         spec = importlib.machinery.ModuleSpec(
             name=fullname,
-            loader=AccelLoader(real_spec, patch, self.accelerator.exclude),
+            loader=AccelLoader(real_spec, transform, self.accelerator.exclude),
             origin=real_spec.origin,
             loader_state=real_spec.loader_state,
             is_package=real_spec.submodule_search_locations is not None,
@@ -227,7 +232,7 @@ class Accelerator:
     """
 
     exclude: Callable[[str], bool]
-    patches: dict[str, PatchType]
+    transforms: dict[str, ModuleTransform]
     _lock: RLock
     _installed: bool
 
@@ -239,7 +244,7 @@ def __init__(
         else:
             self.exclude = frozenset(exclude or ()).__contains__
 
-        self.patches = {}
+        self.transforms = {}
         self._lock = RLock()
         self._installed = False
 
@@ -257,22 +262,38 @@ def enabled(self) -> bool:
         # minimizing the changes needed if/when we want to add that feature.
         return self.installed
 
-    def register(self, name: str, patch: PatchType):
-        """Register a new patch for a module.
+    def register(
+        self,
+        name: str,
+        override: LazyNamespace | None = None,
+        patch: LazyNamespace | None = None,
+    ):
+        """Register a new override or patch for a module.
+
+        Overrides are only visible to non-excluded modules, and don't mutate
+        the original module. Patches apply everywhere and do mutate the original
+        module.
+
+        For example, if `sklearn` is in the `exclude` list for the accelerator,
+        an override for `sklearn.linear_models.LinearRegression` won't be
+        visible to consumers within `sklearn` itself (they'll still get the
+        original `LinearRegression`). In contrast, a `patch` will be visible
+        everywhere, and including for consumers within `sklearn` itself.
 
         Parameters
         ----------
         name : str
             The name of the unaccelerated module to patch.
-        patch : mapping, callable, or str
+        override, patch : mapping, callable, or str
             May be one of the following:
 
-            - A mapping of attributes to override in the accelerated module.
+            - A mapping of attributes to override/patch in the accelerated
+              module.
             - A callable taking the original module and returning a mapping of
-              attributes to override in the accelerated module.
-            - A string name of a module to import containing overrides. All names
-              specified in `__all__` in the module will be used to override
-              attributes in the accelerated module.
+              attributes to override/patch in the accelerated module.
+            - A string name of a module to import containing overrides or
+              patches. All names specified in `__all__` in the module will be
+              used to override/patch attributes in the accelerated module.
 
         Examples
         --------
@@ -287,10 +308,10 @@ def register(self, name: str, patch: PatchType):
         >>> accel.register("foobar", "fast.foobar")
         """
         assert not self._installed
-        assert name not in self.patches
-        self.patches[name] = patch
+        assert name not in self.transforms
+        self.transforms[name] = ModuleTransform(override=override, patch=patch)
 
-    def _maybe_patch(self, name: str) -> None:
+    def _maybe_transform(self, name: str) -> None:
         if (module := sys.modules.get(name)) is None:
             # Not imported yet, import system will load patch lazily later
             return
@@ -312,14 +333,15 @@ def _maybe_patch(self, name: str) -> None:
         if getattr(parent, child_name, None) is module:
             setattr(parent, child_name, accelerated)
 
-        # 4. Apply the patch
-        patch_module(accelerated, self.patches[name])
+        # 4. Apply module transforms
+        self.transforms[name].apply(accelerated)
 
     def install(self) -> None:
         """Install the accelerator.
 
-        This installs the import hooks to intercept future imports of accelerated modules. It also
-        patches previous imports of these modules in a best-effort approach.
+        This installs the import hooks to intercept future imports of
+        accelerated modules. It also wraps/overrides previous imports of these
+        modules in a best-effort approach.
         """
         with self._lock:
             if self._installed:
@@ -328,8 +350,8 @@ def install(self) -> None:
             # Install the import hook. This handles patching any modules imported later.
             sys.meta_path.insert(0, AccelFinder(self))
 
-            # Patch any modules that are already imported.
-            for name in self.patches:
-                self._maybe_patch(name)
+            # Wrap any modules that are already imported.
+            for name in self.transforms:
+                self._maybe_transform(name)
 
             self._installed = True
diff --git a/python/cuml/cuml/accel/core.py b/python/cuml/cuml/accel/core.py
index 9acd6385f8..5a0360501d 100644
--- a/python/cuml/cuml/accel/core.py
+++ b/python/cuml/cuml/accel/core.py
@@ -75,7 +75,7 @@ def debug(self, msg: str) -> None:
 logger = Logger()
 
 
-ACCELERATED_MODULES = [
+_OVERRIDES = {
     "hdbscan",
     "sklearn.cluster",
     "sklearn.covariance",
@@ -88,7 +88,13 @@ def debug(self, msg: str) -> None:
     "sklearn.preprocessing",
     "sklearn.svm",
     "umap",
-]
+}
+
+_PATCHES = {
+    "sklearn.utils",
+}
+
+ACCELERATED_MODULES = sorted(_OVERRIDES.union(_PATCHES))
 
 
 def _exclude_from_acceleration(module: str) -> bool:
@@ -107,7 +113,15 @@ def _exclude_from_acceleration(module: str) -> bool:
 
 ACCEL = Accelerator(exclude=_exclude_from_acceleration)
 for module in ACCELERATED_MODULES:
-    ACCEL.register(module, f"cuml.accel._wrappers.{module}")
+    ACCEL.register(
+        module,
+        override=(
+            f"cuml.accel._overrides.{module}" if module in _OVERRIDES else None
+        ),
+        patch=(
+            f"cuml.accel._patches.{module}" if module in _PATCHES else None
+        ),
+    )
 
 
 def _is_concurrent_managed_access_supported():
diff --git a/python/cuml/cuml/accel/estimator_proxy.py b/python/cuml/cuml/accel/estimator_proxy.py
index 752aa92aeb..c2e711b26e 100644
--- a/python/cuml/cuml/accel/estimator_proxy.py
+++ b/python/cuml/cuml/accel/estimator_proxy.py
@@ -105,7 +105,7 @@ class ProxyBase(BaseEstimator):
     to the GPU with the rules above. If this method raises a ``UnsupportedOnGPU`` error
     then the proxy will fallback to CPU.
 
-    See the definitions in ``cuml.accel._wrappers.linear_model`` for examples.
+    See the definitions in ``cuml.accel._overrides.linear_model`` for examples.
     """
 
     # A set of attribute names that aren't supported by `cuml.accel`.
diff --git a/python/cuml/cuml/accel/pytest_plugin.py b/python/cuml/cuml/accel/pytest_plugin.py
index 72f8aec2ee..86df5e7ec6 100644
--- a/python/cuml/cuml/accel/pytest_plugin.py
+++ b/python/cuml/cuml/accel/pytest_plugin.py
@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -10,7 +10,6 @@
 
 from packaging.requirements import Requirement
 
-from cuml.accel._sklearn_patch import apply_sklearn_patches
 from cuml.accel.core import install
 
 
@@ -32,16 +31,7 @@ class UnmatchedXfailTests(UserWarning):
 def pytest_load_initial_conftests(early_config, parser, args):
     # https://docs.pytest.org/en/7.1.x/reference/\
     # reference.html#pytest.hookspec.pytest_load_initial_conftests
-
-    # Apply sklearn patches BEFORE installing cuml.accel to prevent duplicates
-    apply_sklearn_patches()
-
-    try:
-        install()
-    except RuntimeError:
-        raise RuntimeError(
-            "An existing plugin has already loaded sklearn. Interposing failed."
-        )
+    install()
 
 
 def pytest_addoption(parser):
diff --git a/python/cuml/cuml_accel_tests/test_accelerator.py b/python/cuml/cuml_accel_tests/test_accelerator.py
index 9236418d10..7b13ce5b8d 100644
--- a/python/cuml/cuml_accel_tests/test_accelerator.py
+++ b/python/cuml/cuml_accel_tests/test_accelerator.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 
 from __future__ import annotations
@@ -164,10 +164,10 @@ def fizz():
     assert mod.fizzbuzz() == "fizzbuzz"
 
 
-def test_accelerator_import_in_patch(mockmod):
-    """Check that imports of the original module work fine within a patch"""
+def test_accelerator_import_in_override(mockmod):
+    """Check that imports of the original module work fine within a override"""
 
-    def patch(module):
+    def override(module):
         # Same as `from {mockmod}.utils import fizz`
         fizz = importlib.import_module(f"{mockmod}.utils").fizz
         assert fizz is module.fizz
@@ -175,7 +175,7 @@ def patch(module):
         return {"fizz": lambda: fizz().upper()}
 
     accel = Accelerator()
-    accel.register(f"{mockmod}.utils", patch)
+    accel.register(f"{mockmod}.utils", override)
     accel.install()
 
     mod = importlib.import_module(mockmod)
@@ -201,6 +201,22 @@ def fizz():
     assert mod.utils.fizz is fizz
 
 
+def test_accelerator_module_patch(mockmod):
+    def fizz():
+        return "FIZZ"
+
+    accel = Accelerator()
+    accel.register(f"{mockmod}.utils", patch={"fizz": fizz})
+    accel.install()
+
+    mod = importlib.import_module(mockmod)
+    # Patch applied to original module
+    assert mod.utils._accel_module.fizz is fizz
+    assert mod.utils.fizz() == "FIZZ"
+    assert mod.fizz() == "FIZZ"
+    assert mod.fizzbuzz() == "FIZZbuzz"
+
+
 def test_accel_module(mockmod):
     orig_mod = importlib.import_module(mockmod)
 

From fc6bfcf5f44bc2ac8104b7eba4be11b5724ffee8 Mon Sep 17 00:00:00 2001
From: Jim Crist-Harif <jcristharif@gmail.com>
Date: Fri, 27 Feb 2026 13:32:26 -0600
Subject: [PATCH 2/4] Optimize `Pipeline` data transfer in cuml.accel

This adds an optimization to `Pipeline` when running under cuml.accel.
If a pipeline is composed of accelerated estimators (or only starts with
unaccelerated estimators), then the accelerated estimators will now pass
data on device rather than doing a series of device<->host transfers.
---
 .../cuml/accel/_patches/sklearn/pipeline.py   | 108 +++++++++++
 python/cuml/cuml/accel/core.py                |   1 +
 python/cuml/cuml/accel/estimator_proxy.py     |  13 ++
 python/cuml/cuml_accel_tests/test_pipeline.py | 169 +++++++++++++++++-
 4 files changed, 290 insertions(+), 1 deletion(-)
 create mode 100644 python/cuml/cuml/accel/_patches/sklearn/pipeline.py

diff --git a/python/cuml/cuml/accel/_patches/sklearn/pipeline.py b/python/cuml/cuml/accel/_patches/sklearn/pipeline.py
new file mode 100644
index 0000000000..629f0ea86a
--- /dev/null
+++ b/python/cuml/cuml/accel/_patches/sklearn/pipeline.py
@@ -0,0 +1,108 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+import functools
+import inspect
+
+import cupy as cp
+from cupyx.scipy.sparse import issparse as is_cp_sparse
+from sklearn.pipeline import Pipeline
+from sklearn.utils.metaestimators import available_if
+
+from cuml.accel.estimator_proxy import is_proxy
+from cuml.internals.global_settings import GlobalSettings
+from cuml.internals.outputs import using_output_type
+
+__all__ = ("Pipeline",)
+
+
+def get_output_type(pipeline, reverse=False):
+    """Determine the output type to use for a pipeline method.
+
+    If all cupy-producing steps are only followed by cupy-consuming steps after
+    them, then we may enable cupy outputs. Otherwise we need to fallback to
+    numpy. A pipeline like ``numpy -> numpy -> cupy -> cupy`` is fine to enable
+    cupy for, but ``numpy -> cupy -> numpy -> cupy`` isn't.
+    """
+
+    def flat_steps(pipeline):
+        """Iterate over steps potentially nested pipelines"""
+        for name, step in (
+            reversed(pipeline.steps) if reverse else pipeline.steps
+        ):
+            if step in (None, "passthrough"):
+                continue
+            if isinstance(step, Pipeline):
+                yield from flat_steps(step)
+            else:
+                yield step
+
+    step_iter = flat_steps(pipeline)
+
+    # Skip over any non-cupy producing steps
+    for step in step_iter:
+        if is_proxy(step):
+            break
+    else:
+        # No steps produce cupy
+        return "numpy"
+
+    # If any remaining steps don't consume cupy, we need to fallback to numpy
+    for step in step_iter:
+        if not is_proxy(step):
+            return "numpy"
+    # Tail of the pipeline supports cupy, we can use cupy
+    return "cupy"
+
+
+def patch_method(name):
+    """Patch a sklearn Pipeline method to reduce device<->host transfers."""
+    orig_method = inspect.getattr_static(Pipeline, name)
+    # Unwrap @available_if decorated methods
+    if (check := getattr(orig_method, "check", None)) is not None:
+        orig_method = orig_method.fn
+
+    # `inverse_transform` processes steps in reverse
+    reverse = name == "inverse_transform"
+
+    @functools.wraps(orig_method)
+    def method(self, *args, **kwargs):
+        if name.startswith("fit"):
+            # Validate hyperparameters first if a fit call
+            self._validate_params()
+
+        # Run the original method within the proper output type context
+        with using_output_type(get_output_type(self, reverse=reverse)):
+            out = orig_method(self, *args, **kwargs)
+
+        # Transform output to numpy/scipy.sparse if requested
+        if GlobalSettings().output_type in (None, "numpy") and (
+            isinstance(out, cp.ndarray) or is_cp_sparse(out)
+        ):
+            out = out.get()
+
+        return out
+
+    # Rewrap @available_if decorated methods
+    if check is not None:
+        method = available_if(check)(method)
+
+    setattr(Pipeline, name, method)
+
+
+# These methods run pipeline operations in a series, and need to be
+# patched to reduce data movement between a series of accelerated
+# estimators.
+for method_name in [
+    "decision_function",
+    "fit",
+    "fit_predict",
+    "fit_transform",
+    "inverse_transform",
+    "predict",
+    "predict_log_proba",
+    "predict_proba",
+    "score",
+    "score_samples",
+    "transform",
+]:
+    patch_method(method_name)
diff --git a/python/cuml/cuml/accel/core.py b/python/cuml/cuml/accel/core.py
index 5a0360501d..ffc896fd5a 100644
--- a/python/cuml/cuml/accel/core.py
+++ b/python/cuml/cuml/accel/core.py
@@ -91,6 +91,7 @@ def debug(self, msg: str) -> None:
 }
 
 _PATCHES = {
+    "sklearn.pipeline",
     "sklearn.utils",
 }
 
diff --git a/python/cuml/cuml/accel/estimator_proxy.py b/python/cuml/cuml/accel/estimator_proxy.py
index c2e711b26e..1048832926 100644
--- a/python/cuml/cuml/accel/estimator_proxy.py
+++ b/python/cuml/cuml/accel/estimator_proxy.py
@@ -6,7 +6,9 @@
 import functools
 from typing import Any
 
+import cupy as cp
 import sklearn
+from cupyx.scipy.sparse import issparse as is_cp_sparse
 from packaging.version import Version
 from sklearn.base import (
     BaseEstimator,
@@ -31,6 +33,11 @@ def is_proxy(instance_or_class) -> bool:
     return issubclass(cls, ProxyBase)
 
 
+def ensure_host(x):
+    """Convert any cupy/cupyx.scipy.sparse inputs to their host equivalents"""
+    return x.get() if (isinstance(x, cp.ndarray) or is_cp_sparse(x)) else x
+
+
 class _ReconstructProxy:
     """A function for reconstructing serialized estimators.
 
@@ -318,6 +325,12 @@ def _call_method(self, method: str, *args: Any, **kwargs: Any) -> Any:
 
         # Failed to run on GPU, fallback to CPU
         self._sync_attrs_to_cpu()
+        # Ensure the arguments are on host for the CPU fallback. This is _usually_
+        # already True, but in certain cases (a pipeline with optimized data transfer)
+        # we may need to migrate. In those cases the inputs will only ever be
+        # cupy/cupyx.scipy.sparse objects, so that's all we need to handle here.
+        args = [ensure_host(a) for a in args]
+        kwargs = {k: ensure_host(v) for k, v in kwargs.items()}
         with profilers.track_cpu_call(
             qualname, reason=reason or "Estimator not fit on GPU"
         ):
diff --git a/python/cuml/cuml_accel_tests/test_pipeline.py b/python/cuml/cuml_accel_tests/test_pipeline.py
index ea6dd71039..3175c2637f 100644
--- a/python/cuml/cuml_accel_tests/test_pipeline.py
+++ b/python/cuml/cuml_accel_tests/test_pipeline.py
@@ -1,10 +1,14 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
+import types
 
+import cupy as cp
+import numpy as np
 import pytest
 import scipy as sp
+from sklearn.base import BaseEstimator
 from sklearn.cluster import DBSCAN, KMeans
 from sklearn.datasets import make_classification, make_regression
 from sklearn.decomposition import PCA, TruncatedSVD
@@ -22,9 +26,55 @@
     NearestNeighbors,
 )
 from sklearn.pipeline import Pipeline, make_pipeline
+from sklearn.preprocessing import StandardScaler
 from umap import UMAP
 
 
+class MockMethod:
+    """A simple mock for method types, properly handling method binding"""
+
+    def __init__(self, method):
+        self._method = method
+
+    def __get__(self, obj, objtype=None):
+        if obj is None:
+            return self
+        return types.MethodType(self, obj)
+
+    def __call__(self, *args, **kwargs):
+        self.args = args[1:]  # drop self
+        self.kwargs = kwargs
+        return self._method(*args, **kwargs)
+
+
+@pytest.fixture
+def patch_methods(monkeypatch):
+    """A fixture for patching one or more methods on a class"""
+
+    def patch(cls, *methods):
+        for method in methods:
+            monkeypatch.setattr(cls, method, MockMethod(getattr(cls, method)))
+
+    return patch
+
+
+class HostTransformer(BaseEstimator):
+    """A no-op host-only transformer"""
+
+    def fit(self, X, y=None):
+        assert isinstance(X, np.ndarray)
+        self.n_features_in_ = X.shape[1]
+        return self
+
+    def transform(self, X):
+        assert isinstance(X, np.ndarray)
+        return X
+
+    def inverse_transform(self, X):
+        assert isinstance(X, np.ndarray)
+        return X
+
+
 @pytest.fixture
 def classification_data():
     # Create a synthetic dataset for binary classification
@@ -151,3 +201,120 @@ def test_pipeline_adding_none_value_as_labels(classification_data):
 
     pipeline = make_pipeline(TruncatedSVD(n_components=20))
     pipeline.fit_transform(X_train)
+
+
+@pytest.mark.parametrize(
+    "order, enabled",
+    [
+        ("host", False),
+        ("host-host", False),
+        ("device-host", False),
+        ("device", True),
+        ("device-device", True),
+        ("host-device", True),
+    ],
+)
+@pytest.mark.parametrize("nested", [False, True])
+def test_pipeline_data_transfer(
+    order, enabled, nested, regression_data, patch_methods
+):
+    patch_methods(Ridge, "fit", "predict")
+    X_train, X_test, y_train, y_test = regression_data
+    xp = cp if enabled else np
+
+    steps = [
+        StandardScaler() if step == "device" else HostTransformer()
+        for step in order.split("-")
+    ]
+    if nested:
+        pipeline = make_pipeline(make_pipeline(*steps), Ridge())
+    else:
+        pipeline = make_pipeline(*steps, Ridge())
+
+    pipeline.fit(X_train, y_train)
+    assert isinstance(Ridge.fit.args[0], xp.ndarray)
+    out = pipeline.predict(X_test)
+    assert isinstance(Ridge.predict.args[0], xp.ndarray)
+    # User-facing output is always numpy
+    assert isinstance(out, np.ndarray)
+
+
+@pytest.mark.parametrize(
+    "pipeline, scaler, pca",
+    [
+        (
+            make_pipeline(StandardScaler(), PCA()),
+            (False, True),
+            (True, False),
+        ),
+        (
+            make_pipeline(HostTransformer(), StandardScaler(), PCA()),
+            (False, False),
+            (True, False),
+        ),
+        (
+            make_pipeline(StandardScaler(), HostTransformer(), PCA()),
+            (False, False),
+            (False, False),
+        ),
+        (
+            make_pipeline(StandardScaler(), PCA(), HostTransformer()),
+            (False, True),
+            (False, False),
+        ),
+        (
+            make_pipeline(
+                make_pipeline(HostTransformer(), StandardScaler()), PCA()
+            ),
+            (False, False),
+            (True, False),
+        ),
+        (
+            make_pipeline(
+                make_pipeline(StandardScaler(), PCA()), HostTransformer()
+            ),
+            (False, True),
+            (True, False),
+        ),
+    ],
+)
+def test_pipeline_transform_data_transfer(
+    pipeline, scaler, pca, regression_data, patch_methods
+):
+    patch_methods(StandardScaler, "transform", "inverse_transform")
+    patch_methods(PCA, "transform", "inverse_transform")
+    X = regression_data[0]
+
+    pipeline.fit(X)
+
+    def on_device(method, enabled):
+        xp = cp if enabled else np
+        return isinstance(method.args[0], xp.ndarray)
+
+    out = pipeline.transform(X)
+    assert isinstance(out, np.ndarray)
+    assert on_device(PCA.transform, pca[0])
+    assert on_device(StandardScaler.transform, scaler[0])
+
+    out = pipeline.inverse_transform(X)
+    assert isinstance(out, np.ndarray)
+    assert on_device(PCA.inverse_transform, pca[1])
+    assert on_device(StandardScaler.inverse_transform, scaler[1])
+
+
+def test_pipeline_data_transfer_with_host_fallback(
+    regression_data, patch_methods
+):
+    """Intermediates passed on device, but step falls back to CPU for other reasons.
+
+    Smoketests that the proxy converts device->host before fallback is called."""
+    patch_methods(Ridge, "fit", "predict")
+    X_train, X_test, y_train, y_test = regression_data
+
+    pipeline = make_pipeline(StandardScaler(), Ridge(positive=True))
+    pipeline.fit(X_train, y_train)
+    assert isinstance(Ridge.fit.args[0], cp.ndarray)
+    out = pipeline.predict(X_test)
+    assert isinstance(Ridge.predict.args[0], cp.ndarray)
+    # User-facing output is always numpy
+    assert isinstance(out, np.ndarray)

From fbaa95c6207ae024c666b8d813389f6b5394184c Mon Sep 17 00:00:00 2001
From: Jim Crist-Harif <jcristharif@gmail.com>
Date: Mon, 2 Mar 2026 10:43:53 -0600
Subject: [PATCH 3/4] Fix `set_output` in pipelines

---
 python/cuml/cuml/accel/estimator_proxy.py     | 11 +++++++++--
 python/cuml/cuml_accel_tests/test_pipeline.py |  9 +++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/python/cuml/cuml/accel/estimator_proxy.py b/python/cuml/cuml/accel/estimator_proxy.py
index 1048832926..a8fd90ac2a 100644
--- a/python/cuml/cuml/accel/estimator_proxy.py
+++ b/python/cuml/cuml/accel/estimator_proxy.py
@@ -15,7 +15,10 @@
     ClassNamePrefixFeaturesOutMixin,
     OneToOneFeatureMixin,
 )
-from sklearn.utils._set_output import _wrap_data_with_container
+from sklearn.utils._set_output import (
+    _get_output_config,
+    _wrap_data_with_container,
+)
 
 from cuml.accel import profilers
 from cuml.accel.core import logger
@@ -266,7 +269,11 @@ def _call_gpu_method(self, method: str, *args: Any, **kwargs: Any) -> Any:
         out = gpu_func(*args, **kwargs)
 
         if method in ("transform", "fit_transform"):
-            # Ensure transform result is properly wrapped for `set_output`
+            # Properly wrap output of transform following `set_output` config.
+            # For non-default configuration, we need to ensure the output is
+            # on host before forwarding.
+            if _get_output_config("transform", self)["dense"] != "default":
+                out = ensure_host(out)
             out = _wrap_data_with_container("transform", out, args[0], self)
 
         return self if out is self._gpu else out
diff --git a/python/cuml/cuml_accel_tests/test_pipeline.py b/python/cuml/cuml_accel_tests/test_pipeline.py
index 3175c2637f..f169f41076 100644
--- a/python/cuml/cuml_accel_tests/test_pipeline.py
+++ b/python/cuml/cuml_accel_tests/test_pipeline.py
@@ -6,6 +6,7 @@
 
 import cupy as cp
 import numpy as np
+import pandas as pd
 import pytest
 import scipy as sp
 from sklearn.base import BaseEstimator
@@ -318,3 +319,11 @@ def test_pipeline_data_transfer_with_host_fallback(
     assert isinstance(Ridge.predict.args[0], cp.ndarray)
     # User-facing output is always numpy
     assert isinstance(out, np.ndarray)
+
+
+def test_pipeline_set_output():
+    X, _ = make_regression(random_state=42)
+    X2 = make_pipeline(
+        StandardScaler().set_output(transform="pandas")
+    ).fit_transform(X)
+    assert isinstance(X2, pd.DataFrame)

From 5130dadfd6572cc2185fd5e227ae038c4e199080 Mon Sep 17 00:00:00 2001
From: Jim Crist-Harif <jcristharif@gmail.com>
Date: Mon, 2 Mar 2026 11:02:59 -0600
Subject: [PATCH 4/4] Only optimize data transfer for *transform* methods

Some methods may return dtypes not supported on device, and others would
never be chained in a context where this pipeline optimization would
take effect. Easier to only enabled device arrays for *transform*
methods than optionally gate the inverse.
---
 python/cuml/cuml/accel/estimator_proxy.py     | 18 +++++++++++++-----
 python/cuml/cuml_accel_tests/test_pipeline.py | 15 +++++++++++++++
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/python/cuml/cuml/accel/estimator_proxy.py b/python/cuml/cuml/accel/estimator_proxy.py
index a8fd90ac2a..6161eda48f 100644
--- a/python/cuml/cuml/accel/estimator_proxy.py
+++ b/python/cuml/cuml/accel/estimator_proxy.py
@@ -23,6 +23,7 @@
 from cuml.accel import profilers
 from cuml.accel.core import logger
 from cuml.internals.interop import UnsupportedOnGPU, is_fitted
+from cuml.internals.outputs import using_output_type
 
 SKLEARN_18 = Version(sklearn.__version__) >= Version("1.8.0.dev0")
 
@@ -266,14 +267,21 @@ def _call_gpu_method(self, method: str, *args: Any, **kwargs: Any) -> Any:
             if (gpu_func := getattr(self._gpu, method, None)) is None:
                 raise UnsupportedOnGPU("Method is not implemented in cuml")
 
-        out = gpu_func(*args, **kwargs)
+        # Only transform/fit_transform/inverse_transform with default
+        # set_output config may return device arrays (to support optimized
+        # pipeline data transfers). All other methods must return on host.
+        may_return_on_device = (
+            method in ("transform", "fit_transform", "inverse_transform")
+            and _get_output_config("transform", self)["dense"] == "default"
+        )
+        if may_return_on_device:
+            out = gpu_func(*args, **kwargs)
+        else:
+            with using_output_type("numpy"):
+                out = gpu_func(*args, **kwargs)
 
         if method in ("transform", "fit_transform"):
             # Properly wrap output of transform following `set_output` config.
-            # For non-default configuration, we need to ensure the output is
-            # on host before forwarding.
-            if _get_output_config("transform", self)["dense"] != "default":
-                out = ensure_host(out)
             out = _wrap_data_with_container("transform", out, args[0], self)
 
         return self if out is self._gpu else out
diff --git a/python/cuml/cuml_accel_tests/test_pipeline.py b/python/cuml/cuml_accel_tests/test_pipeline.py
index f169f41076..1715648877 100644
--- a/python/cuml/cuml_accel_tests/test_pipeline.py
+++ b/python/cuml/cuml_accel_tests/test_pipeline.py
@@ -327,3 +327,18 @@ def test_pipeline_set_output():
         StandardScaler().set_output(transform="pandas")
     ).fit_transform(X)
     assert isinstance(X2, pd.DataFrame)
+
+
+def test_pipeline_classifier_predict_non_numeric_labels(patch_methods):
+    X, y = make_classification(random_state=42, n_classes=2)
+    y = np.array(["a", "b"]).take(y)
+
+    patch_methods(LogisticRegression, "fit", "predict")
+
+    pipeline = make_pipeline(StandardScaler(), LogisticRegression())
+    pipeline.fit(X, y)
+    assert isinstance(LogisticRegression.fit.args[0], cp.ndarray)
+    out = pipeline.predict(X)
+    assert isinstance(LogisticRegression.predict.args[0], cp.ndarray)
+    # User-facing output is always numpy
+    assert isinstance(out, np.ndarray)