Execute LLM-generated code

We have built a code-executor that can be used to execute LLM-generated code. This is useful for many situations:

  1. You want to include error-free code in your training code. This method is used in Open Thoughts.

  2. LLM generates some code to generate visualization etc.

  3. Agents and tool-use.

Here is a simple example of code execution in action:

from bespokelabs import curator
from datasets import Dataset

class HelloExecutor(curator.CodeExecutor):     
	def code(self, row):
		return """location = input();print(f"Hello {location}")"""     
		
	def code_input(self, row):        
		return row['location']     
	
	def code_output(self, row, execution_output):        
		row['output'] = execution_output.stdout 
		return row
		
locations = Dataset.from_list([{'location': 'New York'},{'location': 'Tokyo'}])

hello_executor = HelloExecutor() 

print(hello_executor(locations).to_pandas())

The inherited class contains three methods:

  1. code: This is the method that returns the piece of code to be run. This is usually part of the row (you can use curator.LLM to generate this code).

  2. code_input: This is optional, but can return a json that represents values to be passed to input() in the code.

  3. code_output: This is where you parse the output of the execution.

Features:

  1. Full caching and automatic recovery: Similar to Curator.LLM's caching feature, code executor also has inbuilt caching and automatic recovery. Any interrupted runs can be fully recovered and no computation is lost.

  2. Multiple code execution backends: We offer four backends: multiprocessing, docker, ray and E2B. These backends specify different locations where your code can be executed. You can easily switch the backend with a simple parameter change. For example, the hello world example can be run using the ray backend by simply initializing it with `HelloExecutor(backend=ray)`

  3. Progress monitoring using Rich Console:

Backends

We offer four backends for running your code:

  1. Multiprocessing: This is the default backend. This runs code locally and is therefore the least safe option, but is useful for quick execution as it does not require any dependencies.

  2. Docker: It is safer option than multiprocessing.

  3. Ray: If you have a ray cluster, you can use it by setting CodeExecutor(backend="ray"). This is useful when your code can take a long time to run.

  4. E2B: Code can also be run using e2b.dev. Use CodeExecutor(backend="e2b") .

Backend Setup and configuration options

Multiprocessing Backend:

This doesn't require any additional setup. You can configure backend params while initializing as follows:

```python
hello_executor = HelloExecutor(
    backend_params = {    
        "max_requests_per_minute": 1000
    }
)
```        

You can also configure execution parameters:

```python

output = hello_executor(dataset, 
    execution_params = {
      "timeout": 120 # in seconds  
      "memory_limit": 1024 ** 3 
   }
)
```

Docker

With docker, code can be executed in a secure containerized environment. You need docker installed and python's docker client installed on your machine:

  1. pip install docker

  2. In your terminal, run `docker pull python:3.11-slim`

  3. Run the HelloExecutor example with HelloExecutor(backend=docker)

With docker, you can specify a custom docker image to execute your code snippets:

```python
hello_executor = HelloExecutor(
    backend = "docker",
    backend_params = {    
        "image": "andgineer/matplotlib"
    }
)
```  

Ray

As the size of the dataset grows, it becomes harder to scale code execution requirements on a single machine. In such scenarios, one can use the ray backend.

Simply run pip install ray to install the dependencies required for ray backend.

You need to separately spin up a ray cluster and enter the base_url (Installation instructions). If base_url is not entered, then a local ray cluster is spun up.

```python
hello_executor = HelloExecutor(
    backend = "ray",
    backend_params = {    
        "base_url": "<url of ray cluster>"
    }
)
```  

E2B

We also add light support for e2b's hosted code execution backends. While not free, they are secure environments similar to docker environments and have more features.

  1. Run pip install e2b-code-interpreter to install the required dependencies.

  2. Create an account on e2b's website, get the API key and add it to your environment variables.

```python
hello_executor = HelloExecutor(
    backend = "e2b",
)
```  

Conclusion

Check out the examples to get started with code executor. If you have any questions, feel free to join our Discord or send us an email.

Last updated