diff --git a/docker_agent_sandbox/sandbox.py b/docker_agent_sandbox/sandbox.py index 014ca64..9a2958e 100644 --- a/docker_agent_sandbox/sandbox.py +++ b/docker_agent_sandbox/sandbox.py @@ -4,20 +4,16 @@ from __future__ import annotations import io import select as _select -import socket as _socket import tarfile -from contextlib import ExitStack -from select import select as _original_select from pathlib import Path from typing import TYPE_CHECKING -from unittest.mock import patch +import docker +import docker.errors +from loguru import logger + _original_poll = getattr(_select, "poll", None) -import docker -import docker.errors -from docker.utils.socket import consume_socket_output, demux_adaptor, frames_iter -from loguru import logger if TYPE_CHECKING: import docker.models.containers @@ -50,6 +46,8 @@ class DockerSandbox: cpu_limit: float = 8, memory_limit: str = "16g", command: str | None = None, + uid: int | None = None, + gid: int | None = None, ) -> None: self.container_name = container_name self._image = image @@ -65,9 +63,21 @@ class DockerSandbox: self._nano_cpus = int(cpu_limit * 1e9) self._memory_limit = memory_limit self._command = command + self._uid = uid + self._gid = gid self._client: docker.DockerClient = docker.from_env() self._container: docker.models.containers.Container | None = None + @property + def _user(self) -> str | None: + """Docker ``user`` spec ("uid" or "uid:gid"), or None to use the image default.""" + if self._uid is None and self._gid is None: + return None + uid = "" if self._uid is None else str(self._uid) + if self._gid is None: + return uid + return f"{uid}:{self._gid}" + def __enter__(self) -> "DockerSandbox": return self @@ -135,9 +145,13 @@ class DockerSandbox: run_kwargs["security_opt"] = self._security_opt run_kwargs["nano_cpus"] = self._nano_cpus run_kwargs["mem_limit"] = self._memory_limit + if self._user is not None: + run_kwargs["user"] = self._user try: - self._container = self._client.containers.run(self._image, self._command, **run_kwargs) + self._container = self._client.containers.run( + self._image, self._command, **run_kwargs + ) except docker.errors.ImageNotFound: raise RuntimeError( f"Image {self._image!r} not found locally. " @@ -247,13 +261,20 @@ class DockerSandbox: # Command execution # ------------------------------------------------------------------ - def exec(self, command: str, timeout: int = 120, kill_grace: int = 5) -> tuple[int, str]: + def exec( + self, command: str, timeout: int = 120, kill_grace: int = 5 + ) -> tuple[int, str]: if self._container is None: return 1, "Sandbox container is not running." wrapped = [ - "timeout", "--kill-after", f"{kill_grace}s", f"{timeout}s", - "bash", "-c", command, + "timeout", + "--kill-after", + f"{kill_grace}s", + f"{timeout}s", + "bash", + "-c", + command, ] kwargs = {} if self._working_dir is not None: