benchmarking: add memory working set to server telemetry

Author: Nishanth29

benchmarking/README.md

@@ -92,7 +92,7 @@ Three flags control the optional post-run measurements described in
   Defaults to the in-cluster service installed by
   [Optional: Prometheus + Grafana](#optional-prometheus--grafana).
 * `--atelet-lag-s`: how long to wait after the run before reading the
-  atelet's snapshot metrics. Defaults to 70.
+  atelet's metrics. Defaults to 70.
 
 Test-specific flags are appended to the same command; see the sections below.
 
@@ -456,6 +456,25 @@ actually did, independent of what the load generator reported.
   sample with no series counts as 0, since the atelet stops exporting when a
   node has no running actor; if the metric never appeared during the run,
   every field is `null`.
+* `working_set.node`, `ateom`, `atelet`: memory working set in GiB from
+  cAdvisor on the nodes hosting a worker pod, per worker pod (every actor on it
+  included), and per atelet on those nodes. Each has a percentile `summary`
+  (plus `max`) over the steady-state window, `count` of series seen in it, and
+  a `timeseries` every 10s over the whole run with the sum (`total_gb`) and
+  largest (`max_gb`). A spike shorter than 10s can be missed.
+* `working_set.actor`, `per_actor`: the same shape for the actors' own working
+  set as the atelet reports it (`ate_actor_stats_memory_working_set_bytes`),
+  per atelet with every template summed, and the average per measured actor
+  (`ate_actor_stats_sampled_actors`) on each atelet; `count` is atelets. As
+  `per_actor` is an average per atelet, its `max` is the highest atelet
+  average, not the largest actor, and its `total_gb` has no meaning. A gVisor
+  actor is read from its sandbox cgroup on the host, a microVM actor from
+  inside the guest. As with `active_actors`, suspended actors fall out, but
+  through the telemetry-meter a series that stops being exported keeps its
+  last value for up to 5m. Missing samples are skipped, not counted as 0. The
+  atelet samples every `--actor-stats-poll-interval` (default 1m), so these
+  move in steps of a minute or more. `actor` matches one worker pod's actors
+  only when the node runs one worker pod.
 
 Every distribution reports p50, p90, p95 and p99 over the steady-state window.
 The steady-state window runs from the first to the last Locust sample at 90% or
@@ -469,8 +488,9 @@ every 10s.
 The atelet exports before Prometheus scrapes it, so the harvest waits
 `--atelet-lag-s` seconds (default 70, enough for the OTel SDK's 60s default
 export and a 10s scrape)
-and reads the snapshot and `active_actors` windows half that late. The window
-ends at the last full-load sample, so teardown suspends are left out.
+and reads the snapshot, `active_actors` and actor working set windows half
+that late. The window ends at the last full-load sample, so teardown suspends
+are left out.
 
 `metadata.start_ts` and `end_ts` bound the whole run, which the packing
 `timeseries` covers. `steady_start_ts` and `steady_end_ts` bound the

benchmarking/locust/server_telemetry.py

@@ -16,7 +16,7 @@
 
 Queries Prometheus over [T_start, T_end] and the steady-state window [T_steady, T_end]
 to capture cluster packing, node and pod PSI stalls, snapshot sizes, latencies
-and throughput, and running actors.
+and throughput, running actors and memory working set.
 """
 
 import csv
@@ -547,6 +547,129 @@ def _count_active_actors(
     }
 
 
+# Live worker pods: the rate drops a deleted pod's stale series in ~1 min, not 5.
+WORKER_PODS = (
+    "rate(container_cpu_usage_seconds_total"
+    '{namespace="benchmark-workloads", container="ateom"}[1m]) > 0'
+)
+
+# Only nodes/pods with a live worker; max, not sum: a restart briefly has 2 series.
+WORKING_SET_QUERIES = {
+    "node": (
+        'max by (instance) (container_memory_working_set_bytes{container="node"})'
+        f" and on (instance) count by (instance) ({WORKER_PODS})"
+    ),
+    "ateom": (
+        "max by (pod) (container_memory_working_set_bytes"
+        '{namespace="benchmark-workloads", container="ateom"})'
+        f" and on (pod) ({WORKER_PODS})"
+    ),
+    "atelet": (
+        "max by (pod, instance) (container_memory_working_set_bytes"
+        '{namespace="ate-system", container="atelet"})'
+        f" and on (instance) count by (instance) ({WORKER_PODS})"
+    ),
+}
+
+# Running actors per atelet, all templates summed; `instance` when scraped
+# directly, `exported_instance` via a collector. Not node names, so no worker
+# filter: the atelet stops exporting on a node with no running actor.
+ACTOR_WORKING_SET = (
+    "sum by (instance, exported_instance) (ate_actor_stats_memory_working_set_bytes)"
+)
+ACTOR_COUNT = "sum by (instance, exported_instance) (ate_actor_stats_sampled_actors)"
+WORKING_SET_KEYS = (*WORKING_SET_QUERIES, "actor", "per_actor")
+
+BYTES_PER_GIB = 2**30
+
+
+def _read_series(
+    prom_url: str, query: str, start_ts: int, end_ts: int, shift: int = 0
+) -> dict[tuple, dict[int, float]]:
+    """Samples per series every 10s, read `shift` late and moved back by it."""
+    res = query_prometheus_range(
+        prom_url, query, start_ts + shift, end_ts + shift, step="10s"
+    )
+    out: dict[tuple, dict[int, float]] = {}
+    for series in res:
+        samples = out.setdefault(tuple(sorted(series.get("metric", {}).items())), {})
+        for pt in series.get("values", []):
+            try:
+                t, val = int(pt[0]) - shift, float(pt[1])
+            except (ValueError, IndexError, TypeError):
+                continue
+            if not math.isnan(val) and not math.isinf(val):
+                samples[t] = val
+    return out
+
+
+def _summarize_gb(
+    series: dict[tuple, dict[int, float]],
+    steady_start_ts: int,
+    steady_end_ts: int,
+) -> dict[str, Any]:
+    """Steady percentiles, series seen in steady, and a 10s sum/max timeseries."""
+    by_ts: dict[int, list[float]] = {}
+    steady_gb: list[float] = []
+    steady_series = 0
+    for samples in series.values():
+        in_steady = False
+        for t, val in samples.items():
+            gb = val / BYTES_PER_GIB
+            by_ts.setdefault(t, []).append(gb)
+            if steady_start_ts <= t <= steady_end_ts:
+                steady_gb.append(gb)
+                in_steady = True
+        steady_series += in_steady
+    return {
+        "summary": compute_percentiles(steady_gb),
+        "count": steady_series if by_ts else None,
+        "timeseries": [
+            {
+                "timestamp": t,
+                "total_gb": round(sum(v), 4),
+                "max_gb": round(max(v), 4),
+            }
+            for t, v in sorted(by_ts.items())
+        ],
+    }
+
+
+def _measure_working_set(
+    prom_url: str,
+    start_ts: int,
+    end_ts: int,
+    steady_start_ts: int,
+    steady_end_ts: int,
+    lag_s: int = 0,
+) -> dict[str, Any]:
+    """Working set GiB per node, ateom, atelet and actor; null, not 0, with no series.
+
+    `ateom` is the whole worker pod, every actor on it included. `actor` is the
+    atelet's sum over the actors it measured, per atelet; `per_actor` divides it
+    by their count. The atelet samples about once a minute and exports late, so
+    both are read `lag_s // 2` late and moved back onto the run's 10s grid.
+    """
+    def summarize(series: dict[tuple, dict[int, float]]) -> dict[str, Any]:
+        return _summarize_gb(series, steady_start_ts, steady_end_ts)
+
+    out: dict[str, Any] = {
+        key: summarize(_read_series(prom_url, query, start_ts, end_ts))
+        for key, query in WORKING_SET_QUERIES.items()
+    }
+    shift = lag_s // 2
+    actor = _read_series(prom_url, ACTOR_WORKING_SET, start_ts, end_ts, shift)
+    count = _read_series(prom_url, ACTOR_COUNT, start_ts, end_ts, shift)
+    out["actor"] = summarize(actor)
+    # Same atelet and timestamp; a missing or zero count is skipped.
+    out["per_actor"] = summarize({
+        key: {t: v / count[key][t] for t, v in samples.items()
+              if count.get(key, {}).get(t)}
+        for key, samples in actor.items()
+    })
+    return out
+
+
 def harvest_server_telemetry(
     prom_url: str,
     start_ts: int,
@@ -557,9 +680,9 @@ def harvest_server_telemetry(
 ) -> dict[str, Any]:
     """Harvests the ground truth metric streams from Prometheus.
 
-    Only the atelet-exported snapshot and active actor blocks are read late, by
-    up to `lag_s`; packing and PSI are scraped directly and read over the run
-    window as is.
+    Only the atelet-exported snapshot, active actor and actor working set
+    blocks are read late, by up to `lag_s`; packing, PSI and the cAdvisor
+    working set are scraped directly and read over the run window as is.
     """
     steady_end = end_ts if steady_end_ts is None else steady_end_ts
     return {
@@ -578,6 +701,9 @@ def harvest_server_telemetry(
         "active_actors": _count_active_actors(
             prom_url, start_ts, end_ts, steady_start_ts, steady_end, lag_s
         ),
+        "working_set": _measure_working_set(
+            prom_url, start_ts, end_ts, steady_start_ts, steady_end, lag_s
+        ),
     }
 
 
@@ -674,6 +800,11 @@ def log(msg: str) -> None:
     active = telemetry.get("active_actors", {}).get("summary", {})
     for p in ps:
         measurements[f"active_actors_{p}"] = active.get(p)
+    working_set = telemetry.get("working_set", {})
+    for block in WORKING_SET_KEYS:
+        ws = working_set.get(block, {}).get("summary", {})
+        for p in (*ps, "max"):
+            measurements[f"working_set_{block}_{p}_gb"] = ws.get(p)
 
     jsonl_row = {
         "timestamp": data_ts,

benchmarking/locust/unit_tests/test_server_telemetry.py

@@ -40,8 +40,8 @@
           "end_ts": 105, "steady_start_ts": 100}
 
 
-def ranges(packing, node=(), pod=(), active=()):
-    """A query_prometheus_range fake answering packing, PSI and active actors.
+def ranges(packing, node=(), pod=(), active=(), ws=(), actor_ws=()):
+    """A query_prometheus_range fake: packing, PSI, active actors, working set.
 
     Pod PSI only answers the pod-slice selector, so a query that loses the id
     regex (and would also sum the pause container) gets nothing back. Node PSI
@@ -56,6 +56,11 @@ def fake(_url, query, *_args, **_kwargs):
             return packing
         if "ate_actor_stats_sampled_actors" in query:
             return list(active)
+        if "ate_actor_stats_memory_working_set_bytes" in query:
+            return list(actor_ws)
+        # Before the PSI branch: the node working set query has container="node".
+        if "container_memory_working_set_bytes" in query:
+            return list(ws)
         if 'container="node"' in query:
             return list(node) if worker_nodes in query and pod_slice in query else []
         return list(pod) if pod_slice in query else []
@@ -244,7 +249,11 @@ def test_packing_and_checkpoint_math(self, mock_instant, mock_range):
                  "values": [[100, "0.5"], [105, "0.5"]]} for p in "ab"]
         mock_range.side_effect = ranges(
             [partial, full, idle], node=quiet, pod=pods,
-            active=[{"metric": {"instance": "atelet-a"}, "values": [[135, "3"]]}])
+            active=[{"metric": {"instance": "atelet-a"}, "values": [[135, "3"]]}],
+            ws=[{"metric": {"instance": "node-a"},
+                 "values": [[100, str(2**30)], [105, str(2**30)]]}],
+            actor_ws=[{"metric": {"instance": "atelet-a"},
+                       "values": [[135, str(6 * 2**30)]]}])
         # The default 70s lag reads the snapshot window at 135 and 140.
         mock_instant.side_effect = snapshot_prom(
             "11.5",
@@ -278,14 +287,14 @@ def test_packing_and_checkpoint_math(self, mock_instant, mock_range):
 
         sleep.assert_called_once_with(65)  # until end 105 + 70, from 110
         self.assertEqual(summary["metadata"]["atelet_lag_s"], 70)
-        # Packing and PSI are not shifted; active actors are, like snapshots.
-        windows = {("sampled_actors" in c.args[1], c.args[2:4])
+        # Packing, PSI and cAdvisor are not shifted; atelet stats are, like snapshots.
+        windows = {("ate_actor_stats" in c.args[1], c.args[2:4])
                    for c in mock_range.call_args_list}
         self.assertEqual(windows, {(False, (100, 105)), (True, (135, 140))})
         self.assertEqual(
             set(summary),
             {"metadata", "cluster_packing", "node_psi", "pod_psi", "snapshots",
-             "active_actors"},
+             "active_actors", "working_set"},
         )
         packing = summary["cluster_packing"]
         self.assertEqual(packing["summary"]["p50"], 0.8)  # 3 + 1 busy / 5 workers
@@ -305,7 +314,7 @@ def test_packing_and_checkpoint_math(self, mock_instant, mock_range):
         self.assertEqual(row["metric"], "server_summary")
         m = row["measurements"]
         # Every source answered, so a key wired to a wrong name would read None.
-        self.assertEqual(len(m), 49)
+        self.assertEqual(len(m), 74)
         self.assertEqual([k for k, v in m.items() if v is None], [])
         # Every value a string, so one row's types match every other row's.
         self.assertTrue(all(isinstance(v, str) for v in m.values()))
@@ -317,6 +326,11 @@ def test_packing_and_checkpoint_math(self, mock_instant, mock_range):
         self.assertEqual(m["restore_mean_s"], "0.5")
         self.assertEqual(m["checkpoint_mean_s"], "2.0")
         self.assertEqual(m["active_actors_p99"], "3.0")
+        # Read at 135, moved back to 100; 6 GiB over the 3 actors on atelet-a.
+        actor = summary["working_set"]["actor"]
+        self.assertEqual(actor["timeseries"][0]["timestamp"], 100)
+        self.assertEqual(m["working_set_actor_p50_gb"], "6.0")
+        self.assertEqual(m["working_set_per_actor_p50_gb"], "2.0")
 
     @mock.patch("server_telemetry.query_prometheus_range")
     @mock.patch("server_telemetry.query_prometheus_instant")
@@ -466,6 +480,43 @@ def test_steady_window_bounds_packing_and_psi(self, mock_instant, mock_range):
         self.assertEqual(out["node_psi"]["cpu_stall_pct"]["max"], 0.5)
         self.assertEqual(out["pod_psi"]["cpu_stall_pct"]["max"], 0.5)
 
+    @mock.patch("server_telemetry.query_prometheus_range")
+    def test_working_set_gib_and_steady_window(self, mock_range):
+        gib = 2**30
+        mock_range.side_effect = ranges([], ws=[
+            {"metric": {"instance": "node-a"},
+             "values": [[90, str(8 * gib)], [100, str(gib)], [110, str(3 * gib)]]},
+            {"metric": {"instance": "node-b"}, "values": [[110, str(gib)]]},
+        ])
+        out = server_telemetry._measure_working_set("http://p", 90, 110, 100, 110)
+        node = out["node"]
+        # The 8 GiB sample at 90 is ramp-up: in the timeseries, not the summary.
+        self.assertEqual((node["summary"]["p50"], node["summary"]["max"]), (1.0, 3.0))
+        self.assertEqual(node["count"], 2)
+        self.assertEqual(node["timeseries"], [
+            {"timestamp": 90, "total_gb": 8.0, "max_gb": 8.0},
+            {"timestamp": 100, "total_gb": 1.0, "max_gb": 1.0},
+            {"timestamp": 110, "total_gb": 4.0, "max_gb": 3.0},
+        ])
+        self.assertEqual(set(out), {"node", "ateom", "atelet", "actor", "per_actor"})
+
+    @mock.patch("server_telemetry.query_prometheus_range")
+    def test_working_set_null_not_zero(self, mock_range):
+        mock_range.side_effect = ranges([])
+        out = server_telemetry._measure_working_set("http://p", 90, 110, 100, 110)
+        for block in out.values():
+            self.assertEqual(block, {"summary": NO_PERCENTILES, "count": None,
+                                     "timeseries": []})
+
+    def test_working_set_queries_keep_live_workers(self):
+        q = server_telemetry.WORKING_SET_QUERIES
+        live = "rate(container_cpu_usage_seconds_total"
+        for block in q.values():
+            self.assertTrue(block.startswith("max by ("))  # a restart has 2 series
+            self.assertIn(live, block)
+        self.assertIn('container="node"', q["node"])
+        self.assertIn("and on (pod)", q["ateom"])
+
     @mock.patch("server_telemetry.query_prometheus_instant")
     def test_failed_delta_read_is_unknown(self, mock_instant):
         # A failed delta query is not taken as a series born in the window.