Toy model for imitating O2 software memory usage in K8s

Closes OCTRL-1079

Author: knopers8

hacking/o2-software-model/README.md

@@ -0,0 +1,27 @@
+# O2 software toy model
+
+This directory contains a toy model which pretends to be O2 software in terms of how it uses memory and runs on a K8s cluster.
+The manifest file declares an init container which creates a shared memory segment and two identical `memory-user` containers which write to anonymous and/or shared memory with specified rates.
+One can set memory requests and limits to the containers and observe the behaviour.
+
+The manifest also includes a `mem-guard` container, which monitors anonymous memory usage of `memory-user`s and kills them when they go above a set threshold.
+
+See the manifest and environment variables inside for configuring memory consumption rates and thresholds.
+
+See OCTRL-1079 for more context and info which was understood using this model.
+## Running
+
+1. Add `allocate.py` and `mem_guard.py` scripts to a configmap. This way we can reuse the official python image as-is.
+
+```
+kubectl create configmap allocate-script \
+  --from-file=allocate.py=./allocate.py \
+  --from-file=mem_guard.py=./mem_guard.py \
+  -o yaml --dry-run=client | kubectl apply -f -
+```
+
+2. Apply the manifest
+
+```
+kubectl apply -f mem-allocate-test.yaml
+```

hacking/o2-software-model/allocate.py

@@ -0,0 +1,265 @@
+#!/usr/bin/env python3
+import argparse
+import mmap
+import os
+import sys
+import time
+
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Memory load generator: private allocation + shmem(tmpfs file) mmap read/write."
+    )
+
+    # Private memory
+    p.add_argument("--private-rate-bps", type=int, default=0,
+                   help="Private memory allocation+fill rate in bytes/sec. 0 disables.")
+    p.add_argument("--private-bytes", type=int, default=0,
+                   help="Optional cap for private bytes to allocate. 0 means no cap (grow forever).")
+    p.add_argument("--private-chunk-bytes", type=int, default=4 * 1024 * 1024,
+                   help="Chunk size for private allocations (bytes). Default 4MiB.")
+    p.add_argument("--page-touch-bytes", type=int, default=4096,
+                   help="Stride for touching/filling memory to force residency. Default 4096.")
+
+    # Shmem mmap activity
+    p.add_argument("--shmem-path", type=str, default="",
+                   help="Path to a shmem/tmpfs-backed file (e.g. /dev/shm/segment.bin).")
+    p.add_argument("--shmem-bytes", type=int, default=0,
+                   help="Ensure shmem file size is at least this many bytes (create/truncate). 0 = don't change size.")
+    p.add_argument("--shmem-read-rate-bps", type=int, default=0,
+                   help="Read rate from shmem mapping in bytes/sec. Only reads if > 0.")
+    p.add_argument("--shmem-write-rate-bps", type=int, default=0,
+                   help="Write rate to shmem mapping in bytes/sec. Only writes if > 0.")
+
+    # Loop control
+    p.add_argument("--tick-ms", type=int, default=50,
+                   help="Control loop tick in milliseconds. Default 50ms.")
+    p.add_argument("--report-every-s", type=float, default=2.0,
+                   help="Print status every N seconds. Default 2s.")
+
+    return p.parse_args()
+
+
+def touch_buffer(buf: bytearray, stride: int):
+    """Touch/dirty memory with given stride to force page allocation."""
+    n = len(buf)
+    if n == 0:
+        return
+    v = int(time.time_ns() & 0xFF)
+    for i in range(0, n, stride):
+        buf[i] = (buf[i] + v + (i & 0xFF)) & 0xFF
+
+
+def ensure_file_size(path: str, size: int):
+    fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o666)
+    try:
+        os.ftruncate(fd, size)
+    finally:
+        os.close(fd)
+
+
+def main():
+    args = parse_args()
+
+    for name, v in [
+        ("--private-rate-bps", args.private_rate_bps),
+        ("--shmem-read-rate-bps", args.shmem_read_rate_bps),
+        ("--shmem-write-rate-bps", args.shmem_write_rate_bps),
+    ]:
+        if v < 0:
+            print(f"{name} must be >= 0", file=sys.stderr)
+            return 2
+
+    if args.private_chunk_bytes <= 0:
+        print("--private-chunk-bytes must be > 0", file=sys.stderr)
+        return 2
+    if args.page_touch_bytes <= 0:
+        print("--page-touch-bytes must be > 0", file=sys.stderr)
+        return 2
+
+    tick = args.tick_ms / 1000.0
+    report_every = args.report_every_s
+
+    # Private memory state
+    private_chunks: list[bytearray] = []
+    private_allocated = 0
+
+    # Shmem state
+    shmem_fd = None
+    shmem_mm = None
+    shmem_size = 0
+    shmem_r_off = 0
+    shmem_w_off = 0
+    shmem_read_total = 0
+    shmem_written_total = 0
+
+    need_shmem = (args.shmem_read_rate_bps > 0) or (args.shmem_write_rate_bps > 0)
+    if need_shmem:
+        if not args.shmem_path:
+            print("shmem activity requested but --shmem-path is empty", file=sys.stderr)
+            return 2
+
+        if args.shmem_bytes > 0:
+            try:
+                ensure_file_size(args.shmem_path, args.shmem_bytes)
+            except Exception as e:
+                print(f"Failed to create/size shmem file: {e}", file=sys.stderr)
+                return 2
+
+        # For mmap write, we need RDWR; for read-only, RDONLY is enough
+        open_flags = os.O_RDONLY if args.shmem_write_rate_bps == 0 else os.O_RDWR
+        try:
+            shmem_fd = os.open(args.shmem_path, open_flags)
+            st = os.fstat(shmem_fd)
+            shmem_size = st.st_size
+            if shmem_size <= 0:
+                raise RuntimeError("shmem file size is 0; set --shmem-bytes to something > 0")
+
+            access = mmap.ACCESS_READ if args.shmem_write_rate_bps == 0 else mmap.ACCESS_WRITE
+            shmem_mm = mmap.mmap(shmem_fd, shmem_size, access=access)
+        except Exception as e:
+            if shmem_fd is not None:
+                try:
+                    os.close(shmem_fd)
+                except Exception:
+                    pass
+            print(f"Failed to open/mmap shmem file: {e}", file=sys.stderr)
+            return 2
+
+    last = time.monotonic()
+    last_report = last
+
+    private_budget = 0.0
+    shmem_read_budget = 0.0
+    shmem_write_budget = 0.0
+
+    # A tiny accumulator so reads can't be optimized away
+    checksum = 0
+
+    try:
+        while True:
+            now = time.monotonic()
+            dt = now - last
+            last = now
+
+            private_budget += args.private_rate_bps * dt
+            shmem_read_budget += args.shmem_read_rate_bps * dt
+            shmem_write_budget += args.shmem_write_rate_bps * dt
+
+            # --- Private allocation (allocate + touch) ---
+            if args.private_rate_bps > 0:
+                cap = args.private_bytes
+                while private_budget >= 1:
+                    if cap > 0 and private_allocated >= cap:
+                        private_budget = 0
+                        break
+
+                    chunk = args.private_chunk_bytes
+                    if cap > 0:
+                        remaining = cap - private_allocated
+                        if remaining <= 0:
+                            private_budget = 0
+                            break
+                        if remaining < chunk:
+                            chunk = remaining
+
+                    if private_budget < chunk:
+                        chunk = int(private_budget)
+                        if chunk < 64 * 1024:
+                            break
+
+                    buf = bytearray(chunk)
+                    touch_buffer(buf, args.page_touch_bytes)
+                    private_chunks.append(buf)
+                    private_allocated += chunk
+                    private_budget -= chunk
+
+            # --- Shmem mmap write: touches mapped pages, should increase RES/SHR ---
+            if shmem_mm is not None and args.shmem_write_rate_bps > 0:
+                while shmem_write_budget >= 1:
+                    to_write = int(shmem_write_budget)
+                    if to_write <= 0:
+                        break
+
+                    if shmem_w_off >= shmem_size:
+                        shmem_w_off = 0
+
+                    n = min(to_write, shmem_size - shmem_w_off)
+                    if n <= 0:
+                        shmem_w_off = 0
+                        break
+
+                    # Touch each page once (or more if n is large) to fault pages in.
+                    stride = 4096
+                    v = int(time.time_ns() & 0xFF)
+                    end = shmem_w_off + n
+                    for off in range(shmem_w_off, end, stride):
+                        shmem_mm[off] = (shmem_mm[off] + v) & 0xFF
+
+                    shmem_w_off = end
+                    shmem_written_total += n
+                    shmem_write_budget -= n
+
+                    if shmem_w_off >= shmem_size:
+                        shmem_w_off = 0
+
+            # --- Shmem mmap read: faults pages in (RES/SHR can still rise), no dirtying ---
+            if shmem_mm is not None and args.shmem_read_rate_bps > 0:
+                while shmem_read_budget >= 1:
+                    to_read = int(shmem_read_budget)
+                    if to_read <= 0:
+                        break
+
+                    if shmem_r_off >= shmem_size:
+                        shmem_r_off = 0
+
+                    n = min(to_read, shmem_size - shmem_r_off)
+                    if n <= 0:
+                        shmem_r_off = 0
+                        break
+
+                    stride = 4096
+                    end = shmem_r_off + n
+                    for off in range(shmem_r_off, end, stride):
+                        checksum ^= shmem_mm[off]
+
+                    shmem_r_off = end
+                    shmem_read_total += n
+                    shmem_read_budget -= n
+
+                    if shmem_r_off >= shmem_size:
+                        shmem_r_off = 0
+
+            # --- Reporting ---
+            if now - last_report >= report_every:
+                last_report = now
+                parts = [f"private_allocated={private_allocated}B"]
+                if shmem_mm is not None:
+                    parts.append(f"shmem_size={shmem_size}B")
+                    parts.append(f"shmem_written_total={shmem_written_total}B")
+                    parts.append(f"shmem_read_total={shmem_read_total}B")
+                    parts.append(f"checksum={checksum}")
+                    parts.append(f"shmem_path={args.shmem_path}")
+                print(" ".join(parts), flush=True)
+
+            time.sleep(tick)
+
+    except KeyboardInterrupt:
+        print("\nExiting.", file=sys.stderr)
+    finally:
+        if shmem_mm is not None:
+            try:
+                shmem_mm.close()
+            except Exception:
+                pass
+        if shmem_fd is not None:
+            try:
+                os.close(shmem_fd)
+            except Exception:
+                pass
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())