From 0de4aab4d96aacc7e02682f11f9ea0219c2c8051 Mon Sep 17 00:00:00 2001
From: antonkulaga <antonkulaga@gmail.com>
Date: Tue, 22 Oct 2024 23:29:26 +0200
Subject: [PATCH] added mounts to docker

---
 llm_sandbox/docker.py     |  7 +++++--
 llm_sandbox/micromamba.py | 13 ++++++++-----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py
index 4c66c59..aa100e7 100644
--- 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
index 313ea15..fda2306 100644
--- 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,
@@ -11,17 +11,20 @@ class MicromambaSession(SandboxDockerSession):
     """
 
     def __init__(self, client: Optional[docker.DockerClient] = None,
-        image: Optional[str] = None,
-        dockerfile: Optional[str] = "mambaorg/micromamba:latest",
+        image: Optional[str] = "mambaorg/micromamba:latest",
+        dockerfile: Optional[str] = None,
         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