Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Experimental - Stateful code interpreter #326

Closed
wants to merge 31 commits into from

Conversation

jakubno
Copy link
Member

@jakubno jakubno commented Mar 5, 2024

Stateful code interpreter

The current version of the code interpreter SDK allows to run Python code but each run has its own separate context. That means that subsequent runs can't reference to variables, definitions, etc from past code execution runs.

This is suboptimal for a lot of Python use cases with LLMs. Especially GPT-3.5 and 4 expects it runs in a Jupyter Notebook environment. Even when ones tries to convince it otherwise. In practice, LLMs will generate code blocks which have references to previous code blocks. This becomes an issue if a user wants to execute each code block separately which often is the use case.

This new code interpreter template runs a Jupyter server inside the sandbox, which allows for sharing context between code executions.
Additionally, this new template also partly implements the Jupyter Kernel messaging protocol. This means that, for example, support for plotting charts is now improved and we don't need to do hack-ish solutions like in the current production version of our code interpreter.

Current state

Known limited in features such as:

  • All executions share single kernel

We'll be updating this PR with new releases as we gather more user feedback.

Installation

Currently only Python SDK is supported

Install the experimental release candidate

Python

pip install e2b==0.14.10a12

JavaScript

npm install [email protected]

Examples

Minimal example with the sharing context:

Python

from e2b.templates.stateful_code_interpreter import CodeInterpreterV2

with CodeInterpreterV2() as sandbox:
    sandbox.exec_python("x = 1")

    result = sandbox.exec_python("x+=1; x")
    print(result.output)  # outputs 2
 

JavaScript

import { CodeInterpreterV2 } from 'e2b'

const sandbox = await CodeInterpreterV2.create()
await sandbox.execPython('x = 1')

const result = await sandbox.execPython('x+=1; x')
console.log(result.output) // outputs 2
 
await sandbox.close()

Get charts and any display-able data

Python

import base64
import io

from matplotlib import image as mpimg, pyplot as plt

from e2b.templates.stateful_code_interpreter import CodeInterpreterV2


with CodeInterpreterV2() as sandbox:
    # you can install dependencies in "jupyter notebook style" 
    sandbox.exec_python("!pip install matplotlib")

     # plot random graph
    result = sandbox.exec_python(
        """
    import matplotlib.pyplot as plt
    import numpy as np

    x = np.linspace(0, 20, 100)
    y = np.sin(x)

    plt.plot(x, y)
    plt.show()
    """
    )

    # there's your image
    image = result.display_data[0]["image/png"]

    # example how to show the image / prove it works
    i = base64.b64decode(image)
    i = io.BytesIO(i)
    i = mpimg.imread(i, format='PNG')

    plt.imshow(i, interpolation='nearest')
    plt.show()

JavaScript

import { CodeInterpreterV2 } from 'e2b'

const sandbox = await CodeInterpreterV2.create()


const code =  `
import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(0, 20, 100)
y = np.sin(x)

plt.plot(x, y)
plt.show()
`

// you can install dependencies in "jupyter notebook style" 
await sandbox.execPython('!pip install matplotlib')

const result = await sandbox.execPython(code)

// this contains the image data, you can e.g. save it to file or send to frontend
result.display_data[0]['image/png']
 
await sandbox.close()

Streaming code output

Python

from e2b.templates.stateful_code_interpreter import CodeInterpreterV2

code =  """
import time

print("hello")
time.sleep(5)
print("world")
"""
with CodeInterpreterV2() as sandbox:
    sandbox.exec_python(code, on_stdout=print, on_stderr=print)

JavaScript

import { CodeInterpreterV2 } from 'e2b'

code =  `
import time

print("hello")
time.sleep(5)
print("world")
`

const sandbox = await CodeInterpreterV2.create()

await sandbox.execPython(code, out => console.log(out), outErr => console.error(outErr))

Pre-installed Python packages inside the sandbox

The full and always up-to-date list can be found in the requirements.txt file.

# Jupyter server requirements
jupyter-server==2.13.0
ipykernel==6.29.3
ipython==8.22.2

# Other packages
aiohttp==3.9.3
beautifulsoup4==4.12.3
bokeh==3.3.4
gensim==4.3.2
imageio==2.34.0
joblib==1.3.2
librosa==0.10.1
matplotlib==3.8.3
nltk==3.8.1
numpy==1.26.4
opencv-python==4.9.0.80
openpyxl==3.1.2
pandas==1.5.3
plotly==5.19.0
pytest==8.1.0
python-docx==1.1.0
pytz==2024.1
requests==2.26.0
scikit-image==0.22.0
scikit-learn==1.4.1.post1
scipy==1.12.0
seaborn==0.13.2
soundfile==0.12.1
spacy==3.7.4
textblob==0.18.0
tornado==6.4
urllib3==1.26.7
xarray==2024.2.0
xlrd==2.0.1

Custom template using Code Interpreter

If you're using a custom sandbox template, you currently need to go through a bit of a manual setup:

  1. Copy jupyter_server_config.py and start-up.sh from this PR
  2. Make sure you have curl and jq in your Dockerfile or just add
RUN apt-get update && apt-get install -y --no-install-recommends curl jq
  1. Add following commands in your Dockerfile
# Installs jupyter server and kernel
RUN pip install jupyter-server ipykernel ipython
RUN ipython kernel install --name "python3" --user
# Copes jupyter server config
COPY ./jupyter_server_config.py /home/user/.jupyter/
# Setups jupyter server
COPY ./start-up.sh /home/user/.jupyter/
RUN chmod +x /home/user/.jupyter/start-up.sh
  1. Add the following option -c "/home/user/.jupyter/start-up.sh" to e2b template build command or add this line to your e2b.toml.
start_cmd = "/home/user/.jupyter/start-up.sh"

@jakubno jakubno added the python-rc Python SDK - Release candidate label Mar 5, 2024
Copy link

changeset-bot bot commented Mar 5, 2024

⚠️ No Changeset found

Latest commit: 1491f5a

Merging this PR will not cause a version bump for any packages. If these changes should not result in a new version, you're good to go. If these changes should result in a version bump, you need to add a changeset.

This PR includes no changesets

When changesets are added to this PR, you'll see the packages that this PR includes changesets for and the associated semver types

Click here to learn what changesets are, and how to add one.

Click here if you're a maintainer who wants to add a changeset to this PR

@jakubno jakubno added the js-rc JS SDK - Release candidate label Mar 6, 2024
@jakubno jakubno force-pushed the stateful-code-interpreter branch from aaa11fb to 099811a Compare March 6, 2024 23:27
@jakubno jakubno force-pushed the stateful-code-interpreter branch from 099811a to ea8b13a Compare March 6, 2024 23:28
@mlejva mlejva mentioned this pull request Mar 7, 2024
@mlejva
Copy link
Member

mlejva commented Mar 8, 2024

User feedback we've received so far:

  • Change exec_python to exec_jupyter and support other languages like Bash, Julia, or R
  • Make Code Interpreter stateful across multiple sessions
  • Make the API usable via REST and without the SDK
  • Allow to access existing variables from the Jupyter kernel so our users can then let their users to inspect all this data
  • Make it easy for users to use custom sandbox with the new code interpreter
  • Make it easy for users to manage Python and OS packages/libraries
    • override a version of a pre-installed package/library
    • remove a pre-installed package/library
    • add a new package/library

@mlejva mlejva changed the title Experimental feature - Stateful code interpreter Experimental - Stateful code interpreter Mar 17, 2024
@mlejva
Copy link
Member

mlejva commented Mar 20, 2024

UPDATE: This has been much improved in the latest version

  • Execution times are around 150ms
  • Creation times are around 1.5-2s

Benchmarks

I did a basic benchmarking of the new Code Interpreter 2.0 SDK and compared it to our production SDK. All in JS.

It's not completely fair comparison because there's no execPython in the production SDK but it should give us a sense of direction.

Code Interpreter 2.0 SDK

  • 100 iterations
  • Reported an average in the end

Code

import { CodeInterpreterV2 } from 'e2b'

const iterations = 100;
let createSandboxTime = 0;
let execPythonXEquals1Time = 0;
let execPythonXPlusEquals1Time = 0;
let sandboxCloseTime = 0;

for (let i = 0; i < iterations; i++) {
  console.log('Iteration:', i + 1)
  let startTime = performance.now();
  const sandbox = await CodeInterpreterV2.create();
  createSandboxTime += performance.now() - startTime;

  startTime = performance.now();
  await sandbox.execPython('x = 1');
  execPythonXEquals1Time += performance.now() - startTime;

  startTime = performance.now();
  const result = await sandbox.execPython('x+=1; x');
  execPythonXPlusEquals1Time += performance.now() - startTime;

  startTime = performance.now();
  await sandbox.close();
  sandboxCloseTime += performance.now() - startTime;
}

console.log(`Average Create Sandbox Time: ${createSandboxTime / iterations}ms`);
console.log(`Average Execute Python x = 1 Time: ${execPythonXEquals1Time / iterations}ms`);
console.log(`Average Execute Python x+=1; x Time: ${execPythonXPlusEquals1Time / iterations}ms`);
console.log(`Average Sandbox Close Time: ${sandboxCloseTime / iterations}ms`);

Results

Average Create Sandbox Time: 2830.4465337133406ms
Average Execute Python x = 1 Time: 583.8239275050163ms
Average Execute Python x+=1; x Time: 277.89283413648604ms
Average Sandbox Close Time: 0.19184373140335084ms

Production SDK

  • 100 iterations
  • Reported an average in the end

Code

import * as e2b from 'e2b'

const iterations = 100;
let createSandboxTime = 0;
let execPythonXEquals1Time = 0;
let execPythonXPlusEquals1Time = 0;
let sandboxCloseTime = 0;

for (let i = 0; i < iterations; i++) {
  console.log('Iteration:', i + 1)
  let startTime = performance.now();
  const sandbox = await e2b.Sandbox.create();
  createSandboxTime += performance.now() - startTime;

  startTime = performance.now();
  await sandbox.process.startAndWait('python3 -c "x = 1"');
  execPythonXEquals1Time += performance.now() - startTime;

  startTime = performance.now();
  await sandbox.process.startAndWait('python3 -c "x+=1; print(x)"');
  execPythonXPlusEquals1Time += performance.now() - startTime;

  startTime = performance.now();
  await sandbox.close();
  sandboxCloseTime += performance.now() - startTime;
}

console.log(`Average Sandbox Creation Time: ${createSandboxTime / iterations}ms`);
console.log(`Average Process Start and Wait (x=1) Time: ${execPythonXEquals1Time / iterations}ms`);
console.log(`Average Process Start and Wait (x+=1; print(x)) Time: ${execPythonXPlusEquals1Time / iterations}ms`);
console.log(`Average Sandbox Close Time: ${sandboxCloseTime / iterations}ms`);

Results

Average Sandbox Creation Time: 1773.701356651783ms
Average Process Start and Wait (x=1) Time: 171.03909373521805ms
Average Process Start and Wait (x+=1; print(x)) Time: 134.90086081504822ms
Average Sandbox Close Time: 0.16264041662216186ms

Difference

Sandbox creation time: ~1.1s worse
Run a process: ~300ms worse
Sandbox close time: virtually same

@im-calvin
Copy link

Thanks for the great work, it seems to be working for the most part.

I have a quick bug that might come up in the future. I'm using JS SDK.

When I create the sandbox with no initial cwd, it should default to /home/user source. But when I try and access sandboxv2.cwd it returns undefined. Is this intended behaviour?

@mlejva
Copy link
Member

mlejva commented Mar 27, 2024

Hey @im-calvin thank you for catching this.

We're moving the code interpreter SDK to a separate repository. We'll publish it as a separate package. The rest of the work and all the fixes will be there.

https://github.com/e2b-dev/code-interpreter

@@ -0,0 +1,268 @@
import { Sandbox, SandboxOpts } from '../sandbox'
import { ProcessMessage } from '../sandbox/process'
import { randomBytes } from 'crypto'

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please use web crypto for edge environment support (cf workers)

}

private async getKernelID() {
return await this.filesystem.read('/root/.jupyter/kernel_id')

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
return await this.filesystem.read('/root/.jupyter/kernel_id')
return (await this.filesystem.read('/root/.jupyter/kernel_id')).trim();

There was an extra newline at the end on cloudflare workers durable objects.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @lawrencecchen , all the fixes will be incorporated in this repo https://github.com/e2b-dev/code-interpreter

@mlejva
Copy link
Member

mlejva commented Apr 4, 2024

@lawrencecchen you can start using the new SDK https://github.com/e2b-dev/code-interpreter

JavaScript

npm install @e2b/code-interpreter

Python

pip install e2b-code-interpreter

@jakubno jakubno closed this Apr 5, 2024
@jakubno jakubno deleted the stateful-code-interpreter branch April 8, 2024 19:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
js-rc JS SDK - Release candidate python-rc Python SDK - Release candidate
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants