From 0cabd95016f832bd2255a6638e9f05dab4010611 Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 19:49:07 -0700 Subject: [PATCH 01/10] Replace code execution backends with bespokelabs-sandbox Remove the four built-in execution backends (multiprocessing, docker, ray, e2b) and replace them with a single unified backend powered by bespokelabs-sandbox. This delegates sandbox lifecycle, isolation, and multi-provider support to the sandbox package while keeping curator's orchestration layer (caching, retries, request/response files) intact. - Add SandboxCodeExecutionBackend using bespokelabs.sandbox.Sandbox - bespokelabs-sandbox is a required dep; "local" backend works out of the box - code_execution extra installs docker/ray/e2b SDKs for non-local backends - Preserve legacy docker_image and base_url config params for backwards compat - Tar archives retain directory structure (os.path.relpath, not basename) - Regenerate poetry.lock for the new dependency graph - Update tests and examples for new backend names Co-Authored-By: Claude Opus 4.6 (1M context) --- .../code-execution/chart-generation/charts.py | 11 +- .../code-verification/code_verify.py | 4 +- .../math-animation/execute_code.py | 6 +- poetry.lock | 95 +++-- pyproject.toml | 9 +- .../code_execution_backend/_factory.py | 42 ++- .../code_execution_backend/base_backend.py | 45 --- .../code_execution_backend/docker_backend.py | 344 ------------------ .../code_execution_backend/e2b_backend.py | 78 ---- .../multiprocessing_backend.py | 89 ----- .../code_execution_backend/ray_backend.py | 114 ------ .../code_execution_backend/sandbox_backend.py | 146 ++++++++ .../curator/code_executor/code_executor.py | 4 +- .../curator/code_executor/types.py | 4 +- tests/code_executor/test_backend.py | 52 --- tests/code_executor/test_code_execution.py | 40 +- tests/code_executor/test_data_types.py | 11 +- 17 files changed, 276 insertions(+), 818 deletions(-) delete mode 100644 src/bespokelabs/curator/code_executor/code_execution_backend/docker_backend.py delete mode 100644 src/bespokelabs/curator/code_executor/code_execution_backend/e2b_backend.py delete mode 100644 src/bespokelabs/curator/code_executor/code_execution_backend/multiprocessing_backend.py delete mode 100644 src/bespokelabs/curator/code_executor/code_execution_backend/ray_backend.py create mode 100644 src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py diff --git a/examples/code-execution/chart-generation/charts.py b/examples/code-execution/chart-generation/charts.py index 59df4e276..f125ceec1 100644 --- a/examples/code-execution/chart-generation/charts.py +++ b/examples/code-execution/chart-generation/charts.py @@ -108,16 +108,11 @@ def code_output(self, row, execution_output): tar.extractall(path=temp_dir) # read the chart image - # it is in workspace/chart.png for 'docker' backend and chart.png for 'multiprocessing' backend try: - with open(os.path.join(temp_dir, "workspace/chart.png"), "rb") as f: + with open(os.path.join(temp_dir, "chart.png"), "rb") as f: row["chart_image"] = f.read() except Exception: - try: - with open(os.path.join(temp_dir, "chart.png"), "rb") as f: - row["chart_image"] = f.read() - except Exception: - row["chart_image"] = None + row["chart_image"] = None return row @@ -126,7 +121,7 @@ def code_output(self, row, execution_output): # This section demonstrates how to use the ChartCodeExecutor # Initialize executor with the multiprocessing backend (runs code in a separate process) - executor = ChartCodeExecutor(backend="multiprocessing") # or ray, docker (with matplotlib image) + executor = ChartCodeExecutor(backend="local") # or docker, e2b, modal, daytona # Create a simple dataset with one example dataset = Dataset.from_list([{"id": 1}]) diff --git a/examples/code-execution/code-verification/code_verify.py b/examples/code-execution/code-verification/code_verify.py index 9c5ba8b1e..29cdf9d07 100644 --- a/examples/code-execution/code-verification/code_verify.py +++ b/examples/code-execution/code-verification/code_verify.py @@ -1,4 +1,4 @@ -"""This example demonstrates how to use Curator's code execution functionality with Ray backend for executing code on large datasets. +"""This example demonstrates how to use Curator's code execution functionality for executing code on large datasets. It shows how to verify code solutions against expected outputs. @@ -60,7 +60,7 @@ def code_output(self, row, execution_output): if __name__ == "__main__": - executor = TACOCodeExecutor(backend="ray") # or multiprocessing, docker, e2b + executor = TACOCodeExecutor(backend="local") # or docker, e2b, modal, daytona dataset = load_dataset("bespokelabs/sky-t1-taco-test-rejection-sampled-shreyas") execution_output = executor(dataset["train"].select(range(10))) diff --git a/examples/code-execution/math-animation/execute_code.py b/examples/code-execution/math-animation/execute_code.py index f6feb1438..d20325179 100644 --- a/examples/code-execution/math-animation/execute_code.py +++ b/examples/code-execution/math-animation/execute_code.py @@ -75,8 +75,8 @@ def code_output(self, input: Dict, output) -> Dict: with tarfile.open(fileobj=bytes_io) as tar: tar.extractall(path=temp_dir) - # get the video file - video_file = os.path.join(temp_dir, "workspace/media/videos/temp_animation/720p30/video.mp4") + # get the video file (path relative to sandbox workspace) + video_file = os.path.join(temp_dir, "media/videos/temp_animation/720p30/video.mp4") if os.path.exists(video_file): # Copy the video file to the output dataset input["video"] = open(video_file, "rb").read() @@ -111,7 +111,7 @@ def execute_manim_code(dataset_name="pimpalgaonkar/manim_codes_10k", output_data executor = ManimCodeExecutor( backend="docker", backend_params={ - "docker_image": "manimcommunity/manim:latest", + "image": "manimcommunity/manim:latest", }, ) diff --git a/poetry.lock b/poetry.lock index b7af4dce7..b1dec9678 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,24 +1,5 @@ # This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. -[[package]] -name = "aiodocker" -version = "0.24.0" -description = "A simple Docker HTTP API wrapper written with asyncio and aiohttp." -optional = true -python-versions = ">=3.9.0" -files = [ - {file = "aiodocker-0.24.0-py3-none-any.whl", hash = "sha256:2199b7b01f8ce68f9cabab7910ecb26192f6f3494163f1ccffe527b4c3875689"}, - {file = "aiodocker-0.24.0.tar.gz", hash = "sha256:661a6f9a479951f11f793031dcd5d55337e232c4ceaee69d51ceb885e5f16fac"}, -] - -[package.dependencies] -aiohttp = ">=3.8" - -[package.extras] -ci = ["aiohttp (==3.11.6)", "async-timeout (==5.0.1)", "multidict (==6.1.0)", "yarl (==1.17.2)"] -dev = ["async-timeout (==5.0.1)", "codecov (==2.1.13)", "mypy (==1.13.0)", "packaging (==24.1)", "pre-commit (>=3.5.0)", "pytest (==8.3.3)", "pytest-asyncio (==0.24.0)", "pytest-cov (==6.0.0)", "pytest-sugar (==1.0.0)", "ruff (==0.7.1)", "ruff-lsp (==0.0.58)", "towncrier (==24.8.0)"] -doc = ["alabaster (==1.0.0)", "sphinx (==8.1.3)", "sphinx-autodoc-typehints (==2.4.4)", "sphinxcontrib-asyncio (==0.3.0)"] - [[package]] name = "aiofiles" version = "25.1.0" @@ -350,6 +331,27 @@ files = [ docs = ["furo", "jaraco.packaging (>=9.3)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] testing = ["jaraco.test", "pytest (!=8.0.*)", "pytest (>=6,!=8.1.*)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)"] +[[package]] +name = "bespokelabs-sandbox" +version = "0.1.1" +description = "OpenRouter for Sandboxes - Unified Python API for cloud sandbox providers" +optional = false +python-versions = ">=3.10" +files = [ + {file = "bespokelabs_sandbox-0.1.1-py3-none-any.whl", hash = "sha256:fc4dcfa3f0a62c6dca5b99602edfe6ad5d98a411620a9306ab4cd4d70a1e0a09"}, + {file = "bespokelabs_sandbox-0.1.1.tar.gz", hash = "sha256:dee90da99ba1a24a8a93533b5602a123087952ed7c0ec93725a53a24fa850970"}, +] + +[package.extras] +all = ["daytona (>=0.1.0)", "docker (>=7.0.0)", "e2b-code-interpreter (>=1.0.0)", "modal (>=0.60.0)", "ray (>=2.0.0)", "tensorlake (>=0.1.0)"] +daytona = ["daytona (>=0.1.0)"] +dev = ["pytest (>=8.0)", "ruff (>=0.8.0)"] +docker = ["docker (>=7.0.0)"] +e2b = ["e2b-code-interpreter (>=1.0.0)"] +modal = ["modal (>=0.60.0)"] +ray = ["ray (>=2.0.0)"] +tensorlake = ["tensorlake (>=0.1.0)"] + [[package]] name = "blake3" version = "1.0.8" @@ -1173,6 +1175,28 @@ files = [ {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, ] +[[package]] +name = "docker" +version = "7.1.0" +description = "A Python library for the Docker Engine API." +optional = true +python-versions = ">=3.8" +files = [ + {file = "docker-7.1.0-py3-none-any.whl", hash = "sha256:c96b93b7f0a746f9e77d325bcfb87422a3d8bd4f03136ae8a85b37f1898d5fc0"}, + {file = "docker-7.1.0.tar.gz", hash = "sha256:ad8c70e6e3f8926cb8a92619b832b4ea5299e2831c14284663184e200546fa6c"}, +] + +[package.dependencies] +pywin32 = {version = ">=304", markers = "sys_platform == \"win32\""} +requests = ">=2.26.0" +urllib3 = ">=1.26.0" + +[package.extras] +dev = ["coverage (==7.2.7)", "pytest (==7.4.2)", "pytest-cov (==4.1.0)", "pytest-timeout (==2.1.0)", "ruff (==0.1.8)"] +docs = ["myst-parser (==0.18.0)", "sphinx (==5.1.1)"] +ssh = ["paramiko (>=2.4.3)"] +websockets = ["websocket-client (>=1.3.0)"] + [[package]] name = "docstring-parser" version = "0.17.0" @@ -5073,6 +5097,35 @@ files = [ {file = "pytz-2026.1.post1.tar.gz", hash = "sha256:3378dde6a0c3d26719182142c56e60c7f9af7e968076f31aae569d72a0358ee1"}, ] +[[package]] +name = "pywin32" +version = "311" +description = "Python for Window Extensions" +optional = true +python-versions = "*" +files = [ + {file = "pywin32-311-cp310-cp310-win32.whl", hash = "sha256:d03ff496d2a0cd4a5893504789d4a15399133fe82517455e78bad62efbb7f0a3"}, + {file = "pywin32-311-cp310-cp310-win_amd64.whl", hash = "sha256:797c2772017851984b97180b0bebe4b620bb86328e8a884bb626156295a63b3b"}, + {file = "pywin32-311-cp310-cp310-win_arm64.whl", hash = "sha256:0502d1facf1fed4839a9a51ccbcc63d952cf318f78ffc00a7e78528ac27d7a2b"}, + {file = "pywin32-311-cp311-cp311-win32.whl", hash = "sha256:184eb5e436dea364dcd3d2316d577d625c0351bf237c4e9a5fabbcfa5a58b151"}, + {file = "pywin32-311-cp311-cp311-win_amd64.whl", hash = "sha256:3ce80b34b22b17ccbd937a6e78e7225d80c52f5ab9940fe0506a1a16f3dab503"}, + {file = "pywin32-311-cp311-cp311-win_arm64.whl", hash = "sha256:a733f1388e1a842abb67ffa8e7aad0e70ac519e09b0f6a784e65a136ec7cefd2"}, + {file = "pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31"}, + {file = "pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067"}, + {file = "pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852"}, + {file = "pywin32-311-cp313-cp313-win32.whl", hash = "sha256:f95ba5a847cba10dd8c4d8fefa9f2a6cf283b8b88ed6178fa8a6c1ab16054d0d"}, + {file = "pywin32-311-cp313-cp313-win_amd64.whl", hash = "sha256:718a38f7e5b058e76aee1c56ddd06908116d35147e133427e59a3983f703a20d"}, + {file = "pywin32-311-cp313-cp313-win_arm64.whl", hash = "sha256:7b4075d959648406202d92a2310cb990fea19b535c7f4a78d3f5e10b926eeb8a"}, + {file = "pywin32-311-cp314-cp314-win32.whl", hash = "sha256:b7a2c10b93f8986666d0c803ee19b5990885872a7de910fc460f9b0c2fbf92ee"}, + {file = "pywin32-311-cp314-cp314-win_amd64.whl", hash = "sha256:3aca44c046bd2ed8c90de9cb8427f581c479e594e99b5c0bb19b29c10fd6cb87"}, + {file = "pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42"}, + {file = "pywin32-311-cp38-cp38-win32.whl", hash = "sha256:6c6f2969607b5023b0d9ce2541f8d2cbb01c4f46bc87456017cf63b73f1e2d8c"}, + {file = "pywin32-311-cp38-cp38-win_amd64.whl", hash = "sha256:c8015b09fb9a5e188f83b7b04de91ddca4658cee2ae6f3bc483f0b21a77ef6cd"}, + {file = "pywin32-311-cp39-cp39-win32.whl", hash = "sha256:aba8f82d551a942cb20d4a83413ccbac30790b50efb89a75e4f586ac0bb8056b"}, + {file = "pywin32-311-cp39-cp39-win_amd64.whl", hash = "sha256:e0c4cfb0621281fe40387df582097fd796e80430597cb9944f0ae70447bacd91"}, + {file = "pywin32-311-cp39-cp39-win_arm64.whl", hash = "sha256:62ea666235135fee79bb154e695f3ff67370afefd71bd7fea7512fc70ef31e3d"}, +] + [[package]] name = "pywin32-ctypes" version = "0.2.3" @@ -7450,11 +7503,11 @@ test = ["big-O", "jaraco.functools", "jaraco.itertools", "jaraco.test", "more_it type = ["pytest-mypy"] [extras] -code-execution = ["aiodocker", "e2b-code-interpreter", "ray"] +code-execution = ["docker", "e2b-code-interpreter", "ray"] finetune = ["tinker"] vllm = ["vllm"] [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "0a151c57f4090f3e3a7c6364b262b816b7a369fc4a26f3c6e6ef4b2f7f16377d" +content-hash = "ef2e1cff4c0688dee0e30235027af0c96918a123e6ba3874d40b09dfb62da34e" diff --git a/pyproject.toml b/pyproject.toml index c32b02e25..1e7c9eabf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,9 +35,10 @@ aiofiles = ">=22.0,<26.0" anthropic = "^0.84.0" vertexai = "1.71.1" vllm = { version = "^0.6.3", optional = true } -aiodocker = { version = "^0.24.0", optional = true } -ray = { version = "^2.41.0", optional = true } -e2b-code-interpreter = { version = "1.0.3", optional = true } +bespokelabs-sandbox = ">=0.1.1" +docker = { version = ">=7.0.0", optional = true } +ray = { version = ">=2.0.0", optional = true } +e2b-code-interpreter = { version = ">=1.0.0", optional = true } tinker = { version = "^0.8.0", optional = true, python = ">=3.11" } posthog = "^3.11.0" @@ -45,7 +46,7 @@ mistralai = "^1.5.1" [tool.poetry.extras] vllm = ["vllm"] -code_execution = ["ray", "aiodocker", "e2b-code-interpreter"] +code_execution = ["docker", "ray", "e2b-code-interpreter"] finetune = ["tinker"] [tool.poetry.group.dev.dependencies] diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/_factory.py b/src/bespokelabs/curator/code_executor/code_execution_backend/_factory.py index 33f261c91..5d7d2906e 100644 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/_factory.py +++ b/src/bespokelabs/curator/code_executor/code_execution_backend/_factory.py @@ -1,31 +1,35 @@ +"""Factory for creating code execution backends.""" + from typing import Optional from bespokelabs.curator.code_executor.types import CodeExecutionBackendConfig +# Map old backend names to sandbox backend names +_BACKEND_ALIASES = { + "multiprocessing": "local", +} + class _CodeExecutionBackendFactory: - @classmethod - def create(cls, backend: str, backend_params: Optional[CodeExecutionBackendConfig] = None): - if backend_params is None: - backend_params = {} + """Factory for creating code execution backends.""" - if backend == "multiprocessing": - from bespokelabs.curator.code_executor.code_execution_backend.multiprocessing_backend import MultiprocessingCodeExecutionBackend + @classmethod + def create(cls, backend: str, backend_params: Optional[dict] = None): + """Create a code execution backend. - _code_executor = MultiprocessingCodeExecutionBackend(CodeExecutionBackendConfig(**backend_params)) - elif backend == "docker": - from bespokelabs.curator.code_executor.code_execution_backend.docker_backend import DockerCodeExecutionBackend + Args: + backend: Backend name (e.g. "local", "docker", "e2b", "modal", "daytona") + backend_params: Configuration parameters for the backend - _code_executor = DockerCodeExecutionBackend(CodeExecutionBackendConfig(**backend_params)) - elif backend == "ray": - from bespokelabs.curator.code_executor.code_execution_backend.ray_backend import RayCodeExecutionBackend + Returns: + SandboxCodeExecutionBackend instance + """ + if backend_params is None: + backend_params = {} - _code_executor = RayCodeExecutionBackend(CodeExecutionBackendConfig(**backend_params)) - elif backend == "e2b": - from bespokelabs.curator.code_executor.code_execution_backend.e2b_backend import E2BCodeExecutionBackend + backend_name = _BACKEND_ALIASES.get(backend, backend) + config = CodeExecutionBackendConfig(**backend_params) - _code_executor = E2BCodeExecutionBackend(CodeExecutionBackendConfig(**backend_params)) - else: - raise ValueError(f"Unsupported backend: {backend}") + from bespokelabs.curator.code_executor.code_execution_backend.sandbox_backend import SandboxCodeExecutionBackend - return _code_executor + return SandboxCodeExecutionBackend(config, backend_name=backend_name) diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py index e38779798..7eed3b5b3 100644 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py +++ b/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py @@ -2,11 +2,9 @@ import asyncio import glob -import io import json import os import resource -import tarfile import time from abc import abstractmethod from typing import TYPE_CHECKING, List, Optional @@ -657,46 +655,3 @@ async def append_execution_response(self, data: dict, filename: str) -> None: await f.write(json_string + "\n") logger.debug(f"Successfully appended response to {filename}") - @classmethod - def _create_temp_file(cls, content: str, execution_directory: str) -> str: - """Create a temporary file with the given content. - - Args: - content: Content to write to temp file - execution_directory: Directory to create the temp file in - - Returns: - Path to the created temp file - """ - # create execution directory if it doesn't exist - os.makedirs(execution_directory, exist_ok=True) - - temp_file_path = os.path.join(execution_directory, "program.py") - with open(temp_file_path, "w", encoding="utf-8") as temp_file: - temp_file.write(content) - return temp_file_path - - @classmethod - def _get_created_files(cls, program_dir: str) -> bytes: - """Get any files created during code execution. - - Args: - program_dir: Directory containing the executed program and created files - - Returns: - Bytes containing a zip archive of any created files - """ - # Create a BytesIO object to store the zip data - tar_buffer = io.BytesIO() - - # Create a zip archive containing all files in program_dir - with tarfile.open(fileobj=tar_buffer, mode="w") as tar: - # Add all files in program_dir to the archive - for filename in os.listdir(program_dir): - file_path = os.path.join(program_dir, filename) - if os.path.isfile(file_path): - tar.add(file_path, arcname=filename) - - # Get the bytes from the buffer - tar_buffer.seek(0) - return str(tar_buffer.getvalue()) diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/docker_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/docker_backend.py deleted file mode 100644 index 313d15d65..000000000 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/docker_backend.py +++ /dev/null @@ -1,344 +0,0 @@ -"""Docker Code Execution Backend.""" - -import asyncio -import io -import multiprocessing -import os -import tarfile -import tempfile -from concurrent.futures import ProcessPoolExecutor -from functools import partial - -import docker -from docker.errors import DockerException - -from bespokelabs.curator.code_executor.code_execution_backend.base_backend import BaseCodeExecutionBackend -from bespokelabs.curator.code_executor.types import CodeAPIRequest, CodeExecutionOutput, CodeExecutionRequestParams, CodeExecutionResponse -from bespokelabs.curator.log import logger - - -class DockerCodeExecutionBackend(BaseCodeExecutionBackend): - """Docker-based code execution backend.""" - - PYTHON_IMAGE = "python:3.11-slim" - WORKSPACE_DIR = "/workspace" - - def __init__(self, config): - """Initialize the backend.""" - super().__init__(config) - self.config = config - self.PYTHON_IMAGE = config.docker_image if config.docker_image else self.PYTHON_IMAGE - self.client = docker.from_env() - - # Initialize process pool for parallel execution - self.max_workers = getattr(config, "max_workers", multiprocessing.cpu_count()) - self.process_pool = ProcessPoolExecutor(max_workers=self.max_workers) - - logger.debug(f"Initialized Docker backend with {self.max_workers} workers") - - def _ensure_image(self): - """Ensure required Docker image is pulled.""" - try: - self.client.images.get(self.PYTHON_IMAGE) - except docker.errors.ImageNotFound: - logger.info(f"Pulling {self.PYTHON_IMAGE} image...") - self.client.images.pull(self.PYTHON_IMAGE) - except DockerException as e: - raise RuntimeError(f"Error ensuring Docker image: {str(e)}") from e - - async def execute_request(self, request: CodeAPIRequest) -> CodeExecutionResponse: - """Execute a single request in a Docker container using a process pool.""" - return await self.execute_standard_input_request( - request.execution_request.code, request.execution_request.code_input, request.execution_request.execution_params - ) - - @classmethod - def _create_temp_file(cls, content: str) -> str: - """Create a temporary file with the given content.""" - with tempfile.NamedTemporaryFile(delete=False, mode="w", encoding="utf-8") as temp_file: - temp_file.write(content) - return temp_file.name - - def _create_tar_archive(self, data: bytes, filename: str) -> io.BytesIO: - """Create a tar archive containing a single file.""" - tarstream = io.BytesIO() - with tarfile.TarFile(fileobj=tarstream, mode="w") as tar: - tarinfo = tarfile.TarInfo(name=filename) - tarinfo.size = len(data) - tar.addfile(tarinfo, io.BytesIO(data)) - tarstream.seek(0) - return tarstream - - def _get_container_file(self, container, file_path: str) -> str: - """Get content of a file from a Docker container.""" - try: - bits, _ = container.get_archive(file_path) - bio = io.BytesIO() - for chunk in bits: - bio.write(chunk) - bio.seek(0) - - with tarfile.open(fileobj=bio) as tar: - member = tar.next() - f = tar.extractfile(member) - if f: - return f.read().decode("utf-8") - return "" - - except docker.errors.NotFound: - return "" - - def _setup_volume_with_files(self, volume, code_file: str, input_file: str): - """Set up a volume with code and input files.""" - setup_container = self.client.containers.create( - self.PYTHON_IMAGE, command=["sleep", "5"], volumes={volume.name: {"bind": self.WORKSPACE_DIR, "mode": "rw"}} - ) - - try: - setup_container.start() - for src, dest in [(code_file, "program.py"), (input_file, "input.txt")]: - with open(src, "rb") as f: - data = f.read() - tarstream = self._create_tar_archive(data, dest) - setup_container.put_archive(self.WORKSPACE_DIR, tarstream) - finally: - setup_container.stop() - setup_container.remove() - - def _get_created_files(self, container: docker.models.containers.Container, directory: str) -> str: - """Get any files created during code execution.""" - # Get all files in the directory as a tar archive - bits, _ = container.get_archive(directory, encode_stream=True) - bio = io.BytesIO() - for chunk in bits: - bio.write(chunk) - bio.seek(0) - return str(bio.getvalue()) - - def _create_execution_container(self, volume_name: str) -> docker.models.containers.Container: - """Create the main execution container.""" - # TODO: Add a name to the container and volume so that they are unique for a each row - return self.client.containers.create( - self.PYTHON_IMAGE, - command=[ # noqa: E501 - f"bash -c python {self.WORKSPACE_DIR}/program.py < {self.WORKSPACE_DIR}/input.txt > {self.WORKSPACE_DIR}/output.txt 2> {self.WORKSPACE_DIR}/error.txt; echo $? > {self.WORKSPACE_DIR}/exit_code.txt" # noqa: E501 - ], - volumes={volume_name: {"bind": self.WORKSPACE_DIR, "mode": "rw"}}, - working_dir=self.WORKSPACE_DIR, - entrypoint=["/bin/sh", "-c"], - user="root", - ) - - def _handle_execution_result(self, container, status_code: int) -> CodeExecutionOutput: - """Handle the execution result and create appropriate output.""" - stdout = self._get_container_file(container, f"{self.WORKSPACE_DIR}/output.txt") - stderr = self._get_container_file(container, f"{self.WORKSPACE_DIR}/error.txt") - files = self._get_created_files(container, self.WORKSPACE_DIR) - - if status_code == 0: - return CodeExecutionOutput( - message="success", - stdout=stdout, - stderr=stderr, - files=files, - ) - - # Create a more descriptive error message - error_message = f"Program exited with status code {status_code}" - if stderr: - error_message = f"{error_message}\n\nError details:\n{stderr}" - - return CodeExecutionOutput( - message="error", - error=error_message, - stdout=stdout, - stderr=stderr, - files=files, - ) - - async def execute_standard_input_request(self, code: str, code_input: str, execution_params: CodeExecutionRequestParams) -> CodeExecutionOutput: - """Execute code in Docker container using a worker from the process pool.""" - # Create a function that can be executed in a separate process - execution_func = partial( - self._execute_in_process, code=code, code_input=code_input, timeout=120, python_image=self.PYTHON_IMAGE, workspace_dir=self.WORKSPACE_DIR - ) - - # Submit the task to the process pool - loop = asyncio.get_running_loop() - result = await loop.run_in_executor(self.process_pool, execution_func) - - return result - - @staticmethod - def _execute_in_process(code: str, code_input: str, timeout: int, python_image: str, workspace_dir: str) -> CodeExecutionOutput: - """Static method to execute code in a separate process.""" - try: - # Create a new Docker client in this process - client = docker.from_env() - - # Create temporary files - with ( - tempfile.NamedTemporaryFile(delete=False, mode="w", encoding="utf-8") as code_file, - tempfile.NamedTemporaryFile(delete=False, mode="w", encoding="utf-8") as input_file, - ): - code_file.write(code) - code_file_path = code_file.name - - input_file.write(code_input) - input_file_path = input_file.name - - # Ensure image exists - try: - client.images.get(python_image) - except docker.errors.ImageNotFound: - client.images.pull(python_image) - except DockerException as e: - raise RuntimeError(f"Error ensuring Docker image: {str(e)}") from e - - container = None - volume = None - - try: - volume = client.volumes.create() - - # Set up volume with files - setup_container = client.containers.create(python_image, command=["sleep", "5"], volumes={volume.name: {"bind": workspace_dir, "mode": "rw"}}) - - try: - setup_container.start() - for src, dest in [(code_file_path, "program.py"), (input_file_path, "input.txt")]: - with open(src, "rb") as f: - data = f.read() - tarstream = DockerCodeExecutionBackend._create_tar_archive_static(data, dest) - setup_container.put_archive(workspace_dir, tarstream) - finally: - setup_container.stop() - setup_container.remove() - - # Create and run execution container - cmd = ( - f"echo 'DEBUG: Starting container execution' && " - f"echo 'DEBUG: Workspace contents:' && ls -la {workspace_dir} && " - f"python {workspace_dir}/program.py < {workspace_dir}/input.txt > {workspace_dir}/output.txt 2> {workspace_dir}/error.txt; " - f"exit_code=$?; " - f"echo $exit_code > {workspace_dir}/exit_code.txt; " - f"echo 'DEBUG: Execution finished with code:' $exit_code; " - f"exit $exit_code" - ) - - container = client.containers.create( - python_image, - command=[cmd], - volumes={volume.name: {"bind": workspace_dir, "mode": "rw"}}, - working_dir=workspace_dir, - entrypoint=["/bin/bash", "-c"], - user="root", - ) - - container.start() - result = container.wait(timeout=timeout) - - # Handle results - return DockerCodeExecutionBackend._handle_execution_result_static(container, result["StatusCode"], workspace_dir) - - except docker.errors.ContainerError as e: - return CodeExecutionOutput(message="error", error=f"Docker container error: {str(e)}\n\nExit code: {e.exit_status}\n\nStderr: {e.stderr}") - except docker.errors.ImageNotFound as e: - return CodeExecutionOutput(message="error", error=f"Docker image not found: {str(e)}") - except docker.errors.APIError as e: - return CodeExecutionOutput(message="error", error=f"Docker API error: {str(e)}") - except Exception as e: - return CodeExecutionOutput(message="error", error=f"Error during code execution: {str(e)}\n\nType: {type(e).__name__}") - finally: - # Clean up - try: - if container: - logs = container.logs(stdout=True, stderr=True).decode("utf-8", errors="replace") - print(logs) - # container.remove(force=True) - # if volume: - # volume.remove(force=True) - for f in [code_file_path, input_file_path]: - if os.path.exists(f): - os.unlink(f) - client.close() - except Exception: - pass - - except Exception as e: - return CodeExecutionOutput(message="error", error=f"Process execution error: {str(e)}") - - @staticmethod - def _create_tar_archive_static(data: bytes, filename: str) -> io.BytesIO: - """Static version of _create_tar_archive for use in processes.""" - tarstream = io.BytesIO() - with tarfile.TarFile(fileobj=tarstream, mode="w") as tar: - tarinfo = tarfile.TarInfo(name=filename) - tarinfo.size = len(data) - tar.addfile(tarinfo, io.BytesIO(data)) - tarstream.seek(0) - return tarstream - - @staticmethod - def _handle_execution_result_static(container, status_code: int, workspace_dir: str) -> CodeExecutionOutput: - """Static version of _handle_execution_result for use in processes.""" - stdout = DockerCodeExecutionBackend._get_container_file_static(container, f"{workspace_dir}/output.txt") - stderr = DockerCodeExecutionBackend._get_container_file_static(container, f"{workspace_dir}/error.txt") - - # Get created files - try: - bits, _ = container.get_archive(workspace_dir, encode_stream=True) - bio = io.BytesIO() - for chunk in bits: - bio.write(chunk) - bio.seek(0) - files = str(bio.getvalue()) - except Exception: - files = "" - - if status_code == 0: - return CodeExecutionOutput( - message="success", - stdout=stdout, - stderr=stderr, - files=files, - ) - - # Create a more descriptive error message - error_message = f"Program exited with status code {status_code}" - if stderr: - error_message = f"{error_message}\n\nError details:\n{stderr}" - - return CodeExecutionOutput( - message="error", - error=error_message, - stdout=stdout, - stderr=stderr, - files=files, - ) - - @staticmethod - def _get_container_file_static(container, file_path: str) -> str: - """Static version of _get_container_file for use in processes.""" - try: - bits, _ = container.get_archive(file_path) - bio = io.BytesIO() - for chunk in bits: - bio.write(chunk) - bio.seek(0) - - with tarfile.open(fileobj=bio) as tar: - member = tar.next() - f = tar.extractfile(member) - if f: - return f.read().decode("utf-8") - return "" - - except docker.errors.NotFound: - return "" - - async def shutdown(self): - """Cleanup resources including process pool.""" - self.process_pool.shutdown(wait=True) - self.client.close() - logger.debug("Shutting down Docker backend and process pool") diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/e2b_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/e2b_backend.py deleted file mode 100644 index 3faeacf2e..000000000 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/e2b_backend.py +++ /dev/null @@ -1,78 +0,0 @@ -"""E2B Code Execution Backend. - -See https://e2b.dev/ for more information. -""" - -import logging - -from bespokelabs.curator.code_executor.code_execution_backend.base_backend import BaseCodeExecutionBackend -from bespokelabs.curator.code_executor.types import CodeAPIRequest, CodeExecutionOutput, CodeExecutionRequest - -logger = logging.getLogger(__name__) - - -class E2BCodeExecutionBackend(BaseCodeExecutionBackend): - """E2B Code Execution Backend.""" - - def __init__(self, config): - """Initialize the backend.""" - super().__init__(config) - self.config = config - try: - from e2b_code_interpreter import AsyncSandbox - except ModuleNotFoundError as e: - raise ModuleNotFoundError("""Please install 'e2b' extra to use E2BExecutor""") from e - self.Sandbox = AsyncSandbox # Store the class instead of instance - - async def execute_request(self, request: CodeAPIRequest) -> CodeExecutionOutput: - """Execute a single request.""" - self.sandbox = await self.Sandbox.create() # Initialize sandbox for this request - try: - return await self.execute_standard_input_request(request.execution_request) - finally: - await self.sandbox.kill() - - async def execute_standard_input_request(self, request: CodeExecutionRequest) -> CodeExecutionOutput: - """Execute code with E2B sandbox. - - Args: - request: CodeExecutionRequest - - Returns: - CodeExecutionOutput: Execution results - """ - try: - # Run program in E2B sandbox - # save code and input to files in the sandbox root directory - await self.sandbox.files.write("/code.py", request.code) - await self.sandbox.files.write("/input.txt", request.code_input) - - result = await self.sandbox.run_code( - "python /code.py < /input.txt > /output.txt 2> /error.txt", - language="bash", - timeout=request.execution_params.timeout, - ) - - if result.error: - return CodeExecutionOutput( - message="error", - error=f"{result.error.name}: {result.error.value}\n{result.error.traceback}", - ) - - return CodeExecutionOutput( - message="success", - stdout=await self.sandbox.files.read("/output.txt"), - stderr=await self.sandbox.files.read("/error.txt"), - # files=await self.sandbox.files.read("/"), - # TODO: e2b doesn't support seamless reading of an entire directory - ) - - except Exception as e: - return CodeExecutionOutput( - message="error", - error=str(e), - ) - - def __del__(self): - """Clean up when object is destroyed.""" - logger.debug("Shutting down E2B backend") diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/multiprocessing_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/multiprocessing_backend.py deleted file mode 100644 index 8018de103..000000000 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/multiprocessing_backend.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Multiprocessing Code Execution Backend.""" - -import asyncio -import os -import subprocess -from concurrent.futures import ProcessPoolExecutor - -from bespokelabs.curator.code_executor.code_execution_backend.base_backend import BaseCodeExecutionBackend -from bespokelabs.curator.code_executor.types import CodeAPIRequest, CodeExecutionOutput, CodeExecutionRequest -from bespokelabs.curator.log import logger - - -class MultiprocessingCodeExecutionBackend(BaseCodeExecutionBackend): - """Multiprocessing Code Execution Backend.""" - - def __init__(self, config): - """Initialize the backend.""" - super().__init__(config) - self.config = config - self.process_pool = ProcessPoolExecutor(max_workers=os.cpu_count()) - logger.debug(f"Initialized multiprocessing backend with {os.cpu_count()} workers") - - async def execute_request(self, request: CodeAPIRequest) -> CodeExecutionOutput: - """Execute a single request.""" - loop = asyncio.get_running_loop() - - return await loop.run_in_executor( - self.process_pool, - self.execute_standard_input_request, - request.execution_request, - ) - - @classmethod - def execute_standard_input_request(cls, request: CodeExecutionRequest) -> CodeExecutionOutput: - """Execute code with function calls and test cases. - - Args: - request: CodeExecutionRequest - - Returns: - CodeExecutionOutput: Execution results - """ - temp_program_path = None - output = None - try: - temp_program_path = cls._create_temp_file(request.code, request.execution_directory) - try: - # Get directory containing the program file - program_dir = os.path.dirname(temp_program_path) - - # Run program from its directory - result = subprocess.run( - ["python", "program.py"], - input=request.code_input, - text=True, - capture_output=True, - timeout=request.execution_params.timeout, - cwd=request.execution_directory, - ) - - output = CodeExecutionOutput( - message="success", - stdout=result.stdout, - stderr=result.stderr, - files=cls._get_created_files(program_dir), - ) - - except subprocess.TimeoutExpired: - output = CodeExecutionOutput( - message="timeout", - error=f"Execution timed out after {request.execution_params.timeout}s", - ) - - except Exception as e: - output = CodeExecutionOutput( - message="error", - error=str(e), - ) - - finally: - if temp_program_path: - os.unlink(temp_program_path) - - return output - - def __del__(self): - """Clean up pool when object is destroyed.""" - self.process_pool.shutdown(wait=True) - logger.debug("Shutting down multiprocessing backend") diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/ray_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/ray_backend.py deleted file mode 100644 index 2fcdb5baa..000000000 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/ray_backend.py +++ /dev/null @@ -1,114 +0,0 @@ -"""Ray-based Code Execution Backend.""" - -import asyncio -import logging -import os -import subprocess - -import ray - -from bespokelabs.curator.code_executor.code_execution_backend.base_backend import BaseCodeExecutionBackend -from bespokelabs.curator.code_executor.types import CodeAPIRequest, CodeExecutionOutput, CodeExecutionRequest - - -@ray.remote -class CodeExecutorActor: - """Ray actor for executing code in isolation.""" - - def __init__(self): - """Initialize the actor.""" - self.temp_files = [] - - def __del__(self): - """Cleanup temporary files on actor shutdown.""" - if hasattr(self, "temp_files") and self.temp_files: - for temp_file in self.temp_files: - try: - if os.path.exists(temp_file): - os.unlink(temp_file) - except Exception: - pass - - @classmethod - def execute_standard_input_request(cls, request: CodeExecutionRequest) -> CodeExecutionOutput: - """Execute code with function calls and test cases. - - Args: - request: CodeExecutionRequest - - Returns: - CodeExecutionOutput: Execution results - """ - temp_program_path = None - output = None - try: - temp_program_path = BaseCodeExecutionBackend._create_temp_file(request.code, request.execution_directory) - try: - # Get directory containing the program file - program_dir = os.path.dirname(temp_program_path) - - # Run program from its directory - result = subprocess.run( - ["python", "program.py"], - input=request.code_input, - text=True, - capture_output=True, - timeout=request.execution_params.timeout, - cwd=request.execution_directory, - ) - - output = CodeExecutionOutput( - message="success", - stdout=result.stdout, - stderr=result.stderr, - files=BaseCodeExecutionBackend._get_created_files(program_dir), - ) - - except subprocess.TimeoutExpired: - output = CodeExecutionOutput( - message="timeout", - error=f"Execution timed out after {request.execution_params.timeout}s", - ) - - except Exception as e: - output = CodeExecutionOutput( - message="error", - error=str(e), - ) - - finally: - if temp_program_path: - os.unlink(temp_program_path) - - return output - - -class RayCodeExecutionBackend(BaseCodeExecutionBackend): - """Ray-based code execution backend.""" - - def __init__(self, config): - """Initialize the backend.""" - super().__init__(config) - self.config = config - self.logger = logging.getLogger(__name__) - - # Initialize Ray if needed - if not ray.is_initialized(): - if self.config.base_url: - ray.init(address=self.config.base_url) - else: - ray.init(num_cpus=os.cpu_count()) - - self.executor = CodeExecutorActor.remote() - - async def execute_request(self, request: CodeAPIRequest) -> CodeExecutionOutput: - """Execute a single request using Ray actors.""" - # Execute request in the remote executor - result = await asyncio.get_event_loop().run_in_executor(None, ray.get, self.executor.execute_standard_input_request.remote(request.execution_request)) - - return result - - async def shutdown(self): - """Cleanup resources when shutting down the backend.""" - if ray.is_initialized(): - ray.shutdown() diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py new file mode 100644 index 000000000..e1273a4d2 --- /dev/null +++ b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py @@ -0,0 +1,146 @@ +"""Sandbox Code Execution Backend using bespokelabs-sandbox.""" + +import asyncio +import io +import os +import tarfile +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from bespokelabs.curator.code_executor.code_execution_backend.base_backend import BaseCodeExecutionBackend +from bespokelabs.curator.code_executor.types import CodeAPIRequest, CodeExecutionOutput +from bespokelabs.curator.log import logger + +WORKSPACE_DIR = "/workspace" + + +class SandboxCodeExecutionBackend(BaseCodeExecutionBackend): + """Code execution backend using bespokelabs-sandbox.""" + + def __init__(self, config, backend_name="local"): + """Initialize the sandbox backend. + + Args: + config: Backend configuration + backend_name: Sandbox backend name (e.g. "local", "docker", "e2b", "modal", "daytona") + """ + super().__init__(config) + self.config = config + self.backend_name = backend_name + self.thread_pool = ThreadPoolExecutor(max_workers=os.cpu_count()) + + # Build sandbox constructor kwargs from config + self.sandbox_kwargs = {} + # Prefer explicit `image`, fall back to legacy `docker_image` + image = config.image or config.docker_image + if image: + self.sandbox_kwargs["image"] = image + if config.base_url: + self.sandbox_kwargs["base_url"] = config.base_url + + logger.debug(f"Initialized sandbox backend with backend={backend_name}") + + @property + def backend(self) -> str: + """Backend property.""" + return self.backend_name + + async def execute_request(self, request: CodeAPIRequest) -> CodeExecutionOutput: + """Execute a single request using a sandbox.""" + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + self.thread_pool, + partial( + _execute_in_sandbox, + code=request.execution_request.code, + code_input=request.execution_request.code_input or "", + timeout=request.execution_request.execution_params.timeout if request.execution_request.execution_params else 10, + backend_name=self.backend_name, + sandbox_kwargs=self.sandbox_kwargs, + ), + ) + + def __del__(self): + """Clean up thread pool when object is destroyed.""" + self.thread_pool.shutdown(wait=True) + logger.debug("Shutting down sandbox backend") + + +def _execute_in_sandbox( + code: str, + code_input: str, + timeout: int, + backend_name: str, + sandbox_kwargs: dict, +) -> CodeExecutionOutput: + """Execute code in a bespokelabs-sandbox. + + This is a module-level function so it can be used with ProcessPoolExecutor if needed. + """ + from bespokelabs.sandbox import Sandbox + + try: + kwargs = {**sandbox_kwargs} + kwargs.setdefault("timeout_secs", timeout + 30) + + with Sandbox(backend_name, **kwargs) as sb: + sb.write_file(f"{WORKSPACE_DIR}/program.py", code) + sb.write_file(f"{WORKSPACE_DIR}/input.txt", code_input) + + result = sb.execute_command("bash", args=["-c", f"cd {WORKSPACE_DIR} && python program.py < input.txt"]) + + files = _collect_sandbox_files(sb) + + if result.exit_code == 0: + return CodeExecutionOutput( + message="success", + stdout=result.stdout, + stderr=result.stderr, + files=files, + ) + + error_message = f"Program exited with status code {result.exit_code}" + if result.stderr: + error_message = f"{error_message}\n\nError details:\n{result.stderr}" + + return CodeExecutionOutput( + message="error", + error=error_message, + stdout=result.stdout, + stderr=result.stderr, + files=files, + ) + + except Exception as e: + return CodeExecutionOutput( + message="error", + error=str(e), + ) + + +def _collect_sandbox_files(sandbox) -> str: + """Collect files created during execution as a tar archive string. + + Preserves directory structure relative to WORKSPACE_DIR so that + nested output paths (e.g. media/videos/…) survive round-tripping. + """ + try: + file_list = sandbox.list_files(WORKSPACE_DIR) + tar_buffer = io.BytesIO() + + with tarfile.open(fileobj=tar_buffer, mode="w") as tar: + for file_info in file_list: + if file_info.is_dir: + continue + rel_path = os.path.relpath(file_info.path, WORKSPACE_DIR) + content = sandbox.read_file(file_info.path) + if isinstance(content, str): + content = content.encode() + tarinfo = tarfile.TarInfo(name=rel_path) + tarinfo.size = len(content) + tar.addfile(tarinfo, io.BytesIO(content)) + + tar_buffer.seek(0) + return str(tar_buffer.getvalue()) + except Exception: + return "" diff --git a/src/bespokelabs/curator/code_executor/code_executor.py b/src/bespokelabs/curator/code_executor/code_executor.py index fc9966052..31b620b18 100644 --- a/src/bespokelabs/curator/code_executor/code_executor.py +++ b/src/bespokelabs/curator/code_executor/code_executor.py @@ -48,13 +48,13 @@ def code_output(self, row: dict, execution_output: Any): def __init__( self, - backend: str = "multiprocessing", + backend: str = "local", backend_params: dict = None, ): """Initialize the CodeExecutor with specified backend and parameters. Args: - backend: The execution backend to use ("multiprocessing" by default) + backend: The sandbox backend to use (e.g. "local", "docker", "e2b", "modal", "daytona") backend_params: Configuration parameters for the backend """ self._code_executor = _CodeExecutionBackendFactory.create(backend=backend, backend_params=backend_params) diff --git a/src/bespokelabs/curator/code_executor/types.py b/src/bespokelabs/curator/code_executor/types.py index 5c3364ce0..28435160c 100644 --- a/src/bespokelabs/curator/code_executor/types.py +++ b/src/bespokelabs/curator/code_executor/types.py @@ -68,5 +68,7 @@ class CodeExecutionBackendConfig(BaseModel): max_requests_per_minute: int = 10000 max_retries: int = 3 seconds_to_pause_on_rate_limit: int = 10 - base_url: Optional[str] = None + image: Optional[str] = None + # Legacy aliases — mapped to sandbox params at runtime docker_image: Optional[str] = None + base_url: Optional[str] = None diff --git a/tests/code_executor/test_backend.py b/tests/code_executor/test_backend.py index 6e395a56d..79adc149e 100644 --- a/tests/code_executor/test_backend.py +++ b/tests/code_executor/test_backend.py @@ -1,9 +1,6 @@ -import ast import asyncio -import io import json import os -import tarfile import tempfile import time from unittest.mock import Mock, patch @@ -77,53 +74,6 @@ def test_backend_property(backend): assert backend.backend == "test" -@pytest.mark.asyncio -async def test_create_temp_file(temp_dir): - content = "print('hello')" - file_path = BaseCodeExecutionBackend._create_temp_file(content, temp_dir) - - assert os.path.exists(file_path) - with open(file_path) as f: - assert f.read() == content - - -@pytest.mark.asyncio -async def test_get_created_files(temp_dir): - # Create some test files - test_files = {"test1.txt": "content1"} - - for path, content in test_files.items(): - full_path = os.path.join(temp_dir, path) - os.makedirs(os.path.dirname(full_path), exist_ok=True) - with open(full_path, "w") as f: - f.write(content) - - # Add program.py which should be included - with open(os.path.join(temp_dir, "program.py"), "w") as f: - f.write("print('test')") - - # Get tar bytes - tar_bytes = BaseCodeExecutionBackend._get_created_files(temp_dir) - - # Read back the tar contents - tar_buffer = io.BytesIO(ast.literal_eval(tar_bytes)) - with tarfile.open(fileobj=tar_buffer, mode="r") as tar: - files = tar.getnames() - # Convert paths to use forward slashes for consistency - files = [f.replace("\\", "/") for f in files] - - # Verify program.py is included (since base implementation includes all files) - assert "program.py" in files - - # Verify other files are included with correct content - for path, content in test_files.items(): - assert path in files - # Extract and read the file content - member = tar.extractfile(path) - assert member is not None - assert member.read().decode() == content - - @pytest.mark.asyncio async def test_append_execution_response(temp_dir): response_file = os.path.join(temp_dir, "responses.jsonl") @@ -152,8 +102,6 @@ async def test_validate_existing_response_file(backend, temp_dir): code_input="test input", original_row={"input": "test"}, execution_directory=temp_dir, - language="python", - timeout=30, ), attempts_left=3, code_formatter=None, diff --git a/tests/code_executor/test_code_execution.py b/tests/code_executor/test_code_execution.py index 773986aa1..962d1672a 100644 --- a/tests/code_executor/test_code_execution.py +++ b/tests/code_executor/test_code_execution.py @@ -1,15 +1,13 @@ -# test with multiprocessing backend +# test with sandbox backend import pytest from bespokelabs import curator -from tests.conftest import importorskip @pytest.mark.asyncio -async def test_simple_code_execution_multiprocessing(): - """Test simple code execution with basic input/output.""" +async def test_simple_code_execution_local(): + """Test simple code execution with local sandbox backend.""" - # Initialize backend class TestCodeExecutor(curator.CodeExecutor): def code(self, row): return """ @@ -24,33 +22,7 @@ def code_output(self, row, exec_output): row["output"] = exec_output.stdout return row - executor = TestCodeExecutor(backend="multiprocessing") - sample_data = [{"input": "Hello World multiprocessing"}] + executor = TestCodeExecutor(backend="local") + sample_data = [{"input": "Hello World local"}] result = executor(sample_data) - assert result[0]["output"] == "You entered: Hello World multiprocessing\n" - - -@pytest.mark.asyncio -async def test_simple_code_execution_ray(): - """Test simple code execution with basic input/output.""" - importorskip("ray") - - # Initialize backend - class TestCodeExecutor(curator.CodeExecutor): - def code(self, row): - return """ -input_value = input() -print(f"You entered: {input_value}") -""" - - def code_input(self, row): - return row["input"] - - def code_output(self, row, exec_output): - row["output"] = exec_output.stdout - return row - - executor = TestCodeExecutor(backend="ray") - sample_data = [{"input": "Hello World ray"}] - result = executor(sample_data) - assert result[0]["output"] == "You entered: Hello World ray\n" + assert result[0]["output"] == "You entered: Hello World local\n" diff --git a/tests/code_executor/test_data_types.py b/tests/code_executor/test_data_types.py index 340b02aed..aece35e54 100644 --- a/tests/code_executor/test_data_types.py +++ b/tests/code_executor/test_data_types.py @@ -101,15 +101,22 @@ def test_code_execution_backend_config(): assert config.max_requests_per_minute == 10000 assert config.max_retries == 3 assert config.seconds_to_pause_on_rate_limit == 10 + assert config.image is None + assert config.docker_image is None assert config.base_url is None # Test custom values - custom_config = CodeExecutionBackendConfig(max_requests_per_minute=5000, max_retries=5, seconds_to_pause_on_rate_limit=20, base_url="http://test.com") + custom_config = CodeExecutionBackendConfig(max_requests_per_minute=5000, max_retries=5, seconds_to_pause_on_rate_limit=20, image="python:3.11-slim") assert custom_config.max_requests_per_minute == 5000 assert custom_config.max_retries == 5 assert custom_config.seconds_to_pause_on_rate_limit == 20 - assert custom_config.base_url == "http://test.com" + assert custom_config.image == "python:3.11-slim" + + # Test legacy docker_image and base_url still accepted + legacy_config = CodeExecutionBackendConfig(docker_image="manimcommunity/manim:latest", base_url="http://ray-cluster:8265") + assert legacy_config.docker_image == "manimcommunity/manim:latest" + assert legacy_config.base_url == "http://ray-cluster:8265" def test_invalid_code_execution_result(): From ca02e0662d506e002942c0d67587f8f42d952920 Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 19:52:45 -0700 Subject: [PATCH 02/10] Pin bespokelabs-sandbox to 0.1.1 Co-Authored-By: Claude Opus 4.6 (1M context) --- poetry.lock | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index b1dec9678..69bbd2af5 100644 --- a/poetry.lock +++ b/poetry.lock @@ -7510,4 +7510,4 @@ vllm = ["vllm"] [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "ef2e1cff4c0688dee0e30235027af0c96918a123e6ba3874d40b09dfb62da34e" +content-hash = "e60aaa3e3492ae58637bd1d5d950a46518ef2b7cd2d2202dae0abd2286acea3d" diff --git a/pyproject.toml b/pyproject.toml index 1e7c9eabf..220ce3e40 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ aiofiles = ">=22.0,<26.0" anthropic = "^0.84.0" vertexai = "1.71.1" vllm = { version = "^0.6.3", optional = true } -bespokelabs-sandbox = ">=0.1.1" +bespokelabs-sandbox = "0.1.1" docker = { version = ">=7.0.0", optional = true } ray = { version = ">=2.0.0", optional = true } e2b-code-interpreter = { version = ">=1.0.0", optional = true } From a06a48ec834eb7e56a76f21ec4dc79db3d366aad Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 19:57:13 -0700 Subject: [PATCH 03/10] Remove redundant docker/ray/e2b deps and code_execution extra Backend SDKs are managed by bespokelabs-sandbox's own extras (e.g. pip install bespokelabs-sandbox[docker]). Co-Authored-By: Claude Opus 4.6 (1M context) --- poetry.lock | 90 +------------------------------------------------- pyproject.toml | 4 --- 2 files changed, 1 insertion(+), 93 deletions(-) diff --git a/poetry.lock b/poetry.lock index 69bbd2af5..b2224fd15 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1175,28 +1175,6 @@ files = [ {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, ] -[[package]] -name = "docker" -version = "7.1.0" -description = "A Python library for the Docker Engine API." -optional = true -python-versions = ">=3.8" -files = [ - {file = "docker-7.1.0-py3-none-any.whl", hash = "sha256:c96b93b7f0a746f9e77d325bcfb87422a3d8bd4f03136ae8a85b37f1898d5fc0"}, - {file = "docker-7.1.0.tar.gz", hash = "sha256:ad8c70e6e3f8926cb8a92619b832b4ea5299e2831c14284663184e200546fa6c"}, -] - -[package.dependencies] -pywin32 = {version = ">=304", markers = "sys_platform == \"win32\""} -requests = ">=2.26.0" -urllib3 = ">=1.26.0" - -[package.extras] -dev = ["coverage (==7.2.7)", "pytest (==7.4.2)", "pytest-cov (==4.1.0)", "pytest-timeout (==2.1.0)", "ruff (==0.1.8)"] -docs = ["myst-parser (==0.18.0)", "sphinx (==5.1.1)"] -ssh = ["paramiko (>=2.4.3)"] -websockets = ["websocket-client (>=1.3.0)"] - [[package]] name = "docstring-parser" version = "0.17.0" @@ -1224,42 +1202,6 @@ files = [ {file = "docutils-0.22.4.tar.gz", hash = "sha256:4db53b1fde9abecbb74d91230d32ab626d94f6badfc575d6db9194a49df29968"}, ] -[[package]] -name = "e2b" -version = "1.11.1" -description = "E2B SDK that give agents cloud environments" -optional = true -python-versions = "<4.0,>=3.9" -files = [ - {file = "e2b-1.11.1-py3-none-any.whl", hash = "sha256:1ecb123873788472731c101939a494ab852cbcce0f913df6f7ecb194ae932130"}, - {file = "e2b-1.11.1.tar.gz", hash = "sha256:7f7b6f238208d0a23353bb0da01f91a924321b57c61b176506862cbc1493ce8c"}, -] - -[package.dependencies] -attrs = ">=23.2.0" -httpcore = ">=1.0.5,<2.0.0" -httpx = ">=0.27.0,<1.0.0" -packaging = ">=24.1" -protobuf = ">=4.21.0" -python-dateutil = ">=2.8.2" -typing-extensions = ">=4.1.0" - -[[package]] -name = "e2b-code-interpreter" -version = "1.0.3" -description = "E2B Code Interpreter - Stateful code execution" -optional = true -python-versions = "<4.0,>=3.8" -files = [ - {file = "e2b_code_interpreter-1.0.3-py3-none-any.whl", hash = "sha256:c638bd4ec1c99d9c4eaac541bc8b15134cf786f6c7c400d979cef96d62e485d8"}, - {file = "e2b_code_interpreter-1.0.3.tar.gz", hash = "sha256:36475acc001b1317ed129d65970fce6a7cc2d50e3fd3e8a13ad5d7d3e0fac237"}, -] - -[package.dependencies] -attrs = ">=21.3.0" -e2b = ">=1.0.4,<2.0.0" -httpx = ">=0.20.0,<1.0.0" - [[package]] name = "einops" version = "0.8.2" @@ -5097,35 +5039,6 @@ files = [ {file = "pytz-2026.1.post1.tar.gz", hash = "sha256:3378dde6a0c3d26719182142c56e60c7f9af7e968076f31aae569d72a0358ee1"}, ] -[[package]] -name = "pywin32" -version = "311" -description = "Python for Window Extensions" -optional = true -python-versions = "*" -files = [ - {file = "pywin32-311-cp310-cp310-win32.whl", hash = "sha256:d03ff496d2a0cd4a5893504789d4a15399133fe82517455e78bad62efbb7f0a3"}, - {file = "pywin32-311-cp310-cp310-win_amd64.whl", hash = "sha256:797c2772017851984b97180b0bebe4b620bb86328e8a884bb626156295a63b3b"}, - {file = "pywin32-311-cp310-cp310-win_arm64.whl", hash = "sha256:0502d1facf1fed4839a9a51ccbcc63d952cf318f78ffc00a7e78528ac27d7a2b"}, - {file = "pywin32-311-cp311-cp311-win32.whl", hash = "sha256:184eb5e436dea364dcd3d2316d577d625c0351bf237c4e9a5fabbcfa5a58b151"}, - {file = "pywin32-311-cp311-cp311-win_amd64.whl", hash = "sha256:3ce80b34b22b17ccbd937a6e78e7225d80c52f5ab9940fe0506a1a16f3dab503"}, - {file = "pywin32-311-cp311-cp311-win_arm64.whl", hash = "sha256:a733f1388e1a842abb67ffa8e7aad0e70ac519e09b0f6a784e65a136ec7cefd2"}, - {file = "pywin32-311-cp312-cp312-win32.whl", hash = "sha256:750ec6e621af2b948540032557b10a2d43b0cee2ae9758c54154d711cc852d31"}, - {file = "pywin32-311-cp312-cp312-win_amd64.whl", hash = "sha256:b8c095edad5c211ff31c05223658e71bf7116daa0ecf3ad85f3201ea3190d067"}, - {file = "pywin32-311-cp312-cp312-win_arm64.whl", hash = "sha256:e286f46a9a39c4a18b319c28f59b61de793654af2f395c102b4f819e584b5852"}, - {file = "pywin32-311-cp313-cp313-win32.whl", hash = "sha256:f95ba5a847cba10dd8c4d8fefa9f2a6cf283b8b88ed6178fa8a6c1ab16054d0d"}, - {file = "pywin32-311-cp313-cp313-win_amd64.whl", hash = "sha256:718a38f7e5b058e76aee1c56ddd06908116d35147e133427e59a3983f703a20d"}, - {file = "pywin32-311-cp313-cp313-win_arm64.whl", hash = "sha256:7b4075d959648406202d92a2310cb990fea19b535c7f4a78d3f5e10b926eeb8a"}, - {file = "pywin32-311-cp314-cp314-win32.whl", hash = "sha256:b7a2c10b93f8986666d0c803ee19b5990885872a7de910fc460f9b0c2fbf92ee"}, - {file = "pywin32-311-cp314-cp314-win_amd64.whl", hash = "sha256:3aca44c046bd2ed8c90de9cb8427f581c479e594e99b5c0bb19b29c10fd6cb87"}, - {file = "pywin32-311-cp314-cp314-win_arm64.whl", hash = "sha256:a508e2d9025764a8270f93111a970e1d0fbfc33f4153b388bb649b7eec4f9b42"}, - {file = "pywin32-311-cp38-cp38-win32.whl", hash = "sha256:6c6f2969607b5023b0d9ce2541f8d2cbb01c4f46bc87456017cf63b73f1e2d8c"}, - {file = "pywin32-311-cp38-cp38-win_amd64.whl", hash = "sha256:c8015b09fb9a5e188f83b7b04de91ddca4658cee2ae6f3bc483f0b21a77ef6cd"}, - {file = "pywin32-311-cp39-cp39-win32.whl", hash = "sha256:aba8f82d551a942cb20d4a83413ccbac30790b50efb89a75e4f586ac0bb8056b"}, - {file = "pywin32-311-cp39-cp39-win_amd64.whl", hash = "sha256:e0c4cfb0621281fe40387df582097fd796e80430597cb9944f0ae70447bacd91"}, - {file = "pywin32-311-cp39-cp39-win_arm64.whl", hash = "sha256:62ea666235135fee79bb154e695f3ff67370afefd71bd7fea7512fc70ef31e3d"}, -] - [[package]] name = "pywin32-ctypes" version = "0.2.3" @@ -7503,11 +7416,10 @@ test = ["big-O", "jaraco.functools", "jaraco.itertools", "jaraco.test", "more_it type = ["pytest-mypy"] [extras] -code-execution = ["docker", "e2b-code-interpreter", "ray"] finetune = ["tinker"] vllm = ["vllm"] [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "e60aaa3e3492ae58637bd1d5d950a46518ef2b7cd2d2202dae0abd2286acea3d" +content-hash = "d0a1158bf9ef12b75985649660bc7179d83c8e25b7c5e3b2b030dd0624e3f329" diff --git a/pyproject.toml b/pyproject.toml index 220ce3e40..7371eb12a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,9 +36,6 @@ anthropic = "^0.84.0" vertexai = "1.71.1" vllm = { version = "^0.6.3", optional = true } bespokelabs-sandbox = "0.1.1" -docker = { version = ">=7.0.0", optional = true } -ray = { version = ">=2.0.0", optional = true } -e2b-code-interpreter = { version = ">=1.0.0", optional = true } tinker = { version = "^0.8.0", optional = true, python = ">=3.11" } posthog = "^3.11.0" @@ -46,7 +43,6 @@ mistralai = "^1.5.1" [tool.poetry.extras] vllm = ["vllm"] -code_execution = ["docker", "ray", "e2b-code-interpreter"] finetune = ["tinker"] [tool.poetry.group.dev.dependencies] From eae8ed71bf70ea5b36f42db87601f64aaadca16d Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 20:10:43 -0700 Subject: [PATCH 04/10] Remove code_execution extra from CI workflow The code_execution extra was removed from pyproject.toml since backend SDKs are now managed by bespokelabs-sandbox's own extras. Co-Authored-By: Claude Opus 4.6 (1M context) --- .github/workflows/lint-and-test.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/lint-and-test.yaml b/.github/workflows/lint-and-test.yaml index 395cabb0b..14f34d452 100644 --- a/.github/workflows/lint-and-test.yaml +++ b/.github/workflows/lint-and-test.yaml @@ -43,7 +43,7 @@ jobs: - name: Install dependencies run: | pip install poetry - poetry install --with dev --extras "vllm code_execution" + poetry install --with dev --extras "vllm" - name: Run ruff run: | From daa3c1d13a8e9df91a28e78ac17ae6f4406f75df Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 20:22:19 -0700 Subject: [PATCH 05/10] Always return message=success for completed executions Non-zero exit codes are valid outcomes (e.g. code-verification workloads). Setting error on non-zero exits caused validate_existing_response_file to discard and re-run those responses on every resume. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../code_execution_backend/sandbox_backend.py | 15 +-------------- 1 file changed, 1 insertion(+), 14 deletions(-) diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py index e1273a4d2..434257c16 100644 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py +++ b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py @@ -91,21 +91,8 @@ def _execute_in_sandbox( files = _collect_sandbox_files(sb) - if result.exit_code == 0: - return CodeExecutionOutput( - message="success", - stdout=result.stdout, - stderr=result.stderr, - files=files, - ) - - error_message = f"Program exited with status code {result.exit_code}" - if result.stderr: - error_message = f"{error_message}\n\nError details:\n{result.stderr}" - return CodeExecutionOutput( - message="error", - error=error_message, + message="success", stdout=result.stdout, stderr=result.stderr, files=files, From 51b87d5c53f40b0ffb8594fd87145925a37ffa1e Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 20:23:36 -0700 Subject: [PATCH 06/10] fix: Auto-format base_backend.py and curator_response.py with ruff Co-Authored-By: Claude Opus 4.6 (1M context) --- .../curator/code_executor/code_execution_backend/base_backend.py | 1 - src/bespokelabs/curator/types/curator_response.py | 1 - 2 files changed, 2 deletions(-) diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py index 7eed3b5b3..ef5c79dd2 100644 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py +++ b/src/bespokelabs/curator/code_executor/code_execution_backend/base_backend.py @@ -654,4 +654,3 @@ async def append_execution_response(self, data: dict, filename: str) -> None: async with aiofiles.open(filename, "a") as f: await f.write(json_string + "\n") logger.debug(f"Successfully appended response to {filename}") - diff --git a/src/bespokelabs/curator/types/curator_response.py b/src/bespokelabs/curator/types/curator_response.py index 718979740..2d219ae38 100644 --- a/src/bespokelabs/curator/types/curator_response.py +++ b/src/bespokelabs/curator/types/curator_response.py @@ -366,4 +366,3 @@ def append(self, other: "CuratorResponse"): self.cost_info.add(other.cost_info) self.request_stats.add(other.request_stats) self.performance_stats.add(other.performance_stats) - From c79105d06c790b733610fc82c10e68fa9ec6ac2f Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 20:53:55 -0700 Subject: [PATCH 07/10] Add unit tests for CuratorResponse and misc to meet 80% coverage - Test all dataclass add/append methods on TokenUsage, CostInfo, RequestStats, PerformanceStats - Test CuratorResponse save/load, to_dict, get_failed_requests, __getattr__ deprecation, update_tracker_stats offline path - Test safe_model_dump normal and fallback paths Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/unittests/test_curator_response.py | 155 +++++++++++++++++++++++ tests/unittests/test_misc.py | 28 ++++ 2 files changed, 183 insertions(+) create mode 100644 tests/unittests/test_curator_response.py create mode 100644 tests/unittests/test_misc.py diff --git a/tests/unittests/test_curator_response.py b/tests/unittests/test_curator_response.py new file mode 100644 index 000000000..b16beb9e8 --- /dev/null +++ b/tests/unittests/test_curator_response.py @@ -0,0 +1,155 @@ +import json +import tempfile +import warnings +from pathlib import Path + +import pytest +from datasets import Dataset + +from bespokelabs.curator.types.curator_response import ( + CostInfo, + CuratorResponse, + PerformanceStats, + RequestStats, + TokenUsage, +) + + +def test_token_usage_add(): + a = TokenUsage(input=10, output=20, total=30) + b = TokenUsage(input=5, output=15, total=20) + a.add(b) + assert a.input == 15 + assert a.output == 35 + assert a.total == 50 + + +def test_cost_info_add(): + a = CostInfo(total_cost=1.0, input_cost_per_million=2.0, output_cost_per_million=3.0, projected_remaining_cost=0.5) + b = CostInfo(total_cost=0.5, input_cost_per_million=1.0, output_cost_per_million=1.5, projected_remaining_cost=0.2) + a.add(b) + assert a.total_cost == 1.5 + assert a.input_cost_per_million == 3.0 + assert a.output_cost_per_million == 4.5 + assert a.projected_remaining_cost == 0.7 + + +def test_cost_info_add_none_costs(): + a = CostInfo(total_cost=1.0, input_cost_per_million=None, output_cost_per_million=None) + b = CostInfo(total_cost=0.5, input_cost_per_million=1.0, output_cost_per_million=1.5) + a.add(b) + assert a.total_cost == 1.5 + assert a.input_cost_per_million is None + assert a.output_cost_per_million is None + + +def test_request_stats_add(): + a = RequestStats(total=10, succeeded=8, failed=1, in_progress=1, cached=2) + b = RequestStats(total=5, succeeded=4, failed=0, in_progress=1, cached=1) + a.add(b) + assert a.total == 15 + assert a.succeeded == 12 + assert a.failed == 1 + assert a.in_progress == 2 + assert a.cached == 3 + + +def test_performance_stats_add(): + a = PerformanceStats(total_time=10.0, requests_per_minute=60.0, input_tokens_per_minute=1000.0, output_tokens_per_minute=500.0, max_concurrent_requests=5) + b = PerformanceStats(total_time=5.0, requests_per_minute=30.0, input_tokens_per_minute=500.0, output_tokens_per_minute=250.0, max_concurrent_requests=3) + a.add(b) + assert a.total_time == 15.0 + assert a.requests_per_minute == 90.0 + assert a.max_concurrent_requests == 8 + + +def test_curator_response_post_init_none_metadata(): + ds = Dataset.from_dict({"a": [1]}) + resp = CuratorResponse(dataset=ds, metadata=None) + assert resp.metadata == {} + + +def test_curator_response_getattr_raises(): + ds = Dataset.from_dict({"a": [1]}) + resp = CuratorResponse(dataset=ds) + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + with pytest.raises(AttributeError): + _ = resp.nonexistent_attr + + +def test_curator_response_update_tracker_stats_offline(): + ds = Dataset.from_dict({"a": [1]}) + resp = CuratorResponse(dataset=ds) + resp.update_tracker_stats(None) + assert resp.token_usage.total == 0 + + +def test_curator_response_get_failed_requests_no_path(): + ds = Dataset.from_dict({"a": [1]}) + resp = CuratorResponse(dataset=ds, failed_requests_path=None) + assert list(resp.get_failed_requests()) == [] + + +def test_curator_response_get_failed_requests_with_file(): + ds = Dataset.from_dict({"a": [1]}) + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + f.write(json.dumps({"error": "timeout"}) + "\n") + f.write(json.dumps({"error": "rate_limit"}) + "\n") + path = Path(f.name) + + resp = CuratorResponse(dataset=ds, failed_requests_path=path) + failed = list(resp.get_failed_requests()) + assert len(failed) == 2 + assert failed[0]["error"] == "timeout" + path.unlink() + + +def test_curator_response_append(): + ds = Dataset.from_dict({"a": [1]}) + resp1 = CuratorResponse( + dataset=ds, + model_name="gpt-4", + token_usage=TokenUsage(input=100, output=200, total=300), + cost_info=CostInfo(total_cost=1.0), + request_stats=RequestStats(total=10, succeeded=9, failed=1), + performance_stats=PerformanceStats(total_time=5.0), + ) + resp2 = CuratorResponse( + dataset=ds, + model_name="gpt-4", + token_usage=TokenUsage(input=50, output=100, total=150), + cost_info=CostInfo(total_cost=0.5), + request_stats=RequestStats(total=5, succeeded=5, failed=0), + performance_stats=PerformanceStats(total_time=3.0), + ) + resp1.append(resp2) + assert resp1.token_usage.total == 450 + assert resp1.cost_info.total_cost == 1.5 + assert resp1.request_stats.total == 15 + assert resp1.performance_stats.total_time == 8.0 + + +def test_curator_response_to_dict(): + ds = Dataset.from_dict({"a": [1, 2]}) + resp = CuratorResponse(dataset=ds, model_name="gpt-4") + d = resp.to_dict() + assert d["dataset"]["size"] == 2 + assert d["model_name"] == "gpt-4" + assert "token_usage" in d + assert "cost_info" in d + + +def test_curator_response_save_and_load(): + ds = Dataset.from_dict({"a": [1, 2]}) + resp = CuratorResponse( + dataset=ds, + model_name="gpt-4", + token_usage=TokenUsage(input=100, output=200, total=300), + ) + with tempfile.TemporaryDirectory() as tmpdir: + resp.save(tmpdir) + loaded = CuratorResponse.load(tmpdir, ds) + assert loaded.model_name == "gpt-4" + assert loaded.token_usage.input == 100 + assert loaded.token_usage.total == 300 diff --git a/tests/unittests/test_misc.py b/tests/unittests/test_misc.py new file mode 100644 index 000000000..0a65e10f4 --- /dev/null +++ b/tests/unittests/test_misc.py @@ -0,0 +1,28 @@ +from bespokelabs.curator.misc import safe_model_dump + + +def test_safe_model_dump_with_pydantic(): + from pydantic import BaseModel + + class MyModel(BaseModel): + name: str + value: int + + model = MyModel(name="test", value=42) + result = safe_model_dump(model) + assert result == {"name": "test", "value": 42} + + +def test_safe_model_dump_fallback(): + + class FakeModel: + def __init__(self): + self.name = "test" + self.value = 42 + + def model_dump(self): + raise TypeError("not a real pydantic model") + + result = safe_model_dump(FakeModel()) + assert result["name"] == "test" + assert result["value"] == 42 From 60d74d7d6f2a0179cacad53c3a3d74cd09ff8eb8 Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 21:07:07 -0700 Subject: [PATCH 08/10] Enforce per-execution timeout via shell timeout command The sandbox lifecycle timeout (timeout_secs) is a safety net, not the per-execution limit. Use `timeout N` in the shell command to kill the python process after exactly N seconds, matching the old subprocess.run timeout behavior. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../code_executor/code_execution_backend/sandbox_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py index 434257c16..25998933f 100644 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py +++ b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py @@ -87,7 +87,7 @@ def _execute_in_sandbox( sb.write_file(f"{WORKSPACE_DIR}/program.py", code) sb.write_file(f"{WORKSPACE_DIR}/input.txt", code_input) - result = sb.execute_command("bash", args=["-c", f"cd {WORKSPACE_DIR} && python program.py < input.txt"]) + result = sb.execute_command("bash", args=["-c", f"cd {WORKSPACE_DIR} && timeout {timeout} python program.py < input.txt"]) files = _collect_sandbox_files(sb) From 993e0b8201c77504c611a1e3ff681050559424a7 Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 21:16:54 -0700 Subject: [PATCH 09/10] fix: Format test_misc.py with ruff Co-Authored-By: Claude Opus 4.6 (1M context) --- tests/unittests/test_misc.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/unittests/test_misc.py b/tests/unittests/test_misc.py index 0a65e10f4..02e314a24 100644 --- a/tests/unittests/test_misc.py +++ b/tests/unittests/test_misc.py @@ -14,7 +14,6 @@ class MyModel(BaseModel): def test_safe_model_dump_fallback(): - class FakeModel: def __init__(self): self.name = "test" From d7dd3f64b1a8c69a1cd340013917921f5f623e7f Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Fri, 27 Mar 2026 21:33:09 -0700 Subject: [PATCH 10/10] Improve sandbox backend error handling and add unit tests - Distinguish exit code 124 (timeout) with message="timeout" - Format non-zero exit code errors with stderr details - Preserve stdout/stderr/files in exception handler - Add unit tests with fake sandbox for all error paths Co-Authored-By: Claude Opus 4.6 (1M context) --- .../code_execution_backend/sandbox_backend.py | 46 ++++++- tests/code_executor/test_sandbox_backend.py | 116 ++++++++++++++++++ 2 files changed, 158 insertions(+), 4 deletions(-) create mode 100644 tests/code_executor/test_sandbox_backend.py diff --git a/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py index 25998933f..3f52c85ae 100644 --- a/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py +++ b/src/bespokelabs/curator/code_executor/code_execution_backend/sandbox_backend.py @@ -79,6 +79,10 @@ def _execute_in_sandbox( """ from bespokelabs.sandbox import Sandbox + sb = None + result = None + files = "" + try: kwargs = {**sandbox_kwargs} kwargs.setdefault("timeout_secs", timeout + 30) @@ -88,23 +92,57 @@ def _execute_in_sandbox( sb.write_file(f"{WORKSPACE_DIR}/input.txt", code_input) result = sb.execute_command("bash", args=["-c", f"cd {WORKSPACE_DIR} && timeout {timeout} python program.py < input.txt"]) - files = _collect_sandbox_files(sb) + stdout = result.stdout + stderr = result.stderr + + if result.exit_code == 0: + return CodeExecutionOutput( + message="success", + stdout=stdout, + stderr=stderr, + files=files, + ) + + if result.exit_code == 124: + return CodeExecutionOutput( + message="timeout", + error=f"Execution timed out after {timeout}s", + stdout=stdout, + stderr=stderr, + files=files, + ) return CodeExecutionOutput( - message="success", - stdout=result.stdout, - stderr=result.stderr, + message="error", + error=_format_exit_code_error(result.exit_code, stderr), + stdout=stdout, + stderr=stderr, files=files, ) except Exception as e: + if sb is not None and not files: + files = _collect_sandbox_files(sb) + return CodeExecutionOutput( message="error", error=str(e), + stdout=getattr(result, "stdout", None), + stderr=getattr(result, "stderr", None), + files=files, ) +def _format_exit_code_error(exit_code: int | None, stderr: str | None) -> str: + """Format a descriptive error message for non-zero exits.""" + status = "unknown" if exit_code is None else str(exit_code) + error_message = f"Program exited with status code {status}" + if stderr: + error_message = f"{error_message}\n\nError details:\n{stderr}" + return error_message + + def _collect_sandbox_files(sandbox) -> str: """Collect files created during execution as a tar archive string. diff --git a/tests/code_executor/test_sandbox_backend.py b/tests/code_executor/test_sandbox_backend.py new file mode 100644 index 000000000..e2aec72e5 --- /dev/null +++ b/tests/code_executor/test_sandbox_backend.py @@ -0,0 +1,116 @@ +import sys +from types import ModuleType, SimpleNamespace + +from bespokelabs.curator.code_executor.code_execution_backend import sandbox_backend + + +def _install_fake_sandbox(monkeypatch, sandbox_cls): + module = ModuleType("bespokelabs.sandbox") + module.Sandbox = sandbox_cls + monkeypatch.setitem(sys.modules, "bespokelabs.sandbox", module) + + +def test_execute_in_sandbox_returns_error_for_nonzero_exit(monkeypatch): + class FakeSandbox: + def __init__(self, backend_name, **kwargs): + self.backend_name = backend_name + self.kwargs = kwargs + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def write_file(self, path, content): + pass + + def execute_command(self, command, args=None): + return SimpleNamespace(exit_code=1, stdout="partial output", stderr="Traceback") + + _install_fake_sandbox(monkeypatch, FakeSandbox) + monkeypatch.setattr(sandbox_backend, "_collect_sandbox_files", lambda sandbox: "files-archive") + + output = sandbox_backend._execute_in_sandbox( + code="print('hi')", + code_input="", + timeout=10, + backend_name="local", + sandbox_kwargs={}, + ) + + assert output.message == "error" + assert output.error == "Program exited with status code 1\n\nError details:\nTraceback" + assert output.stdout == "partial output" + assert output.stderr == "Traceback" + assert output.files == "files-archive" + + +def test_execute_in_sandbox_returns_timeout_for_timeout_exit(monkeypatch): + class FakeSandbox: + def __init__(self, backend_name, **kwargs): + self.backend_name = backend_name + self.kwargs = kwargs + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def write_file(self, path, content): + pass + + def execute_command(self, command, args=None): + return SimpleNamespace(exit_code=124, stdout="", stderr="") + + _install_fake_sandbox(monkeypatch, FakeSandbox) + monkeypatch.setattr(sandbox_backend, "_collect_sandbox_files", lambda sandbox: "files-archive") + + output = sandbox_backend._execute_in_sandbox( + code="print('hi')", + code_input="", + timeout=7, + backend_name="local", + sandbox_kwargs={}, + ) + + assert output.message == "timeout" + assert output.error == "Execution timed out after 7s" + assert output.files == "files-archive" + + +def test_execute_in_sandbox_preserves_files_when_command_raises(monkeypatch): + class FakeSandbox: + def __init__(self, backend_name, **kwargs): + self.backend_name = backend_name + self.kwargs = kwargs + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def write_file(self, path, content): + pass + + def execute_command(self, command, args=None): + raise RuntimeError("sandbox command failed") + + _install_fake_sandbox(monkeypatch, FakeSandbox) + monkeypatch.setattr(sandbox_backend, "_collect_sandbox_files", lambda sandbox: "files-archive") + + output = sandbox_backend._execute_in_sandbox( + code="print('hi')", + code_input="", + timeout=10, + backend_name="local", + sandbox_kwargs={}, + ) + + assert output.message == "error" + assert output.error == "sandbox command failed" + assert output.stdout is None + assert output.stderr is None + assert output.files == "files-archive"