Skip to content

Commit

Permalink
feat: Re-use Docker container for code execution
Browse files Browse the repository at this point in the history
- Create a unique container name based on agent ID
- Check if the container with the name exists, otherwise create a new container
- If the container is not running, start it; otherwise, restart it
- Execute the code in the container
- Return the output of the code execution

This change enables reusing the same container for consecutive code execution commands, allowing for iterative changes to the execution environment.

Note: This change also includes handling the case where the Docker image is not found locally by pulling it from Docker Hub. The image used in this case is "python:3-alpine".
  • Loading branch information
Pwuts committed Oct 31, 2023
1 parent c3569d1 commit c65b71d
Showing 1 changed file with 54 additions and 37 deletions.
91 changes: 54 additions & 37 deletions autogpts/autogpt/autogpt/commands/execute_code.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from tempfile import NamedTemporaryFile

import docker
from docker.errors import DockerException, ImageNotFound
from docker.errors import DockerException, ImageNotFound, NotFound
from docker.models.containers import Container as DockerContainer

from autogpt.agents.agent import Agent
Expand Down Expand Up @@ -135,58 +135,75 @@ def execute_python_file(

logger.debug("AutoGPT is not running in a Docker container")
try:
assert agent.state.agent_id, "Need Agent ID to attach Docker container"

client = docker.from_env()
# You can replace this with the desired Python image/version
# You can find available Python images on Docker Hub:
# https://hub.docker.com/_/python
image_name = "python:3-alpine"
container_is_fresh = False
container_name = f"{agent.state.agent_id}_sandbox"
try:
client.images.get(image_name)
logger.debug(f"Image '{image_name}' found locally")
except ImageNotFound:
logger.info(
f"Image '{image_name}' not found locally, pulling from Docker Hub..."
)
# Use the low-level API to stream the pull response
low_level_client = docker.APIClient()
for line in low_level_client.pull(image_name, stream=True, decode=True):
# Print the status and progress, if available
status = line.get("status")
progress = line.get("progress")
if status and progress:
logger.info(f"{status}: {progress}")
elif status:
logger.info(status)

logger.debug(f"Running {file_path} in a {image_name} container...")
container: DockerContainer = client.containers.run(
image_name,
container: DockerContainer = client.containers.get(container_name) # type: ignore
except NotFound:
try:
client.images.get(image_name)
logger.debug(f"Image '{image_name}' found locally")
except ImageNotFound:
logger.info(
f"Image '{image_name}' not found locally, pulling from Docker Hub..."
)
# Use the low-level API to stream the pull response
low_level_client = docker.APIClient()
for line in low_level_client.pull(image_name, stream=True, decode=True):
# Print the status and progress, if available
status = line.get("status")
progress = line.get("progress")
if status and progress:
logger.info(f"{status}: {progress}")
elif status:
logger.info(status)

logger.debug(f"Creating new {image_name} container...")
container: DockerContainer = client.containers.run(
image_name,
["sleep", "60"], # Max 60 seconds to prevent permanent hangs
volumes={
str(agent.workspace.root): {
"bind": "/workspace",
"mode": "rw",
}
},
working_dir="/workspace",
stderr=True,
stdout=True,
detach=True,
name=container_name,
) # type: ignore
container_is_fresh = True

if not container.status == "running":
container.start()
elif not container_is_fresh:
container.restart()

logger.debug(f"Running {file_path} in container {container.name}...")
exec_result = container.exec_run(
[
"python",
"-B",
file_path.relative_to(agent.workspace.root).as_posix(),
]
+ args,
volumes={
str(agent.workspace.root): {
"bind": "/workspace",
"mode": "rw",
}
},
working_dir="/workspace",
stderr=True,
stdout=True,
detach=True,
) # type: ignore

container.wait()
logs = container.logs().decode("utf-8")
container.remove()
)

# print(f"Execution complete. Output: {output}")
# print(f"Logs: {logs}")
if exec_result.exit_code != 0:
raise CodeExecutionError(exec_result.output.decode("utf-8"))

return logs
return exec_result.output.decode("utf-8")

except DockerException as e:
logger.warn(
Expand Down

0 comments on commit c65b71d

Please sign in to comment.