add gpu workflow and tests

Author: alinelena

.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@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+        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

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(

tests/test_elasticity.py

@@ -7,7 +7,9 @@
 from ase.build import bulk
 from ase.io import read
 import numpy as np
+import pytest
 from pytest import approx
+import torch
 
 from janus_core.calculations.elasticity import Elasticity
 from janus_core.helpers.mlip_calculators import choose_calculator
@@ -17,15 +19,19 @@
 MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model"
 
 
-def test_calc_elasticity(tmp_path):
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_calc_elasticity(tmp_path, device):
     """Test calculating elasticity for Aluminium."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     elasticity_path = tmp_path / "Al-elastic_tensor.dat"
     log_file = tmp_path / "Al-elasticity.log"
     generated_path = tmp_path / "Al-elasticity-generated.extxyz"
     minimized_path = tmp_path / "Al-elasticity-opt.extxyz"
 
     struct = bulk("Al", crystalstructure="fcc")
-    struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH)
+    struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device)
 
     elasticity = Elasticity(
         struct,
@@ -86,15 +92,19 @@ def test_calc_elasticity(tmp_path):
         assert struct.info["strain"] == approx(strain.voigt)
 
 
-def test_no_optimize_no_write_voigt(tmp_path):
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_no_optimize_no_write_voigt(tmp_path, device):
     """Test calculating elasticity for Aluminium without optimization."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     elasticity_path = tmp_path / "Al-elastic_tensor.dat"
     log_file = tmp_path / "Al-elasticity.log"
     generated_path = tmp_path / "Al-elasticity-generated.extxyz"
     minimized_path = tmp_path / "Al-elasticity-opt.extxyz"
 
     struct = bulk("Al", crystalstructure="fcc")
-    struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH)
+    struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device)
 
     elasticity = Elasticity(
         struct,

tests/test_elasticity_cli.py

@@ -6,7 +6,9 @@
 
 from ase.io import read
 import numpy as np
+import pytest
 from pytest import approx
+import torch
 from typer.testing import CliRunner
 import yaml
 
@@ -26,8 +28,12 @@ def test_help():
     assert "Usage: janus elasticity [OPTIONS]" in strip_ansi_codes(result.stdout)
 
 
-def test_elasticity_opt_none(tmp_path):
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_elasticity_opt_none(tmp_path, device):
     """Test calculating the ElasticTensor from the CLI without optimisation."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     with chdir(tmp_path):
         results_dir = Path("./janus_results")
         elasticity_path = results_dir / "NaCl-elastic_tensor.dat"
@@ -44,6 +50,8 @@ def test_elasticity_opt_none(tmp_path):
                 DATA_PATH / "NaCl.cif",
                 "--arch",
                 "mace_mp",
+                "--device",
+                device,
                 "--n-strains",
                 "2",
                 "--no-minimize",
@@ -95,8 +103,12 @@ def test_elasticity_opt_none(tmp_path):
         assert elasticity_summary["config"]["n_strains"] == 2
 
 
-def test_elasticity_opt_all(tmp_path):
+@pytest.mark.parametrize("device", ["cpu", "cuda"])
+def test_elasticity_opt_all(tmp_path, device):
     """Test calculating the ElasticTensor from the command line."""
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA not available")
+
     with chdir(tmp_path):
         results_dir = Path("./janus_results")
         elasticity_path = results_dir / "NaCl-elastic_tensor.dat"
@@ -113,6 +125,8 @@ def test_elasticity_opt_all(tmp_path):
                 DATA_PATH / "NaCl.cif",
                 "--arch",
                 "mace_mp",
+                "--device",
+                device,
                 "--n-strains",
                 "2",
                 "--minimize-all",