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

feat(sandbox): Candidate Implementation of Sandbox Plugin to Support Jupyter #1255

Merged
merged 24 commits into from
Apr 23, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
24 commits
Select commit Hold shift + click to select a range
7c76759
initialize plugin definition
xingyaoww Apr 20, 2024
ba07415
initialize plugin definition
xingyaoww Apr 20, 2024
fdc5109
simplify mixin
xingyaoww Apr 20, 2024
1ada575
further improve plugin mixin
xingyaoww Apr 20, 2024
9bed95f
add cache dir for pip
xingyaoww Apr 20, 2024
abf0fbb
support clean up cache
xingyaoww Apr 20, 2024
bbac1a6
add script for setup jupyter and execution server
xingyaoww Apr 20, 2024
18cf4ba
integrate JupyterRequirement to ssh_box
xingyaoww Apr 20, 2024
b7f2f60
source bashrc at the end of plugin load
xingyaoww Apr 20, 2024
cf33615
add execute_cli that accept code via stdin
xingyaoww Apr 20, 2024
3908ae7
make JUPYTER_EXEC_SERVER_PORT configurable via env var
xingyaoww Apr 20, 2024
fadb609
increase background cmd sleep time
xingyaoww Apr 20, 2024
d3c829c
Update opendevin/sandbox/plugins/mixin.py
xingyaoww Apr 21, 2024
9481c7b
add mixin to base class
xingyaoww Apr 21, 2024
456d9cb
make jupyter requirement a dataclass
xingyaoww Apr 21, 2024
0d788ef
source plugins only when >0 requirements
xingyaoww Apr 21, 2024
a436c8b
add `sandbox_plugins` for each agent & have controller take care of it
xingyaoww Apr 21, 2024
bac492a
update build.sh to make logs available in /opendevin/logs
xingyaoww Apr 21, 2024
3f5030d
switch to use config for lib and cache dir
xingyaoww Apr 21, 2024
a7b4f2e
fix permission issue with /workspace
xingyaoww Apr 21, 2024
6be6ac8
use python to implement execute_cli to avoid stdin escape issue
xingyaoww Apr 21, 2024
c1a4e56
wait until jupyter is avaialble
xingyaoww Apr 22, 2024
76b5cfa
Merge commit '2242702cf94eab7275f2cb148859135018d9b280' into integrat…
xingyaoww Apr 22, 2024
da1b3ba
support plugin via copying instead of mounting
xingyaoww Apr 22, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ check-system:
echo "$(RED)Unsupported system detected. Please use macOS, Linux, or Windows Subsystem for Linux (WSL).$(RESET)"; \
exit 1; \
fi

check-python:
@echo "$(YELLOW)Checking Python installation...$(RESET)"
@if command -v python3.11 > /dev/null; then \
Expand Down Expand Up @@ -218,6 +218,12 @@ setup-config-prompts:
workspace_dir=$${workspace_dir:-$(DEFAULT_WORKSPACE_DIR)}; \
echo "WORKSPACE_BASE=\"$$workspace_dir\"" >> $(CONFIG_FILE).tmp

# Clean up all caches
clean:
@echo "$(YELLOW)Cleaning up caches...$(RESET)"
@rm -rf opendevin/.cache
@echo "$(GREEN)Caches cleaned up successfully.$(RESET)"

# Help
help:
@echo "$(BLUE)Usage: make [target]$(RESET)"
Expand Down
3 changes: 3 additions & 0 deletions agenthub/codeact_agent/codeact_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
)
from opendevin.parse_commands import parse_command_file
from opendevin.state import State
from opendevin.sandbox.plugins import PluginRequirement, JupyterRequirement

COMMAND_DOCS = parse_command_file()
COMMAND_SEGMENT = (
Expand Down Expand Up @@ -69,6 +70,8 @@ class CodeActAgent(Agent):
The agent works by passing the model a list of action-observation pairs and prompting the model to take the next step.
"""

sandbox_plugins: List[PluginRequirement] = [JupyterRequirement()]

def __init__(
self,
llm: LLM,
Expand Down
2 changes: 2 additions & 0 deletions opendevin/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from opendevin.state import State
from opendevin.llm.llm import LLM
from opendevin.exceptions import AgentAlreadyRegisteredError, AgentNotRegisteredError
from opendevin.sandbox.plugins import PluginRequirement


class Agent(ABC):
Expand All @@ -17,6 +18,7 @@ class Agent(ABC):
"""

_registry: Dict[str, Type['Agent']] = {}
sandbox_plugins: List[PluginRequirement] = []

def __init__(
self,
Expand Down
8 changes: 7 additions & 1 deletion opendevin/config.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import os

import argparse
import toml
import pathlib
from dotenv import load_dotenv

from opendevin.schema import ConfigType
Expand All @@ -18,6 +18,7 @@
ConfigType.WORKSPACE_MOUNT_PATH: None,
ConfigType.WORKSPACE_MOUNT_PATH_IN_SANDBOX: '/workspace',
ConfigType.WORKSPACE_MOUNT_REWRITE: None,
ConfigType.CACHE_DIR: os.path.join(os.path.dirname(os.path.abspath(__file__)), '.cache'),
ConfigType.LLM_MODEL: 'gpt-3.5-turbo-1106',
ConfigType.SANDBOX_CONTAINER_IMAGE: 'ghcr.io/opendevin/sandbox',
ConfigType.RUN_AS_DEVIN: 'true',
Expand Down Expand Up @@ -145,3 +146,8 @@ def get(key: str, required: bool = False):
if not value and required:
raise KeyError(f"Please set '{key}' in `config.toml` or `.env`.")
return value


_cache_dir = config.get('CACHE_DIR')
if _cache_dir:
pathlib.Path(_cache_dir).mkdir(parents=True, exist_ok=True)
Comment on lines +151 to +153
Copy link
Collaborator

Choose a reason for hiding this comment

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

I'm not entirely certain of the purpose of cache dir--do we still need this?

Copy link
Collaborator Author

@xingyaoww xingyaoww Apr 22, 2024

Choose a reason for hiding this comment

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

We need this (at least for now) to cache the pip install commands in Jupyter requirement dependency - without this, you need to re-download the whl for every sandbox initialization, which can be way too slow. I still felt it is a bit slow with cache (you still need to spend the CPU hours)... we probably need to figure out a way to cache it in the future further (e.g., save the docker image with the requirement installed and re-use it as much as possible)

4 changes: 4 additions & 0 deletions opendevin/controller/action_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
AgentErrorObservation,
NullObservation,
)
from opendevin.sandbox.plugins import PluginRequirement


class ActionManager:
Expand Down Expand Up @@ -41,6 +42,9 @@ def __init__(
else:
raise ValueError(f'Invalid sandbox type: {sandbox_type}')

def init_sandbox_plugins(self, plugins: List[PluginRequirement]):
self.sandbox.init_plugins(plugins)

async def run_action(self, action: Action, agent_controller) -> Observation:
observation: Observation = NullObservation('')
if not action.executable:
Expand Down
2 changes: 2 additions & 0 deletions opendevin/controller/agent_controller.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@ def __init__(
self.action_manager = ActionManager(self.id, container_image)
self.max_chars = max_chars
self.callbacks = callbacks
# Initialize agent-required plugins for sandbox (if any)
self.action_manager.init_sandbox_plugins(agent.sandbox_plugins)

def update_state_for_step(self, i):
if self.state is None:
Expand Down
33 changes: 33 additions & 0 deletions opendevin/sandbox/docker/exec_box.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
import sys
import time
import uuid
import tarfile
from glob import glob
from collections import namedtuple
from typing import Dict, List, Tuple

Expand Down Expand Up @@ -122,6 +124,37 @@ def run_command(container, command):
return -1, f'Command: "{cmd}" timed out'
return exit_code, logs.decode('utf-8')

def copy_to(self, host_src: str, sandbox_dest: str, recursive: bool = False):
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

Code duplication from DockerSSHBox which is not ideal, but i guess we can leave them temporarily for this PR and plan to deprecate ExecBox in the next few days.

# mkdir -p sandbox_dest if it doesn't exist
exit_code, logs = self.container.exec_run(
['/bin/bash', '-c', f'mkdir -p {sandbox_dest}'],
workdir=SANDBOX_WORKSPACE_DIR,
)
if exit_code != 0:
raise Exception(
f'Failed to create directory {sandbox_dest} in sandbox: {logs}')

if recursive:
assert os.path.isdir(host_src), 'Source must be a directory when recursive is True'
files = glob(host_src + '/**/*', recursive=True)
srcname = os.path.basename(host_src)
tar_filename = os.path.join(os.path.dirname(host_src), srcname + '.tar')
with tarfile.open(tar_filename, mode='w') as tar:
for file in files:
tar.add(file, arcname=os.path.relpath(file, os.path.dirname(host_src)))
else:
assert os.path.isfile(host_src), 'Source must be a file when recursive is False'
srcname = os.path.basename(host_src)
tar_filename = os.path.join(os.path.dirname(host_src), srcname + '.tar')
with tarfile.open(tar_filename, mode='w') as tar:
tar.add(host_src, arcname=srcname)

with open(tar_filename, 'rb') as f:
data = f.read()

self.container.put_archive(os.path.dirname(sandbox_dest), data)
os.remove(tar_filename)

def execute_in_background(self, cmd: str) -> Process:
result = self.container.exec_run(
self.get_exec_cmd(cmd), socket=True, workdir=SANDBOX_WORKSPACE_DIR
Expand Down
19 changes: 19 additions & 0 deletions opendevin/sandbox/docker/local_box.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,25 @@ def execute(self, cmd: str) -> Tuple[int, str]:
except subprocess.TimeoutExpired:
return -1, 'Command timed out'

def copy_to(self, host_src: str, sandbox_dest: str, recursive: bool = False):
# mkdir -p sandbox_dest if it doesn't exist
res = subprocess.run(f'mkdir -p {sandbox_dest}', shell=True, text=True, cwd=config.get('WORKSPACE_BASE'))
if res.returncode != 0:
raise RuntimeError(f'Failed to create directory {sandbox_dest} in sandbox')

if recursive:
res = subprocess.run(
f'cp -r {host_src} {sandbox_dest}', shell=True, text=True, cwd=config.get('WORKSPACE_BASE')
)
if res.returncode != 0:
raise RuntimeError(f'Failed to copy {host_src} to {sandbox_dest} in sandbox')
else:
res = subprocess.run(
f'cp {host_src} {sandbox_dest}', shell=True, text=True, cwd=config.get('WORKSPACE_BASE')
)
if res.returncode != 0:
raise RuntimeError(f'Failed to copy {host_src} to {sandbox_dest} in sandbox')

def execute_in_background(self, cmd: str) -> Process:
process = subprocess.Popen(
cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
Expand Down
68 changes: 63 additions & 5 deletions opendevin/sandbox/docker/ssh_box.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
import sys
import time
import uuid
import tarfile
from glob import glob
from collections import namedtuple
from typing import Dict, List, Tuple, Union

Expand All @@ -15,6 +17,7 @@
from opendevin.sandbox.sandbox import Sandbox
from opendevin.sandbox.process import Process
from opendevin.sandbox.docker.process import DockerProcess
from opendevin.sandbox.plugins.jupyter import JupyterRequirement
xingyaoww marked this conversation as resolved.
Show resolved Hide resolved
from opendevin.schema import ConfigType
from opendevin.utils import find_available_tcp_port
from opendevin.exceptions import SandboxInvalidBackgroundCommandError
Expand Down Expand Up @@ -58,10 +61,10 @@ class DockerSSHBox(Sandbox):
background_commands: Dict[int, Process] = {}

def __init__(
self,
container_image: str | None = None,
timeout: int = 120,
sid: str | None = None,
self,
container_image: str | None = None,
timeout: int = 120,
sid: str | None = None,
):
# Initialize docker client. Throws an exception if Docker is not reachable.
try:
Expand Down Expand Up @@ -137,6 +140,22 @@ def setup_user(self):
)
if exit_code != 0:
raise Exception(f'Failed to set password in sandbox: {logs}')

# chown the home directory
exit_code, logs = self.container.exec_run(
['/bin/bash', '-c', 'chown opendevin:root /home/opendevin'],
workdir=SANDBOX_WORKSPACE_DIR,
)
if exit_code != 0:
raise Exception(
f'Failed to chown home directory for opendevin in sandbox: {logs}')
Comment on lines +144 to +151
Copy link
Collaborator

Choose a reason for hiding this comment

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

Maybe we want to pull this out in to the core Sandbox implementation, as initialize_homedir. Some sandboxes might override it, but most are going to do exactly this.

(TBH there's probably a lot of logic we could deduplicate between SSH, Exec, and Local)

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I agree! I was also thinking heavily about this while implementing this, maybe we should make a BaseClass to keep all the common logic (including this one!)

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

well... just realized this might just work on DockerSSHBox and ExecBox, since in LocalBox we are just re-using the current user.

Copy link
Collaborator

Choose a reason for hiding this comment

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

Hmm good point. I could see putting this logic in the base class, and then having a pass override for the LocalBox that makes it a no-op

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I guess once we deprecate ExecBox, this logic will not be duplicated since DockerSSHBox, LocalBox and E2BBox will likely to be very different from each other in terms of implementation. So maybe we can leave this as is, and then try to deprecate ExecBox in the next few days/weeks?

exit_code, logs = self.container.exec_run(
['/bin/bash', '-c', f'chown opendevin:root {SANDBOX_WORKSPACE_DIR}'],
workdir=SANDBOX_WORKSPACE_DIR,
)
if exit_code != 0:
raise Exception(
f'Failed to chown workspace directory for opendevin in sandbox: {logs}')
else:
exit_code, logs = self.container.exec_run(
# change password for root
Expand Down Expand Up @@ -208,6 +227,37 @@ def execute(self, cmd: str) -> Tuple[int, str]:
exit_code = int(exit_code.lstrip('echo $?').strip())
return exit_code, command_output

def copy_to(self, host_src: str, sandbox_dest: str, recursive: bool = False):
# mkdir -p sandbox_dest if it doesn't exist
exit_code, logs = self.container.exec_run(
['/bin/bash', '-c', f'mkdir -p {sandbox_dest}'],
workdir=SANDBOX_WORKSPACE_DIR,
)
if exit_code != 0:
raise Exception(
f'Failed to create directory {sandbox_dest} in sandbox: {logs}')

if recursive:
assert os.path.isdir(host_src), 'Source must be a directory when recursive is True'
files = glob(host_src + '/**/*', recursive=True)
srcname = os.path.basename(host_src)
tar_filename = os.path.join(os.path.dirname(host_src), srcname + '.tar')
with tarfile.open(tar_filename, mode='w') as tar:
for file in files:
tar.add(file, arcname=os.path.relpath(file, os.path.dirname(host_src)))
else:
assert os.path.isfile(host_src), 'Source must be a file when recursive is False'
srcname = os.path.basename(host_src)
tar_filename = os.path.join(os.path.dirname(host_src), srcname + '.tar')
with tarfile.open(tar_filename, mode='w') as tar:
tar.add(host_src, arcname=srcname)

with open(tar_filename, 'rb') as f:
data = f.read()

self.container.put_archive(os.path.dirname(sandbox_dest), data)
os.remove(tar_filename)

def execute_in_background(self, cmd: str) -> Process:
result = self.container.exec_run(
self.get_exec_cmd(cmd), socket=True, workdir=SANDBOX_WORKSPACE_DIR
Expand Down Expand Up @@ -307,6 +357,11 @@ def restart_docker_container(self):
'bind': SANDBOX_WORKSPACE_DIR,
'mode': 'rw'
},
# mount cache directory to /home/opendevin/.cache for pip cache reuse
config.get('CACHE_DIR'): {
'bind': '/home/opendevin/.cache' if RUN_AS_DEVIN else '/root/.cache',
'mode': 'rw'
},
},
)
logger.info('Container started')
Expand Down Expand Up @@ -355,8 +410,11 @@ def close(self):
logger.info(
"Interactive Docker container started. Type 'exit' or use Ctrl+C to exit.")

# Initialize required plugins
ssh_box.init_plugins([JupyterRequirement()])

bg_cmd = ssh_box.execute_in_background(
"while true; do echo 'dot ' && sleep 1; done"
"while true; do echo 'dot ' && sleep 10; done"
)

sys.stdout.flush()
Expand Down
4 changes: 4 additions & 0 deletions opendevin/sandbox/e2b/sandbox.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,10 @@ def execute(self, cmd: str) -> Tuple[int, str]:
assert process_output.exit_code is not None
return process_output.exit_code, logs_str

def copy_to(self, host_src: str, sandbox_dest: str, recursive: bool = False):
# FIXME
raise NotImplementedError('Copying files to E2B sandbox is not implemented yet')
Copy link
Collaborator Author

Choose a reason for hiding this comment

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

I guess this shouldn't be so hard to implement for E2B since we already have .filesystem ready? cc @mlejva

Copy link
Contributor

@mlejva mlejva Apr 22, 2024

Choose a reason for hiding this comment

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

@xingyaoww yes, shouldn't be hard. It's possible to upload files (up to 100MB) to the sandbox or one can just write to a file

Copy link
Contributor

Choose a reason for hiding this comment

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

@rbren let me quickly implement this before merging


def execute_in_background(self, cmd: str) -> Process:
process = self.sandbox.process.start(cmd)
e2b_process = E2BProcess(process, cmd)
Expand Down
7 changes: 7 additions & 0 deletions opendevin/sandbox/plugins/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,7 @@
from .mixin import PluginMixin
from .requirement import PluginRequirement

# Requirements
from .jupyter import JupyterRequirement

__all__ = ['PluginMixin', 'PluginRequirement', 'JupyterRequirement']
11 changes: 11 additions & 0 deletions opendevin/sandbox/plugins/jupyter/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import os
from dataclasses import dataclass
from opendevin.sandbox.plugins.requirement import PluginRequirement


@dataclass
class JupyterRequirement(PluginRequirement):
name: str = 'jupyter'
host_src: str = os.path.dirname(os.path.abspath(__file__)) # The directory of this file (sandbox/plugins/jupyter)
sandbox_dest: str = '/opendevin/plugins/jupyter'
bash_script_path: str = 'setup.sh'
25 changes: 25 additions & 0 deletions opendevin/sandbox/plugins/jupyter/execute_cli
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#!/usr/bin/env python3
import os
import sys
import time
import requests

# Read the Python code from STDIN
code = sys.stdin.read()

# Set the default kernel ID
kernel_id = 'default'

# try 5 times until success
PORT = os.environ.get('JUPYTER_EXEC_SERVER_PORT')
POST_URL = f'http://localhost:{PORT}/execute'

for i in range(5):
response = requests.post(POST_URL, json={'kernel_id': kernel_id, 'code': code})
# if "500: Internal Server Error" is not in the response, break the loop
if '500: Internal Server Error' not in response.text:
break
time.sleep(1)

# Print the response
print(str(response.text))
Loading