From bc7ad8651b4a83b93a12ddff062df476267dcb77 Mon Sep 17 00:00:00 2001 From: Martin Simonovsky Date: Wed, 13 Nov 2024 10:28:00 +0000 Subject: [PATCH] Add commit_container option --- llm_sandbox/docker.py | 5 ++++- llm_sandbox/session.py | 3 +++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py index 80d73ef..1002c7f 100644 --- a/llm_sandbox/docker.py +++ b/llm_sandbox/docker.py @@ -31,6 +31,7 @@ def __init__( dockerfile: Optional[str] = None, lang: str = SupportedLanguage.PYTHON, keep_template: bool = False, + commit_container: bool = True, verbose: bool = False, mounts: Optional[list[Mount]] = None, ): @@ -41,6 +42,7 @@ def __init__( :param dockerfile: Path to the Dockerfile, if image is not provided :param lang: Language of the code :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 """ super().__init__(lang, verbose) @@ -71,6 +73,7 @@ def __init__( self.container: Optional[Container] = None self.path = None self.keep_template = keep_template + self.commit_container = commit_container self.is_create_template: bool = False self.verbose = verbose self.mounts = mounts @@ -114,7 +117,7 @@ def open(self): def close(self): if self.container: - if isinstance(self.image, Image): + if self.commit_container and isinstance(self.image, Image): self.container.commit(self.image.tags[-1]) self.container.remove(force=True) diff --git a/llm_sandbox/session.py b/llm_sandbox/session.py index 71bf4b4..5e56dfa 100644 --- a/llm_sandbox/session.py +++ b/llm_sandbox/session.py @@ -14,6 +14,7 @@ def __new__( dockerfile: Optional[str] = None, lang: str = SupportedLanguage.PYTHON, keep_template: bool = False, + commit_container: bool = True, verbose: bool = False, use_kubernetes: bool = False, kube_namespace: Optional[str] = "default", @@ -25,6 +26,7 @@ def __new__( :param dockerfile: Path to the Dockerfile, if image is not provided :param lang: Language of the code :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 (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' @@ -46,5 +48,6 @@ def __new__( dockerfile=dockerfile, lang=lang, keep_template=keep_template, + commit_container=commit_container, verbose=verbose, )