ipybox 0.2.0

Python code execution sandbox based on IPython and Docker

ipybox

A secure, lightweight Python code execution sandbox based on IPython and Docker, specifically designed for AI/ML applications and LLM agents.

Overview

This project provides a sandboxed Python code execution environment built on IPython and Docker. It offers:

• Secure Execution: Code runs in Docker containers, preventing unauthorized system access
• Flexible Dependencies: Supports static and runtime dependency management
• Real-time Streaming: Chunked streaming of execution output as it's generated
• Image Support: Handle image outputs from matplotlib and other visualization libraries
• Resource Control: Container lifecycle management and built-in timeout mechanisms
• Reproducible Environment: Consistent execution environment across different systems
• LLM Agent Ready: Ideal for AI applications that need to execute Python code

This project is in early beta, with active development of new features ongoing.

Installation

Python package

pip install ipybox

Container image

Note: Building the container image requires Docker to be installed. Containers will run with the same user and group IDs as the user who built the image, ensuring proper file permissions on mounted host directories.

Default build

To build a container image with default settings and no extra dependencies:

python -m ipybox build

This creates a Docker image tagged as gradion/ipybox with base Python dependencies required for the code execution environment.

Custom build

To create a custom image with additional dependencies for your application, create a dependencies file (e.g., dependencies.txt) following. For example:

pandas = "^2.2"
scikit-learn = "^1.5"
matplotlib = "^3.9"

To build the image with custom tag and dependencies:

python -m ipybox build \
    --tag my-box:v1 \
    --dependencies path/to/dependencies.txt

The dependencies file should list Python packages in Poetry dependency specification format. These will be installed in addition to the base dependencies required for the execution environment. The execution container also supports installing dependencies at runtime.

Usage

The following examples demonstrate how to use the ExecutionContainer and ExecutionClient context managers to execute Python code in an IPython environment running in a Docker container. Runnable scripts of the following code snippets are available in the examples directory.

Basic usage

Here's a simple example that demonstrates how to execute Python code in an execution container. The ExecutionContainer context manager creates and starts a container for code execution, and the ExecutionClient context manager connects to an IPython kernel running in the container. The example below executes the code print('Hello, world!') and prints the output text:

from gradion.executor import ExecutionContainer, ExecutionClient

# Create and start a container for code execution
async with ExecutionContainer(tag="gradion/executor") as container:
    # Create and connect to an IPython kernel
    async with ExecutionClient(host="localhost", port=container.port) as client:
        # Execute Python code and await the result
        result = await client.execute("print('Hello, world!')")
        # Print the execution output text
        print(f"Output: {result.text}")  # Output: Hello, world!

The default image used by ExecutionContainer is gradion/executor. You can specify a custom image with the tag argument like in ExecutionContainer(tag="my-box:v1"), for example.

State management

Code execution within the same client context is stateful i.e. you can reference variables from previous executions. Code executions in different client contexts are isolated from each other:

async with ExecutionContainer() as container:
    async with ExecutionClient(host="localhost", port=container.port) as client_1:
        # Execute code that defines variable x
        result = await client_1.execute("x = 1")
        assert result.text is None

        # Reference variable x defined in previous execution
        result = await client_1.execute("print(x)")
        assert result.text == "1"

    async with ExecutionClient(host="localhost", port=container.port) as client_2:
        # Variable x is not defined in this client context
        try:
            await client_2.execute("print(x)")
        except ExecutionError as e:
            assert e.args[0] == "NameError: name 'x' is not defined"

Output streaming

The execution client supports streaming output as it's generated during code execution:

async with ExecutionContainer() as container:
    async with ExecutionClient(host="localhost", port=container.port) as client:
        # Code that produces output gradually
        code = """
        import time
        for i in range(5):
            print(f"Processing step {i}")
            time.sleep(1)
        """

        # Submit the execution and stream the output
        execution = await client.submit(code)
        async for chunk in execution.stream():
            # Output will be printed gradually:
            print(f"Received output: {chunk}")
            # Received output: Processing step 0
            # Received output: Processing step 1
            # Received output: Processing step 2
            # Received output: Processing step 3
            # Received output: Processing step 4

        # Get the aggregated result
        result = await execution.result()
        # Print the aggregated output text
        print(f"Aggregated output:\n{result.text}")
        # Aggregated output:
        # Processing step 0
        # Processing step 1
        # Processing step 2
        # Processing step 3
        # Processing step 4

The stream() method accepts an optional timeout argument (defaults to 120 seconds). In case of timeout, the execution is automatically terminated by interrupting the kernel.

Installing dependencies at runtime

async with ExecutionContainer() as container:
    async with ExecutionClient(host="localhost", port=container.port) as client:
        # Install the einops package
        await client.execute("!pip install einops")
        # Then you can use it in the following code
        # execution within the same client context
        result = await client.execute("""
            import einops
            print(einops.__version__)
        """)
        print(f"Output: {result.text}")  # Output: 0.8.0

Creating and returning plots

Plots created with matplotlib or other libraries are returned as PIL images:

async with ExecutionContainer() as container:
    async with ExecutionClient(host="localhost", port=container.port) as client:
        execution = await client.submit("""
            !pip install matplotlib

            import matplotlib.pyplot as plt
            import numpy as np
            import time

            x = np.linspace(0, 10, 100)
            plt.figure(figsize=(8, 6))
            plt.plot(x, np.sin(x))
            plt.title('Sine Wave')
            plt.show()

            print("Plot generation complete!")
            """)

        # Stream the output text as it's generated
        async for chunk in execution.stream():
            print(chunk, end="", flush=True)

        # Obtain the execution result
        result = await execution.result()
        assert "Plot generation complete!" in result.text

        # Get created PIL image from the result
        result.images[0].save(f"output.png")

Images are not part of the output stream, but are available as a images list in the result object. The example above saves the created image as output.png:

Bind mounts

You can mount host directories into the container to allow code execution to access external files:

# Map host paths to container paths. Host paths can be absolute or relative.
# Container paths must be relative and are created as subdirectories of the
# /app directory in the container. The /app directory is the container's
# working directory.
binds = {
    "./data": "data",  # Read data from host
    "./output": "output"  # Write results to host
}

# Create a data file on the host
async with aiofiles.open("data/input.txt", "w") as f:
    await f.write("hello world")

async with ExecutionContainer(binds=binds) as container:
    async with ExecutionClient(host="localhost", port=container.port) as client:
        # Read from mounted data directory
        result = await client.execute("""
            with open('data/input.txt') as f:
                data = f.read()

            # Process data...
            processed = data.upper()

            # Write to mounted output directory
            with open('output/result.txt', 'w') as f:
                f.write(processed)
        """)

# Check the result file on the host
async with aiofiles.open("output/result.txt", "r") as f:
    assert await f.read() == "HELLO WORLD"

Environment variables

Environment variables can be passed to the container for configuration or secrets:

env = {
    "API_KEY": "secret-key-123",
    "DEBUG": "1"
}

async with ExecutionContainer(env=env) as container:
    async with ExecutionClient(host="localhost", port=container.port) as client:
        # Access environment variables in executed code
        result = await client.execute("""
            import os

            api_key = os.environ['API_KEY']
            print(f"Using API key: {api_key}")

            debug = bool(int(os.environ.get('DEBUG', '0')))
            if debug:
                print("Debug mode enabled")
        """)
        print(result.text)
        # Using API key: secret-key-123
        # Debug mode enabled