Tweaks

tests/tools/test_memory_monitor.py

@@ -11,6 +11,7 @@
 
 import gc
 import os
+import queue
 import sys
 import time
 from pathlib import Path
@@ -76,7 +77,7 @@ def test_memory_monitor_api(tmpdir):
 
 
 @pytest.mark.skipif(is_windows(), reason="Check on linux only")
-@pytest.mark.parametrize("memory_type", MemoryType.__members__.values())
+@pytest.mark.parametrize("memory_type", (MemoryType.RSS, MemoryType.SYSTEM))
 def test_memory_type(memory_type):
     memory_monitor = MemoryMonitor(memory_type=memory_type).start()
     allocate(BYTES_TO_ALLOCATE_SMALL)
@@ -120,6 +121,16 @@ def test_monitor_memory_for_callable(max_value):
         assert len(time_values) == len(memory_values)
 
 
+@pytest.mark.skipif(is_windows(), reason="Check on linux only")
+def test_empty_logs(tmpdir):
+    memory_monitor = MemoryMonitor().start()
+    memory_monitor.stop()
+    memory_monitor._memory_values_queue = queue.Queue()  # make sure no logs are recorded
+    time_values, memory_values = memory_monitor.get_data()
+    assert len(time_values) == len(memory_values) == 0
+    memory_monitor.save_memory_logs(time_values, memory_values, Path(tmpdir))
+
+
 @pytest.mark.skipif(ISOLATION_RUN_ENV_VAR not in os.environ, reason="Should be run via isolation proxy")
 def test_memory_values_isolated():
     baseline_memory = float(os.environ[BASELINE_MEMORY_VAR]) if BASELINE_MEMORY_VAR in os.environ else None

tools/memory_monitor.py

@@ -10,6 +10,7 @@
 # limitations under the License.
 
 import atexit
+import logging
 import os
 import queue
 import subprocess
@@ -23,6 +24,8 @@
 import matplotlib.pyplot as plt
 import psutil
 
+logger = logging.getLogger("memory_monitor")
+
 
 class MemoryType(Enum):
     RSS = "rss"
@@ -80,6 +83,10 @@ def __init__(
         """
         self.interval = interval
         self.memory_type = memory_type
+        if memory_type == MemoryType.SYSTEM:
+            logger.warning(
+                "Please be aware that MemoryType.SYSTEM is affected by other processes that change RAM availability."
+            )
         self.memory_unit = memory_unit
 
         self._monitoring_thread_should_stop = False
@@ -137,6 +144,8 @@ def get_data(self, memory_from_zero: Optional[bool] = False) -> Tuple[List, List
         to memory values.
         """
         memory_usage_data = list(self._memory_values_queue.queue)
+        if len(memory_usage_data) == 0:
+            return [], []
         time_values, memory_values = tuple(zip(*memory_usage_data))
         time_values = _subtract_first_element(list(time_values))
         if memory_from_zero:
@@ -171,6 +180,9 @@ def save_memory_logs(
         # Save measurements to text file
         log_filepath = save_dir / f"{filename_label}.txt"
         with open(log_filepath, "w") as log_file:
+            if len(time_values) == 0:
+                log_file.write("No measurements recorded.\nPlease make sure logging duration or interval were enough.")
+                return
             for timestamp, memory_usage in zip(time_values, memory_values):
                 log_file.write(f"{timestamp} {memory_usage:.3f}\n")
 
@@ -222,11 +234,19 @@ def _monitor_memory(self):
                 header_line = -3
                 res_column = 5  # Resident Memory Size (KiB): The non-swapped physical memory a task is using.
 
-                res = subprocess.run(
-                    f"top -n 1 -p {os.getpid()}".split(" "),
-                    capture_output=True,
-                    text=True,
-                )
+                try:
+                    res = subprocess.run(
+                        f"top -n 1 -p {os.getpid()}".split(" "),
+                        capture_output=True,
+                        text=True,
+                        check=True,
+                    )
+                except subprocess.CalledProcessError as e:
+                    raise RuntimeError(
+                        f"top command returned non-zero exit code. Can't collect memory values.\n"
+                        f"Make sure top is an available executable name. Possibly, running python "
+                        f"script through terminal may work.\nOriginal exception:\n{e}"
+                    )
                 stdout, _ = res.stdout, res.stderr
                 lines = stdout.split(new_line_delimiter)
                 if len(lines) < abs(header_line):