added mounts to docker

vndee · Oct 23, 2024 · b6069e6 · b6069e6
1 parent fd0e091
commit b6069e6
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py
@@ -7,6 +7,7 @@
 
 from docker.models.images import Image
 from docker.models.containers import Container
+from docker.types import Mount
 from llm_sandbox.utils import (
     image_exists,
     get_libraries_installation_command,
@@ -31,6 +32,7 @@ def __init__(
         lang: str = SupportedLanguage.PYTHON,
         keep_template: bool = False,
         verbose: bool = False,
+        mounts: Optional[list[Mount]] = None
     ):
         """
         Create a new sandbox session
@@ -71,7 +73,8 @@ def __init__(
         self.keep_template = keep_template
         self.is_create_template: bool = False
         self.verbose = verbose
-
+        self.mounts = mounts
+
     def open(self):
         warning_str = (
             "Since the `keep_template` flag is set to True the docker image will not be removed after the session ends "
@@ -105,7 +108,7 @@ def open(self):
                 if self.verbose:
                     print(f"Using image {self.image.tags[-1]}")
 
-        self.container = self.client.containers.run(self.image, detach=True, tty=True)
+        self.container = self.client.containers.run(self.image, detach=True, tty=True, mounts=self.mounts)
 
     def close(self):
         if self.container:

diff --git a/llm_sandbox/micromamba.py b/llm_sandbox/micromamba.py
@@ -3,7 +3,7 @@
 from llm_sandbox.session import SandboxDockerSession
 from llm_sandbox.docker import ConsoleOutput
 from llm_sandbox.const import SupportedLanguage
-
+from docker.types import Mount
 class MicromambaSession(SandboxDockerSession):
     """
     SandboxDockerSession does not allow activation of micromamba environment,
@@ -15,13 +15,16 @@ def __init__(self, client: Optional[docker.DockerClient] = None,
         dockerfile: Optional[str] = "mambaorg/micromamba:latest",
         lang: str = SupportedLanguage.PYTHON,
         keep_template: bool = False,
-        verbose: bool = False, environment: str = "base"):
+        verbose: bool = False, 
+        mounts: Optional[list[Mount]] = None,
+        environment: str = "base"):
         super().__init__(client=client,
                          image=image,
                          dockerfile=dockerfile,
                          lang=lang,
                          keep_template=keep_template,
-                         verbose=verbose)
+                         verbose=verbose,
+                         mounts=mounts)
         self.environment = environment