From 41b7141ecdd94ec53344677b5c5e98976df82885 Mon Sep 17 00:00:00 2001
From: wben <walterbenitez7@gmail.com>
Date: Sun, 1 Dec 2024 21:03:33 -0300
Subject: [PATCH 1/4] First support for podman

---
 llm_sandbox/podman.py        | 308 +++++++++++++++++++++++++++++++++++
 tests/test_session_podman.py | 155 ++++++++++++++++++
 2 files changed, 463 insertions(+)
 create mode 100644 llm_sandbox/podman.py
 create mode 100644 tests/test_session_podman.py

diff --git a/llm_sandbox/podman.py b/llm_sandbox/podman.py
new file mode 100644
index 0000000..c613fca
--- /dev/null
+++ b/llm_sandbox/podman.py
@@ -0,0 +1,308 @@
+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, ImageNotFound
+from podman.domain.containers import Container
+from podman.domain.images import Image
+from llm_sandbox.utils import (
+    image_exists,
+    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/tests/test_session_podman.py b/tests/test_session_podman.py
new file mode 100644
index 0000000..7860cf1
--- /dev/null
+++ b/tests/test_session_podman.py
@@ -0,0 +1,155 @@
+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()

From 786a815e2d423c86ab531dfb3815f34de561c9b7 Mon Sep 17 00:00:00 2001
From: wben <walterbenitez7@gmail.com>
Date: Mon, 2 Dec 2024 03:42:12 -0300
Subject: [PATCH 2/4] Fixed lint problems

---
 llm_sandbox/podman.py        | 9 ++++-----
 tests/test_session_podman.py | 9 ++++-----
 2 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/llm_sandbox/podman.py b/llm_sandbox/podman.py
index c613fca..12e543e 100644
--- a/llm_sandbox/podman.py
+++ b/llm_sandbox/podman.py
@@ -6,11 +6,10 @@
 from collections.abc import Iterator
 
 from podman import PodmanClient
-from podman.errors import NotFound, ImageNotFound
+from podman.errors import NotFound
 from podman.domain.containers import Container
 from podman.domain.images import Image
 from llm_sandbox.utils import (
-    image_exists,
     get_libraries_installation_command,
     get_code_file_extension,
     get_code_execution_command,
@@ -104,7 +103,9 @@ def open(self):
                     print(f"Using image {self.image.tags[-1]}")
             except NotFound:
                 if self.verbose:
-                    print(f"Image {self.image} not found locally. Attempting to pull...")
+                    print(
+                        f"Image {self.image} not found locally. Attempting to pull..."
+                    )
 
                 try:
                     # Attempt to pull the image
@@ -127,7 +128,6 @@ def open(self):
         )
         self.container.start()
 
-
     def close(self):
         if self.container:
             if self.commit_container and isinstance(self.image, Image):
@@ -304,5 +304,4 @@ def execute_command(
         if self.verbose:
             print(output)
 
-
         return ConsoleOutput(output)
diff --git a/tests/test_session_podman.py b/tests/test_session_podman.py
index 7860cf1..c9cf651 100644
--- a/tests/test_session_podman.py
+++ b/tests/test_session_podman.py
@@ -34,7 +34,7 @@ def test_init_with_invalid_lang(self):
     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"])
@@ -49,7 +49,9 @@ def test_open_with_image(self):
         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")
+        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(
@@ -64,9 +66,6 @@ def test_open_with_image(self):
         # 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

From ffd84f4ec5cea4874288aa0309444f834e45546d Mon Sep 17 00:00:00 2001
From: wben <walterbenitez7@gmail.com>
Date: Mon, 2 Dec 2024 03:50:30 -0300
Subject: [PATCH 3/4] Added podman as dependency

---
 pyproject.toml | 2 ++
 1 file changed, 2 insertions(+)

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]

From a6ea98472edf607cf7eb164e856e7cd05b2282c3 Mon Sep 17 00:00:00 2001
From: wben <walterbenitez7@gmail.com>
Date: Mon, 2 Dec 2024 03:54:16 -0300
Subject: [PATCH 4/4] Added lock file

---
 poetry.lock | 24 ++++++++++++++++++++++--
 1 file changed, 22 insertions(+), 2 deletions(-)

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"