Usage

The two main classes of the ipybox package are ExecutionContainer and ExecutionClient.

Note

Runnable scripts of the source code on this page are available in the examples directory.

Basic usage

For executing code in ipybox you first need to create a Docker container from an ipybox Docker image and then an IPython kernel running in that container. This is done with the ExecutionContainer and the ExecutionClient context managers.

from ipybox import ExecutionClient, ExecutionContainer


async with ExecutionContainer(tag="gradion-ai/ipybox") as container:  # (1)!
    async with ExecutionClient(port=container.port) as client:  # (2)!
        result = await client.execute("print('Hello, world!')")  # (3)!
        print(f"Output: {result.text}")  # (4)!

Create and start a container for code execution
Create and connect to an IPython kernel
Execute Python code and await the result
Output: Hello, world!

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

Note

Instead of letting the ExecutionContainer context manager handle the lifecycle of the container, you can also manually manage the container lifecycle.

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(port=container.port) as client_1:  # (1)!
        result = await client_1.execute("x = 1")  # (2)!
        assert result.text is None
        result = await client_1.execute("print(x)")  # (3)!
        assert result.text == "1"

    async with ExecutionClient(port=container.port) as client_2:  # (4)!
        try:
            await client_2.execute("print(x)")  # (5)!
        except ExecutionError as e:
            assert e.args[0] == "NameError: name 'x' is not defined"

First client context
Execute code that defines variable x
Reference variable x defined in previous execution
Second client context
Variable x is not defined in client_2 context

Output streaming

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

async with ExecutionContainer() as container:
    async with ExecutionClient(port=container.port) as client:
        code = """
        import time
        for i in range(5):
            print(f"Processing step {i}")
            time.sleep(1)
        """  # (1)!

        execution = await client.submit(code)  # (2)!
        print("Streaming output:")
        async for chunk in execution.stream():  # (3)!
            print(f"Received output: {chunk.strip()}")  # (4)!

        result = await execution.result()  # (5)!
        print("\nAggregated output:")
        print(result.text)  # (6)!

Code that produces gradual output
Submit the code for execution
Stream the output

Prints one line per second:

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 output as a single result

Prints the aggregated output:

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(port=container.port) as client:
        execution = await client.submit("!pip install einops")  # (1)!
        async for chunk in execution.stream():  # (2)!
            print(chunk, end="", flush=True)

        result = await client.execute("""
            import einops
            print(einops.__version__)
        """)  # (3)!
        print(f"Output: {result.text}")  # (4)!

Install the einops package using pip

Stream the installation progress. Something like

Collecting einops
Downloading einops-0.8.0-py3-none-any.whl (10.0 kB)
Installing collected packages: einops
Successfully installed einops-0.8.0

Import and use the installed package
Prints Output: 0.8.0

You can also install and use a package within a single execution. There's no need to have two separate executions as done in the example above.

Creating and returning plots

Plots created with matplotlib or other libraries are returned as PIL images. Images are not part of the output stream, but are available as images list in the result object.

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

            import matplotlib.pyplot as plt
            import numpy as np

            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!")
            """)  # (1)!

        async for chunk in execution.stream():  # (2)!
            print(chunk, end="", flush=True)

        result = await execution.result()
        result.images[0].save("sine.png")  # (3)!

Install matplotlib and generate a plot
Stream output text (installation progress and print statement)
Get attached image from execution result and save it as sine.png

Bind mounts

Bind mounts allow executed code to read and write files on the host machine.

await aiofiles.os.makedirs("data", exist_ok=True)
await aiofiles.os.makedirs("output", exist_ok=True)

binds = {  # (1)!
    "./data": "data",  # (2)!
    "./output": "output",  # (3)!
}

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(port=container.port) as client:
        await client.execute("""
            with open('data/input.txt') as f:
                data = f.read()

            processed = data.upper()

            with open('output/result.txt', 'w') as f:
                f.write(processed)
        """)  # (4)!

async with aiofiles.open("output/result.txt", "r") as f:  # (5)!
    result = await f.read()
    assert result == "HELLO WORLD"

Map host paths to container paths.
For reading files from host.
For writing files to host.
Read from mounted data directory, convert to uppercase and write to mounted output directory
Verify the results on host

Environment variables

Environment variables can be set on the container for passing secrets or configuration data, for example.

# Define environment variables for the container
env = {"API_KEY": "secret-key-123", "DEBUG": "1"}  # (1)!

async with ExecutionContainer(env=env) as container:
    async with ExecutionClient(port=container.port) as client:
        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")
        """)  # (2)!
        print(result.text)  # (3)!

Define environment variables for the container
Access environment variables in executed code

Prints

Using API key: secret-key-123
Debug mode enabled

Manual container lifecycle management

Instead of using ExecutionContainer as a context manager, you can also manually run() and kill() the container. This is useful for running the container on a separate host listening to a user-defined host port (e.g. 7777 in the example below).

container = ExecutionContainer(port=7777)  # (1)!
await container.run()  # (2)!
assert container.port == 7777

# do some work ...

await container.kill()  # (3)!

Create an ExecutionContainer instance using a fixed port.
Run the container (detached).
Cleanup.