Repository: desy-ml/cheetah (language: python)

Issue:
Run CI on multiple platforms
To ensure cheetah stays widely compatible, it would be a good idea to run the automatic tests not only for Ubuntu but also on Windows and MacOS. Although it is probably sufficient to run the test for one version of python on the alternative platforms, and not the whole matrix.


Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index eee86420..6d38d919 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -38,6 +38,7 @@
 - Added imports to the code example in the README (see #188) (@jank324)
 - Refactor definitions of physical constants (see #189) (@hespe)
 - Fix the quadrupole strength units in the quadrupole docstring (see #202) (@ansantam)
+- Add CI runs for macOS (arm64) and Windows (see #226) (@cr-xu, @jank324, @hespe)
 
 ## [v0.6.3](https://github.com/desy-ml/cheetah/releases/tag/v0.6.3) (2024-03-28)
 
diff --git a/cheetah/utils/__init__.py b/cheetah/utils/__init__.py
index 4bf5ed6b..06bb63f0 100644
--- a/cheetah/utils/__init__.py
+++ b/cheetah/utils/__init__.py
@@ -1,2 +1,3 @@
+from .device import is_mps_available_and_functional  # noqa: F401
 from .kde import kde_histogram_1d, kde_histogram_2d  # noqa: F401
 from .unique_name_generator import UniqueNameGenerator  # noqa: F401
diff --git a/cheetah/utils/device.py b/cheetah/utils/device.py
new file mode 100644
index 00000000..e6ac629d
--- /dev/null
+++ b/cheetah/utils/device.py
@@ -0,0 +1,13 @@
+import torch
+
+
+def is_mps_available_and_functional():
+    """Check if MPS is available and functional (for GitHub Actions)."""
+    if not torch.backends.mps.is_available():
+        return False
+    try:
+        # Try to allocate a small tensor on the MPS device
+        torch.tensor([1.0], device="mps")
+        return True
+    except RuntimeError:
+        return False


Test patch:
diff --git a/.github/workflows/pytest.yaml b/.github/workflows/pytest.yaml
index 99ca8099..04b6dbdf 100644
--- a/.github/workflows/pytest.yaml
+++ b/.github/workflows/pytest.yaml
@@ -4,10 +4,16 @@ on: [push, pull_request]
 
 jobs:
   build:
-    runs-on: ubuntu-latest
     strategy:
       matrix:
+        os: [ubuntu-latest]
         python-version: ["3.9", "3.10", "3.11", "3.12"]
+        include:
+          - os: macos-latest
+            python-version: "3.12"
+          - os: windows-latest
+            python-version: "3.12"
+    runs-on: ${{ matrix.os }}
 
     steps:
       - uses: actions/checkout@v4
diff --git a/tests/test_bmad_conversion.py b/tests/test_bmad_conversion.py
index ecf91cfe..9dc78849 100644
--- a/tests/test_bmad_conversion.py
+++ b/tests/test_bmad_conversion.py
@@ -2,6 +2,7 @@
 import torch
 
 import cheetah
+from cheetah.utils import is_mps_available_and_functional
 
 
 def test_bmad_tutorial():
@@ -47,7 +48,7 @@ def test_bmad_tutorial():
         pytest.param(
             torch.device("mps"),
             marks=pytest.mark.skipif(
-                not torch.backends.mps.is_available(), reason="MPS not available"
+                not is_mps_available_and_functional(), reason="MPS not available"
             ),
         ),
     ],
diff --git a/tests/test_device_dtype.py b/tests/test_device_dtype.py
index bd7e13a8..67e61205 100644
--- a/tests/test_device_dtype.py
+++ b/tests/test_device_dtype.py
@@ -2,6 +2,7 @@
 import torch
 
 import cheetah
+from cheetah.utils import is_mps_available_and_functional
 
 
 @pytest.mark.parametrize(
@@ -16,7 +17,7 @@
         pytest.param(
             torch.device("mps"),
             marks=pytest.mark.skipif(
-                not torch.backends.mps.is_available(), reason="MPS not available"
+                not is_mps_available_and_functional(), reason="MPS not available"
             ),
         ),
     ],
@@ -79,7 +80,7 @@ def test_change_quadrupole_dtype():
         pytest.param(
             torch.device("mps"),
             marks=pytest.mark.skipif(
-                not torch.backends.mps.is_available(), reason="MPS not available"
+                not is_mps_available_and_functional(), reason="MPS not available"
             ),
         ),
     ],