feat: code sandbox

src/corral/sandbox/__init__.py

@@ -0,0 +1,14 @@
+from corral.sandbox.base import Sandbox
+from corral.sandbox.config import NetworkPolicy, ResourceLimits, SandboxConfig
+from corral.sandbox.result import ExecResult
+from corral.sandbox.tool import create_sandbox, create_sandbox_tools
+
+__all__ = [
+    "ExecResult",
+    "NetworkPolicy",
+    "ResourceLimits",
+    "Sandbox",
+    "SandboxConfig",
+    "create_sandbox",
+    "create_sandbox_tools",
+]

src/corral/sandbox/_capture.py

@@ -0,0 +1,151 @@
+"""Shared output capture and parsing utilities for sandboxed execution.
+
+Extracted from corral.utils.code_tools so both the legacy tool and the
+new sandbox backends can reuse the same logic.
+"""
+
+from __future__ import annotations
+
+import json
+from contextlib import suppress
+
+
+def generate_output_capture_code() -> str:
+    """Generate Python code appended to user code to capture execution results.
+
+    The generated code inspects ``locals()`` for user-defined variables and
+    prints them as a JSON line prefixed with ``EXECUTION_RESULT:``.
+    """
+    return """
+
+import json
+import types
+import sys
+from collections.abc import Iterable
+
+result_ = {}
+
+def is_json_serializable(obj, max_depth=10, current_depth=0):
+    '''Check if object is JSON serializable with depth limit'''
+    if current_depth > max_depth:
+        return False
+
+    try:
+        # Handle basic types
+        if obj is None or isinstance(obj, (bool, int, float, str)):
+            return True
+
+        # Handle sequences (but not strings)
+        if isinstance(obj, (list, tuple)):
+            return all(is_json_serializable(item, max_depth, current_depth + 1) for item in obj)
+
+        # Handle dictionaries
+        if isinstance(obj, dict):
+            return all(
+                isinstance(k, str) and is_json_serializable(v, max_depth, current_depth + 1)
+                for k, v in obj.items()
+            )
+
+        # Quick test with actual JSON serialization for edge cases
+        json.dumps(obj)
+        return True
+    except (TypeError, ValueError, RecursionError, OverflowError):
+        return False
+
+def safe_convert_to_serializable(obj):
+    '''Convert common non-serializable types to serializable ones'''
+    try:
+        import numpy as np
+        # Handle numpy arrays
+        if hasattr(obj, '__module__') and obj.__module__ == 'numpy':
+            if hasattr(obj, 'tolist'):
+                return obj.tolist()
+            elif hasattr(obj, 'item'):
+                return obj.item()
+    except ImportError:
+        pass
+
+    # Handle sets
+    if isinstance(obj, set):
+        return list(obj)
+
+    # Handle other iterables (but not strings/bytes)
+    if hasattr(obj, '__iter__') and not isinstance(obj, (str, bytes, dict)):
+        try:
+            return list(obj)
+        except:
+            pass
+
+    return obj
+
+# Try to capture from preferred variable names first
+preferred_vars = ['output', 'result', 'filtered_data', 'processed_data', 'dataset']
+captured = False
+
+for var_name in preferred_vars:
+    if var_name in locals():
+        var_value = locals()[var_name]
+        converted_value = safe_convert_to_serializable(var_value)
+        if is_json_serializable(converted_value):
+            result_[var_name] = converted_value
+            captured = True
+            break
+
+# Fallback: capture any user-defined variables
+if not captured:
+    excluded = {
+        '__builtins__', '__name__', '__doc__', '__package__', '__loader__',
+        '__spec__', '__annotations__', '__cached__', '__file__',
+        'json', 'types', 'sys', 'Iterable', 'result_', 'preferred_vars',
+        'captured', 'excluded', 'local_vars', 'is_json_serializable',
+        'safe_convert_to_serializable', 'var_name', 'var_value', 'converted_value'
+    }
+
+    local_vars = dict(locals())  # Create snapshot
+
+    # Collect all suitable variables
+    suitable_vars = {}
+    for var_name, var_value in local_vars.items():
+        if (not var_name.startswith('_') and
+            var_name not in excluded and
+            not isinstance(var_value, types.ModuleType) and
+            not callable(var_value)):
+
+            converted_value = safe_convert_to_serializable(var_value)
+            if is_json_serializable(converted_value):
+                suitable_vars[var_name] = converted_value
+
+    # If we have suitable variables, include them all
+    if suitable_vars:
+        result_.update(suitable_vars)
+
+print('EXECUTION_RESULT:', json.dumps(result_))
+"""
+
+
+def parse_execution_output(stdout: str) -> tuple[dict, list[str]]:
+    """Parse subprocess output to extract execution results and regular output.
+
+    Parameters
+    ----------
+    stdout
+        Raw stdout from the subprocess.
+
+    Returns
+    -------
+    tuple[dict, list[str]]
+        ``(execution_result, output_lines)`` where *execution_result* is the
+        parsed JSON dict and *output_lines* are all non-marker lines.
+    """
+    stdout_lines = stdout.strip().split("\n") if stdout.strip() else []
+    execution_result: dict = {}
+    output_lines: list[str] = []
+
+    for line in stdout_lines:
+        if line.startswith("EXECUTION_RESULT:"):
+            with suppress(json.JSONDecodeError):
+                execution_result = json.loads(line[len("EXECUTION_RESULT:") :])
+        else:
+            output_lines.append(line)
+
+    return execution_result, output_lines

src/corral/sandbox/_state.py

@@ -0,0 +1,56 @@
+"""Pickle-based state persistence between sandbox executions."""
+
+RESTORE_TEMPLATE = """\
+import pickle as __pkl__
+from pathlib import Path as __Path__
+__state_path__ = __Path__({state_file!r})
+if __state_path__.exists():
+    with __state_path__.open('rb') as __f__:
+        __saved__ = __pkl__.load(__f__)
+    locals().update(__saved__)
+    del __saved__
+del __pkl__, __Path__, __state_path__
+try:
+    del __f__
+except NameError:
+    pass
+"""
+
+SAVE_TEMPLATE = """\
+import pickle as __pkl__
+import types as __types__
+__save__ = {{}}
+__excluded__ = {{
+    '__builtins__', '__name__', '__doc__', '__package__', '__loader__',
+    '__spec__', '__annotations__', '__cached__', '__file__',
+    '__pkl__', '__types__', '__f__', '__save__', '__excluded__', '__state_path__',
+    '__saved__', '__Path__',
+    'result_', 'is_json_serializable', 'safe_convert_to_serializable',
+    'preferred_vars', 'captured', 'excluded', 'local_vars',
+    'json', 'types', 'sys', 'Iterable',
+    'var_name', 'var_value', 'converted_value', 'suitable_vars',
+}}
+for __k__, __v__ in dict(locals()).items():
+    if (not __k__.startswith('__') and
+        __k__ not in __excluded__ and
+        not isinstance(__v__, __types__.ModuleType) and
+        not callable(__v__)):
+        try:
+            __pkl__.dumps(__v__)
+            __save__[__k__] = __v__
+        except Exception:
+            pass
+with open({state_file!r}, 'wb') as __f__:
+    __pkl__.dump(__save__, __f__)
+del __pkl__, __types__, __save__, __excluded__, __k__, __v__, __f__
+"""
+
+
+def generate_state_restore_code(state_file: str) -> str:
+    """Generate Python code that restores variables from a previous execution."""
+    return RESTORE_TEMPLATE.format(state_file=state_file)
+
+
+def generate_state_save_code(state_file: str) -> str:
+    """Generate Python code that saves user-defined variables for the next execution."""
+    return SAVE_TEMPLATE.format(state_file=state_file)

src/corral/sandbox/base.py

@@ -0,0 +1,86 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from corral.sandbox.config import SandboxConfig
+from corral.sandbox.result import ExecResult
+
+
+class Sandbox(ABC):
+    """Abstract base for sandboxed code execution backends."""
+
+    def __init__(self, config: SandboxConfig):
+        self.config = config
+        self._started = False
+
+    @abstractmethod
+    def start(self) -> None:
+        """Initialize the sandbox.
+
+        Called once per trial, not per execution bcause we want to keep for all subtasks the same sandbox.
+        """
+
+    @abstractmethod
+    def stop(self) -> None:
+        """Tear down the sandbox and clean up all resources."""
+
+    @abstractmethod
+    def execute(
+        self,
+        code: str,
+        timeout: int | None = None,
+        env: dict[str, str] | None = None,
+    ) -> ExecResult:
+        """Execute Python code in the sandbox.
+
+        Parameters
+        ----------
+        code
+            Python source code string to execute.
+        timeout
+            Override for the default timeout in seconds.
+        env
+            Additional environment variables for this execution.
+
+        Returns
+        -------
+        ExecResult
+        """
+
+    @abstractmethod
+    def execute_command(
+        self,
+        command: str,
+        timeout: int | None = None,
+        cwd: str | None = None,
+    ) -> ExecResult:
+        """Execute a shell command in the sandbox.
+
+        Parameters
+        ----------
+        command
+            Shell command string.
+        timeout
+            Override timeout in seconds.
+        cwd
+            Working directory override.
+
+        Returns
+        -------
+        ExecResult
+        """
+
+    @abstractmethod
+    def upload_file(self, local_path: str | Path, sandbox_path: str) -> None:
+        """Copy a file from the host into the sandbox."""
+
+    @abstractmethod
+    def download_file(self, sandbox_path: str, local_path: str | Path) -> None:
+        """Copy a file from the sandbox to the host."""
+
+    def __enter__(self):
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop()
+        return False

src/corral/sandbox/config.py

@@ -0,0 +1,43 @@
+from pydantic import BaseModel, Field
+
+
+class ResourceLimits(BaseModel):
+    """Resource constraints for the sandbox."""
+
+    cpu_count: int = 2
+    memory_mb: int = 2048  # TODO: ml task might require more
+    storage_mb: int = 4096
+    timeout_seconds: int = 300
+    max_output_bytes: int = 60 * 1024  # needs to be defined in the code_tools.py file
+
+
+class NetworkPolicy(BaseModel):
+    """Network access policy for the sandbox."""
+
+    allow_network: bool = False
+    allowed_hosts: list[str] = Field(default_factory=list)
+
+
+class SandboxConfig(BaseModel):
+    """Configuration for a sandbox instance."""
+
+    backend: str = "docker"
+
+    # Docker-specific
+    docker_image: str = "python:3.11-slim"  # TODO: 3.12 or 3.12
+    dockerfile: str | None = None
+
+    # Environment
+    working_dir: str = "/workspace"
+    python_packages: list[str] = Field(default_factory=list)
+    pip_install_timeout: int = 120
+
+    # Resource limits
+    resources: ResourceLimits = Field(default_factory=ResourceLimits)
+
+    # Network
+    network: NetworkPolicy = Field(default_factory=NetworkPolicy)
+
+    # State persistence
+    persistent_state: bool = True
+    state_dir: str | None = None