0.17.1

improved get_device()

Author: bartzbeielstein

RELEASE_NOTES.txt

@@ -1,3 +1,7 @@
+spotpython-0.17.1
+
+- improved get_device function
+
 spotpython-0.17.0
 
 - designs.py and spacefilling.py:

pyproject.toml

@@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "spotpython"
-version = "0.17.0"
+version = "0.17.1"
 authors = [
   { name="T. Bartz-Beielstein", email="tbb@bartzundbartz.de" }
 ]

src/spotpython/utils/device.py

@@ -2,24 +2,43 @@
 
 
 def getDevice(device=None):
-    """Get cpu, gpu or mps device for training.
+    """Get CPU, GPU (CUDA), or MPS device for training.
+
     Args:
         device (str):
-            Device for training. If None or "auto" the device is selected automatically.
+            Device for training. If None or "auto", the device is selected automatically based on availability.
 
     Returns:
         device (str):
             Device for training.
 
+    Raises:
+        ValueError: If the requested device is not recognized or available.
+
     Examples:
         >>> from spotpython.utils.device import getDevice
-        >>> getDevice()
-        'cuda:0'
+            getDevice()
+                'cuda:0'
     """
     if device is None or device == "auto":
+        # Automatically select device
         device = "cpu"
         if torch.cuda.is_available():
             device = "cuda:0"
         elif torch.backends.mps.is_available():
             device = "mps"
+        return device
+
+    # Check the explicit device request
+    if device.startswith("cuda"):
+        if not torch.cuda.is_available():
+            raise ValueError("CUDA device requested but no CUDA device is available.")
+    elif device == "mps":
+        if not torch.backends.mps.is_available():
+            raise ValueError("MPS device requested but MPS is not available.")
+    elif device == "cpu":
+        return "cpu"
+    else:
+        raise ValueError(f"Unrecognized device: {device}. Valid options are 'cpu', 'cuda:x', or 'mps'.")
+
     return device

test/test_get_device.py

@@ -0,0 +1,51 @@
+import pytest
+from unittest.mock import patch
+import torch
+from spotpython.utils.device import getDevice
+
+def test_get_device_auto_cpu():
+    with patch('torch.cuda.is_available') as mock_cuda, patch('torch.backends.mps.is_available') as mock_mps:
+        mock_cuda.return_value = False
+        mock_mps.return_value = False
+        assert getDevice() == 'cpu'
+
+def test_get_device_auto_cuda():
+    with patch('torch.cuda.is_available') as mock_cuda, patch('torch.backends.mps.is_available') as mock_mps:
+        mock_cuda.return_value = True
+        mock_mps.return_value = False
+        assert getDevice() == 'cuda:0'
+
+def test_get_device_auto_mps():
+    with patch('torch.cuda.is_available') as mock_cuda, patch('torch.backends.mps.is_available') as mock_mps:
+        mock_cuda.return_value = False
+        mock_mps.return_value = True
+        assert getDevice() == 'mps'
+
+def test_get_device_explicit_cpu():
+    assert getDevice('cpu') == 'cpu'
+
+def test_get_device_explicit_cuda_available():
+    with patch('torch.cuda.is_available') as mock_cuda:
+        mock_cuda.return_value = True
+        assert getDevice('cuda:0') == 'cuda:0'
+
+def test_get_device_explicit_cuda_unavailable():
+    with patch('torch.cuda.is_available') as mock_cuda:
+        mock_cuda.return_value = False
+        with pytest.raises(ValueError, match="CUDA device requested but no CUDA device is available."):
+            getDevice('cuda:0')
+
+def test_get_device_explicit_mps_available():
+    with patch('torch.backends.mps.is_available') as mock_mps:
+        mock_mps.return_value = True
+        assert getDevice('mps') == 'mps'
+
+def test_get_device_explicit_mps_unavailable():
+    with patch('torch.backends.mps.is_available') as mock_mps:
+        mock_mps.return_value = False
+        with pytest.raises(ValueError, match="MPS device requested but MPS is not available."):
+            getDevice('mps')
+
+def test_get_device_invalid():
+    with pytest.raises(ValueError, match="Unrecognized device: invalid_device. Valid options are 'cpu', 'cuda:x', or 'mps'."):
+        getDevice('invalid_device')