Skip to content

Commit

Permalink
feat: DockerCommandLineCodeExecutor support for additional volume mou…
Browse files Browse the repository at this point in the history
…nts, exposed host ports (#5383)

Add the following additional configuration options to
DockerCommandLineCodeExectutor:

- **extra_volumes** (Optional[Dict[str, Dict[str, str]]], optional): A
dictionary of extra volumes (beyond the work_dir) to mount to the
container. Defaults to None.
- **extra_hosts** (Optional[Dict[str, str]], optional): A dictionary of
host mappings to add to the container. (See Docker docs on extra_hosts)
Defaults to None.
- **init_command** (Optional[str], optional): A shell command to run
before each shell operation execution. Defaults to None. 

## Why are these changes needed?

See linked issue below.

In summary: Enable the agents to:
- work with a richer set of sys admin tools on top of code execution
- add support for a 'project' directory the agents can interact on
that's accessible by bash tools and custom scripts

## Related issue number

Closes #5363

## Checks

- [x] I've included any doc changes needed for
https://microsoft.github.io/autogen/. See
https://microsoft.github.io/autogen/docs/Contribute#documentation to
build and test documentation locally.
- [x] I've added tests (if relevant) corresponding to the changes
introduced in this PR.
- [x] I've made sure all auto checks have passed.
  • Loading branch information
andrejpk authored Feb 11, 2025
1 parent a9db384 commit 540c4fb
Show file tree
Hide file tree
Showing 2 changed files with 68 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from hashlib import sha256
from pathlib import Path
from types import TracebackType
from typing import Any, Callable, ClassVar, List, Optional, ParamSpec, Type, Union
from typing import Any, Callable, ClassVar, Dict, List, Optional, ParamSpec, Type, Union

from autogen_core import CancellationToken
from autogen_core.code_executor import (
Expand Down Expand Up @@ -88,6 +88,13 @@ class DockerCommandLineCodeExecutor(CodeExecutor):
the Python process exits with atext. Defaults to True.
functions (List[Union[FunctionWithRequirements[Any, A], Callable[..., Any]]]): A list of functions that are available to the code executor. Default is an empty list.
functions_module (str, optional): The name of the module that will be created to store the functions. Defaults to "functions".
extra_volumes (Optional[Dict[str, Dict[str, str]]], optional): A dictionary of extra volumes (beyond the work_dir) to mount to the container;
key is host source path and value 'bind' is the container path. See Defaults to None.
Example: extra_volumes = {'/home/user1/': {'bind': '/mnt/vol2', 'mode': 'rw'}, '/var/www': {'bind': '/mnt/vol1', 'mode': 'ro'}}
extra_hosts (Optional[Dict[str, str]], optional): A dictionary of host mappings to add to the container. (See Docker docs on extra_hosts) Defaults to None.
Example: extra_hosts = {"kubernetes.docker.internal": "host-gateway"}
init_command (Optional[str], optional): A shell command to run before each shell operation execution. Defaults to None.
Example: init_command="kubectl config use-context docker-hub"
"""

SUPPORTED_LANGUAGES: ClassVar[List[str]] = [
Expand Down Expand Up @@ -126,6 +133,9 @@ def __init__(
]
] = [],
functions_module: str = "functions",
extra_volumes: Optional[Dict[str, Dict[str, str]]] = None,
extra_hosts: Optional[Dict[str, str]] = None,
init_command: Optional[str] = None,
):
if timeout < 1:
raise ValueError("Timeout must be greater than or equal to 1.")
Expand Down Expand Up @@ -157,6 +167,10 @@ def __init__(

self._functions_module = functions_module
self._functions = functions
self._extra_volumes = extra_volumes if extra_volumes is not None else {}
self._extra_hosts = extra_hosts if extra_hosts is not None else {}
self._init_command = init_command

# Setup could take some time so we intentionally wait for the first code block to do it.
if len(functions) > 0:
self._setup_functions_complete = False
Expand Down Expand Up @@ -354,16 +368,22 @@ async def start(self) -> None:
# Let the docker exception escape if this fails.
await asyncio.to_thread(client.images.pull, self._image)

# Prepare the command (if needed)
shell_command = "/bin/sh"
command = ["-c", f"{(self._init_command)};exec {shell_command}"] if self._init_command else None

self._container = await asyncio.to_thread(
client.containers.create,
self._image,
name=self.container_name,
entrypoint="/bin/sh",
entrypoint=shell_command,
command=command,
tty=True,
detach=True,
auto_remove=self._auto_remove,
volumes={str(self._bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}},
volumes={str(self._bind_dir.resolve()): {"bind": "/workspace", "mode": "rw"}, **self._extra_volumes},
working_dir="/workspace",
extra_hosts=self._extra_hosts,
)
await asyncio.to_thread(self._container.start)

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -164,3 +164,48 @@ async def test_docker_commandline_code_executor_start_stop_context_manager() ->
with tempfile.TemporaryDirectory() as temp_dir:
async with DockerCommandLineCodeExecutor(work_dir=temp_dir) as _exec:
pass


@pytest.mark.asyncio
async def test_docker_commandline_code_executor_extra_args() -> None:
if not docker_tests_enabled():
pytest.skip("Docker tests are disabled")

with tempfile.TemporaryDirectory() as temp_dir:
# Create a file in temp_dir to mount
host_file_path = Path(temp_dir) / "host_file.txt"
host_file_path.write_text("This is a test file.")

container_file_path = "/container/host_file.txt"

extra_volumes = {str(host_file_path): {"bind": container_file_path, "mode": "rw"}}
init_command = "echo 'Initialization command executed' > /workspace/init_command.txt"
extra_hosts = {"example.com": "127.0.0.1"}

async with DockerCommandLineCodeExecutor(
work_dir=temp_dir,
extra_volumes=extra_volumes,
init_command=init_command,
extra_hosts=extra_hosts,
) as executor:
cancellation_token = CancellationToken()

# Verify init_command was executed
init_command_file_path = Path(temp_dir) / "init_command.txt"
assert init_command_file_path.exists()

# Verify extra_hosts
ns_lookup_code_blocks = [
CodeBlock(code="import socket; print(socket.gethostbyname('example.com'))", language="python")
]
ns_lookup_result = await executor.execute_code_blocks(ns_lookup_code_blocks, cancellation_token)
assert ns_lookup_result.exit_code == 0
assert "127.0.0.1" in ns_lookup_result.output

# Verify the file is accessible in the volume mounted in extra_volumes
code_blocks = [
CodeBlock(code=f"with open('{container_file_path}') as f: print(f.read())", language="python")
]
code_result = await executor.execute_code_blocks(code_blocks, cancellation_token)
assert code_result.exit_code == 0
assert "This is a test file." in code_result.output

0 comments on commit 540c4fb

Please sign in to comment.