349 lines
15 KiB
Python
349 lines
15 KiB
Python
"""Session-persistent kernels for REMOTE terminal backends (docker/ssh/modal).
|
|
|
|
Remote backends offer one primitive — ``env.execute(cmd)``, run-to-completion
|
|
— so the three things the local kernel gets from owning a child are rebuilt:
|
|
a detached runner (``nohup ... &``, PID recorded, ``kill -0`` probed per cell);
|
|
a file-based CELL protocol in the kernel dir (``cell_req_NNNNNN.json`` /
|
|
``cell_res_NNNNNN.json``), sibling to the unchanged file-based TOOL-RPC protocol
|
|
(req_/res_) whose host-side ``_rpc_poll_loop`` starts per cell with the calling
|
|
thread's context (= per-cell tool authority); and death detection — a failed
|
|
liveness probe reads as *kernel died: state lost* and the next call respawns,
|
|
never a hung poll (every wait is bounded by the cell timeout).
|
|
|
|
Same invariants as local: owner = approval session key with the ``::child::``
|
|
qualifier (one resolver in tools.code_kernel), same generated tool stubs, same
|
|
output post-processing in the caller. ``reset=true`` kills and respawns. Spawn
|
|
failure fails OPEN to the per-call path with a note.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import atexit
|
|
import json
|
|
import logging
|
|
import secrets
|
|
import shlex
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
from tools.code_kernel import RUNNER_CELL_SOURCE, KernelRegistry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Host poll interval for a cell result file; each poll is one env.execute round-trip
|
|
# (0.1-0.4s on ssh/docker), so this is a floor, not a rate.
|
|
_CELL_POLL_INTERVAL = 0.5
|
|
|
|
# The remote runner: a forever-loop that polls for cell request files, execs them in one
|
|
# persistent namespace, writes response files. Pure files + stdlib only (transport-agnostic);
|
|
# cells and tool-RPC share the kernel dir under distinct prefixes.
|
|
REMOTE_KERNEL_RUNNER_SOURCE = '''\
|
|
"""Auto-generated Hermes REMOTE session-kernel runner (file cell protocol)."""
|
|
import contextlib
|
|
import io
|
|
import json
|
|
import os
|
|
import sys
|
|
import time
|
|
import traceback
|
|
|
|
KDIR = os.environ["HERMES_KERNEL_DIR"]
|
|
CELLS = os.path.join(KDIR, "cells")
|
|
_CAPTURE_LIMIT = {capture_limit}
|
|
IDLE_EXIT_SECONDS = {idle_exit}
|
|
|
|
{cell_source}
|
|
|
|
def main():
|
|
execution_count = 0
|
|
last_activity = time.time()
|
|
while True:
|
|
pending = sorted(
|
|
f for f in os.listdir(CELLS)
|
|
if f.startswith("cell_req_") and f.endswith(".json")
|
|
)
|
|
if not pending:
|
|
if time.time() - last_activity > IDLE_EXIT_SECONDS:
|
|
return # self-reap: nobody is talking to us anymore
|
|
time.sleep(0.2)
|
|
continue
|
|
for name in pending:
|
|
req_path = os.path.join(CELLS, name)
|
|
try:
|
|
with open(req_path, "r", encoding="utf-8") as f:
|
|
request = json.load(f)
|
|
except Exception:
|
|
# Partially-written request (ship in progress): retry next tick.
|
|
continue
|
|
os.remove(req_path)
|
|
last_activity = time.time()
|
|
execution_count += 1
|
|
payload, _ = run_cell(request, execution_count)
|
|
res_name = name.replace("cell_req_", "cell_res_")
|
|
tmp = os.path.join(CELLS, res_name + ".tmp")
|
|
with open(tmp, "w", encoding="utf-8") as f:
|
|
json.dump(payload, f, ensure_ascii=False)
|
|
os.replace(tmp, os.path.join(CELLS, res_name))
|
|
if payload["status"] == "exit":
|
|
return
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
'''
|
|
|
|
|
|
def _sh(env, cmd: str, timeout: int = 15) -> str:
|
|
"""Run *cmd* on the remote from ``/`` and return its output text."""
|
|
result = env.execute(cmd, cwd="/", timeout=timeout)
|
|
return (result.get("output", "") if isinstance(result, dict) else "") or ""
|
|
|
|
|
|
@dataclass
|
|
class RemoteKernel:
|
|
"""Host-side record of one detached remote kernel process."""
|
|
|
|
env: Any
|
|
env_type: str
|
|
kernel_dir: str
|
|
pid: str
|
|
rpc_token: str
|
|
owner: str
|
|
last_used: float = field(default_factory=time.monotonic)
|
|
execution_count: int = 0
|
|
cell_seq: int = 0
|
|
|
|
def sh(self, cmd: str, timeout: int = 15) -> str:
|
|
return _sh(self.env, cmd, timeout)
|
|
|
|
def is_alive(self) -> bool:
|
|
"""Bounded liveness probe: kill -0 through the transport. Any transport
|
|
failure counts as dead — a dropped ssh connection and a dead runner are
|
|
indistinguishable from here, and both have the same correct answer (respawn)."""
|
|
try:
|
|
return "ALIVE" in self.sh(f"kill -0 {shlex.quote(self.pid)} 2>/dev/null && echo ALIVE")
|
|
except Exception:
|
|
return False
|
|
|
|
def kill(self) -> None:
|
|
"""Best-effort kill of the runner and its subprocesses, then rm -rf."""
|
|
q_pid = shlex.quote(self.pid)
|
|
for cmd, failure in (
|
|
# Kill the runner's children if the shell gave it a group, then the PID itself.
|
|
(f"pkill -TERM -P {q_pid} 2>/dev/null; kill {q_pid} 2>/dev/null; true",
|
|
"remote kernel kill failed (transport?)"),
|
|
(f"rm -rf {shlex.quote(self.kernel_dir)}", "remote kernel dir cleanup failed"),
|
|
):
|
|
try:
|
|
self.sh(cmd)
|
|
except Exception:
|
|
logger.debug(failure, exc_info=True)
|
|
|
|
|
|
def _kernel_key(owner: str, env_type: str, task_env_id: str) -> Tuple:
|
|
return (owner, "remote", env_type, task_env_id)
|
|
|
|
|
|
# Registry + lock shared-shape with code_kernel; teardown runs outside the lock.
|
|
_REGISTRY = KernelRegistry(lambda kernel: kernel.kill())
|
|
_REMOTE_KERNELS: Dict[Tuple, RemoteKernel] = _REGISTRY.kernels
|
|
|
|
|
|
def shutdown_all_remote_kernels() -> None:
|
|
_REGISTRY.shutdown()
|
|
|
|
|
|
def shutdown_remote_kernels_for_owner(owner: str) -> None:
|
|
"""Session-boundary disposal — wired to the same clear_session hook as
|
|
local kernels, so /new and session close reap both kinds."""
|
|
if owner:
|
|
_REGISTRY.shutdown(owner)
|
|
|
|
|
|
atexit.register(shutdown_all_remote_kernels)
|
|
|
|
|
|
def _spawn_remote_kernel(env, env_type: str, owner: str, task_env_id: str,
|
|
sandbox_tools: frozenset, *, idle_exit: int) -> Optional[RemoteKernel]:
|
|
"""Start a detached kernel runner on the remote. None on failure (dir removed)."""
|
|
from tools.code_execution_tool import (
|
|
MAX_STDOUT_BYTES, _ship_file_to_remote, _env_temp_dir, generate_hermes_tools_module,
|
|
)
|
|
kernel_dir = f"{_env_temp_dir(env)}/hermes_rkernel_{uuid.uuid4().hex[:12]}"
|
|
q_dir = shlex.quote(kernel_dir)
|
|
kernel = None
|
|
try:
|
|
_sh(env, f"mkdir -p {q_dir}/cells {q_dir}/rpc")
|
|
rpc_token = secrets.token_urlsafe(32)
|
|
_ship_file_to_remote(env, f"{kernel_dir}/kernel_runner.py", REMOTE_KERNEL_RUNNER_SOURCE.format(
|
|
cell_source=RUNNER_CELL_SOURCE, capture_limit=MAX_STDOUT_BYTES, idle_exit=idle_exit))
|
|
_ship_file_to_remote(env, f"{kernel_dir}/hermes_tools.py",
|
|
generate_hermes_tools_module(list(sandbox_tools), transport="file"))
|
|
env_prefix = (f"HERMES_KERNEL_DIR={q_dir} HERMES_RPC_DIR={shlex.quote(kernel_dir + '/rpc')} "
|
|
f"HERMES_RPC_TOKEN={shlex.quote(rpc_token)} PYTHONDONTWRITEBYTECODE=1 PYTHONPATH={q_dir}")
|
|
started = _sh(env, f"cd {q_dir} && nohup env {env_prefix} python3 kernel_runner.py "
|
|
f"> {q_dir}/runner.log 2>&1 & echo PID:$!", timeout=20)
|
|
pid = next((line.strip()[4:].strip() for line in started.splitlines()
|
|
if line.strip().startswith("PID:")), "")
|
|
if not pid.isdigit():
|
|
logger.warning("remote kernel spawn returned no PID: %r", started)
|
|
else:
|
|
candidate = RemoteKernel(env=env, env_type=env_type, kernel_dir=kernel_dir,
|
|
pid=pid, rpc_token=rpc_token, owner=owner)
|
|
if candidate.is_alive():
|
|
kernel = candidate
|
|
else:
|
|
# Died instantly (missing python3 was pre-checked by the caller,
|
|
# so this is unexpected) — surface the runner log.
|
|
try:
|
|
logger.warning("remote kernel died at spawn: %s",
|
|
_sh(env, f"cat {q_dir}/runner.log", timeout=10)[:500])
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
logger.warning("remote kernel spawn failed", exc_info=True)
|
|
if kernel is None:
|
|
try:
|
|
_sh(env, f"rm -rf {q_dir}")
|
|
except Exception:
|
|
pass
|
|
return kernel
|
|
|
|
|
|
def _acquire_remote_kernel(env, env_type: str, owner: str, task_env_id: str,
|
|
sandbox_tools: frozenset, *, reset: bool,
|
|
idle_exit: int) -> Tuple[Optional[RemoteKernel], bool, bool, bool]:
|
|
"""Find/respawn the owner's kernel: (kernel|None, reused, state_reset, state_lost)."""
|
|
key = _kernel_key(owner, env_type, task_env_id)
|
|
state_lost = state_reset = False
|
|
with _REGISTRY.lock:
|
|
kernel = _REMOTE_KERNELS.get(key)
|
|
if kernel is not None and reset:
|
|
_REGISTRY.discard(key, kernel)
|
|
kernel, state_reset = None, True
|
|
if kernel is not None and not kernel.is_alive():
|
|
# Transport drop, container restart, self-reaped on idle, OOM — all
|
|
# the same answer: report the loss, respawn fresh (kill is then only
|
|
# best-effort dir cleanup; the process is already gone).
|
|
_REGISTRY.discard(key, kernel)
|
|
kernel, state_lost = None, True
|
|
reused = kernel is not None
|
|
if kernel is None:
|
|
kernel = _spawn_remote_kernel(env, env_type, owner, task_env_id, sandbox_tools, idle_exit=idle_exit)
|
|
if kernel is not None:
|
|
with _REGISTRY.lock:
|
|
_REMOTE_KERNELS[key] = kernel
|
|
return kernel, reused, state_reset, state_lost
|
|
|
|
|
|
def _run_remote_cell(kernel: RemoteKernel, code: str, timeout: int) -> Tuple[str, Dict[str, Any]]:
|
|
"""Ship one cell request and poll for its result: (cell status, payload)."""
|
|
from tools.code_execution_tool import _ship_file_to_remote
|
|
kernel.cell_seq += 1
|
|
seq = f"{kernel.cell_seq:06d}"
|
|
q_cells = shlex.quote(f"{kernel.kernel_dir}/cells")
|
|
q_res = shlex.quote(f"cell_res_{seq}.json")
|
|
_ship_file_to_remote(kernel.env, f"{kernel.kernel_dir}/cells/cell_req_{seq}.json.tmp",
|
|
json.dumps({"id": seq, "code": code}, ensure_ascii=False))
|
|
kernel.sh(f"mv {q_cells}/cell_req_{seq}.json.tmp {q_cells}/cell_req_{seq}.json", timeout=10)
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
try:
|
|
body = kernel.sh(f"cat {q_cells}/{q_res} 2>/dev/null", timeout=20).strip()
|
|
except Exception:
|
|
# One flaky round-trip is not kernel death; liveness decides.
|
|
body = ""
|
|
if body:
|
|
try:
|
|
payload = json.loads(body)
|
|
status = payload.get("status", "error")
|
|
except ValueError:
|
|
payload, status = {}, "protocol-error"
|
|
kernel.sh(f"rm -f {q_cells}/{q_res}", timeout=10)
|
|
return status, payload
|
|
time.sleep(_CELL_POLL_INTERVAL)
|
|
return "timeout", {}
|
|
|
|
|
|
def execute_in_remote_kernel(
|
|
code: str, *, env, env_type: str, task_env_id: str, sandbox_tools: frozenset,
|
|
timeout: int, max_tool_calls: int, reset: bool, idle_exit: int = 1800,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Run one cell in the owner's remote kernel.
|
|
|
|
Returns the raw cell result dict (caller post-processes output), or ``None`` when no
|
|
kernel could be spawned (caller falls open to per-call). ``state_lost`` / ``state_reset`` /
|
|
``reused`` ride in the ``kernel`` sub-dict, matching the local kernel's result shape.
|
|
"""
|
|
from tools.code_kernel import _resolve_owner
|
|
from tools.code_execution_tool import _rpc_poll_loop
|
|
from tools.thread_context import propagate_context_to_thread
|
|
owner = _resolve_owner(task_env_id)
|
|
kernel, reused, state_reset, state_lost = _acquire_remote_kernel(
|
|
env, env_type, owner, task_env_id, sandbox_tools, reset=reset, idle_exit=idle_exit)
|
|
if kernel is None:
|
|
return None # fail open to per-call
|
|
key = _kernel_key(owner, env_type, task_env_id)
|
|
kernel.last_used = time.monotonic()
|
|
# Clean stale tool-RPC requests from a previous cell before arming this cell's poll loop, so
|
|
# a background thread the last cell leaked cannot smuggle a call into this authority window.
|
|
q_rpc = shlex.quote(kernel.kernel_dir + '/rpc')
|
|
try:
|
|
kernel.sh(f"rm -f {q_rpc}/req_* {q_rpc}/res_*", timeout=10)
|
|
except Exception:
|
|
pass
|
|
tool_call_log: list = []
|
|
tool_call_counter, stop_event = [0], threading.Event()
|
|
# Per-cell RPC thread carrying THIS call's approval/session context — the remote analogue
|
|
# of CellAuthority: authority lives exactly as long as the cell's poll loop.
|
|
rpc_thread = threading.Thread(
|
|
target=propagate_context_to_thread(_rpc_poll_loop),
|
|
args=(env, f"{kernel.kernel_dir}/rpc", task_env_id, tool_call_log, tool_call_counter,
|
|
max_tool_calls, sandbox_tools, stop_event, kernel.rpc_token),
|
|
daemon=True,
|
|
)
|
|
rpc_thread.start()
|
|
cell_status, cell_payload = "no-result", {}
|
|
try:
|
|
cell_status, cell_payload = _run_remote_cell(kernel, code, timeout)
|
|
finally:
|
|
stop_event.set()
|
|
rpc_thread.join(timeout=5)
|
|
kernel_info: Dict[str, Any] = {"reused": reused, "remote": True}
|
|
result: Dict[str, Any] = {
|
|
"status": "error", "stdout": cell_payload.get("stdout", ""),
|
|
"stderr": cell_payload.get("stderr", ""), "traceback": cell_payload.get("traceback", ""),
|
|
"tool_calls_made": tool_call_counter[0], "kernel": kernel_info,
|
|
}
|
|
if cell_status in ("timeout", "protocol-error", "no-result"):
|
|
# No safe way to interrupt one cell in place (same contract as local): kill, report, respawn.
|
|
_REGISTRY.discard(key, kernel)
|
|
if cell_status == "timeout":
|
|
result["status"] = "timeout"
|
|
note = ("Cell timed out; the remote session kernel was killed and "
|
|
"its state was lost. The next call starts a fresh kernel.")
|
|
else:
|
|
note = "Remote kernel protocol failure; kernel killed, state lost."
|
|
kernel_info.update(ended=True, state_lost=True, note=note)
|
|
return result
|
|
if cell_status == "exit":
|
|
_REGISTRY.discard(key, kernel)
|
|
kernel_info["ended"] = True
|
|
kernel.execution_count = int(cell_payload.get("execution_count", 0) or 0)
|
|
kernel_info["execution_count"] = kernel.execution_count
|
|
if cell_status in ("ok", "exit"):
|
|
result["status"] = "success"
|
|
result["stdout_clipped"] = bool(cell_payload.get("stdout_clipped"))
|
|
result["stderr_clipped"] = bool(cell_payload.get("stderr_clipped"))
|
|
if state_reset:
|
|
kernel_info["state_reset"] = True
|
|
if state_lost:
|
|
kernel_info.update(state_lost=True, note=(
|
|
"The previous remote kernel was gone (transport drop, container "
|
|
"restart, or idle self-exit); state from earlier calls was lost and a fresh kernel was started."))
|
|
if cell_status == "error" and result["traceback"]:
|
|
result["error"] = result["traceback"].strip().splitlines()[-1]
|
|
return result
|