Skip to content

Commit

Permalink
feat(forge): Add mount method to FileStorage & execute code in mo…
Browse files Browse the repository at this point in the history 
…unted workspace (#7115)

* Add `FileStorage.mount()` method, which mounts (part of) the workspace to a local path
  * Add `watchdog` library to watch file changes in mount
* Amend `CodeExecutorComponent`
  * Amend `execute_python_file` to execute Python files in a workspace mount
  * Amend `execute_python_code` to create temporary .py file in workspace instead of as a local file
  * Add support for `Path` argument to `filename` parameter on `execute_python_file`

* Fix `test_execute_python_code` (by making it async)
  • Loading branch information
@kcze
kcze committed May 24, 2024
1 parent edcbbbc commit 4e02f7d
Show file tree
Hide file tree
Showing 7 changed files with 272 additions and 142 deletions.
42 changes: 42 additions & 0 deletions autogpt/poetry.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

5 changes: 3 additions & 2 deletions autogpt/tests/integration/test_execute_code.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,8 @@ def test_execute_python_file_args(
assert result == f"{random_args_string}\n"


def test_execute_python_code(
@pytest.mark.asyncio
async def test_execute_python_code(
code_executor_component: CodeExecutorComponent,
random_code: str,
random_string: str,
Expand All @@ -93,7 +94,7 @@ def test_execute_python_code(
if not (is_docker_available() or we_are_running_in_a_docker_container()):
pytest.skip("Docker is not available")

result: str = code_executor_component.execute_python_code(random_code)
result: str = await code_executor_component.execute_python_code(random_code)
assert result.replace("\r", "") == f"Hello {random_string}!\n"


Expand Down
226 changes: 119 additions & 107 deletions forge/forge/components/code_executor/code_executor.py
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
import logging
import os
import random
import shlex
import string
import subprocess
from pathlib import Path
from tempfile import NamedTemporaryFile
from typing import Iterator

import docker
Expand Down Expand Up @@ -97,7 +98,7 @@ def get_commands(self) -> Iterator[Command]:
),
},
)
def execute_python_code(self, code: str) -> str:
async def execute_python_code(self, code: str) -> str:
"""
Create and execute a Python file in a Docker container
and return the STDOUT of the executed code.
Expand All @@ -113,18 +114,19 @@ def execute_python_code(self, code: str) -> str:
str: The STDOUT captured from the code when it ran.
"""

tmp_code_file = NamedTemporaryFile(
"w", dir=self.workspace.root, suffix=".py", encoding="utf-8"
)
tmp_code_file.write(code)
tmp_code_file.flush()
temp_path = ""
while True:
temp_path = f"temp{self._generate_random_string()}.py"
if not self.workspace.exists(temp_path):
break
await self.workspace.write_file(temp_path, code)

try:
return self.execute_python_file(tmp_code_file.name)
return self.execute_python_file(temp_path)
except Exception as e:
raise CommandExecutionError(*e.args)
finally:
tmp_code_file.close()
self.workspace.delete_file(temp_path)

@command(
["execute_python_file"],
Expand All @@ -144,7 +146,7 @@ def execute_python_code(self, code: str) -> str:
),
},
)
def execute_python_file(self, filename: str, args: list[str] | str = []) -> str:
def execute_python_file(self, filename: str | Path, args: list[str] = []) -> str:
"""Execute a Python file in a Docker container and return the output
Args:
Expand All @@ -154,13 +156,7 @@ def execute_python_file(self, filename: str, args: list[str] | str = []) -> str:
Returns:
str: The output of the file
"""
logger.info(
f"Executing python file '{filename}' "
f"in working directory '{self.workspace.root}'"
)

if isinstance(args, str):
args = args.split() # Convert space-separated string to a list
logger.info(f"Executing python file '{filename}'")

if not str(filename).endswith(".py"):
raise InvalidArgumentError("Invalid file type. Only .py files are allowed.")
Expand All @@ -179,98 +175,21 @@ def execute_python_file(self, filename: str, args: list[str] | str = []) -> str:
"App is running in a Docker container; "
f"executing {file_path} directly..."
)
result = subprocess.run(
["python", "-B", str(file_path)] + args,
capture_output=True,
encoding="utf8",
cwd=str(self.workspace.root),
)
if result.returncode == 0:
return result.stdout
else:
raise CodeExecutionError(result.stderr)
with self.workspace.mount() as local_path:
result = subprocess.run(
["python", "-B", str(file_path.relative_to(self.workspace.root))]
+ args,
capture_output=True,
encoding="utf8",
cwd=str(local_path),
)
if result.returncode == 0:
return result.stdout
else:
raise CodeExecutionError(result.stderr)

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

client = docker.from_env()
image_name = "python:3-alpine"
container_is_fresh = False
container_name = f"{self.state.agent_id}_sandbox"
try:
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(self.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(self.workspace.root).as_posix(),
]
+ args,
stderr=True,
stdout=True,
)

if exec_result.exit_code != 0:
raise CodeExecutionError(exec_result.output.decode("utf-8"))

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

except DockerException as e:
logger.warning(
"Could not run the script in a container. "
"If you haven't already, please install Docker: "
"https://docs.docker.com/get-docker/"
)
raise CommandExecutionError(f"Could not run the script in a container: {e}")
return self._run_python_code_in_docker(file_path, args)

def validate_command(self, command_line: str, config: Config) -> tuple[bool, bool]:
"""Check whether a command is allowed and whether it may be executed in a shell.
Expand Down Expand Up @@ -395,3 +314,96 @@ def execute_shell_popen(self, command_line: str) -> str:
os.chdir(current_dir)

return f"Subprocess started with PID:'{str(process.pid)}'"

def _run_python_code_in_docker(self, filename: str | Path, args: list[str]) -> str:
"""Run a Python script in a Docker container"""
file_path = self.workspace.get_path(filename)
try:
assert self.state.agent_id, "Need Agent ID to attach Docker container"

client = docker.from_env()
image_name = "python:3-alpine"
container_is_fresh = False
container_name = f"{self.state.agent_id}_sandbox"
with self.workspace.mount() as local_path:
try:
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(local_path.resolve()): {
"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(self.workspace.root).as_posix(),
]
+ args,
stderr=True,
stdout=True,
)

if exec_result.exit_code != 0:
raise CodeExecutionError(exec_result.output.decode("utf-8"))

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

except DockerException as e:
logger.warning(
"Could not run the script in a container. "
"If you haven't already, please install Docker: "
"https://docs.docker.com/get-docker/"
)
raise CommandExecutionError(f"Could not run the script in a container: {e}")

def _generate_random_string(self, length: int = 8):
# Create a string of all letters and digits
characters = string.ascii_letters + string.digits
# Use random.choices to generate a random string
random_string = "".join(random.choices(characters, k=length))
return random_string
Loading

0 comments on commit 4e02f7d

Please sign in to comment.