More ergonomic mem peak estimation in remove_stripe_fw

Author: mfep

httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/prep/stripe.py

@@ -25,7 +25,7 @@
 import numpy as np
 
 from httomo_backends.cufft import CufftType, cufft_estimate_1d
-from httomolibgpu.prep.stripe import DeviceMemStack, remove_stripe_fw
+from httomolibgpu.prep.stripe import remove_stripe_fw
 
 
 __all__ = [
@@ -60,9 +60,7 @@ def _calc_memory_bytes_for_slices_remove_stripe_fw(
     dtype: np.dtype,
     **kwargs,
 ) -> Tuple[int, int]:
-    mem_stack = DeviceMemStack()
-    remove_stripe_fw(dims_shape, mem_stack=mem_stack, **kwargs)
-    return mem_stack.highwater, 0
+    return remove_stripe_fw(dims_shape, calc_peak_gpu_mem=True, **kwargs), 0
 
 
 def _calc_memory_bytes_remove_all_stripe(

tests/test_httomolibgpu.py

@@ -391,7 +391,7 @@ def test_remove_stripe_fw_memoryhook(slices, ensure_clean_memory):
     )  # the amount of memory in bytes needed for the method according to memoryhook
 
     # now we estimate how much of the total memory required for this data
-    (estimated_memory_bytes, subtract_bytes) = _calc_memory_bytes_remove_stripe_fw(
+    (estimated_memory_bytes, subtract_bytes) = _calc_memory_bytes_for_slices_remove_stripe_fw(
         (slices, dim_x, dim_y), dtype=np.float32()
     )
     estimated_memory_mb = round(estimated_memory_bytes / (1024**2), 2)