Something went wrong. Try again.
small gleam coding and persistent agent daemon with a detachable cli and a canvas ui forked from okami.mom/albedo
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272"""A session's kernel outlives its daemon and its connection.
The kernel runs detached and the daemon connects to its socket directly; itreattaches after a restart (graceful or crash) and after the connection dropsmid-cell, so the namespace, the kernel process, and an in-flight cell's resultall survive. A kernel found dead after a restart has its jobs ended."""
import astimport jsonimport osimport signalimport sqlite3import subprocessfrom contextlib import closingimport timeimport unittest
import harnessfrom harness import alivefrom 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@exclusiveclass 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("<system-note>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("<system-note>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()