From 814ee5dcaa1a619067c7e14269aaad530fd47ab6 Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 20 Jan 2025 11:47:06 +0100 Subject: [PATCH 01/11] Added `exit_code` to `ConsoleOutput` --- llm_sandbox/base.py | 13 +++++++------ llm_sandbox/docker.py | 4 ++-- llm_sandbox/kubernetes.py | 2 +- llm_sandbox/micromamba.py | 2 +- llm_sandbox/podman.py | 4 ++-- tests/test_session.py | 12 ++++++++++++ tests/test_session_podman.py | 13 +++++++++++++ 7 files changed, 38 insertions(+), 12 deletions(-) diff --git a/llm_sandbox/base.py b/llm_sandbox/base.py index 46974dd..7c97c1e 100644 --- a/llm_sandbox/base.py +++ b/llm_sandbox/base.py @@ -3,22 +3,23 @@ class ConsoleOutput: - def __init__(self, text: str): + def __init__(self, exit_code: int, text: str): + self._exit_code = exit_code self._text = text + @property + def exit_code(self): + return self._exit_code + @property def text(self): return self._text def __str__(self): - return f"ConsoleOutput(text={self.text})" + return f"ConsoleOutput(text={self.text}, exit_code={self._exit_code})" class KubernetesConsoleOutput(ConsoleOutput): - def __init__(self, exit_code: int, text: str): - super().__init__(text) - self.exit_code = exit_code - def __str__(self): return f"KubernetesConsoleOutput(text={self.text}, exit_code={self.exit_code})" diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py index 4ed1811..c4fc8ca 100644 --- a/llm_sandbox/docker.py +++ b/llm_sandbox/docker.py @@ -196,7 +196,7 @@ def run(self, code: str, libraries: Optional[List] = None) -> ConsoleOutput: self.copy_to_runtime(code_file, code_dest_file) - output = ConsoleOutput("") + output = ConsoleOutput(0, "") commands = get_code_execution_command(self.lang, code_dest_file) for command in commands: if self.lang == SupportedLanguage.GO: @@ -280,4 +280,4 @@ def execute_command( if self.verbose: print(chunk_str, end="") - return ConsoleOutput(output) + return ConsoleOutput(exit_code, output) diff --git a/llm_sandbox/kubernetes.py b/llm_sandbox/kubernetes.py index 29eef68..1d4fea1 100644 --- a/llm_sandbox/kubernetes.py +++ b/llm_sandbox/kubernetes.py @@ -164,7 +164,7 @@ def run(self, code: str, libraries: Optional[List] = None) -> ConsoleOutput: if output.exit_code != 0: break - return ConsoleOutput(output.text) + return ConsoleOutput(output.exit_code, output.text) def copy_to_runtime(self, src: str, dest: str): if not self.container: diff --git a/llm_sandbox/micromamba.py b/llm_sandbox/micromamba.py index 9228952..3bf236f 100644 --- a/llm_sandbox/micromamba.py +++ b/llm_sandbox/micromamba.py @@ -68,4 +68,4 @@ def execute_command( if self.verbose: print(chunk_str, end="") - return ConsoleOutput(output) + return ConsoleOutput(exit_code, output) diff --git a/llm_sandbox/podman.py b/llm_sandbox/podman.py index 12e543e..f70bba4 100644 --- a/llm_sandbox/podman.py +++ b/llm_sandbox/podman.py @@ -218,7 +218,7 @@ def run(self, code: str, libraries: Optional[List] = None) -> ConsoleOutput: self.copy_to_runtime(code_file, code_dest_file) - output = ConsoleOutput("") + output = ConsoleOutput(0, "") commands = get_code_execution_command(self.lang, code_dest_file) for command in commands: if self.lang == SupportedLanguage.GO: @@ -304,4 +304,4 @@ def execute_command( if self.verbose: print(output) - return ConsoleOutput(output) + return ConsoleOutput(exit_code, output) diff --git a/tests/test_session.py b/tests/test_session.py index 4bc6f21..bbaa754 100644 --- a/tests/test_session.py +++ b/tests/test_session.py @@ -124,6 +124,18 @@ def test_execute_empty_command(self): with self.assertRaises(ValueError): self.session.execute_command("") + def test_execute_failing_command(self): + mock_container = MagicMock() + self.session.container = mock_container + + command = "exit 1" + mock_container.exec_run.return_value = (1, iter([])) + + output = self.session.execute_command(command) + mock_container.exec_run.assert_called_with(command, stream=True, tty=True) + self.assertEqual(output.exit_code, 1) + self.assertEqual(output.text, "") + if __name__ == "__main__": unittest.main() diff --git a/tests/test_session_podman.py b/tests/test_session_podman.py index c9cf651..430b70b 100644 --- a/tests/test_session_podman.py +++ b/tests/test_session_podman.py @@ -149,6 +149,19 @@ def test_execute_empty_command(self): with self.assertRaises(ValueError): self.session.execute_command("") + def test_execute_failing_command(self): + mock_container = MagicMock() + self.session.container = mock_container + + command = "exit 1" + mock_container.exec_run.return_value = (1, iter([])) + + output = self.session.execute_command(command) + + mock_container.exec_run.assert_called_with(command, stream=True, tty=True) + self.assertEqual(output.exit_code, 1) + self.assertEqual(output.text, "") + if __name__ == "__main__": unittest.main() From d62aeded7f2675469e94b8fc8818cdb9eafcb78a Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 20 Jan 2025 12:59:13 +0100 Subject: [PATCH 02/11] Add support to toggle `stream` mode --- llm_sandbox/docker.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py index c4fc8ca..6f301ee 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, + stream: bool = True, container_configs: Optional[dict] = None, ): """ @@ -80,6 +81,7 @@ def __init__( self.is_create_template: bool = False self.verbose = verbose self.mounts = mounts + self.stream = stream self.container_configs = container_configs def open(self): @@ -263,17 +265,20 @@ def execute_command( if workdir: exit_code, exec_log = self.container.exec_run( - command, stream=True, tty=True, workdir=workdir + command, stream=self.stream, tty=True, workdir=workdir ) else: exit_code, exec_log = self.container.exec_run( - command, stream=True, tty=True + command, stream=self.stream, tty=True ) output = "" if self.verbose: print("Output:", end=" ") + if not self.stream: + exec_log = [exec_log] + for chunk in exec_log: chunk_str = chunk.decode("utf-8") output += chunk_str From eb4e7c2ef63970be2bad13a161483bccf8bb498f Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 20 Jan 2025 13:01:23 +0100 Subject: [PATCH 03/11] Added docs and fixed typing --- llm_sandbox/base.py | 2 +- llm_sandbox/docker.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/llm_sandbox/base.py b/llm_sandbox/base.py index 7c97c1e..e7404ad 100644 --- a/llm_sandbox/base.py +++ b/llm_sandbox/base.py @@ -3,7 +3,7 @@ class ConsoleOutput: - def __init__(self, exit_code: int, text: str): + def __init__(self, exit_code: Optional[int], text: str): self._exit_code = exit_code self._text = text diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py index 6f301ee..c266baa 100644 --- a/llm_sandbox/docker.py +++ b/llm_sandbox/docker.py @@ -47,6 +47,7 @@ def __init__( :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 stream: if True, the output will be streamed (enabling this option prevents obtaining an exit code of run command) :param container_configs: Additional configurations for the container, i.e. resources limits (cpu_count, mem_limit), etc. """ super().__init__(lang, verbose) From a6fcc9275f89e4f70399242669165b4608227fb3 Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 20 Jan 2025 13:09:59 +0100 Subject: [PATCH 04/11] Add `stream` for `MicromambaSession` --- llm_sandbox/micromamba.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/llm_sandbox/micromamba.py b/llm_sandbox/micromamba.py index 3bf236f..e3d6f51 100644 --- a/llm_sandbox/micromamba.py +++ b/llm_sandbox/micromamba.py @@ -21,6 +21,7 @@ def __init__( keep_template: bool = False, verbose: bool = False, mounts: Optional[list[Mount]] = None, + stream: bool = True, environment: str = "base", ): super().__init__( @@ -31,6 +32,7 @@ def __init__( keep_template=keep_template, verbose=verbose, mounts=mounts, + stream=stream, ) self.environment = environment @@ -51,17 +53,20 @@ def execute_command( if workdir: exit_code, exec_log = self.container.exec_run( - command, stream=True, tty=True, workdir=workdir + command, stream=self.stream, tty=True, workdir=workdir ) else: exit_code, exec_log = self.container.exec_run( - command, stream=True, tty=True + command, stream=self.stream, tty=True ) output = "" if self.verbose: print("Output:", end=" ") + if not self.stream: + exec_log = [exec_log] + for chunk in exec_log: chunk_str = chunk.decode("utf-8") output += chunk_str From 34620e535db5870814bdf1553d70a14d6c226436 Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 20 Jan 2025 13:11:23 +0100 Subject: [PATCH 05/11] Revert "Add `stream` for `MicromambaSession`" This reverts commit a6fcc9275f89e4f70399242669165b4608227fb3. --- llm_sandbox/micromamba.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/llm_sandbox/micromamba.py b/llm_sandbox/micromamba.py index e3d6f51..3bf236f 100644 --- a/llm_sandbox/micromamba.py +++ b/llm_sandbox/micromamba.py @@ -21,7 +21,6 @@ def __init__( keep_template: bool = False, verbose: bool = False, mounts: Optional[list[Mount]] = None, - stream: bool = True, environment: str = "base", ): super().__init__( @@ -32,7 +31,6 @@ def __init__( keep_template=keep_template, verbose=verbose, mounts=mounts, - stream=stream, ) self.environment = environment @@ -53,20 +51,17 @@ def execute_command( if workdir: exit_code, exec_log = self.container.exec_run( - command, stream=self.stream, tty=True, workdir=workdir + command, stream=True, tty=True, workdir=workdir ) else: exit_code, exec_log = self.container.exec_run( - command, stream=self.stream, tty=True + command, stream=True, tty=True ) output = "" if self.verbose: print("Output:", end=" ") - if not self.stream: - exec_log = [exec_log] - for chunk in exec_log: chunk_str = chunk.decode("utf-8") output += chunk_str From bc488adc6d7de4df1a9b5eb059a7880112e1afc8 Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 27 Jan 2025 11:16:32 +0100 Subject: [PATCH 06/11] Add `pod_manifest` to `SandboxKubernetesSession` --- llm_sandbox/kubernetes.py | 42 +++++++++++++------- tests/test_session_kubernetes.py | 66 ++++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+), 14 deletions(-) create mode 100644 tests/test_session_kubernetes.py diff --git a/llm_sandbox/kubernetes.py b/llm_sandbox/kubernetes.py index 1d4fea1..5c16fab 100644 --- a/llm_sandbox/kubernetes.py +++ b/llm_sandbox/kubernetes.py @@ -27,6 +27,7 @@ def __init__( verbose: bool = False, kube_namespace: Optional[str] = "default", env_vars: Optional[dict] = None, + pod_manifest: Optional[dict] = None, ): """ Create a new sandbox session @@ -37,6 +38,8 @@ def __init__( :param keep_template: if True, the image and container will not be removed after the session ends :param verbose: if True, print messages :param kube_namespace: Kubernetes namespace to use, default is 'default' + :param env_vars: Environment variables to use + :param pod_manifest: Pod manifest to use (ignores other settings: `image`, `kube_namespace` and `env_vars`) """ super().__init__(lang, verbose) if lang not in SupportedLanguageValues: @@ -62,11 +65,34 @@ def __init__( self.keep_template = keep_template self.container = None self.env_vars = env_vars + self.pod_manifest = pod_manifest or self._default_pod_manifest() + self._reconfigure_with_pod_manifest() + + def _reconfigure_with_pod_manifest(self): + self.pod_name = self.pod_manifest.get("metadata", {}).get("name", self.pod_name) + self.kube_namespace = self.pod_manifest.get("metadata", {}).get( + "namespace", self.kube_namespace + ) def open(self): self._create_kubernetes_pod() def _create_kubernetes_pod(self): + self.client.create_namespaced_pod( + namespace=self.kube_namespace, body=self.pod_manifest + ) + + while True: + pod = self.client.read_namespaced_pod( + name=self.pod_name, namespace=self.kube_namespace + ) + if pod.status.phase == "Running": + break + time.sleep(1) + + self.container = self.pod_name + + def _default_pod_manifest(self) -> dict: pod_manifest = { "apiVersion": "v1", "kind": "Pod", @@ -83,22 +109,10 @@ def _create_kubernetes_pod(self): } # Add environment variables if provided if self.env_vars: - pod_manifest["spec"]["containers"][0]["env"] = [ + pod_manifest["spec"]["containers"][0]["env"] = [ # type: ignore[index] {"name": key, "value": value} for key, value in self.env_vars.items() ] - self.client.create_namespaced_pod( - namespace=self.kube_namespace, body=pod_manifest - ) - - while True: - pod = self.client.read_namespaced_pod( - name=self.pod_name, namespace=self.kube_namespace - ) - if pod.status.phase == "Running": - break - time.sleep(1) - - self.container = self.pod_name + return pod_manifest def close(self): self._delete_kubernetes_pod() diff --git a/tests/test_session_kubernetes.py b/tests/test_session_kubernetes.py new file mode 100644 index 0000000..40720e8 --- /dev/null +++ b/tests/test_session_kubernetes.py @@ -0,0 +1,66 @@ +import unittest +from unittest.mock import patch, MagicMock +from llm_sandbox.kubernetes import SandboxKubernetesSession + + +class TestSandboxKubernetesSession(unittest.TestCase): + @patch("kubernetes.config.load_kube_config") + def setUp(self, mock_kube_config): + self.mock_kube_config = MagicMock() + mock_kube_config.return_value = self.mock_kube_config + + self.image = "python:3.9.19-bullseye" + self.dockerfile = None + self.lang = "python" + self.keep_template = False + self.verbose = False + + self.session = SandboxKubernetesSession( + image=self.image, + dockerfile=self.dockerfile, + lang=self.lang, + keep_template=self.keep_template, + verbose=self.verbose, + ) + + def test_with_pod_manifest(self): + pod_manifest = { + "apiVersion": "v1", + "kind": "Pod", + "metadata": { + "name": "test", + "namespace": "test", + "labels": {"app": "sandbox"}, + }, + "spec": { + "containers": [ + { + "name": "sandbox-container", + "image": "test", + "tty": True, + "volumeMounts": { + "name": "tmp", + "mountPath": "/tmp", + }, + } + ], + "volumes": [{"name": "tmp", "emptyDir": {"sizeLimit": "5Gi"}}], + }, + } + self.session = SandboxKubernetesSession( + image=self.image, + dockerfile=self.dockerfile, + lang=self.lang, + keep_template=self.keep_template, + verbose=self.verbose, + pod_manifest=pod_manifest, + ) + + self.session.client = MagicMock() + self.session.client.read_namespaced_pod.return_value.status.phase = "Running" + self.session.open() + + self.session.client.create_namespaced_pod.assert_called_with( + namespace="test", + body=pod_manifest, + ) From 779ac4fb3c358bc0fb46f3ce44158638b0625fb7 Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 27 Jan 2025 11:17:24 +0100 Subject: [PATCH 07/11] Move `_default_pod_manifest` higher --- llm_sandbox/kubernetes.py | 44 +++++++++++++++++++-------------------- 1 file changed, 22 insertions(+), 22 deletions(-) diff --git a/llm_sandbox/kubernetes.py b/llm_sandbox/kubernetes.py index 5c16fab..c499780 100644 --- a/llm_sandbox/kubernetes.py +++ b/llm_sandbox/kubernetes.py @@ -68,6 +68,28 @@ def __init__( self.pod_manifest = pod_manifest or self._default_pod_manifest() self._reconfigure_with_pod_manifest() + def _default_pod_manifest(self) -> dict: + pod_manifest = { + "apiVersion": "v1", + "kind": "Pod", + "metadata": { + "name": self.pod_name, + "namespace": self.kube_namespace, + "labels": {"app": "sandbox"}, + }, + "spec": { + "containers": [ + {"name": "sandbox-container", "image": self.image, "tty": True} + ] + }, + } + # Add environment variables if provided + if self.env_vars: + pod_manifest["spec"]["containers"][0]["env"] = [ # type: ignore[index] + {"name": key, "value": value} for key, value in self.env_vars.items() + ] + return pod_manifest + def _reconfigure_with_pod_manifest(self): self.pod_name = self.pod_manifest.get("metadata", {}).get("name", self.pod_name) self.kube_namespace = self.pod_manifest.get("metadata", {}).get( @@ -92,28 +114,6 @@ def _create_kubernetes_pod(self): self.container = self.pod_name - def _default_pod_manifest(self) -> dict: - pod_manifest = { - "apiVersion": "v1", - "kind": "Pod", - "metadata": { - "name": self.pod_name, - "namespace": self.kube_namespace, - "labels": {"app": "sandbox"}, - }, - "spec": { - "containers": [ - {"name": "sandbox-container", "image": self.image, "tty": True} - ] - }, - } - # Add environment variables if provided - if self.env_vars: - pod_manifest["spec"]["containers"][0]["env"] = [ # type: ignore[index] - {"name": key, "value": value} for key, value in self.env_vars.items() - ] - return pod_manifest - def close(self): self._delete_kubernetes_pod() From c6ee7c18dbbff6c2f91e6b8dda6ae9af4e7c5225 Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Mon, 27 Jan 2025 12:30:17 +0100 Subject: [PATCH 08/11] Fix tests --- tests/test_session_kubernetes.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/test_session_kubernetes.py b/tests/test_session_kubernetes.py index 40720e8..36d429c 100644 --- a/tests/test_session_kubernetes.py +++ b/tests/test_session_kubernetes.py @@ -6,9 +6,6 @@ class TestSandboxKubernetesSession(unittest.TestCase): @patch("kubernetes.config.load_kube_config") def setUp(self, mock_kube_config): - self.mock_kube_config = MagicMock() - mock_kube_config.return_value = self.mock_kube_config - self.image = "python:3.9.19-bullseye" self.dockerfile = None self.lang = "python" @@ -23,7 +20,8 @@ def setUp(self, mock_kube_config): verbose=self.verbose, ) - def test_with_pod_manifest(self): + @patch("kubernetes.config.load_kube_config") + def test_with_pod_manifest(self, mock_kube_config): pod_manifest = { "apiVersion": "v1", "kind": "Pod", From f864e168ee54d305ef272f9dd16ec400aaa936d2 Mon Sep 17 00:00:00 2001 From: Tomasz Wrona Date: Tue, 28 Jan 2025 13:04:45 +0100 Subject: [PATCH 09/11] Store code to execute locally in an unique temporary directory --- llm_sandbox/kubernetes.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/llm_sandbox/kubernetes.py b/llm_sandbox/kubernetes.py index 29eef68..809434a 100644 --- a/llm_sandbox/kubernetes.py +++ b/llm_sandbox/kubernetes.py @@ -1,8 +1,10 @@ import io import os +import tempfile import time import uuid import tarfile +from pathlib import Path from typing import List, Optional from kubernetes import client as k8s_client, config @@ -142,16 +144,18 @@ def run(self, code: str, libraries: Optional[List] = None) -> ConsoleOutput: f"Failed to install library {library}: {output}" ) - code_file = f"/tmp/code.{get_code_file_extension(self.lang)}" + code_file_name = f"code.{get_code_file_extension(self.lang)}" if self.lang == SupportedLanguage.GO: code_dest_file = "/example/code.go" else: - code_dest_file = code_file + code_dest_file = f"/tmp/{code_file_name}" - with open(code_file, "w") as f: - f.write(code) + with tempfile.TemporaryDirectory() as tmp_dir: + code_file = Path(tmp_dir) / code_file_name + with open(code_file, "w") as f: + f.write(code) + self.copy_to_runtime(str(code_file), code_dest_file) - self.copy_to_runtime(code_file, code_dest_file) commands = get_code_execution_command(self.lang, code_dest_file) output = KubernetesConsoleOutput(0, "") From 66eb1fbbf9ea3cf0f64a09451dc27b30928c9609 Mon Sep 17 00:00:00 2001 From: Ricardo Madriz Date: Fri, 31 Jan 2025 11:34:39 -0600 Subject: [PATCH 10/11] Lazily load docker and kubernetes modules --- llm_sandbox/session.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/llm_sandbox/session.py b/llm_sandbox/session.py index 42b20f2..34e1b6f 100644 --- a/llm_sandbox/session.py +++ b/llm_sandbox/session.py @@ -1,15 +1,17 @@ -import docker -from typing import Optional, Union -from kubernetes import client as k8s_client +from __future__ import annotations + +from typing import Optional, Union, TYPE_CHECKING, cast from llm_sandbox.const import SupportedLanguage -from llm_sandbox.docker import SandboxDockerSession -from llm_sandbox.kubernetes import SandboxKubernetesSession + +if TYPE_CHECKING: + import docker + from kubernetes import client as k8s_client class SandboxSession: def __new__( cls, - client: Union[docker.DockerClient, k8s_client.CoreV1Api] = None, + client: Union[docker.DockerClient, k8s_client.CoreV1Api, None] = None, image: Optional[str] = None, dockerfile: Optional[str] = None, lang: str = SupportedLanguage.PYTHON, @@ -34,8 +36,11 @@ def __new__( :param container_configs: Additional configurations for the Docker container, i.e. resources limits (cpu_count, mem_limit), etc. """ if use_kubernetes: + from kubernetes import client as k8s_client + from llm_sandbox.kubernetes import SandboxKubernetesSession + return SandboxKubernetesSession( - client=client, + client=cast(k8s_client.CoreV1Api, client), image=image, dockerfile=dockerfile, lang=lang, @@ -44,8 +49,11 @@ def __new__( kube_namespace=kube_namespace, ) + import docker + from llm_sandbox.docker import SandboxDockerSession + return SandboxDockerSession( - client=client, + client=cast(docker.DockerClient, client), image=image, dockerfile=dockerfile, lang=lang, From 6f61f1b05404d4a9f334e69c6b1de481c1e7080d Mon Sep 17 00:00:00 2001 From: Duy Huynh Date: Sun, 2 Feb 2025 13:59:38 +0700 Subject: [PATCH 11/11] Revert "Lazily load docker and kubernetes modules" --- llm_sandbox/session.py | 24 ++++++++---------------- 1 file changed, 8 insertions(+), 16 deletions(-) diff --git a/llm_sandbox/session.py b/llm_sandbox/session.py index 34e1b6f..42b20f2 100644 --- a/llm_sandbox/session.py +++ b/llm_sandbox/session.py @@ -1,17 +1,15 @@ -from __future__ import annotations - -from typing import Optional, Union, TYPE_CHECKING, cast +import docker +from typing import Optional, Union +from kubernetes import client as k8s_client from llm_sandbox.const import SupportedLanguage - -if TYPE_CHECKING: - import docker - from kubernetes import client as k8s_client +from llm_sandbox.docker import SandboxDockerSession +from llm_sandbox.kubernetes import SandboxKubernetesSession class SandboxSession: def __new__( cls, - client: Union[docker.DockerClient, k8s_client.CoreV1Api, None] = None, + client: Union[docker.DockerClient, k8s_client.CoreV1Api] = None, image: Optional[str] = None, dockerfile: Optional[str] = None, lang: str = SupportedLanguage.PYTHON, @@ -36,11 +34,8 @@ def __new__( :param container_configs: Additional configurations for the Docker container, i.e. resources limits (cpu_count, mem_limit), etc. """ if use_kubernetes: - from kubernetes import client as k8s_client - from llm_sandbox.kubernetes import SandboxKubernetesSession - return SandboxKubernetesSession( - client=cast(k8s_client.CoreV1Api, client), + client=client, image=image, dockerfile=dockerfile, lang=lang, @@ -49,11 +44,8 @@ def __new__( kube_namespace=kube_namespace, ) - import docker - from llm_sandbox.docker import SandboxDockerSession - return SandboxDockerSession( - client=cast(docker.DockerClient, client), + client=client, image=image, dockerfile=dockerfile, lang=lang,