Reduce the number of required iterations by introducing a "good enough" threshold

Author: mfep

httomo/method_wrappers/generic.py

@@ -479,6 +479,8 @@ def _calculate_max_slices_iterative(
         non_slice_dims_shape: Tuple[int, int],
         available_memory: int,
     ) -> int:
+        MEM_RATIO_THRESHOLD = 0.9
+
         def get_mem_bytes(current_slices):
             try:
                 memory_bytes = self._query.calculate_memory_bytes_for_slices(
@@ -501,14 +503,20 @@ def get_mem_bytes(current_slices):
         slices_high = None
         memory_bytes = get_mem_bytes(current_slices)
         if memory_bytes > available_memory:
+            # Found upper limit, continue to binary search
             slices_high = current_slices
         else:
             # linear approximation
             current_slices = int(current_slices * available_memory / memory_bytes)
             while True:
                 memory_bytes = get_mem_bytes(current_slices)
                 if memory_bytes > available_memory:
+                    # Found upper limit, continue to binary search
                     break
+                elif memory_bytes >= available_memory * MEM_RATIO_THRESHOLD:
+                    # This is "good enough", return
+                    return current_slices
+            
                 # If linear approximation is not enough, just double every iteration
                 current_slices *= 2
             slices_high = current_slices
@@ -520,10 +528,10 @@ def get_mem_bytes(current_slices):
             memory_bytes = get_mem_bytes(current_slices)
             if memory_bytes > available_memory:
                 slices_high = current_slices
-            elif memory_bytes < available_memory:
-                slices_low = current_slices
-            else:  # memory_bytes == available_memory
+            elif memory_bytes >= available_memory * MEM_RATIO_THRESHOLD:
                 return current_slices
+            else:
+                slices_low = current_slices
 
         return slices_low
 

tests/method_wrappers/test_generic.py

@@ -769,6 +769,7 @@ def test_method(data):
     check_slices = lambda slices: memcalc_mock(
         dims_shape=(slices, shape[0], shape[1]), dtype=dummy_block.data.dtype
     )
+    threshold = 0.9
     if check_slices(1) > available_memory:
         # If zero slice fits
         assert max_slices == 0
@@ -780,8 +781,8 @@ def test_method(data):
             with pytest.raises(Exception):
                 check_slices(max_slices + 1)
         else:
-            # And one more slice must not fit
-            assert check_slices(max_slices + 1) > available_memory
+            # And one more slice must not fit OR above threshold
+            assert check_slices(max_slices + 1) > available_memory or check_slices(max_slices) >= available_memory * threshold
 
 
 @pytest.mark.cupy