From ca2aa6e0f58b609164afc427a01849404515ba9b Mon Sep 17 00:00:00 2001
From: Naman Gera <namangera15@gmail.com>
Date: Thu, 8 Jun 2023 11:25:57 +0100
Subject: [PATCH] Add gpu-id flag in the cli

---
 httomo/cli.py            | 26 ++++++++++++++++++++++++++
 httomo/globals.py        |  1 +
 httomo/task_runner.py    |  5 ++---
 httomo/wrappers_class.py |  5 ++++-
 tests/test_cli.py        | 13 +++++++++++++
 5 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/httomo/cli.py b/httomo/cli.py
index d29854b12..d20b0c2b0 100644
--- a/httomo/cli.py
+++ b/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,
diff --git a/httomo/globals.py b/httomo/globals.py
index 6602c665e..905ba6f2a 100644
--- a/httomo/globals.py
+++ b/httomo/globals.py
@@ -1,2 +1,3 @@
 run_out_dir = None
 logger = None
+gpu_id = -1
diff --git a/httomo/task_runner.py b/httomo/task_runner.py
index cf1506ece..03363fb60 100644
--- a/httomo/task_runner.py
+++ b/httomo/task_runner.py
@@ -43,7 +43,6 @@
 MAX_SWEEPS = 1
 
 
-
 @dataclass
 class MethodFunc:
     """
@@ -450,9 +449,9 @@ def _get_method_funcs(yaml_config: Path, comm: MPI.Comm) -> List[MethodFunc]:
                 method_func=wrapper_func,
                 wrapper_func=wrapper_method,
                 parameters=method_conf,
-                cpu=True, # get cpu/gpu meta data info from httomolib methods
+                cpu=True,  # get cpu/gpu meta data info from httomolib methods
                 gpu=False,
-                calc_max_slices=None, # call calc_max_slices function in wrappers
+                calc_max_slices=None,  # call calc_max_slices function in wrappers
                 reslice_ahead=False,
                 pattern=Pattern.all,
                 is_loader=False,
diff --git a/httomo/wrappers_class.py b/httomo/wrappers_class.py
index 32f266dee..98e838c66 100644
--- a/httomo/wrappers_class.py
+++ b/httomo/wrappers_class.py
@@ -5,6 +5,7 @@
 import numpy as np
 from mpi4py.MPI import Comm
 
+import httomo.globals
 from httomo.data import mpiutil
 from httomo.utils import Colour, log_once
 
@@ -41,7 +42,9 @@ def __init__(
         self.comm = comm
         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
diff --git a/tests/test_cli.py b/tests/test_cli.py
index ec3f83d05..b7948e6e2 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -64,3 +64,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