"""Spawn a Daytona linux-vm sandbox and mount a bucket with AIMem. Daytona *container* sandboxes do not reliably expose `/dev/fuse` — they run without `CAP_MKNOD`, so `aimem-mount` cannot create the device node either (`mknod` returns `EPERM`), and mounts fail on the hosts that lack it. A `linux-vm` sandbox has full kernel/device access, so `/dev/fuse` is always present. VM sandboxes can only be created from a VM snapshot. Daytona's declarative image builder does not work in the `us-east-1` region or for `linux-vm` sandboxes, so the snapshot is built from a stock OS image (`DAYTONA_IMAGE`, default `ubuntu:24.04`) and aimem is installed into the running sandbox by `aimem-setup.sh` instead of baked into an image layer. The setup script fetches the released binary from the public CDN by default; drop a Linux `aimem` binary in `examples/daytona/bin/` to install a local build instead (the path CI uses). The Daytona API key needs `snapshot:create` scope. If the Daytona Python SDK is not already installed, this script creates ~/.cache/daytona-aimem and installs it there before re-running itself. Required env on the host: AIMEM_BUCKET, AIMEM_ENDPOINT, AIMEM_ACCESS_KEY, AIMEM_SECRET_KEY, DAYTONA_API_KEY. Optional: DAYTONA_TARGET Daytona region (default us-east-1). DAYTONA_IMAGE Base OS image for the snapshot (default ubuntu:24.04). DAYTONA_DELETE_ON_EXIT=1 deletes the sandbox after the smoke check. DAYTONA_AUTO_DELETE_INTERVAL Minutes a continuously-stopped sandbox lives before Daytona auto-deletes it (default 0 = delete as soon as it stops). DAYTONA_SNAPSHOT_MAX_AGE_MIN Only sweep leaked snapshots older than this many minutes, so a concurrent run's fresh snapshot is never reclaimed (default 120). The account has a hard total-disk quota, so a leaked sandbox (a killed job that never reached the cleanup path) holds disk until reaped. Two safeguards keep `create` under quota: every sandbox is created with an auto-delete interval so Daytona reaps it after it auto-stops even when this process is killed, and each run first sweeps its own terminal-state leftovers from earlier runs. """ import importlib import importlib.util import os import shlex import subprocess import sys import time import uuid from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Optional HERE = Path(__file__).parent CONTRACT_DIR = HERE.parent / "agent-worker-runtime" sys.path.insert(0, str(CONTRACT_DIR)) from contract import ( # noqa: E402 RemoteCommandWorkspaceTools, emit_trace, run_contract, without_storage_credentials, ) # Not a fixed /tmp path: /tmp is world-writable, so another local user could # pre-create daytona-aimem/bin/python there and have os.execv below run their # code as this user instead of creating a venv. ~/.cache is owned by (and # only writable by) the invoking user. BOOTSTRAP_VENV = Path(os.environ.get("DAYTONA_AIMEM_VENV", str(Path.home() / ".cache" / "daytona-aimem"))) DAYTONA_PACKAGE = os.environ.get("DAYTONA_SDK_PACKAGE", "daytona") def truthy(value: Optional[str]) -> bool: return value is not None and value.lower() in ("1", "true", "yes", "on") def require_env(*names: str) -> None: missing = [name for name in names if not os.environ.get(name)] if missing: sys.exit("missing required env: " + ", ".join(missing)) def daytona_uid(sandbox: object, attempts: int = 3) -> Optional[str]: """Return the sandbox's effective uid, or None if it cannot be determined. None means "unknown", which is deliberately distinct from a known non-root uid: `exec` can fail transiently while a freshly created sandbox is still settling, and an unknown uid must not cost the caller a working sandbox. """ for attempt in range(1, attempts + 1): try: return str(sandbox.process.exec("id -u", timeout=60).result or "").strip() except Exception as exc: # noqa: BLE001 - transient exec failure print( f" uid check failed (attempt {attempt}/{attempts}): {exc}", file=sys.stderr, ) if attempt < attempts: time.sleep(2 * attempt) return None def ensure_daytona_sdk() -> None: if importlib.util.find_spec("daytona") is not None: return if sys.prefix == sys.base_prefix: python = BOOTSTRAP_VENV / "bin" / "python" if not python.exists(): print(f"creating Python venv at {BOOTSTRAP_VENV}...", file=sys.stderr) subprocess.check_call([sys.executable, "-m", "venv", str(BOOTSTRAP_VENV)]) os.execv(str(python), [str(python), *sys.argv]) print(f"installing Daytona Python SDK ({DAYTONA_PACKAGE})...", file=sys.stderr) subprocess.check_call([sys.executable, "-m", "pip", "install", DAYTONA_PACKAGE]) importlib.invalidate_caches() require_env( "DAYTONA_API_KEY", "AIMEM_BUCKET", "AIMEM_ENDPOINT", "AIMEM_ACCESS_KEY", "AIMEM_SECRET_KEY", ) ensure_daytona_sdk() from daytona import ( # noqa: E402 CreateSandboxFromSnapshotParams, CreateSnapshotParams, Daytona, DaytonaConfig, DaytonaError, DaytonaNotFoundError, DaytonaRateLimitError, DaytonaTimeoutError, Resources, SandboxClass, SessionExecuteRequest, ) MOUNTPOINT = os.environ.get("AIMEM_MOUNTPOINT", "/workspace") MOUNT_WAIT_SECS = int(os.environ.get("MOUNT_WAIT_SECS", "30")) TARGET = os.environ.get("DAYTONA_TARGET", "us-east-1") BASE_IMAGE = os.environ.get("DAYTONA_IMAGE", "ubuntu:24.04") # Snapshots are org-global and named; a fresh per-run snapshot means concurrent # runs never collide and stale binaries are never reused. Snapshots are also # region-scoped, so one created in another region cannot be reused here. run_id = os.environ.get("GITHUB_RUN_ID", "local") run_attempt = os.environ.get("GITHUB_RUN_ATTEMPT", "1") SMOKE_PREFIX = "aimem-smoke-" LABEL_KEY = "aimem-smoke" SNAPSHOT_NAME = os.environ.get( "AIMEM_SNAPSHOT_NAME", f"{SMOKE_PREFIX}{run_id}-{run_attempt}-{uuid.uuid4().hex[:8]}" ) AUTO_DELETE_INTERVAL = int(os.environ.get("DAYTONA_AUTO_DELETE_INTERVAL", "0")) mount_env = { "AIMEM_BUCKET": os.environ["AIMEM_BUCKET"], "AIMEM_ENDPOINT": os.environ["AIMEM_ENDPOINT"], "AIMEM_ACCESS_KEY": os.environ["AIMEM_ACCESS_KEY"], "AIMEM_SECRET_KEY": os.environ["AIMEM_SECRET_KEY"], "AIMEM_MOUNTPOINT": MOUNTPOINT, # aimem uses explicit creds; keep it from ever probing the VM's IMDS. "AWS_EC2_METADATA_DISABLED": "true", } for key in ( "AIMEM_SESSION_TOKEN", "AIMEM_AGENT", "AIMEM_PREFIX", "AIMEM_REGION", "AIMEM_LOCAL", "AIMEM_ALLOW_OTHER", "AIMEM_TLS_CA_FILE", "AIMEM_INSECURE_SKIP_TLS_VERIFY", ): if os.environ.get(key): mount_env[key] = os.environ[key] delete_on_exit = truthy(os.environ.get("DAYTONA_DELETE_ON_EXIT")) def _state(obj: object) -> str: st = getattr(obj, "state", None) return str(getattr(st, "value", st) or "").lower() def sweep_stale_resources(client: Daytona) -> None: """Best-effort reclaim of leaked smoke resources from earlier runs. Deletes only our own leftovers (``aimem-smoke-*`` snapshot name or our label): sandboxes only in a terminal state, snapshots only past an age cutoff — so a concurrent run's live resource is never touched. Candidates are collected across the full listing and deleted afterwards (deleting mid-pagination shifts the offset window and skips entries). Loud on purpose: the log is the only signal that the identity match works against the fleet. Never raises; a sweep failure must not fail the smoke check. """ terminal_sandbox = {"stopped", "archived", "error", "build_failed"} stale_sandboxes = [] try: for sb in client.list(): snapshot_name = getattr(sb, "snapshot", "") or "" labels = getattr(sb, "labels", {}) or {} if not (snapshot_name.startswith(SMOKE_PREFIX) or LABEL_KEY in labels): continue st = _state(sb) print(f"sweep: sandbox id={sb.id} snapshot={snapshot_name} state={st}", file=sys.stderr) if st in terminal_sandbox: stale_sandboxes.append(sb) except Exception as exc: print(f"sweep: listing sandboxes failed: {exc}", file=sys.stderr) # Delete only after the full listing: removing entries mid-iteration shifts # the paginated window and skips others. for sb in stale_sandboxes: try: client.delete(sb) print(f"sweep: deleted sandbox id={sb.id}", file=sys.stderr) except Exception as exc: print(f"sweep: sandbox delete failed id={sb.id}: {exc}", file=sys.stderr) # Snapshots are per-run-unique and the finally already deletes each run's # own, so only killed-job residue accumulates. A concurrent run's snapshot # sits in `active` between create and its sandbox coming up, so state can't # tell "leaked" from "in use" — gate on AGE instead: anything younger than # the cutoff may belong to a live run and is left alone. try: max_age_min = int(os.environ.get("DAYTONA_SNAPSHOT_MAX_AGE_MIN", "120")) except ValueError: max_age_min = 120 max_age = timedelta(minutes=max_age_min) now = datetime.now(timezone.utc) stale_snapshots = [] try: page = 1 while True: res = client.snapshot.list(page=page) for sn in getattr(res, "items", []) or []: name = getattr(sn, "name", "") or "" if not name.startswith(SMOKE_PREFIX) or name == SNAPSHOT_NAME: continue created = getattr(sn, "created_at", None) if created is not None and created.tzinfo is None: created = created.replace(tzinfo=timezone.utc) is_stale = created is not None and (now - created) > max_age print( f"sweep: snapshot name={name} state={_state(sn)} " f"created={created} stale={is_stale}", file=sys.stderr, ) if is_stale: stale_snapshots.append(sn) total_pages = int(getattr(res, "total_pages", 1) or 1) if page >= total_pages: break page += 1 except Exception as exc: print(f"sweep: listing snapshots failed: {exc}", file=sys.stderr) for sn in stale_snapshots: name = getattr(sn, "name", "") try: client.snapshot.delete(sn) print(f"sweep: deleted snapshot name={name}", file=sys.stderr) except Exception as exc: print(f"sweep: snapshot delete failed name={name}: {exc}", file=sys.stderr) daytona = Daytona(DaytonaConfig(target=TARGET)) sandbox = None snapshot = None try: # Reclaim leaked disk from earlier runs before consuming more of the quota. sweep_stale_resources(daytona) # Snapshots are immutable and their name is org-global: if one already exists # under this name (a pinned AIMEM_SNAPSHOT_NAME, or a re-run before the prior # cleanup landed), delete it first — you cannot re-push an image over a name. try: existing = daytona.snapshot.get(SNAPSHOT_NAME) except DaytonaNotFoundError: existing = None # the common case for a fresh per-run name if existing is not None: # A real delete failure must surface here, not be masked until create(). print(f"deleting pre-existing snapshot {SNAPSHOT_NAME}...", file=sys.stderr) daytona.snapshot.delete(existing) print(f"creating linux-vm snapshot {SNAPSHOT_NAME} in {TARGET}...", file=sys.stderr) snapshot = daytona.snapshot.create( CreateSnapshotParams( name=SNAPSHOT_NAME, image=BASE_IMAGE, resources=Resources(cpu=2, memory=4, disk=8), sandbox_class=SandboxClass.LINUX_VM, ), timeout=int(os.environ.get("DAYTONA_SNAPSHOT_TIMEOUT", "900")), ) # `snapshot.create` returns once the snapshot is ACTIVE (built), but the image # can still be propagating to the target region's runners — a sandbox created # in that window fails to start with an "internal error". Retry the create so # it self-heals once propagation completes. create_params = CreateSandboxFromSnapshotParams( snapshot=SNAPSHOT_NAME, os_user="root", auto_stop_interval=10, # Reap the sandbox even if this process is killed before the finally: # Daytona auto-deletes it once it has been stopped for this long # (0 = as soon as it auto-stops). Labelled so the sweep can find it. auto_delete_interval=AUTO_DELETE_INTERVAL, labels={LABEL_KEY: f"{run_id}-{run_attempt}"}, ) create_timeout = int(os.environ.get("DAYTONA_CREATE_TIMEOUT", "900")) attempts = int(os.environ.get("DAYTONA_CREATE_ATTEMPTS", "5")) for attempt in range(1, attempts + 1): print(f"creating sandbox (attempt {attempt}/{attempts})...", file=sys.stderr) try: sandbox = daytona.create(create_params, timeout=create_timeout) # `os_user="root"` above is a request, not a guarantee: Daytona has # been observed to start a linux-vm sandbox as the unprivileged # `daytona` user (uid 1001) instead, which cannot install into # /usr/local/bin. The SDK still sends the field, so this is # server-side and silent — verify before using the sandbox rather # than failing later in aimem-setup.sh. uid = daytona_uid(sandbox) if uid is None: # Unknown, not proven unprivileged: keep the sandbox. aimem-setup.sh # escalates with sudo if it turns out not to be root. print( f"sandbox id={sandbox.id}: could not verify uid; keeping it " "(setup escalates with sudo when not root)", file=sys.stderr, ) break if uid != "0": print( f"sandbox id={sandbox.id} started as uid={uid}, not root " "(os_user=root not honored); discarding it", file=sys.stderr, ) try: daytona.delete(sandbox) except Exception as exc: # noqa: BLE001 - best-effort cleanup print(f" discard failed: {exc}", file=sys.stderr) if attempt == attempts: sys.exit( "sandbox never started as root after " f"{attempts} attempts; os_user=root was not honored" ) wait = min(15 * attempt, 60) print(f" retrying in {wait}s...", file=sys.stderr) time.sleep(wait) continue break except DaytonaError as exc: # Typed exceptions where the SDK provides them; a message match only # for the untyped "still propagating" case surfaced as a validation error. msg = str(exc).lower() transient = isinstance( exc, (DaytonaNotFoundError, DaytonaRateLimitError, DaytonaTimeoutError) ) or any(s in msg for s in ("internal error", "failed to start")) if attempt == attempts or not transient: raise wait = min(15 * attempt, 60) print( f"sandbox create failed ({exc}); snapshot may still be propagating " f"to {TARGET} — retrying in {wait}s...", file=sys.stderr, ) time.sleep(wait) print(f"sandbox id={sandbox.id}", file=sys.stderr) # Declarative build is unavailable here, so install aimem into the running # sandbox: upload the setup script + mount wrapper (and a local binary when # one is staged in ./bin), then run the setup as root. print("installing aimem into sandbox...", file=sys.stderr) sandbox.fs.upload_file(str(HERE / "aimem-setup.sh"), "/tmp/aimem-setup.sh") sandbox.fs.upload_file(str(HERE / "aimem-mount"), "/tmp/aimem-mount") staged = HERE / "bin" / "aimem" if staged.exists(): print("using locally staged aimem binary", file=sys.stderr) sandbox.fs.upload_file(str(staged), "/tmp/aimem-staged") setup = sandbox.process.exec("bash /tmp/aimem-setup.sh", timeout=600) print(setup.result, file=sys.stderr) if setup.exit_code != 0: sys.exit(f"aimem-setup.sh failed (exit {setup.exit_code})") # Process sessions don't inherit the sandbox env, so hand the mount its # config through a root-only file instead. The bytes upload keeps the # credentials off any command line; install it 0600 as root, out of # world-writable /tmp. env_file = "".join(f"export {k}={shlex.quote(v)}\n" for k, v in mount_env.items()) sandbox.fs.upload_file(env_file.encode(), "/tmp/aimem.env") stage = sandbox.process.exec("install -m 600 /tmp/aimem.env /root/aimem.env && rm -f /tmp/aimem.env") if stage.exit_code != 0: sys.exit(f"staging mount env failed (exit {stage.exit_code}): {stage.result}") sandbox.process.create_session("aimem-mount") sandbox.process.execute_session_command( "aimem-mount", SessionExecuteRequest( command=". /root/aimem.env; aimem-mount > /tmp/aimem.log 2>&1", run_async=True, ), ) for i in range(MOUNT_WAIT_SECS): r = sandbox.process.exec(f"mountpoint -q {MOUNTPOINT} && echo MOUNTED || echo NOT_YET") if "MOUNTED" in r.result: print(f"mounted after {i+1}s", file=sys.stderr) break time.sleep(1) else: print(sandbox.process.exec("cat /tmp/aimem.log").result, file=sys.stderr) sys.exit(f"mount did not come up within {MOUNT_WAIT_SECS}s") def daytona_command(command: str, timeout: int) -> tuple[int, str, str]: try: result = sandbox.process.exec(command, timeout=timeout) except DaytonaTimeoutError: return 124, "", "sandbox command timed out" except DaytonaError as exc: status = exc.status_code if exc.status_code is not None else "unknown" return 1, "", f"sandbox command failed ({type(exc).__name__}, status={status})" output = result.result or "" stderr = output if result.exit_code != 0 else "" return result.exit_code, output, stderr tools = RemoteCommandWorkspaceTools( str(sandbox.id), daytona_command, workspace=MOUNTPOINT, ) try: # Provisioning is complete. The agent-worker phase receives Daytona # control-plane access but no AIMem/S3 credentials in its environment. with without_storage_credentials(): run_contract(tools, "from-daytona.txt", "hello from daytona\n") finally: emit_trace(tools) print("DAYTONA EXTERNAL WORKER CONTRACT OK") print(f"ready: sandbox id={sandbox.id} mount={MOUNTPOINT}", file=sys.stderr) finally: if sandbox is not None and delete_on_exit: print(f"deleting sandbox id={sandbox.id}", file=sys.stderr) daytona.delete(sandbox) if snapshot is not None: try: daytona.snapshot.delete(snapshot) except Exception as exc: print(f"snapshot cleanup failed: {exc}", file=sys.stderr)