From c1bdbbc0dfbb2223431a5fb5886af2a49a8c304c Mon Sep 17 00:00:00 2001 From: Yousef Moazzam Date: Tue, 30 Jun 2026 12:18:00 +0100 Subject: [PATCH 1/2] Add optional data shape param to CPU memory estimator This is useful primarily for running the CPU memory estimation outside of pipeline execution (for example, via the `memory-check` CLI command). --- httomo/cli.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/httomo/cli.py b/httomo/cli.py index d001b27f6..ba32c32c6 100644 --- a/httomo/cli.py +++ b/httomo/cli.py @@ -539,7 +539,12 @@ def execute_sweep_run(pipeline: Pipeline, global_comm: MPI.Comm) -> None: ParamSweepRunner(pipeline, global_comm).execute() -def estimate_cpu_memory(in_data_file: Path, pipeline_file: Path, nprocs: int) -> int: +def estimate_cpu_memory( + in_data_file: Path, + pipeline_file: Path, + nprocs: int, + shape: Optional[tuple[int, int, int]] = None, +) -> int: pipeline = generate_pipeline( in_data_file, pipeline_file, False, MPI.COMM_WORLD, PipelineFormat.Yaml ) @@ -551,14 +556,17 @@ def estimate_cpu_memory(in_data_file: Path, pipeline_file: Path, nprocs: int) -> dtype = dataset.dtype full_shape = dataset.shape - preview_config = parse_preview( - config[0]["parameters"].get("preview", None), full_shape - ) - previewed_shape = ( - preview_config.angles.stop - preview_config.angles.start, - preview_config.detector_y.stop - preview_config.detector_y.start, - preview_config.detector_x.stop - preview_config.detector_x.start, - ) + if shape is None: + preview_config = parse_preview( + config[0]["parameters"].get("preview", None), full_shape + ) + previewed_shape = ( + preview_config.angles.stop - preview_config.angles.start, + preview_config.detector_y.stop - preview_config.detector_y.start, + preview_config.detector_x.stop - preview_config.detector_x.start, + ) + else: + previewed_shape = shape section_memory_peak = 0 for idx in range(len(sections)): From d4febacdc69377912db2d4d0b755b80d2d98ef47 Mon Sep 17 00:00:00 2001 From: Yousef Moazzam Date: Tue, 30 Jun 2026 12:39:01 +0100 Subject: [PATCH 2/2] Move GPU-dependent constructor code to separate methods This allows the creation of pipeline objects (which contain method wrapper objects) without requiring a GPU, which is needed for performing CPU memory estimation of GPU pipelines on a machine that isn't the one to be used for data processing (the main case being a separate launcher application). --- httomo/cli.py | 4 ++++ httomo/method_wrappers/generic.py | 14 +++++++++++--- httomo/runner/method_wrapper.py | 14 ++++++++++++++ tests/method_wrappers/test_generic.py | 2 +- 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/httomo/cli.py b/httomo/cli.py index ba32c32c6..0347de599 100644 --- a/httomo/cli.py +++ b/httomo/cli.py @@ -324,6 +324,10 @@ def run( format_enum, ) + for wrapper in pipeline_object: + wrapper.check() + wrapper.prep() + if not does_contain_sweep: execute_high_throughput_run( pipeline_object, diff --git a/httomo/method_wrappers/generic.py b/httomo/method_wrappers/generic.py index a79ba2ae7..f3188c928 100644 --- a/httomo/method_wrappers/generic.py +++ b/httomo/method_wrappers/generic.py @@ -138,13 +138,21 @@ def __init__( self._query.save_result_default() if save_result is None else save_result ) - if self.is_gpu and not gpu_enabled: - raise ValueError("GPU is not available, please use only CPU methods") - self._side_output: Dict[str, Any] = dict() self._gpu_time_info = GpuTimeInfo() + def check(self) -> None: + """ + Raises + ------ + ValueError + If the method is a GPU method and the machine is not GPU-enabled. + """ + if self.is_gpu and not gpu_enabled: + raise ValueError("GPU is not available, please use only CPU methods") + + def prep(self) -> None: if gpu_enabled: _id = httomo.globals.gpu_id self._gpu_id = get_gpu_id() if _id == -1 else _id diff --git a/httomo/runner/method_wrapper.py b/httomo/runner/method_wrapper.py index c5b7e499a..763299f08 100644 --- a/httomo/runner/method_wrapper.py +++ b/httomo/runner/method_wrapper.py @@ -205,3 +205,17 @@ def calculate_max_slices( something persists afterwards. """ ... # pragma: nocover + + def check(self) -> None: + """ + Perform any checks necessary to ensure that the method is capable of executing on the + machine (for example, if it's a GPU method, check that execution is on a machine with a + GPU that is accessible). + """ + ... # pragma: nocover + + def prep(self) -> None: + """ + Perform any preparation needed prior to beginning pipeline execution. + """ + ... # pragma: nocover diff --git a/tests/method_wrappers/test_generic.py b/tests/method_wrappers/test_generic.py index 68c9dffe2..4b38d1265 100644 --- a/tests/method_wrappers/test_generic.py +++ b/tests/method_wrappers/test_generic.py @@ -213,7 +213,7 @@ def fake_method(data): "fake_method", MPI.COMM_WORLD, make_mock_preview_config(mocker), - ) + ).check() assert "GPU is not available" in str(e)