diff --git a/httomo/method_wrappers/generic.py b/httomo/method_wrappers/generic.py index f3188c928..b89d12b02 100644 --- a/httomo/method_wrappers/generic.py +++ b/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 diff --git a/httomo/method_wrappers/save_intermediate.py b/httomo/method_wrappers/save_intermediate.py index 68442e99f..a538f3898 100644 --- a/httomo/method_wrappers/save_intermediate.py +++ b/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: diff --git a/httomo/method_wrappers/stats_calc.py b/httomo/method_wrappers/stats_calc.py index ee4900b53..496d497fc 100644 --- a/httomo/method_wrappers/stats_calc.py +++ b/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) diff --git a/httomo/methods.py b/httomo/methods.py index ec9571c58..523b1c7b3 100644 --- a/httomo/methods.py +++ b/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) diff --git a/httomo/runner/task_runner.py b/httomo/runner/task_runner.py index 6e06d57ee..812bdca48 100644 --- a/httomo/runner/task_runner.py +++ b/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: diff --git a/httomo/transform_layer.py b/httomo/transform_layer.py index 22ab797c7..1a6c4449f 100644 --- a/httomo/transform_layer.py +++ b/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 diff --git a/tests/method_wrappers/test_save_intermediate.py b/tests/method_wrappers/test_save_intermediate.py index 6725ff76d..21cc80df9 100644 --- a/tests/method_wrappers/test_save_intermediate.py +++ b/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( diff --git a/tests/method_wrappers/test_stats_calc.py b/tests/method_wrappers/test_stats_calc.py index 942274d7d..c1148fa7f 100644 --- a/tests/method_wrappers/test_stats_calc.py +++ b/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 diff --git a/tests/test_methods.py b/tests/test_methods.py index f97c159c8..31ebff5a5 100644 --- a/tests/test_methods.py +++ b/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", diff --git a/tests/test_transform_layer.py b/tests/test_transform_layer.py index c459b522b..52b57a341 100644 --- a/tests/test_transform_layer.py +++ b/tests/test_transform_layer.py @@ -327,6 +327,86 @@ def test_insert_data_checker(mocker: MockerFixture, tmp_path: Path): assert pipeline[9].method_name == "save_to_images" +def test_no_data_checker_after_save_intermediate_data( + mocker: MockerFixture, tmp_path: Path +): + comm = MPI.COMM_SELF + repo = make_mock_repo(mocker) + loader = mocker.create_autospec( + LoaderInterface, + instance=True, + ) + pipeline = Pipeline( + loader=loader, + methods=[ + make_test_method( + mocker, + method_name="remove_outlier", + module_path="httomolibgpu.misc.corr", + save_result=False, + task_id="t1", + gpu=True, + ), + make_test_method( + mocker, + method_name="find_center_vo", + module_path="httomolibgpu.recon.rotation", + save_result=False, + task_id="t2", + gpu=True, + ), + make_test_method( + mocker, + method_name="normalize", + module_path="httomolibgpu.prep.normalize", + save_result=False, + task_id="t3", + gpu=True, + ), + make_test_method( + mocker, + method_name="FBP3d_tomobar", + module_path="httomolibgpu.recon.algorithm", + save_result=True, + task_id="t4", + ), + make_test_method( + mocker, + method_name="calculate_stats", + module_path="httomo.methods", + save_result=False, + task_id="t6", + gpu=False, + ), + make_test_method( + mocker, + method_name="rescale_to_int", + module_path="httomolib.misc.rescale", + save_result=False, + task_id="t7", + gpu=False, + ), + make_test_method( + mocker, + method_name="save_to_images", + module_path="httomolib.misc.images", + save_result=False, + task_id="t8", + gpu=False, + ), + ], + ) + trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path) + pipeline = trans.insert_data_checker(pipeline) + + assert len(pipeline) == 10 + assert pipeline[5].method_name == "FBP3d_tomobar" + assert pipeline[6].method_name == "data_checker" + assert pipeline[6].module_path == "httomolib.misc.utils" + assert pipeline[7].method_name == "calculate_stats" + assert pipeline[8].method_name == "rescale_to_int" + + def test_insert_image_save_after_sweep2(mocker: MockerFixture, tmp_path: Path): comm = MPI.COMM_SELF repo = make_mock_repo(mocker) @@ -426,13 +506,13 @@ def test_insert_paganin_not_last_sweep(mocker: MockerFixture, tmp_path: Path): trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path) pipeline = trans.transform(pipeline) - assert len(pipeline) == 10 - assert pipeline[6].method_name == "save_to_images" - assert pipeline[6].task_id == "saveimage_sweep_t3" - assert pipeline[6].config_params["subfolder_name"] == "images_sweep_paganin_filter" + assert len(pipeline) == 11 + assert pipeline[7].method_name == "save_to_images" + assert pipeline[7].task_id == "saveimage_sweep_t3" + assert pipeline[7].config_params["subfolder_name"] == "images_sweep_paganin_filter" assert pipeline[8].method_name == "FBP3d_tomobar" - assert pipeline[9].task_id == "saveimage_sweep_t4" - assert pipeline[9].config_params["subfolder_name"] == "images_sweep_FBP3d_tomobar" + assert pipeline[10].task_id == "saveimage_sweep_t4" + assert pipeline[10].config_params["subfolder_name"] == "images_sweep_FBP3d_tomobar" def test_insert_paganin_is_last_sweep(mocker: MockerFixture, tmp_path: Path): @@ -472,10 +552,10 @@ def test_insert_paganin_is_last_sweep(mocker: MockerFixture, tmp_path: Path): trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path) pipeline = trans.transform(pipeline) - assert len(pipeline) == 7 - assert pipeline[6].method_name == "save_to_images" - assert pipeline[6].task_id == "saveimage_sweep_t3" - assert pipeline[6].config_params["subfolder_name"] == "images_sweep_paganin_filter" + assert len(pipeline) == 8 + assert pipeline[7].method_name == "save_to_images" + assert pipeline[7].task_id == "saveimage_sweep_t3" + assert pipeline[7].config_params["subfolder_name"] == "images_sweep_paganin_filter" def test_insert_denoise_last_after_FBP_sweep(mocker: MockerFixture, tmp_path: Path): @@ -515,7 +595,7 @@ def test_insert_denoise_last_after_FBP_sweep(mocker: MockerFixture, tmp_path: Pa trans = TransformLayer(comm, repo=repo, save_all=False, out_dir=tmp_path) pipeline = trans.transform(pipeline) - assert len(pipeline) == 7 - assert pipeline[6].method_name == "save_to_images" - assert pipeline[6].task_id == "saveimage_sweep_t3" - assert pipeline[6].config_params["subfolder_name"] == "images_sweep_median_filter" + assert len(pipeline) == 8 + assert pipeline[7].method_name == "save_to_images" + assert pipeline[7].task_id == "saveimage_sweep_t3" + assert pipeline[7].config_params["subfolder_name"] == "images_sweep_median_filter"