Skip to content

fix(autogen-ext): preserve Docker GPU requests in config #8140

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
YZJF (YZJF) wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(autogen-ext): preserve Docker GPU requests in config #8140

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions python/packages/autogen-ext/src/autogen_ext/code_executors/docker/_docker_code_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,7 @@ class DockerCommandLineCodeExecutorConfig(BaseModel):
extra_hosts: Dict[str, str] = {}
init_command: Optional[str] = None
delete_tmp_files: bool = False
device_requests: Optional[List[Dict[str, Any]]] = None


class DockerCommandLineCodeExecutor(CodeExecutor, Component[DockerCommandLineCodeExecutorConfig]):
Expand Down Expand Up @@ -588,6 +589,9 @@ def _to_config(self) -> DockerCommandLineCodeExecutorConfig:
extra_hosts=self._extra_hosts,
init_command=self._init_command,
delete_tmp_files=self._delete_tmp_files,
device_requests=[dict(device_request) for device_request in self._device_requests]
if self._device_requests is not None
else None,
)

@classmethod
Expand All @@ -608,4 +612,7 @@ def _from_config(cls, config: DockerCommandLineCodeExecutorConfig) -> Self:
extra_hosts=config.extra_hosts,
init_command=config.init_command,
delete_tmp_files=config.delete_tmp_files,
device_requests=[DeviceRequest(**device_request) for device_request in config.device_requests]
if config.device_requests is not None
else None,
)
6 changes: 5 additions & 1 deletion python/packages/autogen-ext/tests/code_executors/test_docker_commandline_code_executor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from autogen_core import CancellationToken
from autogen_core.code_executor import CodeBlock
from autogen_ext.code_executors.docker import DockerCommandLineCodeExecutor
from docker.types import DeviceRequest


def docker_tests_enabled() -> bool:
Expand Down Expand Up @@ -258,13 +259,16 @@ async def test_docker_commandline_code_executor_extra_args() -> None:
@pytest.mark.asyncio
async def test_docker_commandline_code_executor_serialization() -> None:
with tempfile.TemporaryDirectory() as temp_dir:
executor = DockerCommandLineCodeExecutor(work_dir=temp_dir)
device_request = DeviceRequest(count=-1, capabilities=[["gpu"]])
executor = DockerCommandLineCodeExecutor(work_dir=temp_dir, device_requests=[device_request])

executor_config = executor.dump_component()
loaded_executor = DockerCommandLineCodeExecutor.load_component(executor_config)

assert executor.bind_dir == loaded_executor.bind_dir
assert executor.timeout == loaded_executor.timeout
assert executor_config.config["device_requests"] == [dict(device_request)]
assert loaded_executor._device_requests == [device_request] # type: ignore[reportPrivateUsage]


def test_invalid_timeout() -> None:
Expand Down