Add cpu and gpu markers to tests to be able to explicitly choose them

Signed-off-by: Samet Akcay <[email protected]>

Author: samet-akcay

.github/actions/pytest/action.yaml

@@ -137,9 +137,9 @@ runs:
 
         # Set device-specific pytest arguments
         if [ "${{ inputs.device }}" = "cpu" ]; then
-          DEVICE_ARGS="--markers='not gpu'"
+          DEVICE_ARGS="-m cpu"
         else
-          DEVICE_ARGS=""
+          DEVICE_ARGS="-m 'cpu or gpu'"  # Run all tests on GPU
         fi
 
         # Run pytest
@@ -150,7 +150,7 @@ runs:
           --timeout=${{ inputs.max-test-time }} \
           --verbosity=0 \
           --durations-only \
-          $DEVICE_ARGS
+          ${DEVICE_ARGS}
 
         # Store test result and duration
         exit_code=$?

pyproject.toml

@@ -291,6 +291,10 @@ skips = ["B101"]
 addopts = ["--strict-markers", "--strict-config", "--showlocals", "-ra"]
 testpaths = "tests"
 pythonpath = "src"
+markers = [
+    "gpu: marks tests that require GPU",
+    "cpu: marks tests that can run on CPU only (default)",
+]
 
 
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

tests/conftest.py

@@ -101,3 +101,10 @@ def checkpoint(model_name: str) -> Path:
         return _ckpt_path
 
     return checkpoint
+
+
+def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
+    """Automatically mark tests as 'cpu' unless they're marked as 'gpu'."""
+    for item in items:
+        if not any(marker.name == "gpu" for marker in item.iter_markers()):
+            item.add_marker(pytest.mark.cpu)