diff --git a/httomo/data/param_sweep_store.py b/httomo/data/param_sweep_store.py index 8f0e2f458..a5b645d79 100644 --- a/httomo/data/param_sweep_store.py +++ b/httomo/data/param_sweep_store.py @@ -1,8 +1,13 @@ from __future__ import annotations +from os import PathLike +from pathlib import Path from typing import Literal, Optional, Tuple +from mpi4py import MPI import numpy as np +import h5py +from httomo.runner.dataset_store_backing import DataSetStoreBacking from httomo.runner.auxiliary_data import AuxiliaryData from httomo.sweep_runner.param_sweep_block import ParamSweepBlock from httomo.sweep_runner.param_sweep_store_interfaces import ParamSweepSource @@ -23,6 +28,17 @@ def __init__(self, source: ParamSweepWriter) -> None: ), "Reader should be created from writer with data not `None`" self._data = source._data self._aux_data = source.aux_data + self._h5filename: Optional[Path] = None + self.h5file: Optional[h5py.File] = None + + if source.is_file_based: + self._h5filename = source.filename + self.h5file = h5py.File(source.filename, "r") + self._data = self.h5file["/data"] + else: + self._data = source._data + + source.finalize() @property def no_of_sweeps(self) -> int: @@ -56,18 +72,52 @@ def read_sweep_results(self) -> ParamSweepBlock: slicing_dim=self.extract_dim, ) + def finalize(self): + self._data = None + if self.h5file is not None: + self.h5file.close() + self.h5file = None + if self._h5filename is not None: + self._h5filename.unlink() + self._h5filename = None + + @property + def is_file_based(self) -> bool: + return self._h5filename is not None + + @property + def filename(self) -> Optional[Path]: + return self._h5filename + class ParamSweepWriter: """Write parameter sweep results, concatenating them along the `detector_y` dim""" - def __init__(self, no_of_sweeps: int) -> None: + def __init__( + self, + no_of_sweeps: int, + comm: MPI.Comm, + temppath: PathLike, + backing: DataSetStoreBacking, + ) -> None: self._concat_dim: Literal[1] = 1 self._no_of_sweeps = no_of_sweeps self._no_of_sweeps_written: int = 0 self._single_shape: Optional[Tuple[int, int, int]] = None self._total_shape: Optional[Tuple[int, int, int]] = None - self._data: Optional[np.ndarray] = None + self._data: Optional[np.ndarray | h5py.Dataset] = None self._slices_per_sweep: Optional[int] = None + self._comm = comm + self._backing = backing + self._temppath = temppath + self._h5filename: Optional[Path] = None + self.h5file: Optional[h5py.File] = None + + if backing == DataSetStoreBacking.File: + self._h5filename = ( + Path(self._temppath) / f"sweep-output-rank-{self._comm.rank}.h5" + ) + self.h5file = h5py.File(self._h5filename, "w") @property def no_of_sweeps(self) -> int: @@ -128,8 +178,14 @@ def write_sweep_result(self, block: ParamSweepBlock): self.single_shape[2], ) self._slices_per_sweep = self.single_shape[self.concat_dim] - self._data = np.empty(shape=self.total_shape, dtype=block.data.dtype) self._aux_data = block.aux_data + if self._backing == DataSetStoreBacking.RAM: + self._data = np.empty(shape=self.total_shape, dtype=block.data.dtype) + elif self._backing == DataSetStoreBacking.File: + assert self.h5file is not None + self._data = self.h5file.create_dataset( + "/data", shape=self.total_shape, dtype=block.data.dtype + ) slices = [slice(None, None, 1)] * 3 sweep_res_start = self.no_of_sweeps_written * self.slices_per_sweep @@ -138,5 +194,22 @@ def write_sweep_result(self, block: ParamSweepBlock): sweep_res_start + self.slices_per_sweep, 1, ) + assert ( + self._data is not None + ), "Numpy array or hdf5 file should be created prior to writing sweep result" self._data[slices[0], slices[1], slices[2]] = block.data self.increment_no_of_sweeps_written() + + def finalize(self): + self._data = None + if self.h5file is not None: + self.h5file.close() + self.h5file = None + + @property + def is_file_based(self) -> bool: + return self._backing is DataSetStoreBacking.File + + @property + def filename(self) -> Optional[Path]: + return self._h5filename diff --git a/httomo/sweep_runner/param_sweep_runner.py b/httomo/sweep_runner/param_sweep_runner.py index df24b6bf4..fc20ffa7d 100644 --- a/httomo/sweep_runner/param_sweep_runner.py +++ b/httomo/sweep_runner/param_sweep_runner.py @@ -7,10 +7,12 @@ import numpy as np import httomo -from httomo.data.param_sweep_store import ParamSweepReader, ParamSweepWriter +from httomo import globals +from httomo.data.param_sweep_store import ParamSweepWriter from httomo.method_wrappers.images import ImagesWrapper from httomo.method_wrappers.save_intermediate import SaveIntermediateFilesWrapper from httomo.runner.block_split import BlockSplitter +from httomo.runner.dataset_store_backing import DataSetStoreBacking from httomo.runner.method_wrapper import MethodWrapper from httomo.runner.pipeline import Pipeline from httomo.sweep_runner.param_sweep_block import ParamSweepBlock @@ -21,6 +23,10 @@ from httomo.preview import PreviewConfig, PreviewDimConfig from httomo.runner.dataset_store_interfaces import DataSetSource +PAGANIN_SWEEP_VALUE_THRESHOLD = 50 +"""The number of sweep values for the paganin filter above which the sweep results should be +written to an hdf5 file as opposed to kept in RAM.""" + class ParamSweepRunner: def __init__( @@ -42,6 +48,7 @@ def __init__( self._start_sweep_idx : self._stop_sweep_idx ] self._set_image_saver_params() + self._comm = comm @property def block(self) -> ParamSweepBlock: @@ -212,8 +219,24 @@ def execute_after_sweep(self): def execute_sweep(self): """Execute all param variations of the same method in the sweep""" - writer = ParamSweepWriter(len(self._sweep_values)) method = self._stages.sweep.method + if ( + method.method_name == "paganin_filter" + and len(self._stages.sweep.values) > PAGANIN_SWEEP_VALUE_THRESHOLD + ): + writer = ParamSweepWriter( + len(self._sweep_values), + self._comm, + globals.run_out_dir, + DataSetStoreBacking.File, + ) + else: + writer = ParamSweepWriter( + len(self._sweep_values), + self._comm, + globals.run_out_dir, + DataSetStoreBacking.RAM, + ) log_once(f"Running {method.method_name} ({method.package_name})") log_once(" Beginning parameter sweep") @@ -254,8 +277,10 @@ def execute_sweep(self): log_once(" Finished parameter sweep") - reader = ParamSweepReader(writer) + reader = writer.make_reader() self._block = reader.read_sweep_results() + if reader.is_file_based: + reader.finalize() def execute(self): """Load input data and execute all stages (before sweep, sweep, after sweep)""" diff --git a/tests/data/test_param_sweep_store.py b/tests/data/test_param_sweep_store.py index 5505fde5b..cede9bad7 100644 --- a/tests/data/test_param_sweep_store.py +++ b/tests/data/test_param_sweep_store.py @@ -1,29 +1,35 @@ +from os import PathLike + import pytest import numpy as np +from mpi4py.MPI import COMM_WORLD from pytest_mock import MockerFixture from httomo.data.param_sweep_store import ParamSweepWriter from httomo.runner.auxiliary_data import AuxiliaryData +from httomo.runner.dataset_store_backing import DataSetStoreBacking from httomo.sweep_runner.param_sweep_block import ParamSweepBlock -def make_param_sweep_writer() -> ParamSweepWriter: +def make_param_sweep_writer(tmp_path: PathLike) -> ParamSweepWriter: NO_OF_SWEEPS = 5 - return ParamSweepWriter(NO_OF_SWEEPS) + return ParamSweepWriter(NO_OF_SWEEPS, COMM_WORLD, tmp_path, DataSetStoreBacking.RAM) -def test_param_sweep_writer_get_no_of_sweeps(): - writer = make_param_sweep_writer() +def test_param_sweep_writer_get_no_of_sweeps(tmp_path: PathLike): + writer = make_param_sweep_writer(tmp_path) assert writer.no_of_sweeps == 5 -def test_param_sweep_writer_get_concat_dim(): - writer = make_param_sweep_writer() +def test_param_sweep_writer_get_concat_dim(tmp_path: PathLike): + writer = make_param_sweep_writer(tmp_path) assert writer.concat_dim == 1 -def test_param_sweep_writer_get_single_shape_before_write_raises_error(): - writer = make_param_sweep_writer() +def test_param_sweep_writer_get_single_shape_before_write_raises_error( + tmp_path: PathLike, +): + writer = make_param_sweep_writer(tmp_path) with pytest.raises(ValueError) as e: writer.single_shape assert ( @@ -31,15 +37,19 @@ def test_param_sweep_writer_get_single_shape_before_write_raises_error(): ) in str(e) -def test_param_sweep_writer_get_slices_per_sweep_before_write_raises_error(): - writer = make_param_sweep_writer() +def test_param_sweep_writer_get_slices_per_sweep_before_write_raises_error( + tmp_path: PathLike, +): + writer = make_param_sweep_writer(tmp_path) with pytest.raises(ValueError) as e: writer.slices_per_sweep assert ("Slices per sweep isn't known until the first write has occurred") in str(e) -def test_param_sweep_writer_get_total_shape_before_write_raises_error(): - writer = make_param_sweep_writer() +def test_param_sweep_writer_get_total_shape_before_write_raises_error( + tmp_path: PathLike, +): + writer = make_param_sweep_writer(tmp_path) with pytest.raises(ValueError) as e: writer.total_shape assert ( @@ -48,8 +58,8 @@ def test_param_sweep_writer_get_total_shape_before_write_raises_error(): ) in str(e) -def test_param_sweep_write_make_reader_errors_if_data_none(): - writer = make_param_sweep_writer() +def test_param_sweep_write_make_reader_errors_if_data_none(tmp_path: PathLike): + writer = make_param_sweep_writer(tmp_path) with pytest.raises(ValueError) as e: writer.make_reader() @@ -57,10 +67,10 @@ def test_param_sweep_write_make_reader_errors_if_data_none(): def test_param_sweep_writer_write_sweep_result_transfers_gpu_arr_to_cpu( - mocker: MockerFixture, + mocker: MockerFixture, tmp_path: PathLike ): SWEEP_RES_SHAPE = (180, 3, 160) - writer = make_param_sweep_writer() + writer = make_param_sweep_writer(tmp_path) block = ParamSweepBlock( data=np.ones(SWEEP_RES_SHAPE, dtype=np.float32), aux_data=AuxiliaryData(angles=np.ones(SWEEP_RES_SHAPE[0], dtype=np.float32)), @@ -77,7 +87,7 @@ def test_param_sweep_writer_write_sweep_result_transfers_gpu_arr_to_cpu( to_cpu_spy.assert_called_once() -def test_param_sweep_writer_reader_write_res_and_read(): +def test_param_sweep_writer_reader_write_res_and_read(tmp_path: PathLike): CONCAT_DIM = 1 SWEEP_RES_SHAPE = (180, 3, 160) NO_OF_SWEEPS = 2 @@ -86,7 +96,9 @@ def test_param_sweep_writer_reader_write_res_and_read(): SWEEP_RES_SHAPE[1] * NO_OF_SWEEPS, SWEEP_RES_SHAPE[2], ) - writer = ParamSweepWriter(NO_OF_SWEEPS) + writer = ParamSweepWriter( + NO_OF_SWEEPS, COMM_WORLD, tmp_path, DataSetStoreBacking.RAM + ) # Define an array that will contain data representing the fake result of two parameter # sweep executions. @@ -131,3 +143,58 @@ def test_param_sweep_writer_reader_write_res_and_read(): np.testing.assert_array_equal( block_with_middle_slices.data[:, 1, :], SWEEP_RESULT_TWO_MID_SLICE ) + + +def test_writer_closes_file_on_finalize(mocker: MockerFixture, tmp_path: PathLike): + CONCAT_DIM = 1 + BLOCK_SHAPE = (180, 3, 160) + block_data = np.arange(np.prod(BLOCK_SHAPE), dtype=np.float32).reshape(BLOCK_SHAPE) + aux_data = AuxiliaryData(np.ones(BLOCK_SHAPE[0], dtype=np.float32)) + writer = ParamSweepWriter( + no_of_sweeps=1, + comm=COMM_WORLD, + temppath=tmp_path, + backing=DataSetStoreBacking.File, + ) + writer.write_sweep_result( + ParamSweepBlock( + data=block_data, + aux_data=aux_data, + slicing_dim=CONCAT_DIM, + ) + ) + mock_fileclose = mocker.patch.object(writer.h5file, "close") + writer.finalize() + mock_fileclose.assert_called_once() + + +def test_making_reader_closes_file_and_deletes(tmp_path: PathLike): + CONCAT_DIM = 1 + BLOCK_SHAPE = (180, 3, 160) + block_data = np.arange(np.prod(BLOCK_SHAPE), dtype=np.float32).reshape(BLOCK_SHAPE) + aux_data = AuxiliaryData(np.ones(BLOCK_SHAPE[0], dtype=np.float32)) + writer = ParamSweepWriter( + no_of_sweeps=1, + comm=COMM_WORLD, + temppath=tmp_path, + backing=DataSetStoreBacking.File, + ) + writer.write_sweep_result( + ParamSweepBlock( + data=block_data, + aux_data=aux_data, + slicing_dim=CONCAT_DIM, + ) + ) + reader = writer.make_reader() + + assert writer.h5file is None + assert writer.filename is not None + assert writer.filename.exists() + assert reader.filename == writer.filename + assert reader.h5file is not None + assert reader.h5file.get("data", None) is not None + + reader.finalize() + + assert not writer.filename.exists()