okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petkernel: keep output once, bound snapshots and replays, drop the outbox scanf479db5

main
branchesmainall refs →
9.4 KiB290 linesraw
1"""Measure one kernel's resident memory at boot and under scaled workloads.
2
3Run in nix develop:
4  python test/manual/kernel_memory_benchmark.py [--scale 1 10] [--output /tmp/kernel-memory.json]
5
6Boots the kernel over stdio with the daemon's default module list, samples RSS
7after `ready`, then runs each workload at every scale and reports the RSS delta
8and the kernel's own accounting of what it retains (captures, finished results,
9definitions, outbox). A workload whose delta grows with scale on a path that
10should be bounded is the finding; the absolute numbers are the record.
11"""
12
13import argparse
14import json
15from pathlib import Path
16import platform
17import select
18import struct
19import subprocess
20import sys
21import tempfile
22import time
23
24ROOT = Path(__file__).resolve().parents[2]
25KERNEL = ROOT / "priv" / "python" / "albedo_kernel.py"  # --kernel-dir points elsewhere
26sys.path.insert(0, str(ROOT / "test"))
27import scratch  # noqa: E402
28
29scratch.claim("kernel-memory")
30MODULES = [
31    "run",
32    "work",
33    "mail",
34    "agents",
35    "paperclips",
36    "files",
37    "memory",
38    "commands",
39    "remote",
40    "browser",
41    "view",
42    "skills",
43    "webhooks",
44]
45MAX_FRAME = 8 * 1024 * 1024
46# Plugins that ask the host something at boot get the shape they expect.
47HOST_ANSWERS = {
48    "agents.self": {"id": "bench", "name": "bench", "depth": 0},
49    "commands.list": [],
50}
51
52ACCOUNTING = """
53import sys as _sys, gc as _gc, json as _json
54_k = _sys.modules["__main__"]
55_gc.collect()
56_caps = list(_k.ARCHIVES.values())
57_fin = list(_k.FINISHED.values())
58_result = {
59    "captures": len(_caps),
60    "capture_bytes": sum(
61        _sys.getsizeof(c.data)
62        + _sys.getsizeof(getattr(c, "_tail", None) or b"")
63        + _sys.getsizeof(getattr(c, "tail_data", b""))
64        + _sys.getsizeof(getattr(c, "raw_data", b""))
65        for c in _caps
66    ),
67    "finished": len(_fin),
68    "finished_bytes": sum(len(str(d.get("output", ""))) + sum(map(len, d.get("images", []))) for d in _fin),
69    "definitions": len(_k.DEFINITIONS),
70    "definition_bytes": sum(map(len, _k.DEFINITIONS.values())),
71    "outbox_entries": len(getattr(getattr(_k.LINK, "outbox", None), "entries", ())),
72    "outbox_bytes": getattr(getattr(_k.LINK, "outbox", None), "size", 0),
73    "modules": len(_sys.modules),
74    "gc_objects": len(_gc.get_objects()),
75}
76print(_json.dumps(_result))
77"""
78
79
80def rss_kib(pid):
81    out = subprocess.run(
82        ["ps", "-o", "rss=", "-p", str(pid)], capture_output=True, text=True
83    )
84    return int(out.stdout.strip() or 0)
85
86
87class Kernel:
88    def __init__(self, modules, cwd, kernel=KERNEL):
89        self.process = subprocess.Popen(
90            [sys.executable, "-u", str(kernel), json.dumps(modules)],
91            stdin=subprocess.PIPE,
92            stdout=subprocess.PIPE,
93            stderr=subprocess.DEVNULL,
94            cwd=cwd,
95            bufsize=0,
96        )
97        assert self.process.stdin is not None and self.process.stdout is not None
98        self.stdin = self.process.stdin
99        self.stdout = self.process.stdout
100        self.buffered = []
101        self.cells = 0
102
103    def send(self, message):
104        data = json.dumps(message).encode()
105        self.stdin.write(struct.pack(">I", len(data)) + data)
106        self.stdin.flush()
107
108    def read_exactly(self, size):
109        data = b""
110        while len(data) < size:
111            chunk = self.stdout.read(size - len(data))
112            if not chunk:
113                raise AssertionError(
114                    f"kernel closed its control stream; last frames: {self.buffered[-3:]}"
115                )
116            data += chunk
117        return data
118
119    def recv(self, timeout):
120        ready, _, _ = select.select([self.stdout], [], [], timeout)
121        if not ready:
122            return None
123        size = struct.unpack(">I", self.read_exactly(4))[0]
124        assert size <= MAX_FRAME
125        return json.loads(self.read_exactly(size))
126
127    def wait_for(self, predicate, timeout=120.0):
128        deadline = time.monotonic() + timeout
129        while time.monotonic() < deadline:
130            for index, frame in enumerate(self.buffered):
131                if predicate(frame):
132                    return self.buffered.pop(index)
133            frame = self.recv(max(0.05, deadline - time.monotonic()))
134            if frame is None:
135                continue
136            if frame.get("type") == "call":
137                # Answer host calls the way an idle session would: accept and forget.
138                value = HOST_ANSWERS.get(frame["method"])
139                self.send(
140                    {
141                        "type": "reply",
142                        "id": frame["id"],
143                        "value": {"ok": True, "value": value},
144                    }
145                )
146                continue
147            self.buffered.append(frame)
148        raise AssertionError(
149            f"no matching frame; buffered: {[f.get('type') for f in self.buffered]}"
150        )
151
152    def execute(self, code, timeout_ms=600_000):
153        self.cells += 1
154        cell = f"bench-{self.cells}"
155        self.send(
156            {"type": "execute", "id": cell, "code": code, "timeout_ms": timeout_ms}
157        )
158        done = self.wait_for(lambda f: f.get("type") == "done" and f.get("id") == cell)
159        if done["status"] != "ok":
160            raise RuntimeError(f"cell {cell} failed: {done['output'][-2000:]}")
161        return done
162
163    def accounting(self):
164        return json.loads(self.execute(ACCOUNTING)["output"])
165
166    def close(self):
167        self.send({"type": "shutdown"})
168        self.process.wait(timeout=10)
169        self.stdin.close()
170        self.stdout.close()
171
172
173def prints(kernel, scale, workspace):
174    kernel.execute(
175        f"for i in range({20_000 * scale}):\n"
176        "    print('line', i, 'of output that is fairly ordinary')"
177    )
178
179
180def big_print(kernel, scale, workspace):
181    kernel.execute(f"print('x' * {2 * 1024 * 1024 * scale})")
182
183
184def jobs_output(kernel, scale, workspace):
185    kernel.execute(
186        f"for _ in range({scale}):\n"
187        f"    batch = [run('head', '-c', '{256 * 1024}', '/dev/zero') for _ in range(8)]\n"
188        "    for j in batch:\n        await j\n"
189    )
190
191
192def cells(kernel, scale, workspace):
193    for i in range(50 * scale):
194        kernel.execute(f"x{i} = {i}\nprint('cell', x{i})")
195
196
197def defs(kernel, scale, workspace):
198    for i in range(50 * scale):
199        kernel.execute(f"def helper_{i}(a, b):\n    return a + b + {i}\n")
200
201
202def audit_opens(kernel, scale, workspace):
203    kernel.execute(
204        "import os\n"
205        f"paths = [os.path.join({workspace!r}, f'f{{i}}.txt') for i in range({2000 * scale})]\n"
206        "for p in paths:\n    open(p, 'w').close()\n"
207        "for p in paths:\n    open(p).close()\n"
208        "for p in paths:\n    os.unlink(p)\n"
209    )
210
211
212# Each runs one workload in a fresh kernel; `scale` multiplies the size that grows.
213WORKLOADS = {
214    f.__name__: f for f in (prints, big_print, jobs_output, cells, defs, audit_opens)
215}
216
217
218def main():
219    parser = argparse.ArgumentParser(
220        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
221    )
222    parser.add_argument("--scale", type=int, nargs="+", default=[1, 10])
223    parser.add_argument("--output", type=Path, default=None)
224    parser.add_argument("--only", nargs="*", default=None, help="workload names to run")
225    parser.add_argument(
226        "--kernel-dir",
227        type=Path,
228        default=KERNEL.parent,
229        help="a priv/python to boot instead of this checkout's, for a baseline",
230    )
231    args = parser.parse_args()
232    kernel_path = args.kernel_dir / "albedo_kernel.py"
233
234    report = {
235        "kernel": str(kernel_path),
236        "platform": platform.platform(),
237        "python": platform.python_version(),
238        "modules": MODULES,
239        "boot": {},
240        "workloads": [],
241    }
242    workspace = tempfile.mkdtemp(prefix="albedo-kernel-bench-")
243    kernel = Kernel(MODULES, workspace, kernel_path)
244    started = time.perf_counter()
245    kernel.wait_for(lambda f: f.get("type") == "ready")
246    boot_seconds = time.perf_counter() - started
247    time.sleep(0.2)
248    boot_rss = rss_kib(kernel.process.pid)
249    boot_accounting = kernel.accounting()
250    report["boot"] = {
251        "seconds": round(boot_seconds, 3),
252        "rss_kib": boot_rss,
253        **boot_accounting,
254    }
255    print(json.dumps({"boot": report["boot"]}), flush=True)
256
257    for name, workload in WORKLOADS.items():
258        if args.only is not None and name not in args.only:
259            continue
260        for scale in args.scale:
261            fresh = Kernel(MODULES, workspace, kernel_path)
262            fresh.wait_for(lambda f: f.get("type") == "ready")
263            time.sleep(0.1)
264            before = rss_kib(fresh.process.pid)
265            started = time.perf_counter()
266            workload(fresh, scale, workspace)
267            seconds = time.perf_counter() - started
268            after = rss_kib(fresh.process.pid)
269            accounting = fresh.accounting()
270            settled = rss_kib(fresh.process.pid)
271            fresh.close()
272            row = {
273                "workload": name,
274                "scale": scale,
275                "seconds": round(seconds, 3),
276                "rss_before_kib": before,
277                "rss_after_kib": after,
278                "rss_settled_kib": settled,
279                "rss_delta_kib": after - before,
280                **accounting,
281            }
282            report["workloads"].append(row)
283            print(json.dumps(row), flush=True)
284    kernel.close()
285    if args.output is not None:
286        args.output.write_text(json.dumps(report, indent=2) + "\n")
287
288
289if __name__ == "__main__":
290    main()