diff --git a/httomo/cli.py b/httomo/cli.py index 4775c00a4..9cca815a4 100644 --- a/httomo/cli.py +++ b/httomo/cli.py @@ -562,17 +562,18 @@ def estimate_cpu_memory(in_data_file: Path, pipeline_file: Path, nprocs: int) -> section_memory_peak = 0 for idx in range(len(sections)): + mem, previewed_shape = estimate_section_memory( + nprocs, + 0, + None, + dtype, + previewed_shape, + sections, + idx, + ) section_memory_peak = max( section_memory_peak, - estimate_section_memory( - nprocs, - 0, - None, - dtype, - previewed_shape, - sections, - idx, - ), + mem, ) return section_memory_peak diff --git a/httomo/runner/dataset_store_backing.py b/httomo/runner/dataset_store_backing.py index 200527e96..e8ed2bdf6 100644 --- a/httomo/runner/dataset_store_backing.py +++ b/httomo/runner/dataset_store_backing.py @@ -52,7 +52,7 @@ def calculate_section_output_chunk_shape( for method in section.methods: if method.memory_gpu is None: continue - output_non_slice_dims = method.calculate_output_dims(input_non_slice_dims) + output_non_slice_dims = method.calculate_output_dims(output_non_slice_dims) output_chunk_shape = list(output_non_slice_dims) output_chunk_shape.insert(slicing_dim, chunk_shape[slicing_dim]) @@ -74,7 +74,7 @@ def estimate_section_memory( sections: List[Section], section_idx: int, consider_pinned_memory_pool: bool = False, -) -> int: +) -> tuple[int, tuple[int, int, int]]: # Get chunk shape created by reader of section `n` (the current section) that will account # for padding. This chunk shape is based on the chunk shape written by the writer of # section `n - 1` (the previous section) @@ -150,12 +150,26 @@ def estimate_section_memory( # See https://github.com/cupy/cupy/issues/9813 cupy_transfer_overhead = total_mem + # Calculate the shape of the global data of the output of the section to pass back to the + # caller. + # + # NOTE: The caller will only use this if performing CPU memory estimation of the entire + # pipeline not at runtime of htttomo (ie, the `memory-check` CLI command). If performing + # CPU memory estimation of a section during runtime of httomo, the caller won't use this + # value. + current_section_slicing_dim = _get_slicing_dim(sections[section_idx].pattern) - 1 + output_global_shape = list(output_chunk_shape) + output_global_shape[current_section_slicing_dim] = ( + output_global_shape[current_section_slicing_dim] * nprocs + ) + return ( padded_input_chunk_bytes + output_chunk_bytes + reslice_bytes + cupy_pinned_cpu_pool_memory - + cupy_transfer_overhead + + cupy_transfer_overhead, + output_global_shape, ) @@ -167,7 +181,7 @@ def determine_store_backing( global_shape: Tuple[int, int, int], section_idx: int, ) -> DataSetStoreBacking: - section_memory = estimate_section_memory( + section_memory, _ = estimate_section_memory( comm.size, comm.rank, comm.allgather,