Changed method to measure memory to make it cross-platform

mlte/measurement/memory/local_process_memory_consumption.py

@@ -10,7 +10,8 @@
 import time
 from typing import Any, Dict, Type
 
-from mlte._private.platform import is_macos, is_windows
+import psutil
+
 from mlte.evidence.metadata import EvidenceMetadata
 from mlte.spec.condition import Condition
 from mlte.validation.result import Failure, Success
@@ -168,10 +169,6 @@ def __init__(self, identifier: str):
         :type identifier: str
         """
         super().__init__(self, identifier)
-        if is_windows() or is_macos():
-            raise RuntimeError(
-                f"Measurement {self.metadata.identifier} is not supported on Windows or macOS."
-            )
 
     def __call__(self, pid: int, poll_interval: int = 1) -> MemoryStatistics:
         """
@@ -187,7 +184,7 @@ def __call__(self, pid: int, poll_interval: int = 1) -> MemoryStatistics:
         """
         stats = []
         while True:
-            kb = _get_memory_usage(pid)
+            kb = _get_memory_usage_psutil(pid)
             if kb == 0:
                 break
             stats.append(kb)
@@ -211,7 +208,7 @@ def value(self) -> Type[MemoryStatistics]:
 # -----------------------------------------------------------------------------
 
 
-def _get_memory_usage(pid: int) -> int:
+def _get_memory_usage_pmap(pid: int) -> int:
     """
     Get the current memory usage for the process with `pid`.
 
@@ -245,3 +242,19 @@ def _get_memory_usage(pid: int) -> int:
         raise RuntimeError(
             f"External program needed to get memory usage was not found: {e}"
         )
+
+
+def _get_memory_usage_psutil(pid: int) -> int:
+    """
+    Get the current memory usage for the process with `pid`.
+
+    :param pid: The identifier of the process
+
+    :return: The current memory usage in KB
+    """
+    try:
+        curr_proc = psutil.Process(pid)
+        mem_in_kb = curr_proc.memory_info().rss / 1024
+        return int(mem_in_kb)
+    except psutil.NoSuchProcess:
+        return 0

test/measurement/memory/test_local_process_memory_consumption.py

@@ -10,9 +10,6 @@
 import time
 from typing import Tuple
 
-import pytest
-
-from mlte._private.platform import is_macos, is_nix, is_windows
 from mlte.context.context import Context
 from mlte.evidence.metadata import EvidenceMetadata, Identifier
 from mlte.measurement.memory import (
@@ -39,11 +36,7 @@ def spin_for(seconds: int):
     return prog
 
 
-@pytest.mark.skipif(
-    is_windows() or is_macos(),
-    reason="ProcessLocalCPUUtilization not supported on Windows or macOS.",
-)
-def test_memory_nix_evaluate() -> None:
+def test_memory_evaluate() -> None:
     start = time.time()
 
     p = spin_for(5)
@@ -56,11 +49,7 @@ def test_memory_nix_evaluate() -> None:
     assert int(time.time() - start) >= SPIN_DURATION
 
 
-@pytest.mark.skipif(
-    is_windows() or is_macos(),
-    reason="ProcessLocalCPUUtilization not supported on Windows or macOS.",
-)
-def test_memory_nix_evaluate_async() -> None:
+def test_memory_evaluate_async() -> None:
     start = time.time()
 
     p = spin_for(5)
@@ -74,11 +63,7 @@ def test_memory_nix_evaluate_async() -> None:
     assert int(time.time() - start) >= SPIN_DURATION
 
 
-@pytest.mark.skipif(
-    is_windows() or is_macos(),
-    reason="ProcessLocalCPUUtilization not supported on Windows or macOS.",
-)
-def test_memory_nix_validate_success() -> None:
+def test_memory_validate_success() -> None:
     p = spin_for(5)
 
     m = LocalProcessMemoryConsumption("identifier")
@@ -93,11 +78,7 @@ def test_memory_nix_validate_success() -> None:
     assert vr.metadata.measurement_type, type(MemoryStatistics).__name__
 
 
-@pytest.mark.skipif(
-    is_windows() or is_macos(),
-    reason="ProcessLocalCPUUtilization not supported on Windows or macOS.",
-)
-def test_memory_nix_validate_failure() -> None:
+def test_memory_validate_failure() -> None:
     p = spin_for(5)
 
     m = LocalProcessMemoryConsumption("identifier")
@@ -112,18 +93,6 @@ def test_memory_nix_validate_failure() -> None:
     assert vr.metadata.measurement_type, type(MemoryStatistics).__name__
 
 
-@pytest.mark.skipif(
-    is_nix(), reason="ProcessLocalCPUUtilization not supported on Windows."
-)
-def test_memory_windows_evaluate() -> None:
-    with pytest.raises(RuntimeError):
-        _ = LocalProcessMemoryConsumption("id")
-
-
-@pytest.mark.skipif(
-    is_windows() or is_macos(),
-    reason="LocalProcessCPUUtilization not supported on Windows or macOS.",
-)
 def test_result_save_load(
     store_with_context: Tuple[Store, Context]  # noqa
 ) -> None: