diff --git a/llm_sandbox/podman.py b/llm_sandbox/podman.py new file mode 100644 index 0000000..12e543e --- /dev/null +++ b/llm_sandbox/podman.py @@ -0,0 +1,307 @@ +import io +import os +import tarfile +import tempfile +from typing import List, Optional, Union +from collections.abc import Iterator + +from podman import PodmanClient +from podman.errors import NotFound +from podman.domain.containers import Container +from podman.domain.images import Image +from llm_sandbox.utils import ( + get_libraries_installation_command, + get_code_file_extension, + get_code_execution_command, +) +from llm_sandbox.base import Session, ConsoleOutput +from llm_sandbox.const import ( + SupportedLanguage, + SupportedLanguageValues, + DefaultImage, + NotSupportedLibraryInstallation, +) + + +class SandboxPodmanSession(Session): + def __init__( + self, + client: Optional[PodmanClient] = None, + image: Optional[str] = None, + dockerfile: Optional[str] = None, + lang: str = SupportedLanguage.PYTHON, + keep_template: bool = False, + commit_container: bool = True, + verbose: bool = False, + mounts: Optional[list] = None, + container_configs: Optional[dict] = None, + ): + """ + Create a new sandbox session + :param client: Podman client, if not provided, a new client will be created based on local podman context + :param image: Podman image to use + :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 podman 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 container_configs: Additional configurations for the container, i.e. resources limits (cpu_count, mem_limit), etc. + """ + super().__init__(lang, verbose) + if image and dockerfile: + raise ValueError("Only one of image or dockerfile should be provided") + + if lang not in SupportedLanguageValues: + raise ValueError( + f"Language {lang} is not supported. Must be one of {SupportedLanguageValues}" + ) + + if not image and not dockerfile: + image = DefaultImage.__dict__[lang.upper()] + + self.lang: str = lang + self.client: Optional[PodmanClient] = client or PodmanClient() + self.image: Union[Image, str] = image + self.dockerfile: Optional[str] = dockerfile + 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 + self.container_configs = container_configs + + def open(self): + warning_str = ( + "Since the `keep_template` flag is set to True, the Podman image will not be removed after the session ends " + "and remains for future use." + ) + + # Build image if a Dockerfile is provided + if self.dockerfile: + self.path = os.path.dirname(self.dockerfile) + if self.verbose: + f_str = f"Building Podman image from {self.dockerfile}" + f_str = f"{f_str}\n{warning_str}" if self.keep_template else f_str + print(f_str) + + self.image, _ = self.client.images.build( + path=self.path, + dockerfile=os.path.basename(self.dockerfile), + tag=f"sandbox-{self.lang.lower()}-{os.path.basename(self.path)}", + ) + self.is_create_template = True + + # Check or pull the image + if isinstance(self.image, str): + try: + # Try to get the image locally + self.image = self.client.images.get(self.image) + if self.verbose: + print(f"Using image {self.image.tags[-1]}") + except NotFound: + if self.verbose: + print( + f"Image {self.image} not found locally. Attempting to pull..." + ) + + try: + # Attempt to pull the image + self.image = self.client.images.pull(self.image) + if self.verbose: + print(f"Successfully pulled image {self.image.tags[-1]}") + self.is_create_template = True + except Exception as e: + raise RuntimeError(f"Failed to pull image {self.image}: {e}") + + # Ensure mounts is an iterable (empty list if None) + mounts = self.mounts if self.mounts is not None else [] + + # Create the container + self.container = self.client.containers.create( + image=self.image.id if isinstance(self.image, Image) else self.image, + tty=True, + mounts=mounts, # Use the adjusted mounts + **self.container_configs if self.container_configs else {}, + ) + self.container.start() + + def close(self): + if self.container: + if self.commit_container and isinstance(self.image, Image): + # Extract repository and tag + if self.image.tags: + full_tag = self.image.tags[-1] + if ":" in full_tag: + repository, tag = full_tag.rsplit(":", 1) + else: + repository = full_tag + tag = "latest" + try: + # Commit the container with repository and tag + self.container.commit(repository=repository, tag=tag) + if self.verbose: + print(f"Committed container as image {repository}:{tag}") + except Exception as e: + if self.verbose: + print(f"Failed to commit container: {e}") + raise + + # Stop and remove the container + self.container.stop() + self.container.remove(force=True) + self.container = None + + if self.is_create_template and not self.keep_template: + # check if the image is used by any other container + containers = self.client.containers.list(all=True) + image_id = ( + self.image.id + if isinstance(self.image, Image) + else self.client.images.get(self.image).id + ) + image_in_use = any( + container.image.id == image_id for container in containers + ) + + if not image_in_use: + if isinstance(self.image, str): + self.client.images.remove(self.image) + elif isinstance(self.image, Image): + self.image.remove(force=True) + else: + raise ValueError("Invalid image type") + else: + if self.verbose: + print( + f"Image {self.image.tags[-1]} is in use by other containers. Skipping removal.." + ) + + 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: + if self.lang.upper() in NotSupportedLibraryInstallation: + raise ValueError( + f"Library installation has not been supported for {self.lang} yet!" + ) + 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: + command = get_libraries_installation_command(self.lang, library) + _ = self.execute_command(command, workdir="/example") + else: + for library in libraries: + command = get_libraries_installation_command(self.lang, library) + _ = self.execute_command(command) + with tempfile.TemporaryDirectory() as directory_name: + code_file = os.path.join( + directory_name, f"code.{get_code_file_extension(self.lang)}" + ) + if self.lang == SupportedLanguage.GO: + code_dest_file = "/example/code.go" + else: + code_dest_file = ( + f"/tmp/code.{get_code_file_extension(self.lang)}" # code_file + ) + + with open(code_file, "w") as f: + f.write(code) + + self.copy_to_runtime(code_file, code_dest_file) + + 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") + else: + output = self.execute_command(command) + + return output + + def copy_from_runtime(self, src: str, dest: str): + if not self.container: + raise RuntimeError( + "Session is not open. Please call open() method before copying files." + ) + + if self.verbose: + print(f"Copying {self.container.short_id}:{src} to {dest}..") + + bits, stat = self.container.get_archive(src) + if stat["size"] == 0: + raise FileNotFoundError(f"File {src} not found in the container") + + tarstream = io.BytesIO(b"".join(bits)) + with tarfile.open(fileobj=tarstream, mode="r") as tar: + tar.extractall(os.path.dirname(dest)) + + def copy_to_runtime(self, src: str, dest: str): + if not self.container: + raise RuntimeError( + "Session is not open. Please call open() method before copying files." + ) + + is_created_dir = False + directory = os.path.dirname(dest) + if directory and not self.container.exec_run(f"test -d {directory}")[0] == 0: + self.container.exec_run(f"mkdir -p {directory}") + is_created_dir = True + + if self.verbose: + if is_created_dir: + print(f"Creating directory {self.container.short_id}:{directory}") + print(f"Copying {src} to {self.container.short_id}:{dest}..") + + tarstream = io.BytesIO() + with tarfile.open(fileobj=tarstream, mode="w") as tar: + tar.add(src, arcname=os.path.basename(src)) + + tarstream.seek(0) + self.container.put_archive(os.path.dirname(dest), tarstream) + + def execute_command( + self, command: Optional[str], workdir: Optional[str] = None + ) -> ConsoleOutput: + if not command: + raise ValueError("Command cannot be empty") + + if not self.container: + raise RuntimeError( + "Session is not open. Please call open() method before executing commands." + ) + + if self.verbose: + print(f"Executing command: {command}") + + if workdir: + exit_code, exec_log = self.container.exec_run( + command, stream=True, tty=True, workdir=workdir + ) + else: + exit_code, exec_log = self.container.exec_run( + command, stream=True, tty=True + ) + if isinstance(exec_log, Iterator): + output = "" + for chunk in exec_log: + chunk_str = chunk.decode("utf-8") + output += chunk_str + if self.verbose: + print(chunk_str, end="") + else: + output = exec_log.decode("utf-8") + + if self.verbose: + print(output) + + return ConsoleOutput(output) diff --git a/poetry.lock b/poetry.lock index 48860fe..cea5a2a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand. [[package]] name = "aiohappyeyeballs" @@ -1981,6 +1981,25 @@ files = [ dev = ["pre-commit", "tox"] testing = ["pytest", "pytest-benchmark"] +[[package]] +name = "podman" +version = "5.2.0" +description = "Bindings for Podman RESTful API" +optional = false +python-versions = ">=3.6" +files = [ + {file = "podman-5.2.0-py3-none-any.whl", hash = "sha256:e42e907b44af3e8578ac3bd4838040f7dd6b4af091f787e51462b979746703eb"}, + {file = "podman-5.2.0.tar.gz", hash = "sha256:6a064a3e9dbfc4ee0ee0c51604164e1f58d9d00b10f702c3e570651aee722ec7"}, +] + +[package.dependencies] +requests = ">=2.24" +tomli = {version = ">=1.2.3", markers = "python_version < \"3.11\""} +urllib3 = "*" + +[package.extras] +progress-bar = ["rich (>=12.5.1)"] + [[package]] name = "pre-commit" version = "3.8.0" @@ -3115,8 +3134,9 @@ propcache = ">=0.2.0" docker = ["docker"] k8s = ["kubernetes"] minimal = [] +podman = ["podman"] [metadata] lock-version = "2.0" python-versions = "^3.10" -content-hash = "51fe26ceb1289d1fc7cc278dad26c5a25d61441f6ed610b191694658b8643783" +content-hash = "70c2477c2ff3d214cf5e2e78bca6bf375f4aab946bee2fdc02752be749f2b929" diff --git a/pyproject.toml b/pyproject.toml index 99202d7..210a422 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,10 +15,12 @@ packages = [ python = "^3.10" docker = "^7.1.0" kubernetes = "^30.1.0" +podman = "^5.2.0" [tool.poetry.extras] docker = ["docker"] k8s = ["kubernetes"] +podman = ["podman"] minimal = [] [tool.poetry.group.dev.dependencies] diff --git a/tests/test_session_podman.py b/tests/test_session_podman.py new file mode 100644 index 0000000..c9cf651 --- /dev/null +++ b/tests/test_session_podman.py @@ -0,0 +1,154 @@ +import os +import tarfile +import unittest +from io import BytesIO +from unittest.mock import patch, MagicMock +from llm_sandbox.podman import SandboxPodmanSession + + +class TestSandboxSession(unittest.TestCase): + @patch("podman.from_env") + def setUp(self, mock_podman_from_env): + self.mock_podman_client = MagicMock() + mock_podman_from_env.return_value = self.mock_podman_client + + self.image = "python:3.9.19-bullseye" + self.dockerfile = None + self.lang = "python" + self.keep_template = False + self.verbose = False + + self.session = SandboxPodmanSession( + client=self.mock_podman_client, + image=self.image, + dockerfile=self.dockerfile, + lang=self.lang, + keep_template=self.keep_template, + verbose=self.verbose, + ) + + def test_init_with_invalid_lang(self): + with self.assertRaises(ValueError): + SandboxPodmanSession(lang="invalid_language") + + def test_init_with_both_image_and_dockerfile(self): + with self.assertRaises(ValueError): + SandboxPodmanSession(image="some_image", dockerfile="some_dockerfile") + + def test_open_with_image(self): + # Mock the image retrieval + mock_image = MagicMock(tags=["python:3.9.19-bullseye"]) + self.mock_podman_client.images.get.return_value = mock_image + + # Mock the container creation + self.mock_podman_client.containers = MagicMock() # Ensure containers is mocked + mock_container = MagicMock() + self.mock_podman_client.containers.create.return_value = mock_container + + # Call the open method + self.session.open() + + # Assert that `images.get` was called to check if the image exists + self.mock_podman_client.images.get.assert_called_once_with( + "python:3.9.19-bullseye" + ) + + # Assert that `containers.create` was called with the correct parameters + self.mock_podman_client.containers.create.assert_called_once_with( + image=mock_image, # Use the mock returned by `images.get` + tty=True, + mounts=[], # Default value when `self.mounts` is None + ) + + # Assert that `start` was called on the container + mock_container.start.assert_called_once() + + # Assert that the `self.container` attribute was set + self.assertEqual(self.session.container, mock_container) + + def test_close(self): + mock_container = MagicMock() + self.session.container = mock_container + mock_container.commit.return_values = MagicMock(tags=["python:3.9.19-bullseye"]) + + self.session.close() + mock_container.remove.assert_called_once() + self.assertIsNone(self.session.container) + + def test_run_without_open(self): + with self.assertRaises(RuntimeError): + self.session.run("print('Hello')") + + def test_run_with_code(self): + self.session.container = MagicMock() + self.session.execute_command = MagicMock(return_value=(0, "Output")) + + result = self.session.run("print('Hello')") + self.session.execute_command.assert_called() + self.assertEqual(result, (0, "Output")) + + def test_copy_to_runtime(self): + self.session.container = MagicMock() + src = "test.txt" + dest = "/tmp/test.txt" + with open(src, "w") as f: + f.write("test content") + + self.session.copy_to_runtime(src, dest) + self.session.container.put_archive.assert_called() + + os.remove(src) + + @patch("tarfile.open") + def test_copy_from_runtime(self, mock_tarfile_open): + self.session.container = MagicMock() + src = "/tmp/test.txt" + dest = "test.txt" + + # Create a mock tarfile + tarstream = BytesIO() + with tarfile.open(fileobj=tarstream, mode="w") as tar: + tarinfo = tarfile.TarInfo(name=os.path.basename(dest)) + tarinfo.size = len(b"test content") + tar.addfile(tarinfo, BytesIO(b"test content")) + + tarstream.seek(0) + self.session.container.get_archive.return_value = ( + [tarstream.read()], + {"size": tarstream.__sizeof__()}, + ) + + def mock_extractall(path): + with open(dest, "wb") as f: + f.write(b"test content") + + mock_tarfile = MagicMock() + mock_tarfile.extractall.side_effect = mock_extractall + mock_tarfile_open.return_value.__enter__.return_value = mock_tarfile + + self.session.copy_from_runtime(src, dest) + self.assertTrue(os.path.exists(dest)) + + os.remove(dest) + + def test_execute_command(self): + mock_container = MagicMock() + self.session.container = mock_container + + command = "echo 'Hello'" + mock_container.exec_run.return_value = (0, iter([b"Hello\n"])) + + output = self.session.execute_command(command) + + # Assert that exec_run was called with the correct arguments + mock_container.exec_run.assert_called_with(command, stream=True, tty=True) + # Assert that the returned output is as expected + self.assertEqual(output.text, "Hello\n") # Match the `output` attribute + + def test_execute_empty_command(self): + with self.assertRaises(ValueError): + self.session.execute_command("") + + +if __name__ == "__main__": + unittest.main()