Merge pull request #750 from DiamondLightSource/data_reducer_order

corrections around data checker GPU and transfer optimisation for save_intermediate

Author: dkazanc

httomo/method_wrappers/generic.py

@@ -334,7 +334,7 @@ def execute(self, block: T) -> T:
             block = self._run_method(block, args)
             block = self._postprocess_data(block)
 
-        if gpu_enabled:
+        if self.is_gpu:
             self._gpu_time_info.kernel = t.elapsed
         else:
             self._gpu_time_info.kernel = 0

httomo/method_wrappers/save_intermediate.py

@@ -1,6 +1,5 @@
 import os
-import pathlib
-from typing import Any, Dict, Optional, Union
+from typing import Dict, Optional
 import weakref
 from mpi4py.MPI import Comm, MIN
 import httomo
@@ -9,10 +8,9 @@
 from httomo.runner.loader import LoaderInterface
 from httomo.runner.method_wrapper import GpuTimeInfo, MethodWrapper
 from httomo.runner.methods_repository_interface import MethodRepository
-from httomo.utils import catchtime, xp
+from httomo.utils import catchtime, xp, gpu_enabled
 
 import h5py
-import numpy as np
 
 
 class SaveIntermediateFilesWrapper(GenericMethodWrapper):
@@ -30,6 +28,7 @@ def __init__(
         output_mapping: Dict[str, str] = {},
         out_dir: Optional[os.PathLike] = None,
         prev_method: Optional[MethodWrapper] = None,
+        next_method_is_cpu: bool = False,
         loader: Optional[LoaderInterface] = None,
         **kwargs,
     ):
@@ -45,6 +44,7 @@ def __init__(
         assert loader is not None
         self._loader = loader
         assert prev_method is not None
+        self._next_method_is_cpu = next_method_is_cpu
 
         filename = f"{prev_method.task_id}-{prev_method.package_name}-{prev_method.method_name}"
         is_saving_recon = prev_method.module_path.endswith(".algorithm")
@@ -62,21 +62,34 @@ def __init__(
         # make sure file gets closed properly
         weakref.finalize(self, self._file.close)
 
+    def _transfer_data(self, block: T) -> T:
+        if block.is_cpu:
+            return block
+        if self._next_method_is_cpu:
+            # convert the whole (GPU) block to CPU if the next method is CPU
+            self._gpu_time_info = GpuTimeInfo()
+            with catchtime() as t:
+                block.to_cpu()
+            self._gpu_time_info.device2host = t.elapsed
+            return block
+        return block
+
     def execute(self, block: T) -> T:
-        # we overwrite the whole execute method here, as we do not need any of the helper
-        # methods from the Generic Wrapper
-        # What we know:
-        # - we do not transfer the dataset as a whole to CPU - only the data and angles locally
-        # (and never back)
-        # - the user does not insert this method - it's automatic - so no config params are
-        # relevant
-        # - we return just the input as it is
-        self._gpu_time_info = GpuTimeInfo()
-        with catchtime() as t:
-            data = (
-                block.data_unpadded if block.is_cpu else xp.asnumpy(block.data_unpadded)
-            )
-        self._gpu_time_info.device2host += t.elapsed
+        # we overwrite the most of the execute method here
+        # we transfer the data to CPU only if the next method is CPU, otherwise we keep it on GPU
+        # in case if save_intermediate is the last method we also keep the data on GPU
+        block = self._transfer_data(block)
+
+        if self._next_method_is_cpu or not gpu_enabled:
+            data = block.data_unpadded
+        else:
+            # Transfer data to CPU while the main block stays on GPU
+            self._gpu_time_info = GpuTimeInfo()
+
+            with catchtime() as t:
+                data = xp.asnumpy(block.data_unpadded)
+
+            self._gpu_time_info.device2host += t.elapsed
 
         MIN_BLOCK_LEN_PARAM = "minimum_block_length"
         if block.chunk_index_unpadded[block.slicing_dim] == 0 and self.comm.size > 1:

httomo/method_wrappers/stats_calc.py

@@ -1,7 +1,7 @@
 from httomo.block_interfaces import T
 from httomo.method_wrappers.generic import GenericMethodWrapper
 from httomo.runner.methods_repository_interface import MethodRepository
-from httomo.utils import catchtime, log_rank, xp, gpu_enabled
+from httomo.utils import log_rank
 
 
 from mpi4py.MPI import Comm
@@ -51,13 +51,6 @@ def _transfer_data(self, dataset: T) -> T:
         return dataset
 
     def _run_method(self, dataset: T, args: Dict[str, Any]) -> T:
-        # transfer data to GPU if we can / have it available (always faster),
-        # but don't want to fail if we don't have a GPU (underlying method works for both)
-        # and don't touch original dataset
-        if gpu_enabled and dataset.is_cpu:
-            with catchtime() as t:
-                args[self._parameters[0]] = xp.asarray(dataset.data)
-            self._gpu_time_info.host2device += t.elapsed
         ret = self._method(**args)
         return self._process_return_type(ret, dataset)
 

httomo/methods.py

@@ -35,15 +35,6 @@ def calculate_stats(
     Returns:
         tuple[(float, float, float, int)]: (min, max, sum, total_elements)
     """
-
-    # do this whereever the data is at the moment (GPU/CPU)
-    if data.device != "cpu":
-        # GPU
-        data = xp.nan_to_num(data, copy=False, nan=0.0, posinf=0, neginf=0)
-    else:
-        # CPU
-        data = np.nan_to_num(data, copy=False, nan=0.0, posinf=0, neginf=0)
-
     return (float(data.min()), float(data.max()), float(data.sum()), data.size)
 
 

httomo/runner/task_runner.py

@@ -269,15 +269,12 @@ def _setup_source_sink(self, section: Section, idx: int):
     def _execute_section_block(
         self, section: Section, block: DataSetBlock
     ) -> DataSetBlock:
-        if_previous_block_is_on_gpu = False
         convert_gpu_block_to_cpu = False
         if_current_block_is_on_gpu = False
 
         for ind, method in enumerate(section):
             if method.implementation == "gpu_cupy":
                 if_current_block_is_on_gpu = True
-            if method.method_name == "calculate_stats" and if_previous_block_is_on_gpu:
-                if_current_block_is_on_gpu = True
 
             if ind == len(section) - 1 and if_current_block_is_on_gpu:
                 convert_gpu_block_to_cpu = True
@@ -309,7 +306,6 @@ def _execute_section_block(
                     method.gpu_time.device2host,
                 )
 
-            if_previous_block_is_on_gpu = if_current_block_is_on_gpu
         return block
 
     def _get_methods_name_for_snapshot(self, section: Section) -> str:

httomo/transform_layer.py

@@ -1,11 +1,9 @@
 import os
 from typing import Optional
 from httomo.method_wrappers import make_method_wrapper
-from httomo.runner.output_ref import OutputRef
 from httomo.method_wrappers.datareducer import DatareducerWrapper
 from httomo.method_wrappers.generic import GenericMethodWrapper
 from httomo.method_wrappers.images import ImagesWrapper
-from httomo.method_wrappers.stats_calc import StatsCalcWrapper
 from httomo.method_wrappers.save_intermediate import SaveIntermediateFilesWrapper
 from httomo.runner.pipeline import Pipeline
 from mpi4py import MPI
@@ -43,25 +41,29 @@ def transform(self, pipeline: Pipeline) -> Pipeline:
         pipeline = self.insert_data_reducer(pipeline)
         if pipeline_is_sweep:
             pipeline = self.remove_redundant_method_in_sweep(pipeline)
-        pipeline = self.insert_data_checker(pipeline)
 
         if pipeline_is_sweep:
             pipeline = self.insert_save_images_after_sweep(pipeline)
         else:
             pipeline = self.insert_save_methods(pipeline)
 
+        pipeline = self.insert_data_checker(pipeline)
         return pipeline
 
     def insert_save_methods(self, pipeline: Pipeline) -> Pipeline:
         loader = pipeline.loader
         methods = []
-        for m in pipeline:
+        for i, m in enumerate(pipeline):
             methods.append(m)
             if (
                 (m.save_result or self._save_all)
                 and m.method_name not in ["save_to_images", "data_checker"]
                 and "center" not in m.method_name
             ):
+                if i + 1 < len(pipeline):
+                    next_method_is_cpu = pipeline[i + 1].is_cpu
+                else:
+                    next_method_is_cpu = False
                 methods.append(
                     SaveIntermediateFilesWrapper(
                         self._repo,
@@ -71,6 +73,7 @@ def insert_save_methods(self, pipeline: Pipeline) -> Pipeline:
                         save_result=False,
                         loader=loader,
                         prev_method=m,
+                        next_method_is_cpu=next_method_is_cpu,
                         task_id=f"save_{m.task_id}",
                         out_dir=self._out_dir,
                     )
@@ -108,6 +111,7 @@ def insert_data_checker(self, pipeline: Pipeline) -> Pipeline:
                 "calculate_stats",
                 "rescale_to_int",
                 "save_to_images",
+                "save_intermediate_data",
             ]
             if (
                 m.method_name not in exceptions_methods

tests/method_wrappers/test_save_intermediate.py

@@ -130,9 +130,18 @@ def save_intermediate_data(
     assert wrp._file.filename.startswith(str(tmp_path))
 
 
-@pytest.mark.parametrize("gpu", [False, True], ids=["CPU", "GPU"])
+@pytest.mark.parametrize("gpu", [False, True], ids=["data_CPU", "data_GPU"])
+@pytest.mark.parametrize(
+    "next_method_is_cpu",
+    [False, True],
+    ids=["Next Method is GPU", "Next Method is CPU"],
+)
 def test_save_intermediate_leaves_gpu_data(
-    mocker: MockerFixture, dummy_block: DataSetBlock, tmp_path: Path, gpu: bool
+    mocker: MockerFixture,
+    dummy_block: DataSetBlock,
+    tmp_path: Path,
+    gpu: bool,
+    next_method_is_cpu: bool,
 ):
     if gpu and not gpu_enabled:
         pytest.skip("No GPU available")
@@ -179,6 +188,7 @@ def save_intermediate_data(
         loader=loader,
         out_dir=tmp_path,
         prev_method=prev_method,
+        next_method_is_cpu=next_method_is_cpu,
         minimum_block_length=dummy_block.shape[0],
     )
 
@@ -189,7 +199,10 @@ def save_intermediate_data(
     with mock.patch("httomo.globals.FRAMES_PER_CHUNK", FRAMES_PER_CHUNK):
         res = wrp.execute(dummy_block)
 
-    assert res.is_gpu == gpu
+    if next_method_is_cpu and gpu:
+        assert res.is_cpu == True
+    else:
+        assert res.is_gpu == gpu
 
 
 @pytest.mark.parametrize(

tests/method_wrappers/test_stats_calc.py

@@ -118,7 +118,7 @@ def calculate_stats(data, comm):
 
 
 @pytest.mark.parametrize("gpu", [False, True], ids=["CPU-input", "GPU-input"])
-def test_calculate_stats_uses_gpu_if_available(
+def test_calculate_stats_agnostic(
     mocker: MockerFixture, dummy_block: DataSetBlock, gpu: bool
 ):
     if gpu and not gpu_enabled:
@@ -127,10 +127,6 @@ def test_calculate_stats_uses_gpu_if_available(
     class FakeModule:
         def calculate_stats(data, comm):
             # regardless of dataset input, we want device data if gpu enabled
-            if gpu_enabled:
-                assert data.device != "cpu"
-            else:
-                assert data.device == "cpu"
             return (1.2, 3.1, 42.0, 10)
 
     mocker.patch(
@@ -149,5 +145,7 @@ def calculate_stats(data, comm):
         dummy_block.to_gpu()
 
     res = wrp.execute(dummy_block)
-
-    assert res.is_gpu == gpu
+    if gpu is True:
+        assert res.is_gpu is True
+    else:
+        assert res.is_cpu is True

tests/test_methods.py

@@ -21,17 +21,6 @@ def test_calculate_stats_simple():
     assert ret == (-10.0, 19.0, np.sum(data), 30)
 
 
-def test_calculate_stats_with_nan_and_inf():
-    data = np.arange(30, dtype=np.float32).reshape((2, 3, 5)) - 10.0
-    expected = np.sum(data) - data[1, 1, 1] - data[0, 2, 3] - data[1, 2, 3]
-    data[1, 1, 1] = float("inf")
-    data[0, 2, 3] = float("-inf")
-    data[1, 2, 3] = float("nan")
-    ret = calculate_stats(data)
-
-    assert ret == (-10.0, 19.0, expected, 30)
-
-
 @pytest.mark.skipif(
     not gpu_enabled or xp.cuda.runtime.getDeviceCount() == 0,
     reason="skipped as cupy is not available",