Skip to content

makes the selection of slices unified across all ranks for paganin sweep estimator #735

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dkazanc wants to merge 2 commits into
base: main
Choose a base branch
from
Open

makes the selection of slices unified across all ranks for paganin sweep estimator #735

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d4cae75
makes the selection of slices unified across all ranks
dkazanc Jul 1, 2026
9ca7fed
linting
dkazanc Jul 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 27 additions & 3 deletions httomo/sweep_runner/param_sweep_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -16,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_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
Expand Down Expand Up @@ -320,7 +328,16 @@ def _slices_to_fit_memory_Paganin(source: DataSetSource) -> int:
_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]

Expand All @@ -340,5 +357,12 @@ def get_mem_bytes(slices):
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
Loading