Merge pull request #19 from richin13/main

Lazily load docker and kubernetes modules

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,