feat: add managed tool gateway and Nous subscription support
- add managed modal and gateway-backed tool integrations\n- improve CLI setup, auth, and configuration for subscriber flows\n- expand tests and docs for managed tool support
This commit is contained in:
282
tools/environments/managed_modal.py
Normal file
282
tools/environments/managed_modal.py
Normal file
@@ -0,0 +1,282 @@
|
||||
"""Managed Modal environment backed by tool-gateway."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import requests
|
||||
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__)
|
||||
|
||||
|
||||
def _request_timeout_env(name: str, default: float) -> float:
|
||||
try:
|
||||
value = float(os.getenv(name, str(default)))
|
||||
return value if value > 0 else default
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
class ManagedModalEnvironment(BaseEnvironment):
|
||||
"""Gateway-owned Modal sandbox with Hermes-compatible execute/cleanup."""
|
||||
|
||||
_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)
|
||||
|
||||
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 = task_id
|
||||
self._persistent = persistent_filesystem
|
||||
self._image = 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) -> dict:
|
||||
exec_command, sudo_stdin = self._prepare_command(command)
|
||||
|
||||
# When a sudo password is present, inject it via a shell-level pipe
|
||||
# (same approach as the direct ModalEnvironment) since the gateway
|
||||
# cannot pipe subprocess stdin directly.
|
||||
if sudo_stdin is not None:
|
||||
import shlex
|
||||
exec_command = (
|
||||
f"printf '%s\\n' {shlex.quote(sudo_stdin.rstrip())} | {exec_command}"
|
||||
)
|
||||
|
||||
exec_cwd = cwd or self.cwd
|
||||
effective_timeout = timeout or self.timeout
|
||||
exec_id = str(uuid.uuid4())
|
||||
payload: Dict[str, Any] = {
|
||||
"execId": exec_id,
|
||||
"command": exec_command,
|
||||
"cwd": exec_cwd,
|
||||
"timeoutMs": int(effective_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,
|
||||
)
|
||||
except Exception as exc:
|
||||
return {
|
||||
"output": f"Managed Modal exec failed: {exc}",
|
||||
"returncode": 1,
|
||||
}
|
||||
|
||||
if response.status_code >= 400:
|
||||
return {
|
||||
"output": self._format_error("Managed Modal exec failed", response),
|
||||
"returncode": 1,
|
||||
}
|
||||
|
||||
body = response.json()
|
||||
status = body.get("status")
|
||||
if status in {"completed", "failed", "cancelled", "timeout"}:
|
||||
return {
|
||||
"output": body.get("output", ""),
|
||||
"returncode": body.get("returncode", 1),
|
||||
}
|
||||
|
||||
if body.get("execId") != exec_id:
|
||||
return {
|
||||
"output": "Managed Modal exec start did not return the expected exec id",
|
||||
"returncode": 1,
|
||||
}
|
||||
|
||||
poll_interval = 0.25
|
||||
deadline = time.monotonic() + effective_timeout + 10
|
||||
|
||||
while time.monotonic() < deadline:
|
||||
if is_interrupted():
|
||||
self._cancel_exec(exec_id)
|
||||
return {
|
||||
"output": "[Command interrupted - Modal sandbox exec cancelled]",
|
||||
"returncode": 130,
|
||||
}
|
||||
|
||||
try:
|
||||
status_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 {
|
||||
"output": f"Managed Modal exec poll failed: {exc}",
|
||||
"returncode": 1,
|
||||
}
|
||||
|
||||
if status_response.status_code == 404:
|
||||
return {
|
||||
"output": "Managed Modal exec not found",
|
||||
"returncode": 1,
|
||||
}
|
||||
|
||||
if status_response.status_code >= 400:
|
||||
return {
|
||||
"output": self._format_error("Managed Modal exec poll failed", status_response),
|
||||
"returncode": 1,
|
||||
}
|
||||
|
||||
status_body = status_response.json()
|
||||
status = status_body.get("status")
|
||||
if status in {"completed", "failed", "cancelled", "timeout"}:
|
||||
return {
|
||||
"output": status_body.get("output", ""),
|
||||
"returncode": status_body.get("returncode", 1),
|
||||
}
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
self._cancel_exec(exec_id)
|
||||
return {
|
||||
"output": f"Managed Modal exec timed out after {effective_timeout}s",
|
||||
"returncode": 124,
|
||||
}
|
||||
|
||||
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:
|
||||
cpu = self._coerce_number(self._sandbox_kwargs.get("cpu"), 1)
|
||||
memory = self._coerce_number(
|
||||
self._sandbox_kwargs.get("memoryMiB", self._sandbox_kwargs.get("memory")),
|
||||
5120,
|
||||
)
|
||||
disk = self._coerce_number(
|
||||
self._sandbox_kwargs.get("ephemeral_disk", self._sandbox_kwargs.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))
|
||||
|
||||
body = response.json()
|
||||
sandbox_id = body.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",
|
||||
}
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
|
||||
return requests.request(
|
||||
method,
|
||||
f"{self._gateway_origin}{path}",
|
||||
headers=headers,
|
||||
json=json,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
@staticmethod
|
||||
def _coerce_number(value: Any, default: float) -> float:
|
||||
try:
|
||||
if value is None:
|
||||
return default
|
||||
return float(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
@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()
|
||||
if text:
|
||||
return f"{prefix}: {text}"
|
||||
return f"{prefix}: HTTP {response.status_code}"
|
||||
@@ -20,6 +20,7 @@ from tools.interrupt import is_interrupted
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SNAPSHOT_STORE = get_hermes_home() / "modal_snapshots.json"
|
||||
_DIRECT_SNAPSHOT_NAMESPACE = "direct"
|
||||
|
||||
|
||||
def _load_snapshots() -> Dict[str, str]:
|
||||
@@ -38,12 +39,72 @@ def _save_snapshots(data: Dict[str, str]) -> None:
|
||||
_SNAPSHOT_STORE.write_text(json.dumps(data, indent=2))
|
||||
|
||||
|
||||
class _AsyncWorker:
|
||||
"""Background thread with its own event loop for async-safe swe-rex calls.
|
||||
def _direct_snapshot_key(task_id: str) -> str:
|
||||
return f"{_DIRECT_SNAPSHOT_NAMESPACE}:{task_id}"
|
||||
|
||||
Allows sync code to submit async coroutines and block for results,
|
||||
even when called from inside another running event loop (e.g. Atropos).
|
||||
"""
|
||||
|
||||
def _get_snapshot_restore_candidate(task_id: str) -> tuple[str | None, bool]:
|
||||
"""Return a snapshot id for direct Modal restore and whether the key is legacy."""
|
||||
snapshots = _load_snapshots()
|
||||
|
||||
namespaced_key = _direct_snapshot_key(task_id)
|
||||
snapshot_id = snapshots.get(namespaced_key)
|
||||
if isinstance(snapshot_id, str) and snapshot_id:
|
||||
return snapshot_id, False
|
||||
|
||||
legacy_snapshot_id = snapshots.get(task_id)
|
||||
if isinstance(legacy_snapshot_id, str) and legacy_snapshot_id:
|
||||
return legacy_snapshot_id, True
|
||||
|
||||
return None, False
|
||||
|
||||
|
||||
def _store_direct_snapshot(task_id: str, snapshot_id: str) -> None:
|
||||
"""Persist the direct Modal snapshot id under the direct namespace."""
|
||||
snapshots = _load_snapshots()
|
||||
snapshots[_direct_snapshot_key(task_id)] = snapshot_id
|
||||
snapshots.pop(task_id, None)
|
||||
_save_snapshots(snapshots)
|
||||
|
||||
|
||||
def _delete_direct_snapshot(task_id: str, snapshot_id: str | None = None) -> None:
|
||||
"""Remove direct Modal snapshot entries for a task, including legacy keys."""
|
||||
snapshots = _load_snapshots()
|
||||
updated = False
|
||||
|
||||
for key in (_direct_snapshot_key(task_id), task_id):
|
||||
value = snapshots.get(key)
|
||||
if value is None:
|
||||
continue
|
||||
if snapshot_id is None or value == snapshot_id:
|
||||
snapshots.pop(key, None)
|
||||
updated = True
|
||||
|
||||
if updated:
|
||||
_save_snapshots(snapshots)
|
||||
|
||||
|
||||
def _resolve_modal_image(image_spec: Any) -> Any:
|
||||
"""Convert registry references or snapshot ids into Modal image objects."""
|
||||
import modal as _modal
|
||||
|
||||
if not isinstance(image_spec, str):
|
||||
return image_spec
|
||||
|
||||
if image_spec.startswith("im-"):
|
||||
return _modal.Image.from_id(image_spec)
|
||||
|
||||
return _modal.Image.from_registry(
|
||||
image_spec,
|
||||
setup_dockerfile_commands=[
|
||||
"RUN rm -rf /usr/local/lib/python*/site-packages/pip* 2>/dev/null; "
|
||||
"python -m ensurepip --upgrade --default-pip 2>/dev/null || true",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class _AsyncWorker:
|
||||
"""Background thread with its own event loop for async-safe swe-rex calls."""
|
||||
|
||||
def __init__(self):
|
||||
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
||||
@@ -101,42 +162,20 @@ class ModalEnvironment(BaseEnvironment):
|
||||
|
||||
sandbox_kwargs = dict(modal_sandbox_kwargs or {})
|
||||
|
||||
# If persistent, try to restore from a previous snapshot
|
||||
restored_image = None
|
||||
restored_snapshot_id = None
|
||||
restored_from_legacy_key = False
|
||||
if self._persistent:
|
||||
snapshot_id = _load_snapshots().get(self._task_id)
|
||||
if snapshot_id:
|
||||
try:
|
||||
import modal
|
||||
restored_image = modal.Image.from_id(snapshot_id)
|
||||
logger.info("Modal: restoring from snapshot %s", snapshot_id[:20])
|
||||
except Exception as e:
|
||||
logger.warning("Modal: failed to restore snapshot, using base image: %s", e)
|
||||
restored_image = None
|
||||
restored_snapshot_id, restored_from_legacy_key = _get_snapshot_restore_candidate(self._task_id)
|
||||
if restored_snapshot_id:
|
||||
logger.info("Modal: restoring from snapshot %s", restored_snapshot_id[:20])
|
||||
|
||||
effective_image = restored_image if restored_image else image
|
||||
|
||||
# Pre-build a modal.Image with pip fix for Modal's legacy image builder.
|
||||
# Some task images have broken pip; fix via ensurepip before Modal uses it.
|
||||
import modal as _modal
|
||||
if isinstance(effective_image, str):
|
||||
effective_image = _modal.Image.from_registry(
|
||||
effective_image,
|
||||
setup_dockerfile_commands=[
|
||||
"RUN rm -rf /usr/local/lib/python*/site-packages/pip* 2>/dev/null; "
|
||||
"python -m ensurepip --upgrade --default-pip 2>/dev/null || true",
|
||||
],
|
||||
)
|
||||
|
||||
# Start the async worker thread and create the deployment on it
|
||||
# so all gRPC channels are bound to the worker's event loop.
|
||||
self._worker.start()
|
||||
|
||||
from swerex.deployment.modal import ModalDeployment
|
||||
|
||||
async def _create_and_start():
|
||||
async def _create_and_start(image_spec: Any):
|
||||
deployment = ModalDeployment(
|
||||
image=effective_image,
|
||||
image=image_spec,
|
||||
startup_timeout=180.0,
|
||||
runtime_timeout=3600.0,
|
||||
deployment_timeout=3600.0,
|
||||
@@ -146,7 +185,30 @@ class ModalEnvironment(BaseEnvironment):
|
||||
await deployment.start()
|
||||
return deployment
|
||||
|
||||
self._deployment = self._worker.run_coroutine(_create_and_start())
|
||||
try:
|
||||
target_image_spec = restored_snapshot_id or image
|
||||
try:
|
||||
effective_image = _resolve_modal_image(target_image_spec)
|
||||
self._deployment = self._worker.run_coroutine(_create_and_start(effective_image))
|
||||
except Exception as exc:
|
||||
if not restored_snapshot_id:
|
||||
raise
|
||||
|
||||
logger.warning(
|
||||
"Modal: failed to restore snapshot %s, retrying with base image: %s",
|
||||
restored_snapshot_id[:20],
|
||||
exc,
|
||||
)
|
||||
_delete_direct_snapshot(self._task_id, restored_snapshot_id)
|
||||
base_image = _resolve_modal_image(image)
|
||||
self._deployment = self._worker.run_coroutine(_create_and_start(base_image))
|
||||
else:
|
||||
if restored_snapshot_id and restored_from_legacy_key:
|
||||
_store_direct_snapshot(self._task_id, restored_snapshot_id)
|
||||
logger.info("Modal: migrated legacy snapshot entry for task %s", self._task_id)
|
||||
except Exception:
|
||||
self._worker.stop()
|
||||
raise
|
||||
|
||||
def execute(self, command: str, cwd: str = "", *,
|
||||
timeout: int | None = None,
|
||||
@@ -160,7 +222,7 @@ class ModalEnvironment(BaseEnvironment):
|
||||
exec_command, sudo_stdin = self._prepare_command(command)
|
||||
|
||||
# Modal sandboxes execute commands via the Modal SDK and cannot pipe
|
||||
# subprocess stdin directly the way a local Popen can. When a sudo
|
||||
# subprocess stdin directly the way a local Popen can. When a sudo
|
||||
# password is present, use a shell-level pipe from printf so that the
|
||||
# password feeds sudo -S without appearing as an echo argument embedded
|
||||
# in the shell string.
|
||||
@@ -175,7 +237,6 @@ class ModalEnvironment(BaseEnvironment):
|
||||
effective_cwd = cwd or self.cwd
|
||||
effective_timeout = timeout or self.timeout
|
||||
|
||||
# Run in a background thread so we can poll for interrupts
|
||||
result_holder = {"value": None, "error": None}
|
||||
|
||||
def _run():
|
||||
@@ -191,6 +252,7 @@ class ModalEnvironment(BaseEnvironment):
|
||||
merge_output_streams=True,
|
||||
)
|
||||
)
|
||||
|
||||
output = self._worker.run_coroutine(_do_execute())
|
||||
result_holder["value"] = {
|
||||
"output": output.stdout,
|
||||
@@ -227,7 +289,7 @@ class ModalEnvironment(BaseEnvironment):
|
||||
|
||||
if self._persistent:
|
||||
try:
|
||||
sandbox = getattr(self._deployment, '_sandbox', None)
|
||||
sandbox = getattr(self._deployment, "_sandbox", None)
|
||||
if sandbox:
|
||||
async def _snapshot():
|
||||
img = await sandbox.snapshot_filesystem.aio()
|
||||
@@ -239,11 +301,12 @@ class ModalEnvironment(BaseEnvironment):
|
||||
snapshot_id = None
|
||||
|
||||
if snapshot_id:
|
||||
snapshots = _load_snapshots()
|
||||
snapshots[self._task_id] = snapshot_id
|
||||
_save_snapshots(snapshots)
|
||||
logger.info("Modal: saved filesystem snapshot %s for task %s",
|
||||
snapshot_id[:20], self._task_id)
|
||||
_store_direct_snapshot(self._task_id, snapshot_id)
|
||||
logger.info(
|
||||
"Modal: saved filesystem snapshot %s for task %s",
|
||||
snapshot_id[:20],
|
||||
self._task_id,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Modal: filesystem snapshot failed: %s", e)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user