asta-code-execution 0.1.2__py3-none-any.whl

asta_code_execution-0.1.2.dist-info/METADATA

@@ -0,0 +1,9 @@
+Metadata-Version: 2.3
+Name: asta-code-execution
+Version: 0.1.2
+Summary: Add your description here
+Author: Allen Institute for Artificial Intelligence
+Requires-Dist: ipython>=9.9.0
+Requires-Python: >=3.13
+Description-Content-Type: text/markdown
+

asta_code_execution-0.1.2.dist-info/RECORD

@@ -0,0 +1,8 @@
+code_execution/__init__.py,sha256=IeWaNDWSA0QNQDqsek2xXowglHp7NTVRoM4qzJXCCHk,371
+code_execution/executor.py,sha256=xKf2fp6RJx-gMrc0UIZKgLxywCdiaXwrMOlZ_n007b0,3394
+code_execution/ipython_session.py,sha256=Yd-KYheOzJ0BNErEPqJHmFf-C0iplmZ6gHnoXdeH65c,8453
+code_execution/process_backend.py,sha256=xFtiQMIXwJN4Iqas6gCza1UlxNqMzSyCWs0gtz67nEE,9270
+code_execution/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+asta_code_execution-0.1.2.dist-info/WHEEL,sha256=eh7sammvW2TypMMMGKgsM83HyA_3qQ5Lgg3ynoecH3M,79
+asta_code_execution-0.1.2.dist-info/METADATA,sha256=3weUIeEeNGsJhQXiyvEJUJ_vnHsS_0ieN1xyx2Nbxlo,245
+asta_code_execution-0.1.2.dist-info/RECORD,,

asta_code_execution-0.1.2.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.8.24
+Root-Is-Purelib: true
+Tag: py3-none-any

code_execution/__init__.py

@@ -0,0 +1,14 @@
+"""Code Execution."""
+
+from .executor import IPythonBackend, IPythonExecutor, LocalIPythonBackend
+from .ipython_session import ExecutionConfig, IPythonSession
+from .process_backend import ProcessIPythonBackend
+
+__all__ = [
+    "ExecutionConfig",
+    "IPythonBackend",
+    "IPythonExecutor",
+    "IPythonSession",
+    "LocalIPythonBackend",
+    "ProcessIPythonBackend",
+]

code_execution/executor.py

@@ -0,0 +1,105 @@
+"""Execution backends for running IPython cells."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+from contextlib import contextmanager
+from typing import Any, Protocol
+
+from .ipython_session import ExecutionConfig, IPythonSession
+
+
+class IPythonBackend(Protocol):
+    """Protocol for backends that execute IPython cells."""
+
+    def run_cell(
+        self,
+        code_str: str,
+        *,
+        use_subprocess: bool = False,
+        timeout_s: float | None = None,
+        allow_mime: Iterable[str] | None = None,
+        matplotlib_backend: str | None = ExecutionConfig.matplotlib_backend,
+    ) -> dict[str, Any]:
+        """Execute a code cell and return normalized outputs."""
+        ...
+
+
+@contextmanager
+def _execution_context(cwd: str | None, env: dict[str, str] | None) -> Iterable[None]:
+    # Temporarily apply environment and working directory changes for a run.
+    previous_cwd = os.getcwd()
+    previous_env: dict[str, str | None] = {}
+    try:
+        if env:
+            for key, value in env.items():
+                previous_env[key] = os.environ.get(key)
+                os.environ[key] = value
+        if cwd:
+            os.chdir(cwd)
+        yield
+    finally:
+        if cwd:
+            os.chdir(previous_cwd)
+        if env:
+            for key in env:
+                previous_value = previous_env.get(key)
+                if previous_value is None:
+                    os.environ.pop(key, None)
+                else:
+                    os.environ[key] = previous_value
+
+
+class LocalIPythonBackend:
+    """Local backend that executes code via IPythonSession."""
+
+    def __init__(self, *, cwd: str | None = None, env: dict[str, str] | None = None) -> None:
+        """Initialize the backend with optional working directory and env vars."""
+        self._cwd = cwd
+        self._env = env
+
+    def run_cell(
+        self,
+        code_str: str,
+        *,
+        use_subprocess: bool = False,
+        timeout_s: float | None = None,
+        allow_mime: Iterable[str] | None = None,
+        matplotlib_backend: str | None = ExecutionConfig.matplotlib_backend,
+    ) -> dict[str, Any]:
+        """Execute a code cell in a local IPython session."""
+        with _execution_context(self._cwd, self._env):
+            session = IPythonSession(
+                use_subprocess=use_subprocess,
+                timeout_s=timeout_s,
+                allow_mime=allow_mime,
+                matplotlib_backend=matplotlib_backend,
+            )
+            return session.run_cell(code_str)
+
+
+class IPythonExecutor:
+    """Facade that executes IPython cells via a configurable backend."""
+
+    def __init__(self, backend: IPythonBackend) -> None:
+        """Initialize the executor with the provided backend."""
+        self._backend = backend
+
+    def run_cell(
+        self,
+        code_str: str,
+        *,
+        use_subprocess: bool = False,
+        timeout_s: float | None = None,
+        allow_mime: Iterable[str] | None = None,
+        matplotlib_backend: str | None = ExecutionConfig.matplotlib_backend,
+    ) -> dict[str, Any]:
+        """Execute a code cell using the configured backend."""
+        return self._backend.run_cell(
+            code_str,
+            use_subprocess=use_subprocess,
+            timeout_s=timeout_s,
+            allow_mime=allow_mime,
+            matplotlib_backend=matplotlib_backend,
+        )

code_execution/ipython_session.py

@@ -0,0 +1,266 @@
+"""IPython-backed execution helpers with optional subprocess isolation."""
+
+import base64
+import traceback
+from collections.abc import Iterable
+from dataclasses import dataclass
+from multiprocessing import get_context
+from typing import Any, cast
+
+from IPython.core.formatters import DisplayFormatter
+from IPython.core.interactiveshell import InteractiveShell
+from IPython.utils.capture import capture_output
+
+
+@dataclass(frozen=True)
+class ExecutionConfig:
+    """Configuration for controlling how an IPython session executes cells."""
+
+    use_subprocess: bool = False
+    timeout_s: float | None = None
+    allow_mime: frozenset[str] = frozenset(
+        {
+            "text/plain",
+            "text/html",
+            "text/markdown",
+            "text/latex",
+            "image/png",
+            "image/svg+xml",
+            "image/jpeg",
+            "application/json",
+            "application/javascript",
+            "application/pdf",
+        }
+    )
+    matplotlib_backend: str | None = "module://matplotlib_inline.backend_inline"
+
+
+def _normalize_value(value: Any) -> Any:
+    # Ensure outputs are JSON-safe for downstream serialization.
+    if isinstance(value, bytes):
+        return base64.b64encode(value).decode("ascii")
+    if isinstance(value, (str, int, float, bool)) or value is None:
+        return value
+    if isinstance(value, (list, tuple)):
+        return [_normalize_value(item) for item in value]
+    if isinstance(value, dict):
+        return {str(key): _normalize_value(item) for key, item in value.items()}
+    return repr(value)
+
+
+def _normalize_mime_bundle(data: dict[str, Any], allow_mime: frozenset[str]) -> dict[str, Any]:
+    # Filter to a predictable MIME allowlist to avoid leaking unexpected payload types.
+    normalized: dict[str, Any] = {}
+    for mime_type, payload in data.items():
+        if mime_type not in allow_mime:
+            continue
+        normalized[mime_type] = _normalize_value(payload)
+    return normalized
+
+
+def _format_error(exc: BaseException | None) -> dict[str, str] | None:
+    # Preserve traceback details so callers can render useful diagnostics.
+    if exc is None:
+        return None
+    return {
+        "type": type(exc).__name__,
+        "message": str(exc),
+        "traceback": "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
+    }
+
+
+def _run_cell_with_shell(
+    shell: InteractiveShell,
+    code_str: str,
+    allow_mime: frozenset[str],
+) -> dict[str, Any]:
+    # Execute in-process to preserve state between calls when isolation isn't needed.
+    with capture_output() as captured:
+        result = shell.run_cell(code_str)
+
+    error = result.error_before_exec or result.error_in_exec
+    outputs = {
+        "stdout": captured.stdout,
+        "stderr": captured.stderr,
+        "rich_outputs": [],
+        "success": result.success,
+        "error": _format_error(error),
+    }
+
+    for display_obj in captured.outputs:
+        outputs["rich_outputs"].append(_normalize_mime_bundle(display_obj.data, allow_mime))
+
+    return outputs
+
+
+def _configure_display_formatters(
+    shell: InteractiveShell,
+    allow_mime: frozenset[str],
+) -> None:
+    formatter = shell.display_formatter
+    if formatter is None:
+        return
+    display_formatter = cast(DisplayFormatter, formatter)
+    available = [mime for mime in allow_mime if mime in display_formatter.formatters]
+    if not available:
+        return
+    for mime in available:
+        display_formatter.formatters[mime].enabled = True
+    display_formatter.active_types = available
+
+
+def _configure_matplotlib_backend(
+    backend: str | None,
+    allow_mime: frozenset[str],
+) -> None:
+    if not backend:
+        return
+    try:
+        import matplotlib
+        import matplotlib_inline.backend_inline as backend_inline
+    except Exception:
+        return
+
+    try:
+        current = matplotlib.get_backend()
+    except Exception:
+        current = None
+
+    if current == backend:
+        return
+
+    try:
+        # Ensure rich display payloads are emitted in non-interactive contexts.
+        matplotlib.use(backend, force=True)
+        if "matplotlib_inline.backend_inline" in backend:
+            formats: list[str] = []
+            if "image/png" in allow_mime:
+                formats.append("png")
+            if "image/svg+xml" in allow_mime:
+                formats.append("svg")
+            if "image/jpeg" in allow_mime:
+                formats.append("jpeg")
+            if formats:
+                backend_inline.set_matplotlib_formats(*formats)
+    except Exception:
+        # If pyplot or another backend is already loaded, keep running.
+        return
+
+
+def _configure_shell(
+    shell: InteractiveShell,
+    allow_mime: frozenset[str],
+    matplotlib_backend: str | None,
+) -> None:
+    _configure_display_formatters(shell, allow_mime)
+    _configure_matplotlib_backend(matplotlib_backend, allow_mime)
+
+
+def _ensure_pip_available() -> None:
+    # Bootstrap pip so %pip magic can install packages on demand.
+    try:
+        import pip  # noqa: F401
+
+        return
+    except Exception:
+        pass
+
+    try:
+        import ensurepip
+    except Exception:
+        return
+
+    try:
+        ensurepip.bootstrap()
+    except Exception:
+        return
+
+
+def _run_cell_in_subprocess(
+    code_str: str,
+    allow_mime: frozenset[str],
+    matplotlib_backend: str | None,
+    connection,
+) -> None:
+    # Run in a child process so a hard timeout can be enforced without blocking.
+    shell = InteractiveShell.instance()
+    _configure_shell(shell, allow_mime, matplotlib_backend)
+    outputs = _run_cell_with_shell(shell, code_str, allow_mime)
+    connection.send(outputs)
+    connection.close()
+
+
+class IPythonSession:
+    """Run code in an IPython shell with optional subprocess isolation."""
+
+    def __init__(
+        self,
+        *,
+        use_subprocess: bool = False,
+        timeout_s: float | None = None,
+        allow_mime: Iterable[str] | None = None,
+        matplotlib_backend: str | None = ExecutionConfig.matplotlib_backend,
+    ) -> None:
+        """Initialize an IPython execution session with optional isolation settings."""
+        _ensure_pip_available()
+        self._config = ExecutionConfig(
+            use_subprocess=use_subprocess,
+            timeout_s=timeout_s,
+            allow_mime=frozenset(allow_mime or ExecutionConfig.allow_mime),
+            matplotlib_backend=matplotlib_backend,
+        )
+        # Create a shell instance that persists variables between calls.
+        self.shell = InteractiveShell.instance()
+        _configure_shell(self.shell, self._config.allow_mime, self._config.matplotlib_backend)
+
+    def run_cell(self, code_str: str) -> dict[str, Any]:
+        """Run a code cell and return captured outputs, errors, and success state."""
+        if self._config.timeout_s is not None and not self._config.use_subprocess:
+            raise ValueError("timeout_s requires use_subprocess=True to enforce a hard timeout")
+        if not self._config.use_subprocess:
+            return _run_cell_with_shell(self.shell, code_str, self._config.allow_mime)
+
+        ctx = get_context("spawn")
+        parent_conn, child_conn = ctx.Pipe(duplex=False)
+        process = ctx.Process(
+            target=_run_cell_in_subprocess,
+            args=(code_str, self._config.allow_mime, self._config.matplotlib_backend, child_conn),
+        )
+        process.start()
+        child_conn.close()
+        process.join(timeout=self._config.timeout_s)
+
+        if process.is_alive():
+            process.terminate()
+            process.join()
+            parent_conn.close()
+            return {
+                "stdout": "",
+                "stderr": "",
+                "rich_outputs": [],
+                "success": False,
+                "error": {
+                    "type": "TimeoutError",
+                    "message": f"Execution exceeded {self._config.timeout_s} seconds",
+                    "traceback": "",
+                },
+            }
+
+        if parent_conn.poll():
+            result = parent_conn.recv()
+            parent_conn.close()
+            return result
+
+        parent_conn.close()
+
+        return {
+            "stdout": "",
+            "stderr": "",
+            "rich_outputs": [],
+            "success": False,
+            "error": {
+                "type": "RuntimeError",
+                "message": "Subprocess exited without returning a result",
+                "traceback": "",
+            },
+        }

code_execution/process_backend.py

@@ -0,0 +1,272 @@
+"""Subprocess-based IPython execution backend for local isolated runs."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+from .ipython_session import ExecutionConfig
+
+_DEFAULT_SANDBOX_PACKAGES = [
+    "ipython",
+    "numpy",
+    "pandas",
+    "matplotlib",
+    "matplotlib-inline",
+    "seaborn",
+    "scikit-learn",
+    "scipy",
+    "statsmodels",
+]
+
+_SANDBOX_RUNNER = """\
+import json
+import sys
+import traceback
+
+from code_execution.ipython_session import ExecutionConfig, IPythonSession
+
+
+def _format_error(exc: BaseException) -> dict[str, str]:
+    return {
+        "type": type(exc).__name__,
+        "message": str(exc),
+        "traceback": "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
+    }
+
+
+def main() -> None:
+    try:
+        payload = json.load(sys.stdin)
+        allow_mime = payload.get("allow_mime")
+        session = IPythonSession(
+            use_subprocess=payload.get("use_subprocess", False),
+            timeout_s=payload.get("timeout_s"),
+            allow_mime=allow_mime,
+            matplotlib_backend=payload.get(
+                "matplotlib_backend",
+                ExecutionConfig.matplotlib_backend,
+            ),
+        )
+        result = session.run_cell(payload["code_str"])
+    except BaseException as exc:
+        result = {
+            "stdout": "",
+            "stderr": "",
+            "rich_outputs": [],
+            "success": False,
+            "error": _format_error(exc),
+        }
+    json.dump(result, sys.stdout)
+
+
+if __name__ == "__main__":
+    main()
+"""
+
+
+def _find_uv() -> str:
+    """Return path to the ``uv`` binary, or raise if not found."""
+    uv = shutil.which("uv")
+    if uv is None:
+        raise RuntimeError(
+            "uv is required for ProcessIPythonBackend but was not found on PATH"
+        )
+    return uv
+
+
+# TODO: Transplant this implementation onto asta_sandbox.SandboxBase (making run_cell
+# async via asyncio.to_thread) and move it into the asta-sandbox library as a first-class
+# local process backend alongside InProcessExecutor and ModalEphemeralExecutor. Once done,
+# _ProcessBackendAdapter in agents.py can be removed and ProcessIPythonBackend can be used
+# directly wherever a SandboxBase is expected.
+class ProcessIPythonBackend:
+    """Backend that executes IPython cells in isolated subprocesses.
+
+    Each backend instance lazily creates a *base sandbox venv* with a curated
+    set of packages (mirroring the Modal sandbox image).  Individual
+    ``run_cell`` invocations get a per-process temporary directory so that any
+    ``uv pip install`` executed by experiment code does **not** mutate the
+    base venv.
+    """
+
+    def __init__(
+        self,
+        *,
+        cwd: str | None = None,
+        env: dict[str, str] | None = None,
+        packages: list[str] | None = None,
+        sandbox_venv_path: str | None = None,
+    ) -> None:
+        """Initialize the backend with optional working directory and env vars."""
+        self._cwd = cwd
+        self._env = env
+        self._packages = packages if packages is not None else list(_DEFAULT_SANDBOX_PACKAGES)
+        self._sandbox_venv_path: Path | None = (
+            Path(sandbox_venv_path) if sandbox_venv_path else None
+        )
+        self._sandbox_ready = False
+        self._owned_venv_dir: tempfile.TemporaryDirectory[str] | None = None
+
+    # -- lazy venv setup -----------------------------------------------------
+
+    def _ensure_sandbox(self) -> Path:
+        """Create the base sandbox venv if it doesn't already exist."""
+        if self._sandbox_ready and self._sandbox_venv_path is not None:
+            return self._sandbox_venv_path
+
+        uv = _find_uv()
+
+        if self._sandbox_venv_path is None:
+            self._owned_venv_dir = tempfile.TemporaryDirectory(prefix="process_sandbox_")
+            self._sandbox_venv_path = Path(self._owned_venv_dir.name)
+
+        venv_path = self._sandbox_venv_path
+        python_bin = venv_path / "bin" / "python"
+
+        if not python_bin.exists():
+            subprocess.run(
+                [uv, "venv", str(venv_path), "--seed", "--python", sys.executable],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+
+        # Install curated packages
+        if self._packages:
+            subprocess.run(
+                [uv, "pip", "install", "--python", str(python_bin), *self._packages],
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+
+        # Install code_execution package (editable) so the sandbox runner works
+        code_exec_pkg = Path(__file__).resolve().parent.parent.parent  # packages/code_execution
+        subprocess.run(
+            [uv, "pip", "install", "--python", str(python_bin), "-e", str(code_exec_pkg)],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+        self._sandbox_ready = True
+        return venv_path
+
+    @property
+    def sandbox_python(self) -> str:
+        """Return the path to the sandbox venv's Python interpreter."""
+        venv = self._ensure_sandbox()
+        return str(venv / "bin" / "python")
+
+    # -- cell execution ------------------------------------------------------
+
+    def run_cell(
+        self,
+        code_str: str,
+        *,
+        use_subprocess: bool = False,
+        timeout_s: float | None = None,
+        allow_mime: Iterable[str] | None = None,
+        matplotlib_backend: str | None = ExecutionConfig.matplotlib_backend,
+    ) -> dict[str, Any]:
+        """Execute a code cell in an isolated subprocess and return normalized outputs."""
+        venv_path = self._ensure_sandbox()
+        python_bin = str(venv_path / "bin" / "python")
+
+        payload = {
+            "code_str": code_str,
+            "use_subprocess": use_subprocess,
+            "timeout_s": None,
+            "allow_mime": list(allow_mime) if allow_mime else None,
+            "matplotlib_backend": matplotlib_backend,
+        }
+
+        child_env = os.environ.copy()
+        if self._env:
+            child_env.update(self._env)
+
+        # Per-cell temp directory for any packages installed by experiment code.
+        # An `install-package` script on PATH installs into this directory
+        # via `uv pip install --target`, keeping the base sandbox venv clean.
+        cell_tmp = tempfile.mkdtemp(prefix="cell_pkgs_")
+        try:
+            cell_install_dir = os.path.join(cell_tmp, "lib")
+            os.makedirs(cell_install_dir, exist_ok=True)
+
+            existing_pp = child_env.get("PYTHONPATH", "")
+            child_env["PYTHONPATH"] = (
+                f"{cell_install_dir}:{existing_pp}" if existing_pp else cell_install_dir
+            )
+
+            # Create an install-package script that experiment code calls
+            # to install packages into the per-process temp directory.
+            script_dir = os.path.join(cell_tmp, "bin")
+            os.makedirs(script_dir, exist_ok=True)
+            real_uv = _find_uv()
+            script_path = os.path.join(script_dir, "install-package")
+            with open(script_path, "w") as f:
+                f.write(f"""#!/bin/sh
+# Install a Python package into the per-process temp directory.
+exec "{real_uv}" pip install --target "{cell_install_dir}" --quiet "$@"
+""")
+            os.chmod(script_path, 0o755)
+
+            venv_bin = str(venv_path / "bin")
+            child_env["PATH"] = f"{script_dir}:{venv_bin}:{child_env.get('PATH', '')}"
+            child_env["VIRTUAL_ENV"] = str(venv_path)
+
+            try:
+                proc = subprocess.run(
+                    [python_bin, "-c", _SANDBOX_RUNNER],
+                    input=json.dumps(payload),
+                    capture_output=True,
+                    text=True,
+                    timeout=timeout_s,
+                    cwd=self._cwd,
+                    env=child_env,
+                )
+            except subprocess.TimeoutExpired:
+                return {
+                    "stdout": "",
+                    "stderr": "",
+                    "rich_outputs": [],
+                    "success": False,
+                    "error": {
+                        "type": "TimeoutError",
+                        "message": f"Process execution timed out after {timeout_s}s",
+                        "traceback": "",
+                    },
+                }
+        finally:
+            shutil.rmtree(cell_tmp, ignore_errors=True)
+
+        stdout = proc.stdout or ""
+        stderr = proc.stderr or ""
+
+        try:
+            result = json.loads(stdout) if stdout.strip() else {}
+        except json.JSONDecodeError:
+            return {
+                "stdout": "",
+                "stderr": stderr,
+                "rich_outputs": [],
+                "success": False,
+                "error": {
+                    "type": "RuntimeError",
+                    "message": "Subprocess output was not valid JSON.",
+                    "traceback": stdout,
+                },
+            }
+
+        if stderr:
+            result.setdefault("stderr", "")
+            result["stderr"] = f"{result['stderr']}{stderr}"
+        return result