Code Execution API Reference

Code Execution.

ExecutionConfig dataclass

Configuration for controlling how an IPython session executes cells.

Source code in packages/code_execution/src/code_execution/ipython_session.py

@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"

IPythonBackend

Bases: Protocol

Protocol for backends that execute IPython cells.

Source code in packages/code_execution/src/code_execution/executor.py

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."""
        ...

run_cell(code_str, *, use_subprocess=False, timeout_s=None, allow_mime=None, matplotlib_backend=ExecutionConfig.matplotlib_backend)

Execute a code cell and return normalized outputs.

Source code in packages/code_execution/src/code_execution/executor.py

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."""
    ...

IPythonExecutor

Facade that executes IPython cells via a configurable backend.

Source code in packages/code_execution/src/code_execution/executor.py

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,
        )

__init__(backend)

Initialize the executor with the provided backend.

Source code in packages/code_execution/src/code_execution/executor.py

def __init__(self, backend: IPythonBackend) -> None:
    """Initialize the executor with the provided backend."""
    self._backend = backend

run_cell(code_str, *, use_subprocess=False, timeout_s=None, allow_mime=None, matplotlib_backend=ExecutionConfig.matplotlib_backend)

Execute a code cell using the configured backend.

Source code in packages/code_execution/src/code_execution/executor.py

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,
    )

IPythonSession

Run code in an IPython shell with optional subprocess isolation.

Source code in packages/code_execution/src/code_execution/ipython_session.py

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()

        # Read the result before waiting for the child to exit. The child sends
        # its outputs as a single pickled frame, which runs to megabytes once
        # figures are attached, while a pipe buffers only ~64KB. Joining first
        # would leave the child blocked in send() with nobody draining the pipe,
        # and the parent blocked in join() waiting for an exit that send() is
        # preventing -- a deadlock that only shows up on cells big enough to
        # overflow the buffer. recv() is what drains it, so it has to come first.
        timeout_s = self._config.timeout_s
        deadline = None if timeout_s is None else monotonic() + timeout_s
        result: dict[str, Any] | None = None
        try:
            if parent_conn.poll(_remaining(deadline)):
                result = parent_conn.recv()
        except EOFError:
            # Child exited without sending; reported below as a missing result.
            pass

        process.join(timeout=_remaining(deadline))
        if process.is_alive():
            process.terminate()
            process.join()
        parent_conn.close()

        if result is not None:
            return result

        if deadline is not None and monotonic() >= deadline:
            return {
                "stdout": "",
                "stderr": "",
                "rich_outputs": [],
                "success": False,
                "error": {
                    "type": "TimeoutError",
                    "message": f"Execution exceeded {timeout_s} seconds",
                    "traceback": "",
                },
            }

        return {
            "stdout": "",
            "stderr": "",
            "rich_outputs": [],
            "success": False,
            "error": {
                "type": "RuntimeError",
                "message": "Subprocess exited without returning a result",
                "traceback": "",
            },
        }

__init__(*, use_subprocess=False, timeout_s=None, allow_mime=None, matplotlib_backend=ExecutionConfig.matplotlib_backend)

Initialize an IPython execution session with optional isolation settings.

Source code in packages/code_execution/src/code_execution/ipython_session.py

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)

run_cell(code_str)

Run a code cell and return captured outputs, errors, and success state.

Source code in packages/code_execution/src/code_execution/ipython_session.py

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()

    # Read the result before waiting for the child to exit. The child sends
    # its outputs as a single pickled frame, which runs to megabytes once
    # figures are attached, while a pipe buffers only ~64KB. Joining first
    # would leave the child blocked in send() with nobody draining the pipe,
    # and the parent blocked in join() waiting for an exit that send() is
    # preventing -- a deadlock that only shows up on cells big enough to
    # overflow the buffer. recv() is what drains it, so it has to come first.
    timeout_s = self._config.timeout_s
    deadline = None if timeout_s is None else monotonic() + timeout_s
    result: dict[str, Any] | None = None
    try:
        if parent_conn.poll(_remaining(deadline)):
            result = parent_conn.recv()
    except EOFError:
        # Child exited without sending; reported below as a missing result.
        pass

    process.join(timeout=_remaining(deadline))
    if process.is_alive():
        process.terminate()
        process.join()
    parent_conn.close()

    if result is not None:
        return result

    if deadline is not None and monotonic() >= deadline:
        return {
            "stdout": "",
            "stderr": "",
            "rich_outputs": [],
            "success": False,
            "error": {
                "type": "TimeoutError",
                "message": f"Execution exceeded {timeout_s} seconds",
                "traceback": "",
            },
        }

    return {
        "stdout": "",
        "stderr": "",
        "rich_outputs": [],
        "success": False,
        "error": {
            "type": "RuntimeError",
            "message": "Subprocess exited without returning a result",
            "traceback": "",
        },
    }

LocalIPythonBackend

Local backend that executes code via IPythonSession.

Source code in packages/code_execution/src/code_execution/executor.py

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)

__init__(*, cwd=None, env=None)

Initialize the backend with optional working directory and env vars.

Source code in packages/code_execution/src/code_execution/executor.py

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

run_cell(code_str, *, use_subprocess=False, timeout_s=None, allow_mime=None, matplotlib_backend=ExecutionConfig.matplotlib_backend)

Execute a code cell in a local IPython session.

Source code in packages/code_execution/src/code_execution/executor.py

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)

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.

Source code in packages/code_execution/src/code_execution/process_backend.py

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 so the sandbox runner can import it.
        code_exec_pkg = Path(__file__).resolve().parent.parent.parent
        if (code_exec_pkg / "pyproject.toml").is_file():
            # Running from local workspace
            install_spec = ["-e", str(code_exec_pkg)]
        else:
            # Running from pip-installed cli
            version = importlib.metadata.version("asta-code-execution")
            install_spec = [f"asta-code-execution=={version}"]
        subprocess.run(
            [uv, "pip", "install", "--python", str(python_bin), *install_spec],
            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

sandbox_python property

Return the path to the sandbox venv's Python interpreter.

__init__(*, cwd=None, env=None, packages=None, sandbox_venv_path=None)

Initialize the backend with optional working directory and env vars.

Source code in packages/code_execution/src/code_execution/process_backend.py

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

run_cell(code_str, *, use_subprocess=False, timeout_s=None, allow_mime=None, matplotlib_backend=ExecutionConfig.matplotlib_backend)

Execute a code cell in an isolated subprocess and return normalized outputs.

Source code in packages/code_execution/src/code_execution/process_backend.py

    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

For Modal-based sandbox execution, see docs/autodiscovery_modal/sandbox_backend.md.