diff --git a/docker_agent_sandbox/sandbox.py b/docker_agent_sandbox/sandbox.py index 877087c..009a699 100644 --- a/docker_agent_sandbox/sandbox.py +++ b/docker_agent_sandbox/sandbox.py @@ -44,6 +44,8 @@ class DockerSandbox: cap_drop: list[str] | None = None, cap_add: list[str] | None = None, security_opt: list[str] | None = None, + cpu_limit: float = 8, + memory_limit: str = "16g", ) -> None: self.container_name = container_name self._image = image @@ -56,6 +58,8 @@ class DockerSandbox: self._cap_drop = cap_drop self._cap_add = cap_add self._security_opt = security_opt + self._nano_cpus = int(cpu_limit * 1e9) + self._memory_limit = memory_limit self._client: docker.DockerClient = docker.from_env() self._container: docker.models.containers.Container | None = None @@ -124,6 +128,8 @@ class DockerSandbox: run_kwargs["cap_add"] = self._cap_add if self._security_opt is not None: run_kwargs["security_opt"] = self._security_opt + run_kwargs["nano_cpus"] = self._nano_cpus + run_kwargs["mem_limit"] = self._memory_limit try: self._container = self._client.containers.run(self._image, **run_kwargs)