refactor(code-execution): one checked-execute helper, complete launch command

Gate follow-ups on the remote lockdown:
- _execute_checked(env, cmd, what, **kw) in code_execution_rpc replaces
  the three copy-pasted "execute, raise if returncode != 0" blocks
  (per-call setup, kernel dir setup, file ship via
  _remote_write(check=True)). env.execute always returns a dict, so the
  isinstance/(r or {}) guards go; the error carries command output only,
  never the payload.
- _ship_env_file_and_launch_prefix returned a half-built "( ... && "
  that both callers had to close; a caller that dropped the ")" or
  composed it differently would lose the load-bearing subshell. It now
  takes the launch command, builds the shared env map (RPC dir, token,
  PYTHONDONTWRITEBYTECODE, routed TZ) itself, and returns the complete
  command; the kernel passes only HERMES_KERNEL_DIR/PYTHONPATH, which
  drops its duplicated TZ block and lazy hermes_time import.
- _private_dirs_cmd(root, *subdirs): every caller spelled each path
  twice for mkdir and chmod.
- _run_remote_cell publishes the cell request with one atomic
  _remote_write instead of ship-to-.tmp then a separate unchecked mv,
  saving a backend round-trip per cell.
This commit is contained in:
kshitijk4poor
2026-09-26 22:40:07 +05:30
committed by kshitij
parent 19cf343c74
commit 1a748cc85a
4 changed files with 72 additions and 72 deletions

View File

@@ -207,11 +207,10 @@ 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 hermes_time import get_timezone_name
from tools.code_execution_rpc import _private_dirs_cmd
from tools.code_execution_rpc import _execute_checked, _private_dirs_cmd
from tools.code_execution_tool import (
MAX_STDOUT_BYTES, _ship_file_to_remote, _env_temp_dir,
_ship_env_file_and_launch_prefix, generate_hermes_tools_module,
_ship_env_file_and_launch, generate_hermes_tools_module,
)
kernel_dir = f"{_env_temp_dir(env)}/hermes_rkernel_{uuid.uuid4().hex[:12]}"
q_dir = shlex.quote(kernel_dir)
@@ -221,37 +220,26 @@ def _spawn_remote_kernel(env, env_type: str, owner: str, task_env_id: str,
# the RPC token (in req files), tool results, and cell code/output.
# Fail closed on setup failure rather than ship secrets into a dir that
# stayed permissive.
setup = env.execute(
_private_dirs_cmd([f"{kernel_dir}/cells", f"{kernel_dir}/rpc"],
[kernel_dir, f"{kernel_dir}/cells", f"{kernel_dir}/rpc"]),
cwd="/", timeout=15)
if not isinstance(setup, dict) or setup.get("returncode", 1) != 0:
raise RuntimeError(
f"remote kernel dir setup failed: {(setup or {}).get('output', setup)!r}")
_execute_checked(env, _private_dirs_cmd(kernel_dir, f"{kernel_dir}/cells",
f"{kernel_dir}/rpc"),
"remote kernel dir setup", timeout=15)
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_map = {"HERMES_KERNEL_DIR": kernel_dir,
"HERMES_RPC_DIR": f"{kernel_dir}/rpc",
"HERMES_RPC_TOKEN": rpc_token,
"PYTHONDONTWRITEBYTECODE": "1",
"PYTHONPATH": kernel_dir}
tz = get_timezone_name() # routed profile's timezone, matching the per-call path
if tz:
env_map["TZ"] = tz
launch_prefix = _ship_env_file_and_launch_prefix(
env, kernel_dir, "kernel.env", env_map)
# kernel.env is removed after sourcing: the runner's env keeps the
# values, so the token file need not sit at rest for the kernel's
# lifetime. runner.log is pre-created 600 so the launch redirect never
# lands at the remote's default umask. The inner `&` stays inside the
# subshell where `$!` resolves to the runner pid.
started = _sh(env, f"{launch_prefix} rm -f ./kernel.env && "
f"touch runner.log && chmod 600 runner.log && "
f"{{ nohup python3 kernel_runner.py > runner.log 2>&1 & "
f'echo "PID:$!"; }} )', timeout=20)
launch_cmd = _ship_env_file_and_launch(
env, kernel_dir, "kernel.env",
"rm -f ./kernel.env && touch runner.log && chmod 600 runner.log && "
'{ nohup python3 kernel_runner.py > runner.log 2>&1 & echo "PID:$!"; }',
rpc_dir=f"{kernel_dir}/rpc", rpc_token=rpc_token,
HERMES_KERNEL_DIR=kernel_dir, PYTHONPATH=kernel_dir)
started = _sh(env, launch_cmd, timeout=20)
pid = next((line.strip()[4:].strip() for line in started.splitlines()
if line.strip().startswith("PID:")), "")
if not pid.isdigit():
@@ -317,9 +305,9 @@ def _run_remote_cell(kernel: RemoteKernel, code: str, timeout: int) -> Tuple[str
kernel.cell_seq += 1
seq = f"{kernel.cell_seq:06d}"
q_cells, q_res = shlex.quote(f"{kernel.kernel_dir}/cells"), 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)
# One round-trip: tmp write + rename publishes the request atomically.
_ship_file_to_remote(kernel.env, f"{kernel.kernel_dir}/cells/cell_req_{seq}.json",
json.dumps({"id": seq, "code": code}, ensure_ascii=False), atomic=True)
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
try: