diff --git a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py index 05924e186643..1ac058a9680f 100644 --- a/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py +++ b/python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py @@ -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 ( @@ -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]] = [ @@ -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.") @@ -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 @@ -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) diff --git a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py index ad74236008fa..6c65835d183d 100644 --- a/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py +++ b/python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py @@ -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