Tests cuda

.github/workflows/bluesky.yml

@@ -0,0 +1,58 @@
+name: gpu-ci
+
+on: [push, pull_request]
+
+jobs:
+
+  tests:
+    runs-on: [self-hosted, gpu]
+    if: github.repository == 'stfc/janus-core'
+    timeout-minutes: 60
+    strategy:
+      matrix:
+        python-version: ["3.10","3.11","3.12"]
+
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install "all" dependencies
+        run: uv sync --extra chgnet --extra dpa3 --extra d3 --extra grace --extra mace --extra orb --extra upet --extra plumed
+
+      - name: Install PLUMED
+        uses: Iximiel/install-plumed@v1
+        id: plumed
+        continue-on-error: true
+
+      - name: Set environment variable based on plumed success
+        run: |
+          if [ "${{ steps.plumed.outcome }}" = "success" ]; then
+            echo "PLUMED_KERNEL=${{ steps.plumed.outputs.plumed_prefix }}/lib/libplumedKernel.dylib" >> $GITHUB_ENV
+          fi
+
+      - name: Download extra models
+        run: |
+          python3 tests/models/extra_models.py tests/models/extra
+
+      - name: Run test suite
+        env:
+          # show timings of tests
+          PYTEST_ADDOPTS: "--durations=0"
+        run: uv run --no-sync pytest -k cuda
+
+      - name: Install updated e3nn dependencies
+        run: |
+          uv sync --extra mattersim --extra fairchem --extra sevennet --extra nequip --extra d3
+          uv pip install --reinstall pynvml
+          uv pip install "fairchem-core[torch-extras]" --no-build-isolation
+
+      - name: Run test suite for updated e3nn dependencies
+        env:
+          # show timings of tests
+          PYTEST_ADDOPTS: "--durations=0"
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py -k cuda

.github/workflows/ci.yml

@@ -46,7 +46,7 @@ jobs:
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run --no-sync pytest --cov janus_core --cov-append .
+        run: uv run --no-sync pytest --cov janus_core --cov-append . -k "not cuda"
 
       - name: Install updated e3nn dependencies
         run: |
@@ -60,7 +60,7 @@ jobs:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py -k "not cuda"
 
       - name: Report coverage to Coveralls
         uses: coverallsapp/github-action@v2

.github/workflows/mac.yml

@@ -42,7 +42,7 @@ jobs:
         env:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
-        run: uv run --no-sync pytest
+        run: uv run --no-sync pytest -k "not cuda"
 
       - name: Install updated e3nn dependencies
         run: |
@@ -55,4 +55,4 @@ jobs:
           # show timings of tests
           PYTEST_ADDOPTS: "--durations=0"
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
-        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py
+        run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py -k "not cuda"

janus_core/helpers/log.py

@@ -5,6 +5,7 @@
 import json
 import logging
 from typing import Literal
+from warnings import warn
 
 from codecarbon import OfflineEmissionsTracker
 from codecarbon.output import LoggerOutput
@@ -187,6 +188,7 @@ def config_tracker(
     OfflineEmissionsTracker | None
         Configured offline codecarbon tracker, if logger is specified.
     """
+    tracker = None
     if janus_logger and track_carbon:
         carbon_logger = LoggerOutput(janus_logger)
         tracker = OfflineEmissionsTracker(
@@ -206,12 +208,11 @@ def config_tracker(
             carbon_logger.removeHandler(carbon_logger.handlers[0])
 
         if not hasattr(tracker, "_emissions"):
-            raise ValueError(
+            warn(
                 "Carbon tracker has not been configured correctly. Please try "
-                "reconfiguring, or disable the tracker."
+                "reconfiguring, or disable the tracker.",
+                stacklevel=2,
             )
-
-    else:
-        tracker = None
+            tracker = None
 
     return tracker

pyproject.toml

@@ -168,10 +168,15 @@ select = [
     "N",
     # isort
     "UP",
+    # mccabe complexity
+    "C901",
 ]
 [tool.ruff.lint.per-file-ignores]
 'containers/*.py' = ['D100', 'E501']
 
+[tool.ruff.lint.mccabe]
+max-complexity = 5
+
 [tool.ruff.lint.isort]
 force-sort-within-sections = true
 required-imports = ["from __future__ import annotations"]

tests/data/mlip_fine_tune.yml

@@ -32,3 +32,4 @@ keep_checkpoints: False
 save_cpu: True
 weight_decay: 1e-8
 eval_interval: 2
+plot: False

tests/data/mlip_train.yml

@@ -44,3 +44,4 @@ clip_grad: 100
 keep_checkpoints: False
 keep_isolated_atoms: True
 save_cpu: True
+plot: False

tests/test_descriptors.py

@@ -7,6 +7,7 @@
 from ase import Atoms
 from ase.io import read
 import pytest
+import torch
 
 from janus_core.calculations.descriptors import Descriptors
 from janus_core.calculations.single_point import SinglePoint
@@ -17,11 +18,15 @@
 MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model"
 
 
-def test_calc_descriptors(tmp_path):
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_calc_descriptors(tmp_path, device):
     """Test calculating equation of state from ASE atoms object."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     struct = read(DATA_PATH / "NaCl.cif")
     log_file = tmp_path / "descriptors.log"
-    struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH)
+    struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device)
 
     descriptors = Descriptors(
         struct,
@@ -41,13 +46,18 @@ def test_calc_descriptors(tmp_path):
     )
 
 
-def test_calc_per_element(tmp_path):
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_calc_per_element(tmp_path, device):
     """Test calculating descriptors for each element from SinglePoint object."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     log_file = tmp_path / "descriptors.log"
     single_point = SinglePoint(
         struct=DATA_PATH / "NaCl.cif",
         arch="mace",
         model=MODEL_PATH,
+        device=device,
     )
 
     descriptors = Descriptors(
@@ -68,14 +78,19 @@ def test_calc_per_element(tmp_path):
     assert atoms.info["mace_Na_descriptor"] == pytest.approx(-0.0020374985791535563)
 
 
-def test_logging(tmp_path):
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_logging(tmp_path, device):
     """Test attaching logger to Descriptors and emissions are saved to info."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     log_file = tmp_path / "descriptors.log"
 
     single_point = SinglePoint(
         struct=DATA_PATH / "NaCl.cif",
         arch="mace_mp",
         model=MODEL_PATH,
+        device=device,
     )
 
     descriptors = Descriptors(
@@ -92,12 +107,17 @@ def test_logging(tmp_path):
     assert single_point.struct.info["emissions"] > 0
 
 
-def test_dispersion():
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_dispersion(device):
     """Test using mace_mp with dispersion."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     single_point = SinglePoint(
         struct=DATA_PATH / "NaCl.cif",
         arch="mace_mp",
         calc_kwargs={"dispersion": False},
+        device=device,
     )
 
     descriptors = Descriptors(
@@ -110,6 +130,7 @@ def test_dispersion():
         struct=DATA_PATH / "NaCl.cif",
         arch="mace_mp",
         calc_kwargs={"dispersion": True},
+        device=device,
     )
 
     descriptors_disp = Descriptors(
@@ -118,17 +139,24 @@ def test_dispersion():
     )
     descriptors_disp.run()
 
-    assert (
-        descriptors_disp.struct.info["mace_mp_d3_descriptor"]
-        == descriptors.struct.info["mace_mp_descriptor"]
+    assert descriptors_disp.struct.info["mace_mp_d3_descriptor"] == pytest.approx(
+        descriptors.struct.info["mace_mp_descriptor"]
     )
 
 
-def test_not_implemented_error():
-    """Test correct error raised if descriptors not implemented."""
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_not_implemented_error(device):
+    """Test correct implemented error raised if descriptors not installed."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
+    from tests.utils import skip_extras
+
+    skip_extras("chgnet")
     single_point = SinglePoint(
         struct=DATA_PATH / "NaCl.cif",
         arch="chgnet",
+        device=device,
     )
     with pytest.raises(NotImplementedError):
         Descriptors(