From 062bb53ad49ee1ab5d1e76cc731be3a7bb3b7826 Mon Sep 17 00:00:00 2001
From: Igor Tatarnikov <61896994+IgorTatarnikov@users.noreply.github.com>
Date: Fri, 31 May 2024 12:32:22 +0100
Subject: [PATCH] Force cpu usage on macos-latest runners (#114)

* Force cpu usage on macos-latest runners

* Moved torch_device setting to root conftest

* Go back to setting keras global state in brainglobe-workflows

* Force backend env variable to be set with a fixture before importing keras

* Add input that should cause an error only on macos-latest

* Run set_device_arm_macos_ci on all darwin platforms

* Try using the CI variable instead

* Pass GITHUB_* env variables via tox

* Moved passenv to the bottom of testenv to match other pyproject files

* Use force_cpu function from cellfinder

* Removed manual setting of GITHUB_ACTIONS
---
 pyproject.toml    |  3 +++
 tests/conftest.py | 24 ++++++++++++++++++++++++
 2 files changed, 27 insertions(+)
 create mode 100644 tests/conftest.py

diff --git a/pyproject.toml b/pyproject.toml
index aae15f17..e91ef78b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -147,4 +147,7 @@ commands =
 description =
     Run tests
     coredev: Run tests with the development version of cellfinder
+passenv =
+    CI
+    GITHUB_*
 """
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 00000000..b61cb636
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,24 @@
+import os
+
+import pytest
+import torch
+from cellfinder.core.tools.system import force_cpu
+
+
+@pytest.fixture(scope="session", autouse=True)
+def set_backend_torch():
+    os.environ["KERAS_BACKEND"] = "torch"
+
+
+@pytest.fixture(scope="session", autouse=True)
+def set_device_arm_macos_ci():
+    """
+    Ensure that the device is set to CPU when running on arm based macOS
+    GitHub runners. This is to avoid the following error:
+    https://discuss.pytorch.org/t/mps-back-end-out-of-memory-on-github-action/189773/5
+    """
+    if (
+        os.getenv("GITHUB_ACTIONS") == "true"
+        and torch.backends.mps.is_available()
+    ):
+        force_cpu()