diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index a176bd3..13dd2a8 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -36,12 +36,17 @@ jobs: python-version: "3.12" - name: Install the project + # The docs group carries `ezmsg-learn[all]`: autodoc imports every + # documented module, including the ones behind the torch/sklearn extras. run: uv sync --only-group docs - name: Build documentation + # --no-sync so the build uses exactly the environment installed above. + # Without it, `uv run` re-syncs to the default groups, which quietly + # decides the docs dependencies for us. run: | cd docs - uv run make html + uv run --no-sync make html - name: Add .nojekyll file run: touch docs/build/html/.nojekyll diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 61d1a04..fea84cb 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -47,3 +47,34 @@ jobs: run: uv run pytest tests/integration -v --tb=long env: PYTHONFAULTHANDLER: 1 + + # The backends are extras; this job proves a base install stays free of them + # and that the modules behind an extra say which extra they need. + minimal-install: + name: Minimal install (no extras) + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v6 + with: + python-version: "3.13" + + - name: Install the project without extras + # Deliberately not `uv sync`: the dev dependency-groups pull in [all]. + run: | + uv venv + uv pip install . pytest + + - name: Assert no backend is installed + run: | + .venv/bin/python -c " + import importlib.util + present = [m for m in ('torch', 'sklearn', 'river', 'pandas') if importlib.util.find_spec(m)] + assert not present, f'base install pulled in {present}' + " + + - name: Run the optional-dependency guards + run: .venv/bin/python -m pytest tests/unit/test_optional_deps.py -v diff --git a/README.md b/README.md index 2c96fd5..0e2a911 100644 --- a/README.md +++ b/README.md @@ -4,31 +4,39 @@ This repository contains a Python package with modules for machine learning (ML) > If you are only interested in offline analysis without concern for reproducibility in online applications, then you should probably look elsewhere. -Processing units include dimensionality reduction, linear regression, and classification that can be initialized with known weights, or adapted on-the-fly with incoming (labeled) data. Machine-learning code depends on `river`, `scikit-learn`, `numpy`, and `torch`. +Processing units include dimensionality reduction, linear regression, and classification that can be initialized with known weights, or adapted on-the-fly with incoming (labeled) data. ## Installation -Install from PyPI: +The base install is NumPy-only. The machine-learning backends are optional extras, so a deployment that uses only the lightweight processors does not pay for a PyTorch or scikit-learn install: ```bash -pip install ezmsg-learn +pip install ezmsg-learn # numpy-only processors +pip install "ezmsg-learn[sklearn]" # + pandas, river, scikit-learn +pip install "ezmsg-learn[torch]" # + torch +pip install "ezmsg-learn[all]" # everything ``` Or install the latest development version: ```bash -pip install git+https://github.com/ezmsg-org/ezmsg-learn@dev +pip install "git+https://github.com/ezmsg-org/ezmsg-learn@dev#egg=ezmsg-learn[all]" ``` +Importing a module whose backend is not installed raises an `ImportError` naming the extra to install. + ## Dependencies -- `ezmsg` -- `ezmsg-baseproc` -- `ezmsg-sigproc` -- `numpy` -- `scipy` -- `scikit-learn` -- `river` +Base (`pip install ezmsg-learn`) — `ezmsg`, `ezmsg-baseproc`, `ezmsg-sigproc`, `numpy`, `scipy`, `array-api-compat`. + +| Extra | Adds | Covers | +| --- | --- | --- | +| _(none)_ | — | `process.ssr`, `process.flatten`, `process.seqseqsampler`, `process.refit_kalman`, `model.cca`, `model.refit_kalman` | +| `sklearn` | `pandas`, `river`, `scikit-learn` | `process.adaptive_linear_regressor`, `process.linear_regressor`, `process.sgd`, `process.slda`, `process.sklearn`, `dim_reduce.*` | +| `torch` | `torch` | `process.base`, `process.torch`, `process.rnn`, `process.transformer`, `process.mlp_old`, `model.mlp`, `model.rnn`, `model.transformer` | +| `all` | both of the above | everything, including all `collection.sample_adapt_regressor` backends | + +`collection.sample_adapt_regressor` imports its backend lazily, so it needs only the extra for the `model_type` in use — and none at all for `model_type="kalman"`. ## Development diff --git a/docs/source/index.rst b/docs/source/index.rst index e240d38..531e8c9 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -26,23 +26,68 @@ Most modules support both: Installation ------------ -Install directly from GitHub: +The base install is NumPy-only; the machine-learning backends are optional +extras, so a deployment that uses only the lightweight processors does not pay +for a PyTorch or scikit-learn install: .. code-block:: bash - pip install git+https://github.com/ezmsg-org/ezmsg-learn + pip install ezmsg-learn # numpy-only processors + pip install "ezmsg-learn[sklearn]" # + pandas, river, scikit-learn + pip install "ezmsg-learn[torch]" # + torch + pip install "ezmsg-learn[all]" # everything + +Or install directly from GitHub: + +.. code-block:: bash + + pip install "git+https://github.com/ezmsg-org/ezmsg-learn#egg=ezmsg-learn[all]" + +Importing a module whose backend is not installed raises an ``ImportError`` +naming the extra to install. Dependencies ^^^^^^^^^^^^ -This package requires: +The base install requires: * ``ezmsg`` - Core ezmsg framework +* ``ezmsg-baseproc`` - Processor base classes * ``ezmsg-sigproc`` - Signal processing extensions * ``numpy`` - Numerical computing -* ``scikit-learn`` - Machine learning utilities -* ``torch`` - Deep learning framework -* ``river`` - Online machine learning +* ``scipy`` - Scientific computing +* ``array-api-compat`` - Array API portability layer + +Optional extras +^^^^^^^^^^^^^^^ + +.. list-table:: + :header-rows: 1 + :widths: 15 25 60 + + * - Extra + - Adds + - Covers + * - *(none)* + - — + - ``process.ssr``, ``process.flatten``, ``process.seqseqsampler``, + ``process.refit_kalman``, ``model.cca``, ``model.refit_kalman`` + * - ``sklearn`` + - ``pandas``, ``river``, ``scikit-learn`` + - ``process.adaptive_linear_regressor``, ``process.linear_regressor``, + ``process.sgd``, ``process.slda``, ``process.sklearn``, ``dim_reduce.*`` + * - ``torch`` + - ``torch`` + - ``process.base``, ``process.torch``, ``process.rnn``, + ``process.transformer``, ``process.mlp_old``, ``model.mlp``, + ``model.rnn``, ``model.transformer`` + * - ``all`` + - both of the above + - everything, including all ``collection.sample_adapt_regressor`` backends + +:mod:`ezmsg.learn.collection.sample_adapt_regressor` imports its backend +lazily, so it needs only the extra for the ``model_type`` in use — and none at +all for ``model_type="kalman"``. Quick Start ----------- diff --git a/pyproject.toml b/pyproject.toml index 42b1b3c..a3779fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,11 +12,26 @@ dependencies = [ "ezmsg>=3.9.0", "ezmsg-baseproc>=1.7.0", "ezmsg-sigproc>=3.0.0", + "numpy", + "scipy", + "array-api-compat", +] + +[project.optional-dependencies] +# The heavyweight backends are opt-in so that the pure-numpy processors +# (process.ssr, process.flatten, process.refit_kalman, ...) can be installed on +# constrained hosts without pulling in a full PyTorch/sklearn stack. +sklearn = [ "pandas>=2.2", "river>=0.22.0", "scikit-learn>=1.6.0", +] +torch = [ "torch>=2.6.0", ] +all = [ + "ezmsg-learn[sklearn,torch]", +] [dependency-groups] dev = [ @@ -28,11 +43,17 @@ lint = [ "ruff==0.16.2", ] test = [ + # The backends are extras, so `uv sync` would not install them; the test + # suite covers all of them. + "ezmsg-learn[all]", "ezmsg-simbiophys>=1.8.0", "hmmlearn>=0.3.3", "pytest>=8.4.1", ] docs = [ + # autodoc imports every documented module, so the docs build needs the + # project and all of its backends -- not just Sphinx. + "ezmsg-learn[all]", "sphinx>=8.1.3", "pydata-sphinx-theme", "sphinx_autodoc_typehints>=3.0.0", diff --git a/src/ezmsg/learn/_optional.py b/src/ezmsg/learn/_optional.py new file mode 100644 index 0000000..5a81ed7 --- /dev/null +++ b/src/ezmsg/learn/_optional.py @@ -0,0 +1,25 @@ +"""Reporting for the optional backend dependencies. + +``ezmsg-learn`` ships its heavyweight backends as extras (``ezmsg-learn[torch]`` +and ``ezmsg-learn[sklearn]``) so the pure-numpy processors can be installed +without them. Modules belonging to an extra still import their backend eagerly; +they only wrap the import so that a missing install reports which extra to +install instead of a bare ``ModuleNotFoundError: No module named 'torch'``. +""" + + +def missing_extra(extra: str, module: str) -> ImportError: + """Build the error raised when a module's optional backend is not installed. + + Args: + extra: Name of the extra that provides the backend, e.g. ``"torch"``. + module: Importing module, normally passed as ``__name__``. + + Returns: + The :class:`ImportError` to raise; chain it off the original with + ``raise missing_extra(...) from exc``. + """ + return ImportError( + f"{module} requires the optional '{extra}' dependencies of ezmsg-learn, which are not installed. " + f'Install them with: pip install "ezmsg-learn[{extra}]"' + ) diff --git a/src/ezmsg/learn/collection/sample_adapt_regressor.py b/src/ezmsg/learn/collection/sample_adapt_regressor.py index 7611a61..b6eb4bf 100644 --- a/src/ezmsg/learn/collection/sample_adapt_regressor.py +++ b/src/ezmsg/learn/collection/sample_adapt_regressor.py @@ -1,3 +1,13 @@ +"""Decode collection wired around a single regressor engine. + +.. note:: + The regressor backends live behind optional extras. ``model_type="mlp"`` + needs ``ezmsg-learn[torch]`` and the River/sklearn model types need + ``ezmsg-learn[sklearn]``; ``model_type="kalman"`` needs neither. The backend + modules are imported only once a backend is selected, so an install carrying + just one extra can still build the collection for that backend. +""" + from dataclasses import field import ezmsg.core as ez @@ -13,17 +23,12 @@ from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace -from ezmsg.learn.process.adaptive_linear_regressor import ( - AdaptiveLinearRegressorSettings, - AdaptiveLinearRegressorUnit, -) from ezmsg.learn.process.flatten import Flatten, FlattenSettings from ezmsg.learn.process.refit_kalman import ( RefitKalmanFilterSettings, RefitKalmanFilterUnit, ) from ezmsg.learn.process.seqseqsampler import SeqSeqSamplerSettings, SeqSeqSamplerUnit -from ezmsg.learn.process.torch import TorchModelSettings, TorchModelUnit from ezmsg.learn.util import AdaptiveLinearRegressor #: Default torch model class used when ``model_type == "mlp"``. @@ -181,16 +186,47 @@ def _build_regressor_unit(settings: SampleAdaptRegressorSettings): """Factory: construct the single regressor unit for ``settings.model_type``. Returns ``(unit, backend)`` where ``backend`` is ``"linear"`` (River/sklearn - via :class:`AdaptiveLinearRegressorUnit`), ``"torch"`` (mlp), or ``"kalman"``. + via ``AdaptiveLinearRegressorUnit``), ``"torch"`` (mlp), or ``"kalman"``. """ backend = _model_backend(settings.model_type) if backend == "torch": + from ezmsg.learn.process.torch import TorchModelUnit + return TorchModelUnit(), backend if backend == "kalman": return RefitKalmanFilterUnit(), backend + + from ezmsg.learn.process.adaptive_linear_regressor import AdaptiveLinearRegressorUnit + return AdaptiveLinearRegressorUnit(), backend +def _build_regressor_settings(backend: str, settings: SampleAdaptRegressorSettings): + """Translate the collection settings into the selected backend's settings.""" + if backend == "torch": + from ezmsg.learn.process.torch import TorchModelSettings + + return TorchModelSettings( + model_class=settings.model_class, + checkpoint_path=settings.model_path, + model_kwargs=dict(settings.model_kwargs), + device=settings.device, + ) + if backend == "kalman": + return RefitKalmanFilterSettings( + checkpoint_path=settings.model_path, + steady_state=settings.steady_state, + ) + + from ezmsg.learn.process.adaptive_linear_regressor import AdaptiveLinearRegressorSettings + + return AdaptiveLinearRegressorSettings( + model_type=settings.model_type, + settings_path=settings.model_path, + model_kwargs=settings.model_kwargs, + ) + + def build_sample_adapt_regressor( settings: SampleAdaptRegressorSettings, ) -> ez.Collection: @@ -227,30 +263,7 @@ class SampleAdaptRegressor(ez.Collection): ADAPTER = DecodeOutputAdapter() def configure(self) -> None: - if backend == "linear": - self.REGRESSOR.apply_settings( - AdaptiveLinearRegressorSettings( - model_type=self.SETTINGS.model_type, - settings_path=self.SETTINGS.model_path, - model_kwargs=self.SETTINGS.model_kwargs, - ) - ) - elif backend == "torch": - self.REGRESSOR.apply_settings( - TorchModelSettings( - model_class=self.SETTINGS.model_class, - checkpoint_path=self.SETTINGS.model_path, - model_kwargs=dict(self.SETTINGS.model_kwargs), - device=self.SETTINGS.device, - ) - ) - else: - self.REGRESSOR.apply_settings( - RefitKalmanFilterSettings( - checkpoint_path=self.SETTINGS.model_path, - steady_state=self.SETTINGS.steady_state, - ) - ) + self.REGRESSOR.apply_settings(_build_regressor_settings(backend, self.SETTINGS)) if use_window: self.WINDOW.apply_settings( diff --git a/src/ezmsg/learn/dim_reduce/adaptive_decomp.py b/src/ezmsg/learn/dim_reduce/adaptive_decomp.py index 667efbd..7c97787 100644 --- a/src/ezmsg/learn/dim_reduce/adaptive_decomp.py +++ b/src/ezmsg/learn/dim_reduce/adaptive_decomp.py @@ -19,7 +19,13 @@ processor_state, ) from ezmsg.util.messages.axisarray import AxisArray, replace -from sklearn.decomposition import IncrementalPCA, MiniBatchNMF + +from .._optional import missing_extra + +try: + from sklearn.decomposition import IncrementalPCA, MiniBatchNMF +except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc class AdaptiveDecompSettings(ez.Settings): diff --git a/src/ezmsg/learn/model/mlp.py b/src/ezmsg/learn/model/mlp.py index 7164b1a..84501f0 100644 --- a/src/ezmsg/learn/model/mlp.py +++ b/src/ezmsg/learn/model/mlp.py @@ -1,5 +1,10 @@ -import torch -import torch.nn +from .._optional import missing_extra + +try: + import torch + import torch.nn +except ImportError as exc: + raise missing_extra("torch", __name__) from exc class MLP(torch.nn.Module): diff --git a/src/ezmsg/learn/model/mlp_old.py b/src/ezmsg/learn/model/mlp_old.py index e5c0bce..6dc1085 100644 --- a/src/ezmsg/learn/model/mlp_old.py +++ b/src/ezmsg/learn/model/mlp_old.py @@ -1,5 +1,10 @@ -import torch -import torch.nn +from .._optional import missing_extra + +try: + import torch + import torch.nn +except ImportError as exc: + raise missing_extra("torch", __name__) from exc class MLP(torch.nn.Sequential): diff --git a/src/ezmsg/learn/model/rnn.py b/src/ezmsg/learn/model/rnn.py index 4b5e429..92df929 100644 --- a/src/ezmsg/learn/model/rnn.py +++ b/src/ezmsg/learn/model/rnn.py @@ -1,6 +1,11 @@ from typing import Optional -import torch +from .._optional import missing_extra + +try: + import torch +except ImportError as exc: + raise missing_extra("torch", __name__) from exc class RNNModel(torch.nn.Module): diff --git a/src/ezmsg/learn/model/transformer.py b/src/ezmsg/learn/model/transformer.py index 935a774..983b0be 100644 --- a/src/ezmsg/learn/model/transformer.py +++ b/src/ezmsg/learn/model/transformer.py @@ -1,6 +1,11 @@ from typing import Optional -import torch +from .._optional import missing_extra + +try: + import torch +except ImportError as exc: + raise missing_extra("torch", __name__) from exc class TransformerModel(torch.nn.Module): diff --git a/src/ezmsg/learn/process/adaptive_linear_regressor.py b/src/ezmsg/learn/process/adaptive_linear_regressor.py index 7965027..3605458 100644 --- a/src/ezmsg/learn/process/adaptive_linear_regressor.py +++ b/src/ezmsg/learn/process/adaptive_linear_regressor.py @@ -13,10 +13,6 @@ import ezmsg.core as ez import numpy as np -import pandas as pd -import river.linear_model -import river.optim -import sklearn.base from array_api_compat import get_namespace, is_numpy_array from ezmsg.baseproc import ( BaseAdaptiveTransformer, @@ -25,8 +21,17 @@ ) from ezmsg.util.messages.axisarray import AxisArray, replace +from .._optional import missing_extra from ..util import AdaptiveLinearRegressor, RegressorType, get_regressor +try: + import pandas as pd + import river.linear_model + import river.optim + import sklearn.base +except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc + class AdaptiveLinearRegressorSettings(ez.Settings): model_type: AdaptiveLinearRegressor = AdaptiveLinearRegressor.LINEAR diff --git a/src/ezmsg/learn/process/base.py b/src/ezmsg/learn/process/base.py index f2277c3..fdb55ec 100644 --- a/src/ezmsg/learn/process/base.py +++ b/src/ezmsg/learn/process/base.py @@ -4,7 +4,13 @@ from pathlib import Path import ezmsg.core as ez -import torch + +from .._optional import missing_extra + +try: + import torch +except ImportError as exc: + raise missing_extra("torch", __name__) from exc class ModelInitMixin: diff --git a/src/ezmsg/learn/process/linear_regressor.py b/src/ezmsg/learn/process/linear_regressor.py index aac7004..2286dd7 100644 --- a/src/ezmsg/learn/process/linear_regressor.py +++ b/src/ezmsg/learn/process/linear_regressor.py @@ -8,10 +8,15 @@ processor_state, ) from ezmsg.util.messages.axisarray import AxisArray, replace -from sklearn.linear_model._base import LinearModel +from .._optional import missing_extra from ..util import RegressorType, StaticLinearRegressor, get_regressor +try: + from sklearn.linear_model._base import LinearModel +except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc + class LinearRegressorSettings(ez.Settings): model_type: StaticLinearRegressor = StaticLinearRegressor.LINEAR diff --git a/src/ezmsg/learn/process/mlp_old.py b/src/ezmsg/learn/process/mlp_old.py index b8c87d3..78d1360 100644 --- a/src/ezmsg/learn/process/mlp_old.py +++ b/src/ezmsg/learn/process/mlp_old.py @@ -2,8 +2,6 @@ import ezmsg.core as ez import numpy as np -import torch -import torch.nn from ezmsg.baseproc import ( BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit, @@ -12,8 +10,15 @@ from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace +from .._optional import missing_extra from ..model.mlp_old import MLP +try: + import torch + import torch.nn +except ImportError as exc: + raise missing_extra("torch", __name__) from exc + class MLPSettings(ez.Settings): hidden_channels: list[int] diff --git a/src/ezmsg/learn/process/rnn.py b/src/ezmsg/learn/process/rnn.py index eb87e22..cf1f8f1 100644 --- a/src/ezmsg/learn/process/rnn.py +++ b/src/ezmsg/learn/process/rnn.py @@ -2,12 +2,12 @@ import ezmsg.core as ez import numpy as np -import torch from ezmsg.baseproc import BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit from ezmsg.baseproc.util.profile import profile_subpub from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace +from .._optional import missing_extra from .base import ModelInitMixin from .torch import ( TorchModelSettings, @@ -15,6 +15,11 @@ TorchProcessorMixin, ) +try: + import torch +except ImportError as exc: + raise missing_extra("torch", __name__) from exc + class RNNSettings(TorchModelSettings): model_class: str = "ezmsg.learn.model.rnn.RNNModel" diff --git a/src/ezmsg/learn/process/sgd.py b/src/ezmsg/learn/process/sgd.py index baedab5..44d67c5 100644 --- a/src/ezmsg/learn/process/sgd.py +++ b/src/ezmsg/learn/process/sgd.py @@ -9,11 +9,16 @@ ) from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace -from sklearn.exceptions import NotFittedError -from sklearn.linear_model import SGDClassifier +from .._optional import missing_extra from ..util import ClassifierMessage +try: + from sklearn.exceptions import NotFittedError + from sklearn.linear_model import SGDClassifier +except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc + class SGDDecoderSettings(ez.Settings): alpha: float = 1e-5 diff --git a/src/ezmsg/learn/process/sklearn.py b/src/ezmsg/learn/process/sklearn.py index 5473f43..9d72270 100644 --- a/src/ezmsg/learn/process/sklearn.py +++ b/src/ezmsg/learn/process/sklearn.py @@ -5,7 +5,6 @@ import ezmsg.core as ez import numpy as np -import pandas as pd from ezmsg.baseproc import ( BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit, @@ -14,6 +13,13 @@ from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace +from .._optional import missing_extra + +try: + import pandas as pd +except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc + class PredictMethod(str, Enum): PREDICT = "predict" diff --git a/src/ezmsg/learn/process/slda.py b/src/ezmsg/learn/process/slda.py index 91a24b9..8db090e 100644 --- a/src/ezmsg/learn/process/slda.py +++ b/src/ezmsg/learn/process/slda.py @@ -19,10 +19,15 @@ ) from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace -from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA +from .._optional import missing_extra from ..util import ClassifierMessage +try: + from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA +except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc + class SLDASettings(ez.Settings): settings_path: str diff --git a/src/ezmsg/learn/process/torch.py b/src/ezmsg/learn/process/torch.py index 0001dba..0efeb3f 100644 --- a/src/ezmsg/learn/process/torch.py +++ b/src/ezmsg/learn/process/torch.py @@ -3,7 +3,6 @@ import ezmsg.core as ez import numpy as np -import torch from ezmsg.baseproc import ( BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit, @@ -15,8 +14,14 @@ from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace +from .._optional import missing_extra from .base import ModelInitMixin +try: + import torch +except ImportError as exc: + raise missing_extra("torch", __name__) from exc + class TorchSimpleSettings(ez.Settings): model_class: str diff --git a/src/ezmsg/learn/process/transformer.py b/src/ezmsg/learn/process/transformer.py index dd71b1d..eea3f22 100644 --- a/src/ezmsg/learn/process/transformer.py +++ b/src/ezmsg/learn/process/transformer.py @@ -1,12 +1,12 @@ import typing import ezmsg.core as ez -import torch from ezmsg.baseproc import BaseAdaptiveTransformer, BaseAdaptiveTransformerUnit from ezmsg.baseproc.util.profile import profile_subpub from ezmsg.util.messages.axisarray import AxisArray from ezmsg.util.messages.util import replace +from .._optional import missing_extra from .base import ModelInitMixin from .torch import ( TorchModelSettings, @@ -14,6 +14,11 @@ TorchProcessorMixin, ) +try: + import torch +except ImportError as exc: + raise missing_extra("torch", __name__) from exc + class TransformerSettings(TorchModelSettings): model_class: str = "ezmsg.learn.model.transformer.TransformerModel" diff --git a/src/ezmsg/learn/util.py b/src/ezmsg/learn/util.py index 3460568..f0f9b53 100644 --- a/src/ezmsg/learn/util.py +++ b/src/ezmsg/learn/util.py @@ -2,10 +2,10 @@ from dataclasses import dataclass, field from enum import Enum -import river.linear_model -import sklearn.linear_model from ezmsg.util.messages.axisarray import AxisArray +from ._optional import missing_extra + # from sklearn.neural_network import MLPClassifier @@ -27,13 +27,44 @@ class StaticLinearRegressor(str, Enum): RIDGE = "ridge" -ADAPTIVE_REGRESSORS = { - AdaptiveLinearRegressor.LINEAR: river.linear_model.LinearRegression, - AdaptiveLinearRegressor.LOGISTIC: river.linear_model.LogisticRegression, - AdaptiveLinearRegressor.SGD: sklearn.linear_model.SGDRegressor, - AdaptiveLinearRegressor.PAR: sklearn.linear_model.PassiveAggressiveRegressor, - # AdaptiveLinearRegressor.MLP: MLPClassifier, -} +# The registries below are built on demand so that the enums and +# :class:`ClassifierMessage` -- which need nothing beyond ezmsg -- stay importable +# without the ``sklearn`` extra installed. +def _adaptive_regressors() -> dict: + try: + import river.linear_model + import sklearn.linear_model + except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc + + return { + AdaptiveLinearRegressor.LINEAR: river.linear_model.LinearRegression, + AdaptiveLinearRegressor.LOGISTIC: river.linear_model.LogisticRegression, + AdaptiveLinearRegressor.SGD: sklearn.linear_model.SGDRegressor, + AdaptiveLinearRegressor.PAR: sklearn.linear_model.PassiveAggressiveRegressor, + # AdaptiveLinearRegressor.MLP: MLPClassifier, + } + + +def _static_regressors() -> dict: + try: + import sklearn.linear_model + except ImportError as exc: + raise missing_extra("sklearn", __name__) from exc + + return { + StaticLinearRegressor.LINEAR: sklearn.linear_model.LinearRegression, + StaticLinearRegressor.RIDGE: sklearn.linear_model.Ridge, + } + + +def __getattr__(name: str) -> dict: + """Resolve the ``*_REGRESSORS`` registries lazily (:pep:`562`).""" + if name == "ADAPTIVE_REGRESSORS": + return _adaptive_regressors() + if name == "STATIC_REGRESSORS": + return _static_regressors() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") # Function to get a regressor by type and name @@ -47,21 +78,15 @@ def get_regressor( if regressor_type == RegressorType.ADAPTIVE: if isinstance(regressor_name, str): regressor_name = AdaptiveLinearRegressor(regressor_name) - return ADAPTIVE_REGRESSORS[regressor_name] + return _adaptive_regressors()[regressor_name] elif regressor_type == RegressorType.STATIC: if isinstance(regressor_name, str): regressor_name = StaticLinearRegressor(regressor_name) - return STATIC_REGRESSORS[regressor_name] + return _static_regressors()[regressor_name] else: raise ValueError(f"Unknown regressor type: {regressor_type}") -STATIC_REGRESSORS = { - StaticLinearRegressor.LINEAR: sklearn.linear_model.LinearRegression, - StaticLinearRegressor.RIDGE: sklearn.linear_model.Ridge, -} - - @dataclass class ClassifierMessage(AxisArray): labels: list[str] = field(default_factory=list) diff --git a/tests/unit/test_optional_deps.py b/tests/unit/test_optional_deps.py new file mode 100644 index 0000000..57de3e9 --- /dev/null +++ b/tests/unit/test_optional_deps.py @@ -0,0 +1,138 @@ +"""Guards for the optional-backend split. + +``torch`` and the sklearn stack ship as extras (``ezmsg-learn[torch]`` / +``ezmsg-learn[sklearn]``) so that the pure-numpy processors can be installed on +constrained hosts. Two properties keep that promise: + +1. The base-install modules must not reach a backend, even transitively. This is + checked in a subprocess, since the test session itself imports everything. +2. A module belonging to an extra must fail with a message naming that extra. + Only checkable on an install *without* the extra, so those tests skip in the + full dev environment and run in CI's minimal-install job. +""" + +import importlib +import importlib.util +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + +BACKEND_MODULES = ("torch", "sklearn", "river", "pandas") + +SRC = Path(__file__).resolve().parents[2] / "src" + +#: Importable with no extras installed. +LIGHT_MODULES = [ + "ezmsg.learn.util", + "ezmsg.learn.process.ssr", + "ezmsg.learn.process.flatten", + "ezmsg.learn.process.seqseqsampler", + "ezmsg.learn.process.refit_kalman", + "ezmsg.learn.model.cca", + "ezmsg.learn.model.refit_kalman", + "ezmsg.learn.linear_model.cca", + "ezmsg.learn.collection.sample_adapt_regressor", +] + +TORCH_MODULES = [ + "ezmsg.learn.model.mlp", + "ezmsg.learn.model.mlp_old", + "ezmsg.learn.model.rnn", + "ezmsg.learn.model.transformer", + "ezmsg.learn.nlin_model.mlp", + "ezmsg.learn.process.base", + "ezmsg.learn.process.torch", + "ezmsg.learn.process.rnn", + "ezmsg.learn.process.transformer", + "ezmsg.learn.process.mlp_old", +] + +SKLEARN_MODULES = [ + "ezmsg.learn.dim_reduce.adaptive_decomp", + "ezmsg.learn.dim_reduce.incremental_decomp", + "ezmsg.learn.process.adaptive_linear_regressor", + "ezmsg.learn.process.linear_regressor", + "ezmsg.learn.process.sgd", + "ezmsg.learn.process.slda", + "ezmsg.learn.process.sklearn", + "ezmsg.learn.linear_model.adaptive_linear_regressor", + "ezmsg.learn.linear_model.linear_regressor", + "ezmsg.learn.linear_model.sgd", + "ezmsg.learn.linear_model.slda", +] + +HAS_TORCH = importlib.util.find_spec("torch") is not None +HAS_SKLEARN = importlib.util.find_spec("sklearn") is not None + + +def _backends_loaded_by(source: str) -> list[str]: + """Run ``source`` in a fresh interpreter; return the backends it left in ``sys.modules``.""" + script = textwrap.dedent(source) + textwrap.dedent(f""" + import sys + print(" ".join(m for m in {BACKEND_MODULES!r} if m in sys.modules)) + """) + proc = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + env={**os.environ, "PYTHONPATH": str(SRC)}, + ) + assert proc.returncode == 0, f"subprocess failed:\n{proc.stderr}" + return proc.stdout.split() + + +@pytest.mark.parametrize("module", LIGHT_MODULES) +def test_light_module_loads_no_backend(module: str) -> None: + """A base-install module must not pull in torch or the sklearn stack.""" + loaded = _backends_loaded_by(f"import importlib; importlib.import_module({module!r})") + assert not loaded, f"{module} transitively imported {loaded}" + + +def test_kalman_collection_loads_no_backend() -> None: + """``model_type="kalman"`` needs neither extra, so building it must load neither.""" + loaded = _backends_loaded_by(""" + from ezmsg.learn.collection.sample_adapt_regressor import ( + SampleAdaptRegressorSettings, + build_sample_adapt_regressor, + ) + + build_sample_adapt_regressor(SampleAdaptRegressorSettings(model_type="kalman")) + """) + assert not loaded, f"the kalman collection transitively imported {loaded}" + + +@pytest.mark.skipif(HAS_TORCH, reason="torch installed; the missing-extra path is unreachable") +@pytest.mark.parametrize("module", TORCH_MODULES) +def test_torch_module_names_its_extra(module: str) -> None: + with pytest.raises(ImportError, match=r"ezmsg-learn\[torch\]"): + importlib.import_module(module) + + +@pytest.mark.skipif(HAS_SKLEARN, reason="scikit-learn installed; the missing-extra path is unreachable") +@pytest.mark.parametrize("module", SKLEARN_MODULES) +def test_sklearn_module_names_its_extra(module: str) -> None: + with pytest.raises(ImportError, match=r"ezmsg-learn\[sklearn\]"): + importlib.import_module(module) + + +@pytest.mark.skipif(HAS_SKLEARN, reason="scikit-learn installed; the missing-extra path is unreachable") +def test_regressor_registry_names_its_extra() -> None: + """``util`` itself stays importable; only the registries need the extra.""" + from ezmsg.learn import util + + with pytest.raises(ImportError, match=r"ezmsg-learn\[sklearn\]"): + util.get_regressor("adaptive", "linear") + + with pytest.raises(ImportError, match=r"ezmsg-learn\[sklearn\]"): + util.ADAPTIVE_REGRESSORS + + +def test_util_module_getattr_rejects_unknown_names() -> None: + from ezmsg.learn import util + + with pytest.raises(AttributeError): + util.NOT_A_REGISTRY