Merge pull request #51 from mgiammar/update_device_handling

Update ttsim3d version and match device handling behavior

Author: jdickerson95

pyproject.toml

@@ -33,6 +33,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Typing :: Typed",
 ]
 # add your package dependencies here
@@ -50,7 +51,7 @@ dependencies = [
     "torch-fourier-slice>=v0.2.0",
     "torch-fourier-filter>=v0.2.3",
     "torch-so3>=v0.2.0",
-    "ttsim3d>=v0.3.0",
+    "ttsim3d>=v0.4.0",
     "lmfit",
     "zenodo-get",
 ]

src/leopard_em/pydantic_models/config/computational_config.py

@@ -1,44 +1,51 @@
 """Computational configuration for 2DTM."""
 
-from typing import Annotated
+from typing import Annotated, Optional, Union
 
 import torch
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, Field
+
+# Type alias for non-negative integer
+NonNegativeInt = Annotated[int, Field(ge=0)]
 
 
 class ComputationalConfig(BaseModel):
     """Serialization of computational resources allocated for 2DTM.
 
+    NOTE: The field `gpu_ids` is not validated at instantiation past being one of the
+    valid types. For example, if "cuda:0" is specified but no CUDA device is available,
+    the instantiation will succeed, and only upon translating `gpu_ids` to a list of
+    `torch.device` objects will an error be raised. This is done to allow for
+    configuration files to be loaded without requiring the actual hardware to be
+    present at the time of loading.
+
     Attributes
     ----------
-    gpu_ids : list[int]
-        Which GPU(s) to use for computation, defaults to 0 which will use device at
-        index 0. A value of -2 or less corresponds to CPU device. A value of -1 will
-        use all available GPUs.
+    gpu_ids : Optional[Union[int, list[int], str, list[str]]]
+        Field which specifies which GPUs to use for computation. The following types
+        of values are allowed:
+        - A single integer, e.g. 0, which means to use GPU with ID 0.
+        - A list of integers, e.g. [0, 2], which means to use GPUs with IDs 0 and 2.
+        - A device specifier string, e.g. "cuda:0", which means to use GPU with ID 0.
+        - A list of device specifier strings, e.g. ["cuda:0", "cuda:1"], which means to
+          use GPUs with IDs 0 and 1.
+        - The specific string "all" which means to use all available GPUs identified
+          by torch.cuda.device_count().
+        - The specific string "cpu" which means to use CPU.
     num_cpus : int
         Total number of CPUs to use, defaults to 1.
     """
 
-    gpu_ids: int | list[int] = [0]
-    num_cpus: Annotated[int, Field(ge=1)] = 1
-
-    @field_validator("gpu_ids")  # type: ignore
-    def validate_gpu_ids(cls, v):  # pylint: disable=no-self-argument
-        """Validate input value for GPU ids."""
-        if isinstance(v, int):
-            v = [v]
-
-        # Check if -1 appears, it is only value in list
-        if -1 in v and len(v) > 1:
-            raise ValueError(
-                "If -1 (all GPUs) is in the list, it must be the only value."
-            )
-
-        # Check if -2 appears, it is only value in list
-        if -2 in v and len(v) > 1:
-            raise ValueError("If -2 (CPU) is in the list, it must be the only value.")
-
-        return v
+    # Type-hinting here is ensuring non-negative integers, and list of at least one
+    gpu_ids: Optional[
+        Union[
+            str,
+            NonNegativeInt,
+            Annotated[list[NonNegativeInt], Field(min_length=1)],
+            Annotated[list[str], Field(min_length=1)],
+        ]
+    ] = [0]
+    num_cpus: NonNegativeInt = 1
 
     @property
     def gpu_devices(self) -> list[torch.device]:
@@ -48,13 +55,29 @@ def gpu_devices(self) -> list[torch.device]:
         -------
         list[torch.device]
         """
-        # Case where gpu_ids is integer
-        if isinstance(self.gpu_ids, int):
-            self.gpu_ids = [self.gpu_ids]
-
-        if -1 in self.gpu_ids:
+        # Handle special string cases first
+        if self.gpu_ids == "all":
+            if not torch.cuda.is_available():
+                raise ValueError("No CUDA devices available.")
             return [torch.device(f"cuda:{i}") for i in range(torch.cuda.device_count())]
-        if -2 in self.gpu_ids:
+
+        if self.gpu_ids == "cpu":
             return [torch.device("cpu")]
 
-        return [torch.device(f"cuda:{gpu_id}") for gpu_id in self.gpu_ids]
+        # Normalize to list for uniform processing
+        gpu_list = self.gpu_ids if isinstance(self.gpu_ids, list) else [self.gpu_ids]
+
+        # Process each item in the normalized list
+        devices = []
+        for gpu_id in gpu_list:
+            if isinstance(gpu_id, int):
+                devices.append(torch.device(f"cuda:{gpu_id}"))
+            elif isinstance(gpu_id, str):
+                devices.append(torch.device(gpu_id))
+            else:
+                raise TypeError(
+                    f"Invalid type for gpu_ids element: {type(gpu_id)}. "
+                    "Expected int or str."
+                )
+
+        return devices

src/leopard_em/pydantic_models/managers/optimize_template_manager.py

@@ -71,7 +71,7 @@ def make_backend_core_function_kwargs(
             Whether to use refined angles or not. Defaults to True.
         """
         # simulate template volume
-        template = self.simulator.run(gpu_ids=self.computational_config.gpu_ids)
+        template = self.simulator.run(device=self.computational_config.gpu_ids)
 
         # The set of "best" euler angles from match template search
         # Check if refined angles exist, otherwise use the original angles

tests/pydantic_models/test_computational_config.py

@@ -1,7 +1,6 @@
 """Tests for the ComputationalConfig model"""
 
 import pytest
-import torch
 from pydantic import ValidationError
 
 from leopard_em.pydantic_models.config import ComputationalConfig
@@ -18,122 +17,14 @@ def test_default_values():
     assert config.num_cpus == 1
 
 
-def test_single_gpu_id():
-    """
-    Test that a single integer gpu_id is converted to a list.
-
-    Verifies that when passing a single integer as gpu_ids, it's converted to a list
-    containing that integer.
-    """
-    config = ComputationalConfig(gpu_ids=1)
-    assert config.gpu_ids == [1]
-
-
-def test_multiple_gpu_ids():
-    """
-    Test that multiple gpu_ids are correctly stored as a list.
-
-    Verifies that when passing a list of gpu_ids, they are correctly stored in the
-    config.
-    """
-    config = ComputationalConfig(gpu_ids=[0, 1, 2])
-    assert config.gpu_ids == [0, 1, 2]
-
-
-def test_all_gpus():
-    """
-    Test the special value -1 for gpu_ids.
-
-    Verifies that when passing -1 as gpu_ids, it's stored correctly, which indicates
-    using all available GPUs.
-    """
-    config = ComputationalConfig(gpu_ids=-1)
-    assert config.gpu_ids == [-1]
-
-
-def test_cpu_only():
-    """
-    Test the special value -2 for gpu_ids.
-
-    Verifies that when passing -2 as gpu_ids, it's stored correctly, which indicates
-    using CPU only.
-    """
-    config = ComputationalConfig(gpu_ids=-2)
-    assert config.gpu_ids == [-2]
-
-
 def test_invalid_gpu_ids():
     """
-    Test that invalid combinations of gpu_ids raise errors.
-
-    Verifies that special values -1 and -2 cannot be combined with other gpu IDs.
-    """
-    with pytest.raises(ValueError, match="If -1"):
-        ComputationalConfig(gpu_ids=[-1, 0])
-
-    with pytest.raises(ValueError, match="If -2"):
-        ComputationalConfig(gpu_ids=[-2, 0])
-
-
-def test_num_cpus():
-    """
-    Test that num_cpus is correctly stored.
+    Test invalid gpu_ids values.
 
-    Verifies that the specified number of CPUs is correctly stored in the config.
-    """
-    config = ComputationalConfig(num_cpus=4)
-    assert config.num_cpus == 4
-
-
-def test_invalid_num_cpus():
-    """
-    Test that invalid num_cpus values raise errors.
-
-    Verifies that specifying zero or negative values for num_cpus raises a
-    ValidationError.
+    Verifies that a ValidationError is raised when invalid gpu_ids are provided.
     """
     with pytest.raises(ValidationError):
-        ComputationalConfig(num_cpus=0)
-
-
-def test_gpu_devices_single_gpu():
-    """
-    Test the gpu_devices property for a single GPU configuration.
-
-    Verifies that the correct torch.device object is created for a single GPU.
-    """
-    config = ComputationalConfig(gpu_ids=0)
-    assert config.gpu_devices == [torch.device("cuda:0")]
-
-
-def test_gpu_devices_multiple_gpus():
-    """
-    Test the gpu_devices property for multiple GPUs configuration.
+        ComputationalConfig(gpu_ids=[-1])  # Negative GPU ID is invalid
 
-    Verifies that the correct torch.device objects are created for multiple GPUs.
-    """
-    config = ComputationalConfig(gpu_ids=[0, 1])
-    assert config.gpu_devices == [torch.device("cuda:0"), torch.device("cuda:1")]
-
-
-def test_gpu_devices_all_gpus(monkeypatch):
-    """
-    Test the gpu_devices property when using all available GPUs.
-
-    Uses monkeypatch to set a fixed number of GPUs for testing, and verifies
-    that the correct device objects are created for all available GPUs.
-    """
-    monkeypatch.setattr(torch.cuda, "device_count", lambda: 2)
-    config = ComputationalConfig(gpu_ids=-1)
-    assert config.gpu_devices == [torch.device("cuda:0"), torch.device("cuda:1")]
-
-
-def test_gpu_devices_cpu():
-    """
-    Test the gpu_devices property in CPU-only mode.
-
-    Verifies that the correct torch.device object for CPU is created when using
-    CPU-only mode.
-    """
-    config = ComputationalConfig(gpu_ids=-2)
-    assert config.gpu_devices == [torch.device("cpu")]
+    with pytest.raises(ValidationError):
+        ComputationalConfig(gpu_ids=[])  # Empty list