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/_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/_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..ffc896fd5a 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,14 @@ def debug(self, msg: str) -> None: "sklearn.preprocessing", "sklearn.svm", "umap", -] +} + +_PATCHES = { + "sklearn.pipeline", + "sklearn.utils", +} + +ACCELERATED_MODULES = sorted(_OVERRIDES.union(_PATCHES)) def _exclude_from_acceleration(module: str) -> bool: @@ -107,7 +114,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..6161eda48f 100644 --- a/python/cuml/cuml/accel/estimator_proxy.py +++ b/python/cuml/cuml/accel/estimator_proxy.py @@ -6,18 +6,24 @@ 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, 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 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") @@ -31,6 +37,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. @@ -105,7 +116,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`. @@ -256,10 +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"): - # Ensure transform result is properly wrapped for `set_output` + # Properly wrap output of transform following `set_output` config. out = _wrap_data_with_container("transform", out, args[0], self) return self if out is self._gpu else out @@ -318,6 +340,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/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) diff --git a/python/cuml/cuml_accel_tests/test_pipeline.py b/python/cuml/cuml_accel_tests/test_pipeline.py index ea6dd71039..1715648877 100644 --- a/python/cuml/cuml_accel_tests/test_pipeline.py +++ b/python/cuml/cuml_accel_tests/test_pipeline.py @@ -1,10 +1,15 @@ # -# 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 pandas as pd 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 +27,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 +202,143 @@ 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) + + +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) + + +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)