Merge branch 'main' into memory

httomo/cli.py

@@ -66,6 +66,12 @@ def check(yaml_config: Path, in_data: Path = None):
     default=1,
     help=" The number of the CPU cores per process.",
 )
+@click.option(
+    "--gpu-id",
+    type=click.INT,
+    default=-1,
+    help="The GPU ID of the device to use.",
+)
 @click.option(
     "--save-all",
     is_flag=True,
@@ -99,6 +105,7 @@ def run(
     dimension: int,
     pad: int,
     ncore: int,
+    gpu_id: int,
     save_all: bool,
     file_based_reslice: bool,
     reslice_dir: Path,
@@ -120,6 +127,25 @@ def run(
         # Copy YAML pipeline file to output directory
         copy(yaml_config, httomo.globals.run_out_dir)
 
+    # try to access the GPU with the ID given
+    try:
+        import cupy as cp
+
+        gpu_count = cp.cuda.runtime.getDeviceCount()
+
+        if gpu_id != -1:
+            if gpu_id not in range(0, gpu_count):
+                raise ValueError(
+                    f"GPU Device not available for access. Use a GPU ID in the range: 0 to {gpu_count} (exclusive)"
+                )
+
+            cp.cuda.Device(gpu_id).use()
+
+        httomo.globals.gpu_id = gpu_id
+
+    except ImportError:
+        pass  # silently pass and run the CPU pipeline
+
     return run_tasks(
         in_file,
         yaml_config,

httomo/globals.py

@@ -1,2 +1,3 @@
 run_out_dir = None
 logger = None
+gpu_id = -1

httomo/wrappers_class.py

@@ -2,12 +2,12 @@
 import numpy as np
 import inspect
 from inspect import Parameter, signature
+import httomo.globals
 from httomo.utils import Colour, log_exception, log_once, gpu_enabled, xp
 from httomo.data import mpiutil
 
 from mpi4py.MPI import Comm
 
-
 def _gpumem_cleanup():
     """cleans up GPU memory and also the FFT plan cache"""
     if gpu_enabled:
@@ -28,7 +28,9 @@ def __init__(
         self.dict_params: Dict[str, Any] = {}
         if gpu_enabled:
             self.num_GPUs = xp.cuda.runtime.getDeviceCount()
-            self.gpu_id = mpiutil.local_rank % self.num_GPUs
+            _id = httomo.globals.gpu_id
+            # if gpu-id was specified in the CLI, use that
+            self.gpu_id = mpiutil.local_rank % self.num_GPUs if _id == -1 else _id
 
     def _transfer_data(self, *args) -> Union[tuple, xp.ndarray, np.ndarray]:
         """Transfer the data between the host and device for the GPU-enabled method

tests/test_cli.py

@@ -59,3 +59,16 @@ def test_cli_pass_output_folder(
     ]
     subprocess.check_output(cmd)
     assert Path(custom_output_dir, "user.log").exists()
+
+
+@pytest.mark.cupy
+def test_cli_pass_gpu_id(cmd, standard_data, standard_loader, output_folder):
+    cmd.insert(7, standard_data)
+    cmd.insert(8, standard_loader)
+    cmd.insert(4, "--gpu-id")
+    cmd.insert(5, "10")
+
+    result = subprocess.run(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
+    )
+    assert "GPU Device not available for access." in result.stderr