From d4cae75c876930eec991f275d2238c79d7d54177 Mon Sep 17 00:00:00 2001 From: algol Date: Wed, 1 Jul 2026 12:19:45 +0100 Subject: [PATCH 1/2] makes the selection of slices unified across all ranks --- httomo/sweep_runner/param_sweep_runner.py | 25 +++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/httomo/sweep_runner/param_sweep_runner.py b/httomo/sweep_runner/param_sweep_runner.py index 3d1d61f93..2a58a6660 100644 --- a/httomo/sweep_runner/param_sweep_runner.py +++ b/httomo/sweep_runner/param_sweep_runner.py @@ -2,6 +2,8 @@ from typing import Any, Dict, List, Optional, Tuple import tqdm + +from mpi4py import MPI from mpi4py.MPI import Comm import numpy as np @@ -16,7 +18,7 @@ from httomo.sweep_runner.param_sweep_block import ParamSweepBlock from httomo.sweep_runner.side_output_manager import SideOutputManager from httomo.sweep_runner.stages import NonSweepStage, Stages, SweepStage -from httomo.utils import catchtime, log_exception, log_once, search_max_slices_iterative +from httomo.utils import catchtime, log_exception, log_rank, log_once, search_max_slices_iterative from httomo.runner.gpu_utils import get_available_gpu_memory, gpumem_cleanup from httomo.preview import PreviewConfig, PreviewDimConfig from httomo.runner.dataset_store_interfaces import DataSetSource @@ -319,8 +321,16 @@ def _slices_to_fit_memory_Paganin(source: DataSetSource) -> int: from httomo_backends.methods_database.packages.backends.httomolibgpu.supporting_funcs.prep.phase import ( _calc_memory_bytes_for_slices_paganin_filter, ) + gpumem_cleanup() available_memory = get_available_gpu_memory(10.0) + available_memory_in_GB = round(available_memory / (1024**3), 2) + memory_str = ( + f"The amount of the available GPU memory is {available_memory_in_GB} GB" + ) + comm = MPI.COMM_WORLD + log_rank(memory_str, comm=comm) + angles_total = source.aux_data.angles_length det_X_length = source.chunk_shape[2] @@ -339,6 +349,13 @@ def get_mem_bytes(slices): return 2**64 finally: gpumem_cleanup() - - gpumem_cleanup() - return search_max_slices_iterative(available_memory, get_mem_bytes) + + slices_max = 3 + if comm.rank == 0: + # assuming here that the GPU devices have the same amount of memory + slices_max = search_max_slices_iterative(available_memory, get_mem_bytes) + comm.Barrier() + # Broadcast + if comm.size > 1: + slices_max = comm.bcast(slices_max, root=0) + return slices_max \ No newline at end of file From 9ca7fed888e3392f0165e835c0fca5763620de18 Mon Sep 17 00:00:00 2001 From: algol Date: Wed, 1 Jul 2026 12:21:30 +0100 Subject: [PATCH 2/2] linting --- httomo/sweep_runner/param_sweep_runner.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/httomo/sweep_runner/param_sweep_runner.py b/httomo/sweep_runner/param_sweep_runner.py index 2a58a6660..6f39c8310 100644 --- a/httomo/sweep_runner/param_sweep_runner.py +++ b/httomo/sweep_runner/param_sweep_runner.py @@ -18,7 +18,13 @@ from httomo.sweep_runner.param_sweep_block import ParamSweepBlock from httomo.sweep_runner.side_output_manager import SideOutputManager from httomo.sweep_runner.stages import NonSweepStage, Stages, SweepStage -from httomo.utils import catchtime, log_exception, log_rank, log_once, search_max_slices_iterative +from httomo.utils import ( + catchtime, + log_exception, + log_rank, + log_once, + search_max_slices_iterative, +) from httomo.runner.gpu_utils import get_available_gpu_memory, gpumem_cleanup from httomo.preview import PreviewConfig, PreviewDimConfig from httomo.runner.dataset_store_interfaces import DataSetSource @@ -321,16 +327,17 @@ def _slices_to_fit_memory_Paganin(source: DataSetSource) -> int: from httomo_backends.methods_database.packages.backends.httomolibgpu.supporting_funcs.prep.phase import ( _calc_memory_bytes_for_slices_paganin_filter, ) + gpumem_cleanup() available_memory = get_available_gpu_memory(10.0) available_memory_in_GB = round(available_memory / (1024**3), 2) memory_str = ( - f"The amount of the available GPU memory is {available_memory_in_GB} GB" - ) + f"The amount of the available GPU memory is {available_memory_in_GB} GB" + ) comm = MPI.COMM_WORLD log_rank(memory_str, comm=comm) - + angles_total = source.aux_data.angles_length det_X_length = source.chunk_shape[2] @@ -349,7 +356,7 @@ def get_mem_bytes(slices): return 2**64 finally: gpumem_cleanup() - + slices_max = 3 if comm.rank == 0: # assuming here that the GPU devices have the same amount of memory @@ -358,4 +365,4 @@ def get_mem_bytes(slices): # Broadcast if comm.size > 1: slices_max = comm.bcast(slices_max, root=0) - return slices_max \ No newline at end of file + return slices_max