diff --git a/README.md b/README.md index 878175d..935603f 100644 --- a/README.md +++ b/README.md @@ -152,7 +152,8 @@ Here is a list of things you can do to contribute: - [x] Add support for Go. - [ ] Add support for Ruby. - [x] Add remote Docker host support. -- [ ] Add remote Kubernetes cluster support. +- [x] Add remote Kubernetes cluster support. +- [ ] Commit the last container state to the image before closing kubernetes session. - [ ] Release version 1.0.0. ### License diff --git a/examples/code_runner_docker.py b/examples/code_runner_docker.py index edeb0aa..96b2f9e 100644 --- a/examples/code_runner_docker.py +++ b/examples/code_runner_docker.py @@ -4,16 +4,16 @@ def run_python_code(): with SandboxSession(lang="python", keep_template=True, verbose=True) as session: output = session.run("print('Hello, World!')") - print(output) + print(output.text) output = session.run( "import numpy as np\nprint(np.random.rand())", libraries=["numpy"] ) - print(output) + print(output.text) session.execute_command("pip install pandas") output = session.run("import pandas as pd\nprint(pd.__version__)") - print(output) + print(output.text) session.copy_to_runtime("README.md", "/sandbox/data.csv") @@ -29,13 +29,13 @@ def run_java_code(): } """, ) - print(output) + print(output.text) def run_javascript_code(): with SandboxSession(lang="javascript", keep_template=True, verbose=True) as session: output = session.run("console.log('Hello, World!')") - print(output) + print(output.text) output = session.run( """ @@ -45,7 +45,7 @@ def run_javascript_code(): """, libraries=["axios"], ) - print(output) + print(output.text) def run_cpp_code(): @@ -59,7 +59,7 @@ def run_cpp_code(): } """, ) - print(output) + print(output.text) output = session.run( """ @@ -75,7 +75,7 @@ def run_cpp_code(): } """, ) - print(output) + print(output.text) # run with libraries output = session.run( @@ -95,7 +95,7 @@ def run_cpp_code(): """, libraries=["libstdc++"], ) - print(output) + print(output.text) def run_go_code(): @@ -109,7 +109,7 @@ def run_go_code(): } """, ) - print(output) + print(output.text) # run with libraries output = session.run( @@ -136,7 +136,7 @@ def run_go_code(): """, libraries=["github.com/spyzhov/ajson"], ) - print(output) + print(output.text) if __name__ == "__main__": diff --git a/examples/code_runner_k8s.py b/examples/code_runner_k8s.py index 06743c0..411e1c1 100644 --- a/examples/code_runner_k8s.py +++ b/examples/code_runner_k8s.py @@ -2,21 +2,157 @@ def run_python_code(): - with SandboxSession(lang="python", keep_template=True, verbose=True, use_kubernetes=True) as session: + with SandboxSession( + lang="python", keep_template=True, verbose=True, use_kubernetes=True + ) as session: output = session.run("print('Hello, World!')") - print(output) + print(output.text) output = session.run( "import numpy as np\nprint(np.random.rand())", libraries=["numpy"] ) - print(output) + print(output.text) session.execute_command("pip install pandas") output = session.run("import pandas as pd\nprint(pd.__version__)") - print(output) + print(output.text) session.copy_to_runtime("README.md", "/sandbox/data.csv") +def run_java_code(): + with SandboxSession( + lang="java", keep_template=True, verbose=True, use_kubernetes=True + ) as session: + output = session.run( + """ + public class Main { + public static void main(String[] args) { + System.out.println("Hello, World!"); + } + } + """, + ) + print(output.text) + + +def run_javascript_code(): + with SandboxSession( + lang="javascript", keep_template=True, verbose=True, use_kubernetes=True + ) as session: + output = session.run("console.log('Hello, World!')") + print(output.text) + + # TODO: Fix this + # output = session.run( + # """ + # const axios = require('axios'); + # axios.get('https://jsonplaceholder.typicode.com/posts/1') + # .then(response => console.log(response.data)); + # """, + # libraries=["axios"], + # ) + # print(output.text) + + +def run_cpp_code(): + with SandboxSession( + lang="cpp", keep_template=True, verbose=True, use_kubernetes=True + ) as session: + output = session.run( + """ + #include + int main() { + std::cout << "Hello, World!" << std::endl; + return 0; + } + """, + ) + print(output.text) + + output = session.run( + """ + #include + #include + int main() { + std::vector v = {1, 2, 3, 4, 5}; + for (int i : v) { + std::cout << i << " "; + } + std::cout << std::endl; + return 0; + } + """, + ) + print(output.text) + + # run with libraries + output = session.run( + """ + #include + #include + #include + int main() { + std::vector v = {1, 2, 3, 4, 5}; + std::reverse(v.begin(), v.end()); + for (int i : v) { + std::cout << i << " "; + } + std::cout << std::endl; + return 0; + } + """, + libraries=["libstdc++"], + ) + print(output.text) + + +def run_go_code(): + with SandboxSession( + lang="go", keep_template=True, verbose=True, use_kubernetes=True + ) as session: + output = session.run( + """ + package main + import "fmt" + func main() { + fmt.Println("Hello, World!") + } + """, + ) + print(output.text) + + # run with libraries + output = session.run( + """ + package main + import ( + "fmt" + "github.com/spyzhov/ajson" + ) + func main() { + fmt.Println("Hello, World!") + json := []byte(`{"price": 100}`) + + root, _ := ajson.Unmarshal(json) + nodes, _ := root.JSONPath("$..price") + for _, node := range nodes { + node.SetNumeric(node.MustNumeric() * 1.25) + node.Parent().AppendObject("currency", ajson.StringNode("", "EUR")) + } + result, _ := ajson.Marshal(root) + + fmt.Printf("%s", result) + } + """, + libraries=["github.com/spyzhov/ajson"], + ) + print(output.text) + + if __name__ == "__main__": - run_python_code() \ No newline at end of file + # run_python_code() + # run_java_code() + run_javascript_code() + # run_cpp_code() + # run_go_code() diff --git a/llm_sandbox/base.py b/llm_sandbox/base.py index 55ce9c0..46974dd 100644 --- a/llm_sandbox/base.py +++ b/llm_sandbox/base.py @@ -2,6 +2,27 @@ from typing import Optional, List +class ConsoleOutput: + def __init__(self, text: str): + self._text = text + + @property + def text(self): + return self._text + + def __str__(self): + return f"ConsoleOutput(text={self.text})" + + +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})" + + class Session(ABC): def __init__(self, lang: str, verbose: bool = True, *args, **kwargs): self.lang = lang @@ -17,7 +38,7 @@ def close(self): raise NotImplementedError @abstractmethod - def run(self, code: str, libraries: Optional[List] = None): + def run(self, code: str, libraries: Optional[List] = None) -> ConsoleOutput: raise NotImplementedError @abstractmethod diff --git a/llm_sandbox/docker.py b/llm_sandbox/docker.py index f320a21..9586638 100644 --- a/llm_sandbox/docker.py +++ b/llm_sandbox/docker.py @@ -12,7 +12,7 @@ get_code_file_extension, get_code_execution_command, ) -from llm_sandbox.base import Session +from llm_sandbox.base import Session, ConsoleOutput from llm_sandbox.const import ( SupportedLanguage, SupportedLanguageValues, @@ -137,7 +137,7 @@ def close(self): f"Image {self.image.tags[-1]} is in use by other containers. Skipping removal.." ) - def run(self, code: str, libraries: Optional[List] = None): + def run(self, code: str, libraries: Optional[List] = None) -> ConsoleOutput: if not self.container: raise RuntimeError( "Session is not open. Please call open() method before running code." @@ -151,16 +151,16 @@ def run(self, code: str, libraries: Optional[List] = None): if self.lang == SupportedLanguage.GO: self.execute_command("mkdir -p /example") - self.execute_command("go mod init example", worKdir="/example") - self.execute_command("go mod tidy", worKdir="/example") + self.execute_command("go mod init example", workdir="/example") + self.execute_command("go mod tidy", workdir="/example") - for lib in libraries: - command = get_libraries_installation_command(self.lang, lib) - self.execute_command(command, worKdir="/example") + for library in libraries: + command = get_libraries_installation_command(self.lang, library) + _ = self.execute_command(command, workdir="/example") else: - for lib in libraries: - command = get_libraries_installation_command(self.lang, lib) - self.execute_command(command) + for library in libraries: + command = get_libraries_installation_command(self.lang, library) + _ = self.execute_command(command) code_file = f"/tmp/code.{get_code_file_extension(self.lang)}" if self.lang == SupportedLanguage.GO: @@ -173,11 +173,11 @@ def run(self, code: str, libraries: Optional[List] = None): self.copy_to_runtime(code_file, code_dest_file) - output = "" + output = ConsoleOutput("") commands = get_code_execution_command(self.lang, code_dest_file) for command in commands: if self.lang == SupportedLanguage.GO: - output = self.execute_command(command, worKdir="/example") + output = self.execute_command(command, workdir="/example") else: output = self.execute_command(command) @@ -224,7 +224,9 @@ def copy_to_runtime(self, src: str, dest: str): tarstream.seek(0) self.container.put_archive(os.path.dirname(dest), tarstream) - def execute_command(self, command: Optional[str], worKdir: Optional[str] = None): + def execute_command( + self, command: Optional[str], workdir: Optional[str] = None + ) -> ConsoleOutput: if not command: raise ValueError("Command cannot be empty") @@ -236,12 +238,14 @@ def execute_command(self, command: Optional[str], worKdir: Optional[str] = None) if self.verbose: print(f"Executing command: {command}") - if worKdir: - _, exec_log = self.container.exec_run( - command, stream=True, tty=True, workdir=worKdir + if workdir: + exit_code, exec_log = self.container.exec_run( + command, stream=True, tty=True, workdir=workdir ) else: - _, exec_log = self.container.exec_run(command, stream=True, tty=True) + exit_code, exec_log = self.container.exec_run( + command, stream=True, tty=True + ) output = "" if self.verbose: @@ -253,4 +257,4 @@ def execute_command(self, command: Optional[str], worKdir: Optional[str] = None) if self.verbose: print(chunk_str, end="") - return output + return ConsoleOutput(output) diff --git a/llm_sandbox/kubernetes.py b/llm_sandbox/kubernetes.py index 789a5d4..5174586 100644 --- a/llm_sandbox/kubernetes.py +++ b/llm_sandbox/kubernetes.py @@ -1,9 +1,13 @@ +import io +import os import time +import uuid +import tarfile from typing import List, Optional from kubernetes import client as k8s_client, config from kubernetes.stream import stream -from llm_sandbox.base import Session +from llm_sandbox.base import Session, ConsoleOutput, KubernetesConsoleOutput from llm_sandbox.utils import ( get_libraries_installation_command, get_code_file_extension, @@ -49,7 +53,7 @@ def __init__( self.image = image self.kube_namespace = kube_namespace - self.pod_name = f"sandbox-{lang.lower()}" + self.pod_name = f"sandbox-{lang.lower()}-{uuid.uuid4().hex}" self.keep_template = keep_template self.container = None @@ -89,36 +93,67 @@ def close(self): self._delete_kubernetes_pod() def _delete_kubernetes_pod(self): - if not self.keep_template: - self.client.delete_namespaced_pod( - name=self.pod_name, - namespace=self.kube_namespace, - body=k8s_client.V1DeleteOptions(), - ) + self.client.delete_namespaced_pod( + name=self.pod_name, + namespace=self.kube_namespace, + body=k8s_client.V1DeleteOptions(), + ) - def run(self, code: str, libraries: Optional[List] = None): + def run(self, code: str, libraries: Optional[List] = None) -> ConsoleOutput: if not self.container: raise RuntimeError( "Session is not open. Please call open() method before running code." ) if libraries: - command = get_libraries_installation_command(self.lang, libraries) - self.execute_command(command) + if self.lang == SupportedLanguage.GO: + self.execute_command("mkdir -p /example") + self.execute_command("go mod init example", workdir="/example") + self.execute_command("go mod tidy", workdir="/example") + + for library in libraries: + install_command = get_libraries_installation_command( + self.lang, library + ) + output = self.execute_command(install_command, workdir="/example") + if output.exit_code != 0: + raise RuntimeError( + f"Failed to install library {library}: {output}" + ) + else: + for library in libraries: + install_command = get_libraries_installation_command( + self.lang, library + ) + output = self.execute_command(install_command) + if output.exit_code != 0: + raise RuntimeError( + f"Failed to install library {library}: {output}" + ) code_file = f"/tmp/code.{get_code_file_extension(self.lang)}" + if self.lang == SupportedLanguage.GO: + code_dest_file = "/example/code.go" + else: + code_dest_file = code_file + with open(code_file, "w") as f: f.write(code) - self.copy_to_runtime(code_file, code_file) - commands = get_code_execution_command(self.lang, code_file) + self.copy_to_runtime(code_file, code_dest_file) + commands = get_code_execution_command(self.lang, code_dest_file) - output = "" + output = KubernetesConsoleOutput(0, "") for command in commands: - exit_code, output = self.execute_command(command) - if exit_code != 0: + if self.lang == SupportedLanguage.GO: + output = self.execute_command(command, workdir="/example") + else: + output = self.execute_command(command) + + if output.exit_code != 0: break - return exit_code, output + + return ConsoleOutput(output.text) def copy_to_runtime(self, src: str, dest: str): if not self.container: @@ -126,11 +161,25 @@ def copy_to_runtime(self, src: str, dest: str): "Session is not open. Please call open() method before copying files." ) + start_time = time.time() if self.verbose: print(f"Copying {src} to {self.container}:{dest}..") + dest_dir = os.path.dirname(dest) + dest_file = os.path.basename(dest) + + if dest_dir: + self.execute_command(f"mkdir -p {dest_dir}") + with open(src, "rb") as f: - exec_command = ["tar", "xvf", "-", "-C", dest] + tarstream = io.BytesIO() + with tarfile.open(fileobj=tarstream, mode="w") as tar: + tarinfo = tarfile.TarInfo(name=dest_file) + tarinfo.size = os.path.getsize(src) + tar.addfile(tarinfo, f) + tarstream.seek(0) + + exec_command = ["tar", "xvf", "-", "-C", dest_dir] resp = stream( self.client.connect_get_namespaced_pod_exec, self.container, @@ -148,9 +197,15 @@ def copy_to_runtime(self, src: str, dest: str): print(resp.read_stdout()) if resp.peek_stderr(): print(resp.read_stderr()) - resp.write_stdin(f.read()) + resp.write_stdin(tarstream.read(4096)) resp.close() + end_time = time.time() + if self.verbose: + print( + f"Copied {src} to {self.container}:{dest} in {end_time - start_time:.2f} seconds" + ) + def copy_from_runtime(self, src: str, dest: str): if not self.container: raise RuntimeError( @@ -180,7 +235,9 @@ def copy_from_runtime(self, src: str, dest: str): if resp.peek_stderr(): print(resp.read_stderr()) - def execute_command(self, command: str): + def execute_command( + self, command: str, workdir: Optional[str] = None + ) -> KubernetesConsoleOutput: if not self.container: raise RuntimeError( "Session is not open. Please call open() method before executing commands." @@ -189,7 +246,11 @@ def execute_command(self, command: str): if self.verbose: print(f"Executing command: {command}") - exec_command = ["/bin/sh", "-c", command] + if workdir: + exec_command = ["sh", "-c", f"cd {workdir} && {command}"] + else: + exec_command = ["/bin/sh", "-c", command] + resp = stream( self.client.connect_get_namespaced_pod_exec, self.container, @@ -199,7 +260,25 @@ def execute_command(self, command: str): stdin=False, stdout=True, tty=False, + _preload_content=False, ) - output = resp.read_stdout() - exit_code = 0 if resp.returncode is None else resp.returncode - return exit_code, output + + output = "" + if self.verbose: + print("Output:", end=" ") + + while resp.is_open(): + resp.update(timeout=1) + if resp.peek_stdout(): + chunk = resp.read_stdout() + output += chunk + if self.verbose: + print(chunk, end="") + if resp.peek_stderr(): + chunk = resp.read_stderr() + output += chunk + if self.verbose: + print(chunk, end="") + + exit_code = resp.returncode + return KubernetesConsoleOutput(exit_code, output)