fix(aggregation): Fix optional dependencies (#329)

* Change tests.yml to add a tests-default-install job
* Rename jobs in tests.yml
* Add try/except around nash_mtl and cagrad import in aggregation/__init__.py
* Add _check_dependencies.py with _OptionalDepsNotInstalledError and check_dependencies_are_installed
* Add call to check_dependencies_are_installed in cagrad.py and nash_mtl.py
* Ignore E402 error in flake8 config

Author: ValerianRey

.github/workflows/tests.yml

@@ -7,7 +7,8 @@ on:
     - cron: '41 16 * * *'  # Every day at 16:41 UTC (to avoid high load at exact hour values).
 
 jobs:
-  Testing:
+  tests-full-install:
+    name: Run tests with full install
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false  # Ensure matrix jobs keep running even if one fails
@@ -30,6 +31,29 @@ jobs:
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
 
+  tests-default-install:
+    name: Run (most) tests with default install
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up PDM
+        uses: pdm-project/setup-pdm@v4
+        with:
+          python-version: '3.13'
+      - name: Install default (without any option) and test dependencies
+        run: pdm install --group test --frozen-lockfile
+      - name: Run unit and doc tests with coverage report
+        run: |
+          pdm run pytest tests/unit tests/doc \
+          --ignore tests/unit/aggregation/test_cagrad.py \
+          --ignore tests/unit/aggregation/test_nash_mtl.py \
+          --ignore tests/doc/test_aggregation.py \
+          --cov=src --cov-report=xml
+      - name: Upload results to Codecov
+        uses: codecov/codecov-action@v4
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+
   build-doc:
     name: Build doc
     runs-on: ubuntu-latest

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
     hooks:
     -   id: flake8  # Check style and syntax. Does not modify code, issues have to be solved manually.
         args: [
-            '--ignore=E501,E203,W503',  # Ignore line length problems, space after colon problems, line break occurring before a binary operator problems.
+            '--ignore=E501,E203,W503,E402',  # Ignore line length problems, space after colon problems, line break occurring before a binary operator problems, module level import not at top of file problems.
             '--per-file-ignores=*/__init__.py:F401',  # Ignore module imported but unused problems in __init__.py files.
         ]
 

src/torchjd/aggregation/__init__.py

@@ -1,6 +1,6 @@
+from ._check_dependencies import _OptionalDepsNotInstalledError
 from .aligned_mtl import AlignedMTL
 from .bases import Aggregator
-from .cagrad import CAGrad
 from .config import ConFIG
 from .constant import Constant
 from .dualproj import DualProj
@@ -9,9 +9,18 @@
 from .krum import Krum
 from .mean import Mean
 from .mgda import MGDA
-from .nash_mtl import NashMTL
 from .pcgrad import PCGrad
 from .random import Random
 from .sum import Sum
 from .trimmed_mean import TrimmedMean
 from .upgrad import UPGrad
+
+try:
+    from .cagrad import CAGrad
+except _OptionalDepsNotInstalledError:  # The required dependencies are not installed
+    pass
+
+try:
+    from .nash_mtl import NashMTL
+except _OptionalDepsNotInstalledError:  # The required dependencies are not installed
+    pass

src/torchjd/aggregation/_check_dependencies.py

@@ -0,0 +1,16 @@
+from importlib.util import find_spec
+
+
+class _OptionalDepsNotInstalledError(ModuleNotFoundError):
+    pass
+
+
+def check_dependencies_are_installed(dependency_names: list[str]) -> None:
+    """
+    Check that the required list of dependencies are installed.
+
+    This can be useful for Aggregators whose dependencies are optional when installing torchjd.
+    """
+
+    if any(find_spec(name) is None for name in dependency_names):
+        raise _OptionalDepsNotInstalledError()

src/torchjd/aggregation/cagrad.py

@@ -1,3 +1,7 @@
+from ._check_dependencies import check_dependencies_are_installed  # noqa
+
+check_dependencies_are_installed(["cvxpy", "clarabel"])
+
 import cvxpy as cp
 import numpy as np
 import torch
@@ -31,8 +35,10 @@ class CAGrad(_WeightedAggregator):
         tensor([0.1835, 1.2041, 1.2041])
 
     .. note::
-        This aggregator has dependencies that are not included by default when installing
-        ``torchjd``. To install them, use ``pip install torchjd[cagrad]``.
+        This aggregator is not installed by default. When not installed, trying to import it should
+        result in the following error:
+        ``ImportError: cannot import name 'CAGrad' from 'torchjd.aggregation'``.
+        To install it, use ``pip install torchjd[cagrad]``.
     """
 
     def __init__(self, c: float, norm_eps: float = 0.0001):

src/torchjd/aggregation/nash_mtl.py

@@ -23,6 +23,9 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+from ._check_dependencies import check_dependencies_are_installed
+
+check_dependencies_are_installed(["cvxpy", "ecos"])
 
 import cvxpy as cp
 import numpy as np
@@ -61,8 +64,10 @@ class NashMTL(_WeightedAggregator):
         tensor([0.0542, 0.7061, 0.7061])
 
     .. note::
-        This aggregator has dependencies that are not included by default when installing
-        ``torchjd``. To install them, use ``pip install torchjd[nash_mtl]``.
+        This aggregator is not installed by default. When not installed, trying to import it should
+        result in the following error:
+        ``ImportError: cannot import name 'NashMTL' from 'torchjd.aggregation'``.
+        To install it, use ``pip install torchjd[nash_mtl]``.
 
     .. warning::
         This implementation was adapted from the `official implementation