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
@@ -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
@@ -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(

0 comments on commit c65b71d

Please sign in to comment.