"""A session's kernel outlives its daemon and its connection. The kernel runs detached and the daemon connects to its socket directly; it reattaches after a restart (graceful or crash) and after the connection drops mid-cell, so the namespace, the kernel process, and an in-flight cell's result all survive. A kernel found dead after a restart has its jobs ended. """ import ast import json import os import signal import sqlite3 import subprocess from contextlib import closing import time import unittest import harness from harness import alive from harness import Albedo, Provider, exclusive, no_spares, python, text def bridges(home): """Pids of the bridges whose kernel lives under this daemon's home.""" listing = subprocess.run( ["ps", "-axo", "pid=,command="], capture_output=True, text=True, check=True ).stdout return [ int(line.split(None, 1)[0]) for line in listing.splitlines() if "albedo_bridge.py" in line and str(home) in line ] # A live job keeps an idle kernel from being released between turns, as a # user's running work does. HOLD = "hold = run('sleep', '120')\n" # exclusive: restarts the daemon @exclusive class KernelDetachTests(unittest.TestCase): def setUp(self): self.cells = [] self.cell_timeout_ms = 60_000 def script(request): if request["input"][-1].get("role") == "user" and self.cells: return python(self.cells.pop(0), timeout_ms=self.cell_timeout_ms) return text("ok") self.provider = Provider(script) self.addCleanup(self.provider.close) self.app = Albedo( self.provider, protocol="responses", prepare=lambda app: app.env.update( ALBEDO_TEST_TIME_SCALE="5", ALBEDO_IDLE_SECONDS="600" ), ) self.app.__enter__() self.addCleanup(self.app.__exit__, None, None, None) self.session = self.app.session() def cell(self, code): """Run one cell through a model turn and return its tool result.""" self.cells.append(code) self.app.prompt(self.session, "run it").close() self.app.idle(self.session) results = [ json.loads(part["value"]) for entry in self.app.history(self.session)["items"] if entry["kind"] == "tool_result" and entry["tool"]["name"] == "python" for part in entry["content"] if part["kind"] == "json" and part["field"] == "result" ] return results[-1] @exclusive def test_variables_survive_graceful_and_crash_restarts(self): first = self.cell(HOLD + "import os\nsurvivor = 41\nos.getpid()") self.app.restart() after_stop = self.cell("survivor += 1\n(survivor, os.getpid())") self.assertEqual(after_stop["value"], f"(42, {first['value']})") self.app.restart(crash=True) after_crash = self.cell("survivor += 1\n(survivor, os.getpid())") self.assertEqual(after_crash["value"], f"(43, {first['value']})") notes = [ part["text"] for entry in self.app.history(self.session)["items"] if entry["kind"] == "note" for part in entry["content"] if part["kind"] == "text" ] self.assertFalse([note for note in notes if "variables are gone" in note]) def test_stopping_with_kernels_ends_them_instead_of_leaving_them(self): kernel = int(self.cell("import os\nsurvivor = 41\nos.getpid()")["value"]) self.app.cli("daemon", "--stop", "--kernels") harness.wait_until( lambda: not alive(kernel), 15, "the kernel outlived daemon --stop --kernels" ) self.app.restart() # A fresh kernel, with the variables the closing session saved. after = self.cell("survivor += 1\n(survivor, os.getpid())") survivor, pid = ast.literal_eval(after["value"]) self.assertEqual(survivor, 42) self.assertNotEqual(pid, kernel) @exclusive def test_a_daemon_booted_on_a_copy_of_the_database_leaves_its_kernels_alone(self): """The copy carries this daemon's kernel records (socket, token, process group); its daemon forgets them instead of attaching.""" first = self.cell(HOLD + "import os\nsurvivor = 41\nos.getpid()") copy = harness.Daemon("copy") with ( closing(sqlite3.connect(self.app.home / "albedo.sqlite")) as source, closing(sqlite3.connect(copy.home / "albedo.sqlite")) as target, ): source.backup(target) copy.boot() self.addCleanup(copy._stop, 15) with closing(sqlite3.connect(copy.home / "albedo.sqlite")) as db: (left,) = db.execute("SELECT count(*) FROM kernel_links").fetchone() self.assertEqual(left, 0) self.assertIn( "kernel records: left 1 kernels of another daemon's home alone", (copy.home / "daemon.log").read_text(), ) after = self.cell("survivor += 1\n(survivor, os.getpid())") self.assertEqual(after["value"], f"(42, {first['value']})") @exclusive def test_a_connection_dropped_mid_cell_still_delivers_the_result(self): first = self.cell(HOLD + "import os\nos.getpid()") self.assertEqual(bridges(self.app.home), [], "a local kernel kept a bridge") # The kernel drops the daemon's connection from its side a second in. slept = self.cell( "import socket, sys, threading\n" "link = sys.modules['__main__'].LINK\n" "threading.Timer(1, link.connection.shutdown, [socket.SHUT_RDWR]).start()\n" "threading.Event().wait(3)\nprint('woke')\nos.getpid()" ) self.assertEqual(slept["status"], "ok", slept) self.assertEqual(slept["output"].strip(), "woke") self.assertEqual(slept["value"], first["value"]) self.assertEqual(bridges(self.app.home), []) def test_concurrent_jobs_keep_running_and_deadlines_count_wall_time(self): # Legacy admission settings must not throttle new kernels. self.app.daemon.env.update( ALBEDO_MAX_LOCAL_JOBS="1", ALBEDO_JOB_GRACE_SECONDS="0.2" ) self.app.restart() program = "import time; time.sleep(5.5); print('progress', flush=True); time.sleep(60)" # The three jobs each use a seven-second wall timeout; give the cell # sixty real seconds after its fivefold timer scaling for spawn and drain. self.cell_timeout_ms = 300_000 result = self.cell( "import sys\n" f"workers = [run(sys.executable, '-c', {program!r}, timeout=7) for _ in range(3)]\n" "for worker in workers:\n await worker\n" "[(worker.timed_out, worker.duration, worker.tail().startswith('progress')) for worker in workers]" ) self.assertEqual(result["status"], "ok", result) outcomes = ast.literal_eval(result["value"]) self.assertEqual(len(outcomes), 3) for timed_out, duration, progressed in outcomes: self.assertTrue(timed_out) self.assertTrue(progressed) self.assertGreaterEqual(duration, 7) self.assertLess(duration, 10) def last_prompt(self): """The newest model request's user text: where a notice would ride.""" request = next( record["request"] for record in reversed(self.provider.requests) if any( item.get("role") == "user" for item in record["request"].get("input", []) ) ) return [ item["content"] for item in request["input"] if item.get("role") == "user" ][-1] def test_a_restart_tells_the_model_what_the_kernel_is_still_running(self): started = self.cell( "import asyncio, os\nj = run('sleep', '120')\n" "await asyncio.sleep(0.3)\n(os.getpid(), j.id)" ) kernel, job = ( part.strip(" '\"") for part in started["value"].strip("()").split(",") ) self.app.restart() # The next turn runs on the reattached kernel, and its request says # the kernel carried on and names the job still running, so a live # handle is not mistaken for one the restart killed. self.cells.append("import os\nos.getpid()") self.app.prompt(self.session, "run it").close() self.app.idle(self.session) prompt = self.last_prompt() self.assertIn("The python kernel reattached", prompt) self.assertIn("Jobs still running: " + job, prompt) self.assertIn("sleep 120", prompt) self.assertEqual(self.cell("import os\nos.getpid()")["value"], str(int(kernel))) # The notice rides the first request after the reattach, once. self.cells.append("import os\nos.getpid()") self.app.prompt(self.session, "run it").close() self.app.idle(self.session) self.assertNotIn("The python kernel reattached", self.last_prompt()) def test_a_job_ending_after_a_restart_wakes_its_idle_session(self): started = self.cell("j = run('sleep', '4')\nj.id") job = started["value"].strip("'\"") requests = len(self.provider.requests) # The restart leaves the session idle and unloaded; only its kernel # comes back, and the job's end must still reach the model. self.app.restart() deadline = time.monotonic() + 30 while len(self.provider.requests) == requests and time.monotonic() < deadline: time.sleep(0.1) self.assertGreater(len(self.provider.requests), requests, "no wake") self.app.idle(self.session) self.assertIn(job, self.last_prompt()) def test_a_kernel_killed_while_the_daemon_was_down_has_its_jobs_ended(self): started = self.cell( "import asyncio, os\nj = run('sleep', '300')\nawait asyncio.sleep(0.5)\n(os.getpid(), j.process.pid, j.id)" ) kernel, job, handle = ast.literal_eval(started["value"]) # Graceful stop leaves the kernel running; it dies hard before the # next daemon reaches it, and its job's group outlives it. self.app.restart(prepare=lambda _: os.kill(kernel, signal.SIGKILL)) deadline = time.monotonic() + 15 while alive(job) and time.monotonic() < deadline: time.sleep(0.1) self.assertFalse(alive(job), "the dead kernel's job is still running") # The new kernel's first request tells the model what died with the # old one, from the job record the daemon kept. self.cells.append("import os\nos.getpid()") self.app.prompt(self.session, "run it").close() self.app.idle(self.session) prompt = self.last_prompt() self.assertIn("The python kernel restarted", prompt) self.assertIn(f"Jobs stopped by the restart: {handle} (sleep 300)", prompt) # no spares: with a pool, an idle kernel with nothing running is released # before the restart, so only a machine without spares reattaches one @no_spares def test_a_restart_tells_the_model_when_nothing_was_running(self): self.cell("survivor = 1") self.app.restart() self.cells.append("survivor") self.app.prompt(self.session, "run it").close() self.app.idle(self.session) self.assertIn( "The python kernel reattached; its variables and jobs carried on." " Nothing was running.", self.last_prompt(), ) if __name__ == "__main__": unittest.main()