From 272c262f76386c1994f628e5f0e7e0e12cae0299 Mon Sep 17 00:00:00 2001
From: Duy Huynh <vndee.huynh@gmail.com>
Date: Mon, 18 Nov 2024 12:17:17 +0700
Subject: [PATCH] feat: add `container_configs` argument for docker session

---
 llm_sandbox/docker.py  | 10 +++++++++-
 llm_sandbox/session.py |  3 +++
 pyproject.toml         |  2 +-
 3 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py
index 1002c7f..607029f 100644
--- a/llm_sandbox/docker.py
+++ b/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,
         )
 
     def close(self):
diff --git a/llm_sandbox/session.py b/llm_sandbox/session.py
index 5e56dfa..42b20f2 100644
--- a/llm_sandbox/session.py
+++ b/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,
         )
diff --git a/pyproject.toml b/pyproject.toml
index 1dfe5db..30d2989 100644
--- a/pyproject.toml
+++ b/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 <vndee.huynh@gmail.com>"]
 license = "MIT"