Merge pull request #7 from vndee/feat/container-configs-arguments-for…

…-docker-session

feat: add `container_configs` argument for docker session

llm_sandbox/docker.py

@@ -34,6 +34,7 @@ def __init__(
         commit_container: bool = True,
         verbose: bool = False,
         mounts: Optional[list[Mount]] = None,
+        container_configs: Optional[dict] = None,
     ):
         """
         Create a new sandbox session
@@ -44,6 +45,8 @@ def __init__(
         :param keep_template: if True, the image and container will not be removed after the session ends
         :param commit_container: if True, the Docker container will be commited after the session ends
         :param verbose: if True, print messages
+        :param mounts: List of mounts to be mounted to the container
+        :param container_configs: Additional configurations for the container, i.e. resources limits (cpu_count, mem_limit), etc.
         """
         super().__init__(lang, verbose)
         if image and dockerfile:
@@ -77,6 +80,7 @@ def __init__(
         self.is_create_template: bool = False
         self.verbose = verbose
         self.mounts = mounts
+        self.container_configs = container_configs
 
     def open(self):
         warning_str = (
@@ -112,7 +116,11 @@ def open(self):
                     print(f"Using image {self.image.tags[-1]}")
 
         self.container = self.client.containers.run(
-            self.image, detach=True, tty=True, mounts=self.mounts
+            self.image,
+            detach=True,
+            tty=True,
+            mounts=self.mounts,
+            **self.container_configs if self.container_configs else {},
         )
 
     def close(self):

llm_sandbox/session.py

@@ -18,6 +18,7 @@ def __new__(
         verbose: bool = False,
         use_kubernetes: bool = False,
         kube_namespace: Optional[str] = "default",
+        container_configs: Optional[dict] = None,
     ):
         """
         Create a new sandbox session
@@ -30,6 +31,7 @@ def __new__(
         :param verbose: if True, print messages (default is True)
         :param use_kubernetes: if True, use Kubernetes instead of Docker (default is False)
         :param kube_namespace: Kubernetes namespace to use (only if 'use_kubernetes' is True), default is 'default'
+        :param container_configs: Additional configurations for the Docker container, i.e. resources limits (cpu_count, mem_limit), etc.
         """
         if use_kubernetes:
             return SandboxKubernetesSession(
@@ -50,4 +52,5 @@ def __new__(
             keep_template=keep_template,
             commit_container=commit_container,
             verbose=verbose,
+            container_configs=container_configs,
         )

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "llm-sandbox"
-version = "0.1.2-alpha"
+version = "0.1.4"
 description = "Lightweight and portable LLM sandbox runtime (code interpreter) Python library"
 authors = ["Duy Huynh <[email protected]>"]
 license = "MIT"