scikit-learn · betatim · Mar 29, 2022 · Apr 1, 2022 · Apr 4, 2022 · Apr 4, 2022
diff --git a/doc/computing.rst b/doc/computing.rst
@@ -14,3 +14,4 @@ Computing with scikit-learn
     computing/scaling_strategies
     computing/computational_performance
     computing/parallelism
+    computing/engine
diff --git a/doc/computing/engine.rst b/doc/computing/engine.rst
@@ -0,0 +1,29 @@
+.. Places parent toc into the sidebar
+
+:parenttoc: True
+
+.. _engine:
+
+Computation Engines (experimental)
+==================================
+
+**This API is experiment** which means that it is subject to change without
+any backward compatibility guarantees.
+
+TODO: explain goals here
+
+Activating an engine
+--------------------
+
+TODO: installing third party engine provider packages
+
+TODO: how to list installed engines
+
+TODO: how to install a plugin
+
+Writing a new engine provider
+-----------------------------
+
+TODO: show engine API of a given estimator.
+
+TODO: give example setup.py with setuptools to define an entrypoint.
diff --git a/doc/whats_new/v1.2.rst b/doc/whats_new/v1.2.rst
@@ -2,6 +2,26 @@
 
 .. currentmodule:: sklearn
 
+
+TODO: move to doc/whats_new/v1.3.rst once it exists in main:
+
+- |Enhancement| Experimential engine API (no backward compatibility guarantees)
+  to allow for external packages to contribute alternative implementations for
+  the core computational routines of some selected scikit-learn estimators.
+
+  Currently, the following estimators allow alternative implementations:
+
+  - :class:`~sklearn.cluster.KMeans` (only for the LLoyd algorithm).
+  - TODO: add more here
+
+  External engine providers include:
+
+  - https://github.com/soda-inria/sklearn-numba-dpex that provided a KMeans
+    engine optimized for OpenCL enabled GPUs.
+  - TODO: add more here
+
+  :pr:`24497` by :user:`ogrisel`, :user:`fcharras`.
+
 .. _changes_1_2:
 
 Version 1.2.0

diff --git a/setup.py b/setup.py
@@ -606,6 +606,7 @@ def setup_package():
         python_requires=python_requires,
         install_requires=min_deps.tag_to_packages["install"],
         package_data={"": ["*.pxd"]},
+        entry_points={"pytest11": ["sklearn_plugin_testing = sklearn._engine.testing"]},
         zip_safe=False,  # the package can run out of an .egg file
         include_package_data=True,
         extras_require={

diff --git a/sklearn/_config.py b/sklearn/_config.py
@@ -14,6 +14,7 @@
     ),
     "enable_cython_pairwise_dist": True,
     "array_api_dispatch": False,
+    "engine_provider": (),
     "transform_output": "default",
 }
 _threadlocal = threading.local()
@@ -53,6 +54,7 @@ def set_config(
     pairwise_dist_chunk_size=None,
     enable_cython_pairwise_dist=None,
     array_api_dispatch=None,
+    engine_provider=None,
     transform_output=None,
 ):
     """Set global scikit-learn configuration
@@ -122,6 +124,15 @@ def set_config(
 
         .. versionadded:: 1.2
 
+    engine_provider : str or sequence of str, default=None
+        Enable computational engine implementation provided by third party
+        packages to leverage specific hardware platforms using frameworks or
+        libraries outside of the usual scikit-learn project dependencies.
+
+        See the :ref:`User Guide <engine>` for more details.
+
+        .. versionadded:: 1.3
+
     transform_output : str, default=None
         Configure output of `transform` and `fit_transform`.
 
@@ -155,6 +166,8 @@ def set_config(
         local_config["enable_cython_pairwise_dist"] = enable_cython_pairwise_dist
     if array_api_dispatch is not None:
         local_config["array_api_dispatch"] = array_api_dispatch
+    if engine_provider is not None:
+        local_config["engine_provider"] = engine_provider
     if transform_output is not None:
         local_config["transform_output"] = transform_output
 
@@ -169,6 +182,7 @@ def config_context(
     pairwise_dist_chunk_size=None,
     enable_cython_pairwise_dist=None,
     array_api_dispatch=None,
+    engine_provider=None,
     transform_output=None,
 ):
     """Context manager for global scikit-learn configuration.
@@ -237,6 +251,15 @@ def config_context(
 
         .. versionadded:: 1.2
 
+    engine_provider : str or sequence of str, default=None
+        Enable computational engine implementation provided by third party
+        packages to leverage specific hardware platforms using frameworks or
+        libraries outside of the usual scikit-learn project dependencies.
+
+        See the :ref:`User Guide <engine>` for more details.
+
+        .. versionadded:: 1.3
+
     transform_output : str, default=None
         Configure output of `transform` and `fit_transform`.
 
@@ -285,6 +308,7 @@ def config_context(
         pairwise_dist_chunk_size=pairwise_dist_chunk_size,
         enable_cython_pairwise_dist=enable_cython_pairwise_dist,
         array_api_dispatch=array_api_dispatch,
+        engine_provider=engine_provider,
         transform_output=transform_output,
     )
 

diff --git a/sklearn/_engine/__init__.py b/sklearn/_engine/__init__.py
@@ -0,0 +1,4 @@
+from .base import get_engine_classes, list_engine_provider_names
+
+
+__all__ = ["get_engine_classes", "list_engine_provider_names"]
diff --git a/sklearn/_engine/base.py b/sklearn/_engine/base.py
@@ -0,0 +1,115 @@
+from importlib.metadata import entry_points
+from importlib import import_module
+from functools import lru_cache
+import warnings
+
+from .._config import get_config
+
+SKLEARN_ENGINES_ENTRY_POINT = "sklearn_engines"
+
+
+class EngineSpec:
+
+    __slots__ = ["name", "provider_name", "module_name", "engine_qualname"]
+
+    def __init__(self, name, provider_name, module_name, engine_qualname):
+        self.name = name
+        self.provider_name = provider_name
+        self.module_name = module_name
+        self.engine_qualname = engine_qualname
+
+    def get_engine_class(self):
+        engine = import_module(self.module_name)
+        for attr in self.engine_qualname.split("."):
+            engine = getattr(engine, attr)
+        return engine
+
+
+def _parse_entry_point(entry_point):
+    module_name, engine_qualname = entry_point.value.split(":")
+    provider_name = next(iter(module_name.split(".", 1)))
+    return EngineSpec(entry_point.name, provider_name, module_name, engine_qualname)
+
+
+@lru_cache
+def _parse_entry_points(provider_names=None):
+    specs = []
+    all_entry_points = entry_points()
+    if hasattr(all_entry_points, "select"):
+        engine_entry_points = all_entry_points.select(group=SKLEARN_ENGINES_ENTRY_POINT)
+    else:
+        engine_entry_points = all_entry_points.get(SKLEARN_ENGINES_ENTRY_POINT, ())
+    for entry_point in engine_entry_points:
+        try:
+            spec = _parse_entry_point(entry_point)
+            if provider_names is not None and spec.provider_name not in provider_names:
+                # Skip entry points that do not match the requested provider names.
+                continue
+            specs.append(spec)
+        except Exception as e:
+            # Do not raise an exception in case an invalid package has been
+            # installed in the same Python env as scikit-learn: just warn and
+            # skip.
+            warnings.warn(
+                f"Invalid {SKLEARN_ENGINES_ENTRY_POINT} entry point"
+                f" {entry_point.name} with value {entry_point.value}: {e}"
+            )
+    if provider_names is not None:
+        observed_provider_names = {spec.provider_name for spec in specs}
+        missing_providers = set(provider_names) - observed_provider_names
+        if missing_providers:
+            raise RuntimeError(
+                "Could not find any provider for the"
+                f" {SKLEARN_ENGINES_ENTRY_POINT} entry point with name(s):"
+                f" {', '.join(repr(p) for p in sorted(missing_providers))}"
+            )
+    return specs
+
+
+def list_engine_provider_names():
+    """Find the list of sklearn_engine provider names
+
+    This function only inspects the metadata and should trigger any module import.
+    """
+    return sorted({spec.provider_name for spec in _parse_entry_points()})
+
+
+def _get_engine_classes(engine_name, provider_names, engine_specs, default):
+    specs_by_provider = {}
+    for spec in engine_specs:
+        if spec.name != engine_name:
+            continue
+        specs_by_provider.setdefault(spec.provider_name, spec)
+
+    for provider_name in provider_names:
+        spec = specs_by_provider.get(provider_name)
+        if spec is not None:
+            # XXX: should we return an instance or the class itself?
+            yield spec.provider_name, spec.get_engine_class()
+
+    yield "default", default
+
+
+def get_engine_classes(engine_name, default, verbose=False):
+    provider_names = get_config()["engine_provider"]
+    if isinstance(provider_names, str):
+        provider_names = (provider_names,)
+    elif not isinstance(provider_names, tuple):
+        # Make sure the provider names are a tuple to make it possible for the
+        # lru cache to hash them.
+        provider_names = tuple(provider_names)
+    if not provider_names:
+        yield "default", default
+        return
+    engine_specs = _parse_entry_points(provider_names=provider_names)
+    for provider, engine_class in _get_engine_classes(
+        engine_name=engine_name,
+        provider_names=provider_names,
+        engine_specs=engine_specs,
+        default=default,
+    ):
+        if verbose:
+            print(
+                f"trying engine {engine_class.__module__}.{engine_class.__qualname__} ."
+            )
+        yield provider, engine_class
diff --git a/sklearn/_engine/testing.py b/sklearn/_engine/testing.py
@@ -0,0 +1,38 @@
+from pytest import xfail, hookimpl
+
+from sklearn import config_context
+
+from sklearn.exceptions import NotSupportedByEngineError
+
+
+# TODO: document this pytest plugin + write a tutorial on how to develop a new plugin
+# and explain good practices regarding testing against sklearn test modules.
+def pytest_addoption(parser):
+    group = parser.getgroup("Sklearn plugin testing")
+    group.addoption(
+        "--sklearn-engine-provider",
+        action="store",
+        nargs=1,
+        type=str,
+        help="Name of the an engine provider for sklearn to activate for all tests.",
+    )
+
+
+@hookimpl(hookwrapper=True)
+def pytest_pyfunc_call(pyfuncitem):
+    engine_provider = pyfuncitem.config.getoption("sklearn_engine_provider")
+    if engine_provider is None:
+        yield
+        return
+
+    with config_context(engine_provider=engine_provider):
+        try:
+            outcome = yield
+            outcome.get_result()
+        except NotSupportedByEngineError:
+            xfail(
+                reason=(
+                    "This test cover features that are not supported by the "
+                    f"engine provided by {engine_provider}."
+                )
+            )
diff --git a/sklearn/_engine/tests/__init__.py b/sklearn/_engine/tests/__init__.py
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,4 @@
		from .base import get_engine_classes, list_engine_provider_names


		__all__ = ["get_engine_classes", "list_engine_provider_names"]