Teknium
206 lines
10 KiB
Python
206 lines
10 KiB
Python
"""Managed Modal environment backed by tool-gateway.
|
|
|
|
Deliberately overrides :meth:`BaseEnvironment.execute`: the tool-gateway does command
|
|
preparation, CWD tracking and env-snapshot management server-side, so the base
|
|
``_wrap_command`` / ``_wait_for_process`` / snapshot machinery does not apply.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import requests
|
|
import shlex
|
|
import time
|
|
import uuid
|
|
from typing import Any, Dict, Optional
|
|
|
|
from tools.environments.base import BaseEnvironment
|
|
from tools.interrupt import is_interrupted
|
|
from tools.managed_tool_gateway import resolve_managed_tool_gateway
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_TERMINAL_EXEC_STATUSES = frozenset({"completed", "failed", "cancelled", "timeout"})
|
|
_CLIENT_TIMEOUT_GRACE_SECONDS = 10.0
|
|
|
|
|
|
def _coerce_number(value: Any, default: float) -> float:
|
|
try:
|
|
return default if value is None else float(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def _request_timeout_env(name: str, default: float) -> float:
|
|
value = _coerce_number(os.getenv(name), default)
|
|
return value if value > 0 else default
|
|
|
|
|
|
def _result(output: str, returncode: int = 1) -> dict:
|
|
return {"output": output, "returncode": returncode}
|
|
|
|
|
|
class ManagedModalEnvironment(BaseEnvironment):
|
|
"""Gateway-owned Modal sandbox with Hermes-compatible execute/cleanup."""
|
|
|
|
_stdin_mode = "payload"
|
|
_CONNECT_TIMEOUT_SECONDS = _request_timeout_env("TERMINAL_MANAGED_MODAL_CONNECT_TIMEOUT_SECONDS", 1.0)
|
|
_POLL_READ_TIMEOUT_SECONDS = _request_timeout_env("TERMINAL_MANAGED_MODAL_POLL_READ_TIMEOUT_SECONDS", 5.0)
|
|
_CANCEL_READ_TIMEOUT_SECONDS = _request_timeout_env("TERMINAL_MANAGED_MODAL_CANCEL_READ_TIMEOUT_SECONDS", 5.0)
|
|
|
|
def __init__(self, image: str, cwd: str = "/root", timeout: int = 60,
|
|
modal_sandbox_kwargs: Optional[Dict[str, Any]] = None,
|
|
persistent_filesystem: bool = True, task_id: str = "default"):
|
|
super().__init__(cwd=cwd, timeout=timeout)
|
|
# Managed Modal does not sync or mount host credential files.
|
|
try:
|
|
from tools.credential_files import get_credential_file_mounts
|
|
except Exception:
|
|
get_credential_file_mounts = None
|
|
if get_credential_file_mounts is not None and get_credential_file_mounts():
|
|
raise ValueError(
|
|
"Managed Modal does not support host credential-file passthrough. "
|
|
"Use TERMINAL_MODAL_MODE=direct when skills or config require "
|
|
"credential files inside the sandbox."
|
|
)
|
|
gateway = resolve_managed_tool_gateway("modal")
|
|
if gateway is None:
|
|
raise ValueError("Managed Modal requires a configured tool gateway and Nous user token")
|
|
self._gateway_origin = gateway.gateway_origin.rstrip("/")
|
|
self._nous_user_token = gateway.nous_user_token
|
|
self._task_id, self._persistent, self._image = task_id, persistent_filesystem, image
|
|
self._sandbox_kwargs = dict(modal_sandbox_kwargs or {})
|
|
self._create_idempotency_key = str(uuid.uuid4())
|
|
self._sandbox_id = self._create_sandbox()
|
|
|
|
def execute(self, command: str, cwd: str = "", *, timeout: int | None = None, stdin_data: str | None = None,
|
|
rewrite_compound_background: bool = True, bounded_capture: bool = False) -> dict:
|
|
# Signature parity with BaseEnvironment.execute only: the gateway runs commands
|
|
# explicitly (no shell background rewriting) and returns the remote result in one
|
|
# payload, so streaming-time bounding does not apply (the terminal tool's final
|
|
# truncation still caps it).
|
|
del rewrite_compound_background, bounded_capture
|
|
exec_command, sudo_stdin = self._prepare_command(command)
|
|
if sudo_stdin is not None:
|
|
# Feed sudo via a shell pipe: the transport has no direct stdin piping.
|
|
exec_command = f"printf '%s\\n' {shlex.quote(sudo_stdin.rstrip())} | {exec_command}"
|
|
timeout = timeout or self.timeout
|
|
exec_id = str(uuid.uuid4())
|
|
payload: Dict[str, Any] = {"execId": exec_id, "command": exec_command, "cwd": cwd or self.cwd,
|
|
"timeoutMs": int(timeout * 1000)}
|
|
if stdin_data is not None:
|
|
payload["stdinData"] = stdin_data
|
|
try:
|
|
response = self._request("POST", f"/v1/sandboxes/{self._sandbox_id}/execs", json=payload, timeout=10)
|
|
body = response.json() if response.status_code < 400 else None
|
|
except Exception as exc:
|
|
return _result(f"Managed Modal exec failed: {exc}")
|
|
if body is None:
|
|
return _result(self._format_error("Managed Modal exec failed", response))
|
|
if (final := self._result_from_body(body)) is not None:
|
|
return final
|
|
if body.get("execId") != exec_id:
|
|
return _result("Managed Modal exec start did not return the expected exec id")
|
|
deadline = time.monotonic() + timeout + _CLIENT_TIMEOUT_GRACE_SECONDS
|
|
_now = time.monotonic()
|
|
_activity_state = {"last_touch": _now, "start": _now}
|
|
while True:
|
|
if is_interrupted():
|
|
self._cancel_exec(exec_id)
|
|
return _result("[Command interrupted - Modal sandbox exec cancelled]", 130)
|
|
try:
|
|
if (result := self._poll_exec(exec_id)) is not None:
|
|
return result
|
|
except Exception as exc:
|
|
return _result(f"Managed Modal exec failed: {exc}")
|
|
if time.monotonic() >= deadline:
|
|
self._cancel_exec(exec_id)
|
|
return _result(f"Managed Modal exec timed out after {timeout}s", 124)
|
|
# Periodic activity touch so the gateway knows we're alive (lazy import:
|
|
# tests stub tools.environments.base with only BaseEnvironment)
|
|
try:
|
|
from tools.environments.base import touch_activity_if_due
|
|
touch_activity_if_due(_activity_state, "modal command running")
|
|
except Exception:
|
|
pass
|
|
time.sleep(0.25)
|
|
|
|
@staticmethod
|
|
def _result_from_body(body: dict) -> dict | None:
|
|
"""Final result dict if the exec body reports a terminal status, else ``None``."""
|
|
if body.get("status") in _TERMINAL_EXEC_STATUSES:
|
|
return _result(body.get("output", ""), body.get("returncode", 1))
|
|
|
|
def _poll_exec(self, exec_id: str) -> dict | None:
|
|
try:
|
|
response = self._request("GET", f"/v1/sandboxes/{self._sandbox_id}/execs/{exec_id}",
|
|
timeout=(self._CONNECT_TIMEOUT_SECONDS, self._POLL_READ_TIMEOUT_SECONDS))
|
|
except Exception as exc:
|
|
return _result(f"Managed Modal exec poll failed: {exc}")
|
|
if response.status_code == 404:
|
|
return _result("Managed Modal exec not found")
|
|
if response.status_code >= 400:
|
|
return _result(self._format_error("Managed Modal exec poll failed", response))
|
|
return self._result_from_body(response.json())
|
|
|
|
def _cancel_exec(self, exec_id: str) -> None:
|
|
try:
|
|
self._request("POST", f"/v1/sandboxes/{self._sandbox_id}/execs/{exec_id}/cancel",
|
|
timeout=(self._CONNECT_TIMEOUT_SECONDS, self._CANCEL_READ_TIMEOUT_SECONDS))
|
|
except Exception as exc:
|
|
logger.warning("Managed Modal exec cancel failed: %s", exc)
|
|
|
|
def cleanup(self):
|
|
if not getattr(self, "_sandbox_id", None):
|
|
return
|
|
try:
|
|
self._request("POST", f"/v1/sandboxes/{self._sandbox_id}/terminate",
|
|
json={"snapshotBeforeTerminate": self._persistent}, timeout=60)
|
|
except Exception as exc:
|
|
logger.warning("Managed Modal cleanup failed: %s", exc)
|
|
finally:
|
|
self._sandbox_id = None
|
|
|
|
def _create_sandbox(self) -> str:
|
|
kw = self._sandbox_kwargs
|
|
cpu = _coerce_number(kw.get("cpu"), 1)
|
|
memory = _coerce_number(kw.get("memoryMiB", kw.get("memory")), 5120)
|
|
disk = _coerce_number(kw.get("ephemeral_disk", kw.get("diskMiB")), None)
|
|
create_payload = {
|
|
"image": self._image, "cwd": self.cwd, "cpu": cpu, "memoryMiB": memory, "timeoutMs": 3_600_000,
|
|
"idleTimeoutMs": max(300_000, int(self.timeout * 1000)),
|
|
"persistentFilesystem": self._persistent, "logicalKey": self._task_id,
|
|
}
|
|
if disk is not None:
|
|
create_payload["diskMiB"] = disk
|
|
response = self._request("POST", "/v1/sandboxes", json=create_payload, timeout=60,
|
|
extra_headers={"x-idempotency-key": self._create_idempotency_key})
|
|
if response.status_code >= 400:
|
|
raise RuntimeError(self._format_error("Managed Modal create failed", response))
|
|
sandbox_id = response.json().get("id")
|
|
if not isinstance(sandbox_id, str) or not sandbox_id:
|
|
raise RuntimeError("Managed Modal create did not return a sandbox id")
|
|
return sandbox_id
|
|
|
|
def _request(self, method: str, path: str, *, json: Dict[str, Any] | None = None, timeout: int = 30,
|
|
extra_headers: Dict[str, str] | None = None) -> requests.Response:
|
|
headers = {"Authorization": f"Bearer {self._nous_user_token}", "Content-Type": "application/json",
|
|
**(extra_headers or {})}
|
|
return requests.request(method, f"{self._gateway_origin}{path}", headers=headers, json=json, timeout=timeout)
|
|
|
|
@staticmethod
|
|
def _format_error(prefix: str, response: requests.Response) -> str:
|
|
try:
|
|
payload = response.json()
|
|
if isinstance(payload, dict):
|
|
message = payload.get("error") or payload.get("message") or payload.get("code")
|
|
if isinstance(message, str) and message:
|
|
return f"{prefix}: {message}"
|
|
return f"{prefix}: {json.dumps(payload, ensure_ascii=False)}"
|
|
except Exception:
|
|
pass
|
|
text = response.text.strip()
|
|
return f"{prefix}: {text}" if text else f"{prefix}: HTTP {response.status_code}"
|