feat: DockerCommandLineCodeExecutor support for additional volume mounts, exposed host ports

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)
 

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