342 lines
14 KiB
Python
342 lines
14 KiB
Python
"""Modal cloud execution environment using the native Modal SDK directly
|
|
(``Sandbox.create()`` + ``Sandbox.exec()``) with persistent snapshots across sessions."""
|
|
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
import itertools
|
|
import logging
|
|
import shlex
|
|
import tarfile
|
|
import threading
|
|
from pathlib import Path
|
|
from typing import Any, Optional
|
|
|
|
from hermes_constants import get_hermes_home
|
|
from tools.environments.base import BaseEnvironment, _ThreadedProcessHandle, _load_json_store, _save_json_store
|
|
from tools.environments.file_sync import (
|
|
FileSyncManager, iter_sync_files, quoted_mkdir_command, quoted_rm_command, unique_parent_dirs,
|
|
)
|
|
from tools.environments.remote_common import bash_argv, ensure_lazy_dep
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_SNAPSHOT_STORE = get_hermes_home() / "modal_snapshots.json"
|
|
_DIRECT_SNAPSHOT_NAMESPACE = "direct"
|
|
|
|
|
|
def _load_snapshots() -> dict:
|
|
return _load_json_store(_SNAPSHOT_STORE)
|
|
|
|
|
|
def _save_snapshots(data: dict) -> None:
|
|
_save_json_store(_SNAPSHOT_STORE, data)
|
|
|
|
|
|
def _direct_snapshot_key(task_id: str) -> str:
|
|
return f"{_DIRECT_SNAPSHOT_NAMESPACE}:{task_id}"
|
|
|
|
|
|
def _get_snapshot_restore_candidate(task_id: str) -> tuple[str | None, bool]:
|
|
"""Return (snapshot_id, from_legacy_key); the namespaced key wins over the legacy bare task id."""
|
|
snapshots = _load_snapshots()
|
|
for key, legacy in ((_direct_snapshot_key(task_id), False), (task_id, True)):
|
|
snapshot_id = snapshots.get(key)
|
|
if isinstance(snapshot_id, str) and snapshot_id:
|
|
return snapshot_id, legacy
|
|
return None, False
|
|
|
|
|
|
def _store_direct_snapshot(task_id: str, snapshot_id: str) -> None:
|
|
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:
|
|
snapshots = _load_snapshots()
|
|
updated = False
|
|
for key in (_direct_snapshot_key(task_id), task_id):
|
|
value = snapshots.get(key)
|
|
if value is not None and (snapshot_id is None or value == snapshot_id):
|
|
snapshots.pop(key, None)
|
|
updated = True
|
|
if updated:
|
|
_save_snapshots(snapshots)
|
|
|
|
|
|
def _ensure_modal_sdk() -> None:
|
|
"""Lazy-install modal on demand. Idempotent — fast no-op once installed."""
|
|
ensure_lazy_dep("terminal.modal")
|
|
|
|
|
|
def _resolve_modal_image(image_spec: Any) -> Any:
|
|
"""Convert registry references or snapshot ids into Modal image objects. Registry images
|
|
get pip repaired (ensurepip) before Modal's bootstrap; ubuntu/debian also get python3."""
|
|
_ensure_modal_sdk()
|
|
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)
|
|
|
|
setup_commands = [
|
|
"RUN rm -rf /usr/local/lib/python*/site-packages/pip* 2>/dev/null; "
|
|
"python -m ensurepip --upgrade --default-pip 2>/dev/null || true",
|
|
]
|
|
if any(base in image_spec.lower() for base in ("ubuntu", "debian")):
|
|
setup_commands.insert(0,
|
|
"RUN apt-get update -qq && apt-get install -y -qq python3 python3-venv > /dev/null 2>&1 || true"
|
|
)
|
|
return _modal.Image.from_registry(image_spec, setup_dockerfile_commands=setup_commands)
|
|
|
|
|
|
async def _stream_stdin(proc, payload: str, chunk_size: int) -> None:
|
|
"""Write ``payload`` to ``proc.stdin`` in ``chunk_size`` pieces, draining after each, then EOF."""
|
|
offset = 0
|
|
while offset < len(payload):
|
|
proc.stdin.write(payload[offset:offset + chunk_size])
|
|
await proc.stdin.drain.aio()
|
|
offset += chunk_size
|
|
proc.stdin.write_eof()
|
|
await proc.stdin.drain.aio()
|
|
|
|
|
|
class _AsyncWorker:
|
|
"""Background thread with its own event loop for async-safe Modal calls."""
|
|
|
|
def __init__(self):
|
|
self._loop: Optional[asyncio.AbstractEventLoop] = None
|
|
self._thread: Optional[threading.Thread] = None
|
|
self._started = threading.Event()
|
|
|
|
def start(self):
|
|
self._thread = threading.Thread(target=self._run_loop, daemon=True)
|
|
self._thread.start()
|
|
self._started.wait(timeout=30)
|
|
|
|
def _run_loop(self):
|
|
self._loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(self._loop)
|
|
self._started.set()
|
|
self._loop.run_forever()
|
|
|
|
def run_coroutine(self, coro, timeout=600):
|
|
from agent.async_utils import safe_schedule_threadsafe
|
|
if self._loop is None or self._loop.is_closed():
|
|
if asyncio.iscoroutine(coro):
|
|
coro.close()
|
|
raise RuntimeError("AsyncWorker loop is not running")
|
|
future = safe_schedule_threadsafe(coro, self._loop)
|
|
if future is None:
|
|
raise RuntimeError("AsyncWorker loop is not running")
|
|
return future.result(timeout=timeout)
|
|
|
|
def stop(self):
|
|
if self._loop and self._loop.is_running():
|
|
self._loop.call_soon_threadsafe(self._loop.stop)
|
|
if self._thread:
|
|
self._thread.join(timeout=10)
|
|
|
|
|
|
class ModalEnvironment(BaseEnvironment):
|
|
"""Modal cloud execution via native Modal sandboxes: spawn-per-call via _ThreadedProcessHandle
|
|
wrapping async SDK calls, cancel_fn wired to sandbox.terminate for interrupt support."""
|
|
|
|
_stdin_mode = "heredoc"
|
|
_snapshot_timeout = 60 # Modal cold starts can be slow
|
|
# Modal SDK stdin buffer limit: the command-router path allows 16 MB but the
|
|
# legacy server path caps at 2 MB, so chunks stay under 2 MB and each is
|
|
# flushed individually via drain().
|
|
_STDIN_CHUNK_SIZE = 1 * 1024 * 1024 # 1 MB — safe for both transport paths
|
|
|
|
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)
|
|
self._persistent = persistent_filesystem
|
|
self._task_id = task_id
|
|
self._sandbox = None
|
|
self._app = None
|
|
self._worker = _AsyncWorker()
|
|
self._sync_manager: FileSyncManager | None = None # initialized after sandbox creation
|
|
sandbox_kwargs = dict(modal_sandbox_kwargs or {})
|
|
restored_snapshot_id = None
|
|
restored_from_legacy_key = False
|
|
if self._persistent:
|
|
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])
|
|
_ensure_modal_sdk()
|
|
import modal as _modal
|
|
cred_mounts = []
|
|
try:
|
|
from tools.credential_files import get_credential_file_mounts, iter_skills_files, iter_cache_files
|
|
|
|
# from_iterable keeps each source lazy so a failure mid-way leaves the earlier mounts in place
|
|
for entry in itertools.chain.from_iterable(
|
|
fn() for fn in (get_credential_file_mounts, iter_skills_files, iter_cache_files)
|
|
):
|
|
cred_mounts.append(
|
|
_modal.Mount.from_local_file(entry["host_path"], remote_path=entry["container_path"])
|
|
)
|
|
except Exception as e:
|
|
logger.debug("Modal: could not load credential file mounts: %s", e)
|
|
self._worker.start()
|
|
|
|
async def _create_sandbox(image_spec: Any):
|
|
app = await _modal.App.lookup.aio("hermes-agent", create_if_missing=True)
|
|
create_kwargs = dict(sandbox_kwargs)
|
|
if cred_mounts:
|
|
create_kwargs["mounts"] = list(create_kwargs.pop("mounts", [])) + cred_mounts
|
|
sandbox = await _modal.Sandbox.create.aio(
|
|
"sleep", "infinity", image=image_spec, app=app,
|
|
timeout=int(create_kwargs.pop("timeout", 3600)), **create_kwargs,
|
|
)
|
|
return app, sandbox
|
|
try:
|
|
try:
|
|
effective_image = _resolve_modal_image(restored_snapshot_id or image)
|
|
self._app, self._sandbox = self._worker.run_coroutine(_create_sandbox(effective_image), timeout=300)
|
|
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)
|
|
self._app, self._sandbox = self._worker.run_coroutine(
|
|
_create_sandbox(_resolve_modal_image(image)), timeout=300)
|
|
else:
|
|
if restored_snapshot_id and restored_from_legacy_key:
|
|
_store_direct_snapshot(self._task_id, restored_snapshot_id)
|
|
except Exception:
|
|
self._worker.stop()
|
|
raise
|
|
logger.info("Modal: sandbox created (task=%s)", self._task_id)
|
|
self._sync_manager = FileSyncManager(
|
|
get_files_fn=lambda: iter_sync_files("/root/.hermes"),
|
|
upload_fn=self._modal_upload, delete_fn=self._modal_delete,
|
|
bulk_upload_fn=self._modal_bulk_upload, bulk_download_fn=self._modal_bulk_download,
|
|
)
|
|
self._sync_manager.sync(force=True)
|
|
self.init_session()
|
|
|
|
def _modal_upload(self, host_path: str, remote_path: str) -> None:
|
|
"""Upload a single file via base64 piped through stdin."""
|
|
b64 = base64.b64encode(Path(host_path).read_bytes()).decode("ascii")
|
|
container_dir = str(Path(remote_path).parent)
|
|
cmd = f"mkdir -p {shlex.quote(container_dir)} && base64 -d > {shlex.quote(remote_path)}"
|
|
|
|
async def _write():
|
|
proc = await self._sandbox.exec.aio("bash", "-c", cmd)
|
|
await _stream_stdin(proc, b64, self._STDIN_CHUNK_SIZE)
|
|
await proc.wait.aio()
|
|
self._worker.run_coroutine(_write(), timeout=30)
|
|
|
|
def _modal_bulk_upload(self, files: list[tuple[str, str]]) -> None:
|
|
"""Upload many files as one in-memory gzipped tar streamed through stdin
|
|
into ``base64 -d | tar xzf -``, avoiding the SDK's 64 KB exec-arg limit."""
|
|
if not files:
|
|
return
|
|
buf = io.BytesIO()
|
|
with tarfile.open(fileobj=buf, mode="w:gz") as tar:
|
|
for host_path, remote_path in files:
|
|
tar.add(host_path, arcname=remote_path.lstrip("/"))
|
|
payload = base64.b64encode(buf.getvalue()).decode("ascii")
|
|
cmd = f"{quoted_mkdir_command(unique_parent_dirs(files))} && base64 -d | tar xzf - -C /"
|
|
|
|
async def _bulk():
|
|
proc = await self._sandbox.exec.aio("bash", "-c", cmd)
|
|
await _stream_stdin(proc, payload, self._STDIN_CHUNK_SIZE)
|
|
exit_code = await proc.wait.aio()
|
|
if exit_code != 0:
|
|
stderr_text = await proc.stderr.read.aio()
|
|
raise RuntimeError(f"Modal bulk upload failed (exit {exit_code}): {stderr_text}")
|
|
self._worker.run_coroutine(_bulk(), timeout=120)
|
|
|
|
def _modal_bulk_download(self, dest: Path) -> None:
|
|
"""Download remote .hermes/ as a tar archive (sandboxes run as root, so /root/.hermes)."""
|
|
async def _download():
|
|
proc = await self._sandbox.exec.aio("bash", "-c", "tar cf - -C / root/.hermes")
|
|
data = await proc.stdout.read.aio()
|
|
exit_code = await proc.wait.aio()
|
|
if exit_code != 0:
|
|
raise RuntimeError(f"Modal bulk download failed (exit {exit_code})")
|
|
return data
|
|
tar_bytes = self._worker.run_coroutine(_download(), timeout=120)
|
|
if isinstance(tar_bytes, str):
|
|
tar_bytes = tar_bytes.encode()
|
|
dest.write_bytes(tar_bytes)
|
|
|
|
def _modal_delete(self, remote_paths: list[str]) -> None:
|
|
"""Batch-delete remote files via exec."""
|
|
rm_cmd = quoted_rm_command(remote_paths)
|
|
|
|
async def _rm():
|
|
proc = await self._sandbox.exec.aio("bash", "-c", rm_cmd)
|
|
await proc.wait.aio()
|
|
self._worker.run_coroutine(_rm(), timeout=15)
|
|
|
|
def _before_execute(self) -> None:
|
|
"""Sync files to sandbox via FileSyncManager (rate-limited internally)."""
|
|
self._sync_manager.sync()
|
|
|
|
def _run_bash(self, cmd_string: str, *, login: bool = False, timeout: int = 120, stdin_data: str | None = None):
|
|
"""Return a _ThreadedProcessHandle wrapping an async Modal sandbox exec."""
|
|
sandbox = self._sandbox
|
|
worker = self._worker
|
|
|
|
def cancel():
|
|
worker.run_coroutine(sandbox.terminate.aio(), timeout=15)
|
|
|
|
def exec_fn() -> tuple[str, int]:
|
|
async def _do():
|
|
process = await sandbox.exec.aio(*bash_argv(cmd_string, login), timeout=timeout)
|
|
stdout = await process.stdout.read.aio()
|
|
stderr = await process.stderr.read.aio()
|
|
exit_code = await process.wait.aio()
|
|
if isinstance(stdout, bytes):
|
|
stdout = stdout.decode("utf-8", errors="replace")
|
|
if isinstance(stderr, bytes):
|
|
stderr = stderr.decode("utf-8", errors="replace")
|
|
output = stdout
|
|
if stderr:
|
|
output = f"{stdout}\n{stderr}" if stdout else stderr
|
|
return output, exit_code
|
|
return worker.run_coroutine(_do(), timeout=timeout + 30)
|
|
return _ThreadedProcessHandle(exec_fn, cancel_fn=cancel)
|
|
|
|
def cleanup(self):
|
|
"""Snapshot the filesystem (if persistent) then stop the sandbox."""
|
|
if self._sandbox is None:
|
|
return
|
|
if self._sync_manager:
|
|
logger.info("Modal: syncing files from sandbox...")
|
|
self._sync_manager.sync_back()
|
|
if self._persistent:
|
|
try:
|
|
async def _snapshot():
|
|
img = await self._sandbox.snapshot_filesystem.aio()
|
|
return img.object_id
|
|
try:
|
|
snapshot_id = self._worker.run_coroutine(_snapshot(), timeout=60)
|
|
except Exception:
|
|
snapshot_id = None
|
|
if snapshot_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)
|
|
try:
|
|
self._worker.run_coroutine(self._sandbox.terminate.aio(), timeout=15)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self._worker.stop()
|
|
self._sandbox = None
|
|
self._app = None
|