Get benchmark node memory stats from /proc (#7856)

Co-authored-by: Amaury Chamayou <amchamay@microsoft.com>

Author: eddyashton

tests/e2e_suite.py

@@ -3,6 +3,7 @@
 
 import infra.e2e_args
 import infra.network
+import infra.proc
 import suite.test_suite as s
 import suite.test_requirements as reqs
 import infra.logging_app as app
@@ -24,6 +25,22 @@ class TestStatus(Enum):
     skipped = auto()
 
 
+def mem_stats(network):
+    mem = {}
+    for node in network.get_joined_nodes():
+        try:
+            pid = node.remote.remote.proc.pid
+            stats = infra.proc.get_proc_memory_stats(pid)
+            if stats is not None:
+                mem[node.local_node_id] = stats
+        except (AttributeError, OSError) as exc:
+            LOG.debug(
+                f"Unable to collect memory stats for node "
+                f"{getattr(node, 'local_node_id', '<unknown>')}: {exc}"
+            )
+    return mem
+
+
 def run(args):
     chosen_suite = []
 
@@ -126,6 +143,7 @@ def filter_fun(x):
                 {
                     "status": status.name,
                     "elapsed (s)": round(test_elapsed, 2),
+                    "memory": mem_stats(new_network),
                 }
             )
 

tests/infra/basicperf.py

@@ -20,6 +20,7 @@
 import ccf.ledger
 import plotext as plt
 import infra.bencher
+import infra.proc
 
 
 def configure_remote_client(args, client_id, client_host, common_dir):
@@ -439,6 +440,15 @@ def run(args):
 
                 perf_label = args.perf_label
 
+                if not args.stop_primary_after_s:
+                    primary, _ = network.find_primary()
+                    mem = infra.proc.get_proc_memory_stats(
+                        primary.remote.remote.proc.pid
+                    )
+                    if mem is not None:
+                        bf = infra.bencher.Bencher()
+                        bf.set_memory(perf_label, mem)
+
                 network.stop_all_nodes()
 
                 agg = []

tests/infra/bencher.py

@@ -6,6 +6,8 @@
 import dataclasses
 from typing import Optional, Union
 
+from loguru import logger as LOG
+
 BENCHER_FILE = "bencher.json"
 
 # See https://bencher.dev/docs/reference/bencher-metric-format/
@@ -76,6 +78,17 @@ def __init__(self):
             with open(BENCHER_FILE, "w+") as bf:
                 json.dump({}, bf)
 
+    def set_memory(self, key: str, proc_stats: dict):
+        LOG.info(
+            f"Memory: RSS={proc_stats['current_rss']}, "
+            f"Peak RSS={proc_stats['peak_rss']}, "
+            f"Virtual={proc_stats['virtual_size']}"
+        )
+        self.set(
+            key,
+            Memory(proc_stats["current_rss"], high_value=proc_stats["peak_rss"]),
+        )
+
     def set(self, key: str, metric: Union[Latency, Throughput, Memory]):
         with open(BENCHER_FILE, "r") as bf:
             data = json.load(bf)

tests/infra/piccolo_driver.py

@@ -15,6 +15,7 @@
 from piccolo import generator
 from piccolo import analyzer
 import infra.bencher
+import infra.proc
 
 
 def get_command_args(args, network, get_command):
@@ -215,6 +216,12 @@ def run(get_command, args):
                         infra.bencher.Throughput(perf_result),
                     )
 
+                primary, _ = network.find_primary()
+                mem = infra.proc.get_proc_memory_stats(primary.remote.remote.proc.pid)
+                if mem is not None:
+                    bf = infra.bencher.Bencher()
+                    bf.set_memory(perf_label, mem)
+
                 for remote_client in clients:
                     remote_client.stop()
 

tests/infra/proc.py

@@ -1,10 +1,39 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the Apache 2.0 License.
 from subprocess import run, Popen, PIPE
+from pathlib import Path
+from typing import Optional, Dict
 
 from loguru import logger as LOG
 
 
+def get_proc_memory_stats(pid: int) -> Optional[Dict[str, int]]:
+    """Read memory statistics for a process from /proc/<pid>/status.
+
+    Returns a dict with keys:
+      - current_rss: current resident set size in bytes
+      - peak_rss: peak resident set size (VmHWM) in bytes
+      - virtual_size: total virtual memory size in bytes
+    Returns None if the process info cannot be read.
+    """
+    try:
+        status_path = Path(f"/proc/{pid}/status")
+        text = status_path.read_text()
+    except (OSError, PermissionError):
+        return None
+
+    fields = {"VmRSS": "current_rss", "VmHWM": "peak_rss", "VmSize": "virtual_size"}
+    result = {}
+    for line in text.splitlines():
+        parts = line.split(":", 1)
+        if len(parts) == 2 and parts[0].strip() in fields:
+            key = fields[parts[0].strip()]
+            # Values in /proc/*/status are in kB
+            value_str = parts[1].strip().split()[0]
+            result[key] = int(value_str) * 1024
+    return result if len(result) == len(fields) else None
+
+
 def ccall(*args, path=None, log_output=True, env=None):
     suffix = f" [cwd: {path}]" if path else ""
     cmd = " ".join(args)

tests/infra/runner.py

@@ -159,6 +159,12 @@ def run(get_command, args):
                         infra.bencher.Throughput(perf_result),
                     )
 
+                primary, _ = network.find_primary()
+                mem = infra.proc.get_proc_memory_stats(primary.remote.remote.proc.pid)
+                if mem is not None:
+                    bf = infra.bencher.Bencher()
+                    bf.set_memory(perf_label, mem)
+
                 for remote_client in clients:
                     remote_client.stop()