Fix progress_utils_test.py to restrict mocking of system-wide modules to the module under test.

This test was failing because sometimes other modules were calling `time.time` or `tracemalloc.get_traced_memory`, stealing the mocked return value.

Also add a warning that running with `tracemalloc` can significantly slow down the code.

PiperOrigin-RevId: 895884076

Author: gonnet

ai_edge_quantizer/utils/progress_utils.py

@@ -15,6 +15,7 @@
 
 """Utility functions to display a progress bar and progress report."""
 
+import logging
 import time
 import tracemalloc
 import tqdm
@@ -55,42 +56,68 @@ def close(self):
 
 
 class ProgressReport:
-  """A class to generate a progress report for the quantization process."""
+  """A class to generate a progress report for the quantization process.
+  
+  If initialized with `trace_memory=True`, it will also track the peak memory
+  use using `tracemalloc`, which may hurt performance and interfere with other
+  code using `tracemalloc` concurrently.
+  """
+  _description: str
+  _trace_memory: bool
+  _start_time: float | None = None
+  _tracemalloc_started_by_me: bool = False
 
-  def __init__(self, description: str = ''):
+  def __init__(self, description: str = '', trace_memory: bool = False):
     self._description = description
-    self._start_time = None
+    self._trace_memory = trace_memory
 
   def capture_progess_start(self):
     self._start_time = time.time()
-    tracemalloc.start()
+    if self._trace_memory:
+      logging.warning(
+          'Progress bar reporting with `trace_memory=True` is enabled which may'
+          ' significantly slow down your computations!'
+      )
+      self._tracemalloc_started_by_me = not tracemalloc.is_tracing()
+      if self._tracemalloc_started_by_me:
+        tracemalloc.start()
+      else:
+        tracemalloc.reset_peak()
+
+  def _capture_progress_end(self) -> int | None:
+    if self._trace_memory:
+      _, mem_peak_bytes = tracemalloc.get_traced_memory()
+      if self._tracemalloc_started_by_me:
+        tracemalloc.stop()
+        self._tracemalloc_started_by_me = False
+      return mem_peak_bytes
 
   def render_report(
       self,
       original_size: int,
       quantized_size: int,
       quantization_ratio: float,
-      memory_peak: float,
       total_time: float,
+      memory_peak: float | None,
   ):
     """Prints out the progress report."""
     print(f'Original model size: {original_size/1024:.2f} KB')
     print(f'Quantized model size: {quantized_size/1024:.2f} KB')
     print(f'Quantization Ratio: {quantization_ratio:.2f}')
     print(f'Total time: {total_time:.2f} seconds')
-    print(f'Memory peak: {memory_peak:.2f} MB')
+    if memory_peak is not None:
+      print(f'Memory peak: {memory_peak/1024/1024:.2f} MB')
 
   def generate_progress_report(self, original_model, quantized_model):
     original_size = len(original_model)
     quantized_size = len(quantized_model)
     quantization_ratio = quantized_size / original_size
     total_time = time.time() - self._start_time
-    _, mem_peak_bytes = tracemalloc.get_traced_memory()
-    mem_peak_mb = mem_peak_bytes / 1024 / 1024
+    mem_peak_bytes = self._capture_progress_end()
     self.render_report(
         original_size,
         quantized_size,
         quantization_ratio,
-        mem_peak_mb,
         total_time,
+        mem_peak_bytes,
     )

ai_edge_quantizer/utils/progress_utils_test.py

@@ -15,74 +15,84 @@
 
 import io
 import sys
-import time
-import tracemalloc
 from unittest import mock
 
-import absl.testing.absltest as absltest
+from absl.testing import absltest
+from absl.testing import parameterized
+import tqdm
+
 from ai_edge_quantizer.utils import progress_utils
 
 
 class ProgressBarTest(absltest.TestCase):
 
-  @mock.patch('tqdm.tqdm')
-  def test_progress_bar_update(self, mock_tqdm):
-    mock_progress_bar_instance = mock_tqdm.return_value
+  def setUp(self):
+    super().setUp()
+    self.mock_tqdm = self.enter_context(
+        mock.patch.object(tqdm, 'tqdm', autospec=True, spec_set=True)
+    )
+
+  def test_progress_bar_update(self):
+    mock_progress_bar_instance = self.mock_tqdm.return_value
     with progress_utils.ProgressBar(total_steps=10) as pb:
       pb.update_single_step()
       pb.update_single_step()
 
-    mock_tqdm.assert_called_once_with(
+    self.mock_tqdm.assert_called_once_with(
         total=10, desc='', leave=True, disable=False
     )
     self.assertEqual(mock_progress_bar_instance.update.call_count, 2)
     mock_progress_bar_instance.update.assert_called_with(1)
     mock_progress_bar_instance.close.assert_called_once()
 
-  @mock.patch('tqdm.tqdm')
-  def test_progress_bar_disable(self, mock_tqdm):
-    mock_progress_bar_instance = mock_tqdm.return_value
+  def test_progress_bar_disable(self):
+    mock_progress_bar_instance = self.mock_tqdm.return_value
     with progress_utils.ProgressBar(total_steps=10, disable=True):
       pass
-    mock_tqdm.assert_called_once_with(
+    self.mock_tqdm.assert_called_once_with(
         total=10, desc='', leave=True, disable=True
     )
     mock_progress_bar_instance.close.assert_called_once()
 
-  @mock.patch('tqdm.tqdm')
-  def test_progress_bar_disappear_on_finish(self, mock_tqdm):
-    mock_progress_bar_instance = mock_tqdm.return_value
+  def test_progress_bar_disappear_on_finish(self):
+    mock_progress_bar_instance = self.mock_tqdm.return_value
     with progress_utils.ProgressBar(total_steps=10, disappear_on_finish=True):
       pass
-    mock_tqdm.assert_called_once_with(
+    self.mock_tqdm.assert_called_once_with(
         total=10, desc='', leave=False, disable=False
     )
     mock_progress_bar_instance.close.assert_called_once()
 
 
-class ProgressReportTest(absltest.TestCase):
+class ProgressReportTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
     self.mock_time = self.enter_context(
-        mock.patch.object(time, 'time', autospec=True)
+        mock.patch.object(progress_utils, 'time', autospec=True, spec_set=True)
     )
-    self.mock_tracemalloc_start = self.enter_context(
-        mock.patch.object(tracemalloc, 'start', autospec=True)
-    )
-    self.mock_tracemalloc_get_traced_memory = self.enter_context(
-        mock.patch.object(tracemalloc, 'get_traced_memory', autospec=True)
-    )
-
-  def test_generate_progress_report(self):
-    self.mock_time.side_effect = [100.0, 105.5]  # Start time, end time.
-    # Mock memory: current=1MB, peak=2MB.
-    self.mock_tracemalloc_get_traced_memory.return_value = (
-        1 * 1024 * 1024,
-        2 * 1024 * 1024,
+    self.mock_tracemalloc = self.enter_context(
+        mock.patch.object(
+            progress_utils, 'tracemalloc', autospec=True, spec_set=True
+        )
     )
 
-    progress_report = progress_utils.ProgressReport()
+  @parameterized.named_parameters(
+      ('trace_memory_enabled', True), ('trace_memory_disabled', False)
+  )
+  def test_generate_progress_report(self, trace_memory: bool):
+    self.mock_time.time.side_effect = [100.0, 105.5]  # Start time, end time.
+
+    if trace_memory:
+      self.mock_tracemalloc.is_tracing.return_value = False
+      self.mock_tracemalloc.start.side_effect = None
+      self.mock_tracemalloc.stop.return_value = None
+      self.mock_tracemalloc.get_traced_memory.return_value = (
+          1 * 1024 * 1024,
+          2 * 1024 * 1024,
+      )
+
+    progress_report = progress_utils.ProgressReport(trace_memory=trace_memory)
     progress_report.capture_progess_start()
 
     original_model = b'\x01' * 2048  # 2KB.
@@ -92,13 +102,17 @@ def test_generate_progress_report(self):
     with mock.patch.object(sys, 'stdout', mock_stdout):
       progress_report.generate_progress_report(original_model, quantized_model)
 
-    self.mock_tracemalloc_start.assert_called_once()
     output = mock_stdout.getvalue()
     self.assertIn('Original model size: 2.00 KB', output)
     self.assertIn('Quantized model size: 1.00 KB', output)
     self.assertIn('Quantization Ratio: 0.50', output)
     self.assertIn('Total time: 5.50 seconds', output)
-    self.assertIn('Memory peak: 2.00 MB', output)
+
+    if trace_memory:
+      self.mock_tracemalloc.is_tracing.assert_called_once_with()
+      self.mock_tracemalloc.start.assert_called_once_with()
+      self.mock_tracemalloc.stop.assert_called_once_with()
+      self.assertIn('Memory peak: 2.00 MB', output)
 
 
 if __name__ == '__main__':