okami.mom/albedo

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

nekomimi.petconnect to local kernels from erlang instead of keeping a python bridge per kernelce3e329

main
5.9 KiB155 linesraw
1"""Idle kernel reaping persists Python state and explains restoration to the model."""
2
3import sqlite3
4import subprocess
5import time
6import unittest
7
8from harness import alive, latest_user
9from harness import Albedo, Provider, exclusive, python, text
10from stream_support import StreamProbe
11
12
13def notes_for(app, session, cursor):
14    return [
15        event["data"]["text"]
16        for event in app.stream_page(session, cursor)["events"]
17        if event.get("type") == "note"
18    ]
19
20
21# Short enough that the test waits seconds, long enough to outlast a turn.
22IDLE_SECONDS = 2
23
24
25def kernels(home):
26    """Pids of the kernels whose run directory is under this daemon's home."""
27    listing = subprocess.run(
28        ["ps", "-axo", "pid=,command="], capture_output=True, text=True, check=True
29    ).stdout
30    return [
31        int(line.split(None, 1)[0])
32        for line in listing.splitlines()
33        if "albedo_kernel.py" in line and f"{home}/run/" in line
34    ]
35
36
37# exclusive: changes the daemon idle timeout and counts all its kernels
38@exclusive
39class IdleReapTests(unittest.TestCase):
40    def setUp(self):
41        def script(request):
42            if request["input"][-1].get("role") == "user":
43                prompt = latest_user(request)
44                if prompt.startswith("remember"):
45                    return python("answer = 7\nrows = [1, 2, 3]")
46                if prompt.startswith("recall"):
47                    return python("print('recalled', answer, rows)")
48            return text("ok")
49
50        provider = Provider(script)
51        self.addCleanup(provider.close)
52        self.provider = provider
53        self.app = Albedo(
54            provider,
55            protocol="responses",
56            prepare=lambda app: app.env.update(
57                ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS), ALBEDO_INSPECT="1"
58            ),
59        )
60        self.app.__enter__()
61        self.addCleanup(self.app.__exit__, None, None, None)
62
63    def turn(self, session, prompt):
64        self.app.prompt(session, prompt).close()
65        self.app.idle(session)
66
67    def test_idle_kernel_reaps_and_restores_variables_with_notices(self):
68        app = self.app
69        existing = set(kernels(app.home))
70        session = app.session()
71        cursor = app.stream_page(session)
72        self.assertEqual(
73            set(kernels(app.home)) - existing,
74            set(),
75            "session must not eagerly open a kernel",
76        )
77        self.turn(session, "remember this")
78        self.assertEqual(len(set(kernels(app.home)) - existing), 1)
79        deadline = time.monotonic() + IDLE_SECONDS * 4
80        while time.monotonic() < deadline and set(kernels(app.home)) - existing:
81            time.sleep(0.1)
82        self.assertEqual(
83            set(kernels(app.home)) - existing, set(), "idle kernel outlived its limit"
84        )
85        notes = notes_for(app, session, cursor)
86        self.assertIn("released", notes[-1])
87        self.assertIn("2 variables saved to disk", notes[-1])
88        self.assertTrue((app.home / "kernels" / f"{session}.state").exists())
89
90        self.turn(session, "recall them")
91        prompt = latest_user(self.provider.requests[-1]["request"])
92        self.assertTrue(prompt.startswith("recall them<system-note>"), prompt)
93        self.assertIn("restored from disk: answer, rows", prompt)
94        self.assertTrue(
95            any(
96                "recalled 7 [1, 2, 3]" in item.get("output", "")
97                for item in self.provider.requests[-1]["request"]["input"]
98            )
99        )
100        notes = notes_for(app, session, cursor)
101        self.assertTrue(
102            any("restored 2 variables from disk" in note for note in notes), notes
103        )
104
105    def test_reattached_kernel_without_a_loaded_session_is_released(self):
106        app = self.app
107        app.restart(prepare=lambda a: a.env.update(ALBEDO_IDLE_SECONDS="600"))
108        session = app.session()
109        self.turn(session, "remember this")
110        with sqlite3.connect(app.home / "albedo.sqlite") as database:
111            (kernel,) = database.execute(
112                "SELECT pid FROM kernel_links WHERE session=?", (session,)
113            ).fetchone()
114        saved = app.home / "kernels" / f"{session}.state"
115        self.assertFalse(saved.exists())
116        # A crash leaves the kernel waiting; the next daemon reattaches it
117        # without loading the session.
118        app.restart(
119            crash=True,
120            prepare=lambda a: a.env.update(ALBEDO_IDLE_SECONDS=str(IDLE_SECONDS)),
121        )
122        deadline = time.monotonic() + IDLE_SECONDS * 5
123        while time.monotonic() < deadline and alive(kernel):
124            time.sleep(0.1)
125        self.assertFalse(alive(kernel), "reattached kernel outlived its limit")
126        self.assertTrue(saved.exists())
127        self.turn(session, "recall them")
128        self.assertTrue(
129            any(
130                "recalled 7 [1, 2, 3]" in item.get("output", "")
131                for item in self.provider.requests[-1]["request"]["input"]
132            )
133        )
134
135    def test_blocked_sweep_admits_one_worker_and_continues_after_owner_death(self):
136        app = self.app
137        held, survivor = app.session(), app.session()
138        self.turn(held, "remember held values")
139        self.turn(survivor, "remember surviving values")
140        probe = StreamProbe(app)
141        result = probe.call("maintenance_json", f'[<<"{held}">>]')
142        self.assertEqual(result["admitted"], 1)
143        self.assertEqual(result["outcome"], "completed")
144        saved = app.home / "kernels" / f"{survivor}.state"
145        deadline = time.monotonic() + IDLE_SECONDS * 4
146        while not saved.exists() and time.monotonic() < deadline:
147            time.sleep(0.1)
148        self.assertTrue(saved.exists(), "one dead owner prevented survivor release")
149        self.turn(survivor, "recall surviving values")
150        self.assertTrue(
151            any(
152                "recalled 7 [1, 2, 3]" in item.get("output", "")
153                for item in self.provider.requests[-1]["request"]["input"]
154            )
155        )