Support subclassing proxy estimators

python/cuml/cuml/accel/estimator_proxy.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 from __future__ import annotations
 
+import abc
 import functools
 from typing import Any
 
@@ -50,7 +51,7 @@ def is_proxy(instance_or_class) -> bool:
         cls = instance_or_class
     else:
         cls = type(instance_or_class)
-    return issubclass(cls, ProxyBase)
+    return isinstance(cls, ProxyBaseMeta) and hasattr(cls, "_cpu_class")
 
 
 def ensure_host(x):
@@ -113,7 +114,59 @@ def __call__(self, cls_path, cpu, load_on_gpu):
 _reconstruct_proxy = _ReconstructProxy()
 
 
-class ProxyBase(BaseEstimator):
+class ProxyBaseMeta(abc.ABCMeta):
+    """A metaclass for `ProxyBase` types.
+
+    Most of the magic of `ProxyBase` lives in `ProxyBase.__init_subclass__`.
+    However, to support subclassing proxy estimators (which may make use of
+    sklearn internals, and thus cannot be actual proxies), we need a way to
+    dynamically modify the bases of a class. We also want these subclasses to
+    identify as subclasses (and instances) of the proxy class, even if the
+    proxy class isn't a true base.
+
+    Unfortunately, metaclasses don't compose as well - the metaclass of a new
+    class must be a (non-strict) subclass of the metaclass of all base classes.
+    As such, if any subclasses introduce a new metaclass to do other magic, the
+    magic of `ProxyBaseMeta` will cause issues. To work around this, we subclass
+    `ProxyBaseMeta` from the most common metaclass in use (`abc.ABCMeta`) so
+    these can at least be mixed.
+    """
+
+    def __new__(cls, name, bases, ns, **kwargs):
+        # If any base classes are ProxyBaseMeta instances _with_ a cpu class
+        # defined, replace them with their CPU class.
+        bases = tuple(
+            getattr(base, "_cpu_class", base)
+            if isinstance(base, ProxyBaseMeta)
+            else base
+            for base in bases
+        )
+        return super().__new__(cls, name, bases, ns, **kwargs)
+
+    def __subclasscheck__(self, subclass):
+        """Check if a class is a subclass"""
+        # Check if it's a true subclass
+        if super().__subclasscheck__(subclass):
+            return True
+        # Check if its a subclass of _cpu_class (if available)
+        if (cpu_class := getattr(self, "_cpu_class", None)) is not None:
+            return cpu_class.__subclasscheck__(subclass)
+        # Not a subclass
+        return False
+
+    def __instancecheck__(self, instance):
+        """Check if an object is an instance."""
+        # Check if it's a true instance
+        if super().__instancecheck__(instance):
+            return True
+        # Check if its an instance of _cpu_class (if available)
+        if (cpu_class := getattr(self, "_cpu_class", None)) is not None:
+            return cpu_class.__instancecheck__(instance)
+        # Not an instance
+        return False
+
+
+class ProxyBase(BaseEstimator, metaclass=ProxyBaseMeta):
     """A base class for defining new Proxy estimators.
 
     Subclasses should define ``_gpu_class``, which must be a subclass of
@@ -804,16 +857,18 @@ class ArrayAPIProxyBase(ProxyBase):
     """
 
     def __init_subclass__(cls, **kwargs):
-        # Programmatically create a new private cuml.Base class that wraps the
-        # sklearn array-api-enabled model in a cuml consistent API.
-        cls._gpu_class = type(
-            cls.__name__,
-            (_ArrayAPIWrapper,),
-            {
-                "_cpu_class_path": cls._cpu_class_path,
-                "_params_from_cpu_override": getattr(
-                    cls, "_params_from_cpu", None
-                ),
-            },
-        )
+        # If _cpu_class_path not defined, skip generation of accelerated class
+        if hasattr(cls, "_cpu_class_path"):
+            # Programmatically create a new private cuml.Base class that wraps the
+            # sklearn array-api-enabled model in a cuml consistent API.
+            cls._gpu_class = type(
+                cls.__name__,
+                (_ArrayAPIWrapper,),
+                {
+                    "_cpu_class_path": cls._cpu_class_path,
+                    "_params_from_cpu_override": getattr(
+                        cls, "_params_from_cpu", None
+                    ),
+                },
+            )
         super().__init_subclass__(**kwargs)

python/cuml/cuml_accel_tests/test_estimator_proxy.py

@@ -31,6 +31,7 @@
 from sklearn.preprocessing import StandardScaler
 
 from cuml.accel import is_proxy
+from cuml.accel.estimator_proxy import ProxyBase
 
 SKLEARN_16 = Version(sklearn.__version__) >= Version("1.6.0")
 SKLEARN_18 = Version(sklearn.__version__) >= Version("1.8.0.dev0")
@@ -877,3 +878,46 @@ def test_array_api_proxy_fallback_older_sklearn():
     model = StandardScaler().fit(X)
     # Ran on CPU
     assert model._gpu is None
+
+
+@pytest.mark.parametrize(
+    "Base",
+    [
+        pytest.param(LinearRegression, id="ProxyBase"),
+        pytest.param(RandomForestRegressor, id="ProxyBase-with-ABCMeta"),
+        pytest.param(StandardScaler, id="ArrayAPIProxyBase"),
+    ],
+)
+def test_subclass_of_proxy_isnt_accelerated(Base):
+    class Sub(Base):
+        pass
+
+    X, y = make_regression(n_samples=200, random_state=42)
+
+    base_model = Base().fit(X, y)
+    sub_model = Sub().fit(X, y)
+
+    # A base class and instance are proxies
+    assert is_proxy(Base)
+    assert is_proxy(base_model)
+
+    # issubclass/isinstance work on true proxy classes/instances
+    assert issubclass(Base, ProxyBase)
+    assert isinstance(base_model, Base)
+    assert isinstance(base_model, ProxyBase)
+
+    # A subclass (and instance) are not proxies
+    assert not is_proxy(Sub)
+    assert not is_proxy(sub_model)
+
+    # A subclass doesn't have any concrete proxy bases in its mro:
+    assert Base not in Sub.__mro__
+    # A subclass does have the original CPU model in its mro:
+    assert Base._cpu_class in Sub.__mro__
+
+    # A subclass still identifies as a subclass
+    assert issubclass(Sub, Base)
+    assert issubclass(Sub, ProxyBase)
+    # A subclass instance still identifies as an instance
+    assert isinstance(sub_model, Base)
+    assert isinstance(sub_model, ProxyBase)

python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-list.yaml

@@ -213,8 +213,6 @@
   - "sklearn.ensemble.tests.test_forest::test_warm_start[RandomForestClassifier]"
   - "sklearn.ensemble.tests.test_forest::test_warm_start[RandomForestRegressor]"
   - "sklearn.ensemble.tests.test_forest::test_warm_start_oob[RandomForestRegressor]"
-  - "sklearn.ensemble.tests.test_stacking::test_stacking_without_n_features_in[make_classification-StackingClassifier-LogisticRegression]"
-  - "sklearn.ensemble.tests.test_stacking::test_stacking_without_n_features_in[make_regression-StackingRegressor-LinearRegression]"
   - "sklearn.ensemble.tests.test_voting::test_predict_on_toy_problem[42]"
   - "sklearn.ensemble.tests.test_voting::test_sample_weight[42]"
   - "sklearn.ensemble.tests.test_voting::test_set_estimator_drop"
@@ -247,8 +245,6 @@
   - "sklearn.linear_model.tests.test_coordinate_descent::test_enet_sample_weight_consistency[42-csr_matrix-False-0.01-False]"
   - "sklearn.linear_model.tests.test_coordinate_descent::test_enet_sample_weight_consistency[42-csr_matrix-False-0.01-True]"
   - "sklearn.linear_model.tests.test_coordinate_descent::test_enet_toy"
-  - "sklearn.linear_model.tests.test_coordinate_descent::test_lassoCV_does_not_set_precompute[False-False]"
-  - "sklearn.linear_model.tests.test_coordinate_descent::test_lassoCV_does_not_set_precompute[auto-False]"
   - "sklearn.linear_model.tests.test_coordinate_descent::test_lasso_alpha_warning"
   - "sklearn.linear_model.tests.test_coordinate_descent::test_lasso_dual_gap"
   - "sklearn.linear_model.tests.test_coordinate_descent::test_lasso_readonly_data"
@@ -787,6 +783,7 @@
 - reason: Tests that fail due to poking at sklearn internals, failures don't indicate bugs in cuml.accel
   marker: cuml_accel_invalid_sklearn_tests
   tests:
+  - "sklearn.model_selection.tests.test_search::test_grid_search_score_method"
   - "sklearn.svm.tests.test_svm::test_gamma_scale"
   - "sklearn.svm.tests.test_svm::test_svc_raises_error_internal_representation"
 - reason: This test asserts a copy hasn't happened, but that's not actually guaranteed by the interface.
@@ -1044,10 +1041,6 @@
   - "sklearn.tests.test_common::test_estimators[TSNE()-check_fit2d_predict1d]"
   - "sklearn.tests.test_common::test_estimators[TSNE()-check_methods_sample_order_invariance]"
   - "sklearn.tests.test_common::test_estimators[TSNE()-check_methods_subset_invariance]"
-- reason: Config dispatch with parallel processing not working correctly with StandardScaler proxy
-  marker: parallel_config
-  tests:
-  - "sklearn.utils.tests.test_parallel::test_dispatch_config_parallel[2]"
 - reason: Calibration temperature scaling differs with cuml.accel in sklearn 1.8
   condition: scikit-learn>=1.8
   tests:
@@ -1108,14 +1101,12 @@
   tests:
   - "sklearn.feature_selection.tests.test_rfe::test_rfe_wrapped_estimator[RFECV-4-importance_getter0]"
   - "sklearn.feature_selection.tests.test_rfe::test_rfe_wrapped_estimator[RFECV-4-regressor_.coef_]"
-  - "sklearn.model_selection.tests.test_search::test_grid_search_no_score"
   - "sklearn.svm.tests.test_sparse::test_linearsvc[lil_array-dok_array]"
   - "sklearn.svm.tests.test_sparse::test_linearsvc[lil_matrix-dok_matrix]"
   - "sklearn.svm.tests.test_sparse::test_linearsvc_iris[csr_array]"
   - "sklearn.svm.tests.test_sparse::test_linearsvc_iris[csr_matrix]"
   - "sklearn.svm.tests.test_sparse::test_sparse_liblinear_intercept_handling"
   - "sklearn.svm.tests.test_svm::test_dense_liblinear_intercept_handling"
-  - "sklearn.tests.test_calibration::test_calibration_default_estimator"
   - "sklearn.tests.test_calibration::test_calibration_multiclass[1-True-sigmoid]"
   - "sklearn.tests.test_common::test_estimators[LinearSVC()-check_classifier_data_not_an_array]"
   - "sklearn.tests.test_common::test_estimators[LinearSVC()-check_dtype_object]"