Enable CPU memory estimation of GPU pipelines to be callable on non-GPU machine

httomo/cli.py

@@ -324,6 +324,10 @@ def run(
         format_enum,
     )
 
+    for wrapper in pipeline_object:
+        wrapper.check()
+        wrapper.prep()
+
     if not does_contain_sweep:
         execute_high_throughput_run(
             pipeline_object,
@@ -539,7 +543,12 @@ def execute_sweep_run(pipeline: Pipeline, global_comm: MPI.Comm) -> None:
     ParamSweepRunner(pipeline, global_comm).execute()
 
 
-def estimate_cpu_memory(in_data_file: Path, pipeline_file: Path, nprocs: int) -> int:
+def estimate_cpu_memory(
+    in_data_file: Path,
+    pipeline_file: Path,
+    nprocs: int,
+    shape: Optional[tuple[int, int, int]] = None,
+) -> int:
     pipeline = generate_pipeline(
         in_data_file, pipeline_file, False, MPI.COMM_WORLD, PipelineFormat.Yaml
     )
@@ -551,14 +560,17 @@ def estimate_cpu_memory(in_data_file: Path, pipeline_file: Path, nprocs: int) ->
         dtype = dataset.dtype
         full_shape = dataset.shape
 
-    preview_config = parse_preview(
-        config[0]["parameters"].get("preview", None), full_shape
-    )
-    previewed_shape = (
-        preview_config.angles.stop - preview_config.angles.start,
-        preview_config.detector_y.stop - preview_config.detector_y.start,
-        preview_config.detector_x.stop - preview_config.detector_x.start,
-    )
+    if shape is None:
+        preview_config = parse_preview(
+            config[0]["parameters"].get("preview", None), full_shape
+        )
+        previewed_shape = (
+            preview_config.angles.stop - preview_config.angles.start,
+            preview_config.detector_y.stop - preview_config.detector_y.start,
+            preview_config.detector_x.stop - preview_config.detector_x.start,
+        )
+    else:
+        previewed_shape = shape
 
     section_memory_peak = 0
     for idx in range(len(sections)):

httomo/method_wrappers/generic.py

@@ -138,13 +138,21 @@ def __init__(
             self._query.save_result_default() if save_result is None else save_result
         )
 
-        if self.is_gpu and not gpu_enabled:
-            raise ValueError("GPU is not available, please use only CPU methods")
-
         self._side_output: Dict[str, Any] = dict()
 
         self._gpu_time_info = GpuTimeInfo()
 
+    def check(self) -> None:
+        """
+        Raises
+        ------
+        ValueError
+            If the method is a GPU method and the machine is not GPU-enabled.
+        """
+        if self.is_gpu and not gpu_enabled:
+            raise ValueError("GPU is not available, please use only CPU methods")
+
+    def prep(self) -> None:
         if gpu_enabled:
             _id = httomo.globals.gpu_id
             self._gpu_id = get_gpu_id() if _id == -1 else _id

httomo/runner/method_wrapper.py

@@ -205,3 +205,17 @@ def calculate_max_slices(
         something persists afterwards.
         """
         ...  # pragma: nocover
+
+    def check(self) -> None:
+        """
+        Perform any checks necessary to ensure that the method is capable of executing on the
+        machine (for example, if it's a GPU method, check that execution is on a machine with a
+        GPU that is accessible).
+        """
+        ...  # pragma: nocover
+
+    def prep(self) -> None:
+        """
+        Perform any preparation needed prior to beginning pipeline execution.
+        """
+        ...  # pragma: nocover

tests/method_wrappers/test_generic.py

@@ -213,7 +213,7 @@ def fake_method(data):
             "fake_method",
             MPI.COMM_WORLD,
             make_mock_preview_config(mocker),
-        )
+        ).check()
 
     assert "GPU is not available" in str(e)