From 43df0e9fc21e2e95821459c62815ff1bfe95081c Mon Sep 17 00:00:00 2001 From: Duy Huynh Date: Sat, 6 Jul 2024 20:27:42 +0700 Subject: [PATCH] feat: add pre-commit as pylint --- .github/workflows/pylint.yaml | 11 +++++ .pre-commit-config.yaml | 26 +++++++++++ Makefile | 25 +++++++++++ llm_sandbox/session.py | 84 +++++++++++++++++++++++++++++++---- llm_sandbox/utils.py | 2 - tests/busybox.Dockerfile | 2 +- 6 files changed, 138 insertions(+), 12 deletions(-) create mode 100644 .github/workflows/pylint.yaml create mode 100644 .pre-commit-config.yaml create mode 100644 Makefile diff --git a/.github/workflows/pylint.yaml b/.github/workflows/pylint.yaml new file mode 100644 index 0000000..6b2f3f8 --- /dev/null +++ b/.github/workflows/pylint.yaml @@ -0,0 +1,11 @@ +name: pylint + +on: [pull_request, push] + +jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: actions/setup-python@v3 + - uses: pre-commit/action@v3.0.0 \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..7271c7d --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,26 @@ +repos: +- repo: https://github.com/pre-commit/pre-commit-hooks + rev: v2.3.0 + hooks: + - id: end-of-file-fixer + - id: trailing-whitespace + - id: check-added-large-files +- repo: https://github.com/psf/black + rev: 23.10.0 + hooks: + - id: black +- repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.7.1 + hooks: + - id: mypy + additional_dependencies: [pydantic] + exclude: tests +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.1.7 + hooks: + - id: ruff + args: [--fix] +- repo: https://github.com/gitleaks/gitleaks + rev: v8.17.0 + hooks: + - id: gitleaks diff --git a/Makefile b/Makefile new file mode 100644 index 0000000..f740f90 --- /dev/null +++ b/Makefile @@ -0,0 +1,25 @@ +PHONY: hello env lint + +RED=\033[0;31m +GREEN=\033[0;32m +YELLOW=\033[0;33m +BLUE=\033[0;34m +MAGENTA=\033[0;35m +CYAN=\033[0;36m +RESET=\033[0m + +hello: + @echo "${MAGENTA}Hello, $$(whoami)!${RESET}" + @echo "${GREEN}Current Time:${RESET}\t\t${YELLOW}$$(date)${RESET}" + @echo "${GREEN}Working Directory:${RESET}\t${YELLOW}$$(pwd)${RESET}" + @echo "${GREEN}Shell:${RESET}\t\t\t${YELLOW}$$(echo $$SHELL)${RESET}" + @echo "${GREEN}Terminal:${RESET}\t\t${YELLOW}$$(echo $$TERM)${RESET}" + +env: + @echo "To activate the Poetry environment, run:" + @echo "source $$(poetry env info --path)/bin/activate" + +lint: + @echo "Running linter..." + @source $$(poetry env info --path)/bin/activate && pre-commit run --all-files + @echo "Done." diff --git a/llm_sandbox/session.py b/llm_sandbox/session.py index 43fe2a9..c8c258b 100644 --- a/llm_sandbox/session.py +++ b/llm_sandbox/session.py @@ -1,25 +1,86 @@ -from typing import List +import os +import docker +from typing import List, Optional, Union -from .const import SupportedLanguage +from docker.models.images import Image +from llm_sandbox.utils import image_exists +from llm_sandbox.const import SupportedLanguage class SandboxSession: - def __init__(self, image: str = None, dockerfile: str = None, lang: str = SupportedLanguage.PYTHON): - # raise error if both the image and dockerfile exist since we only need one + def __init__( + self, + image: Optional[str] = None, + dockerfile: Optional[str] = None, + lang: str = SupportedLanguage.PYTHON, + keep_template: bool = False, + ): + """ + Create a new sandbox session + :param image: Docker 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 + """ if image and dockerfile: raise ValueError("Only one of image or dockerfile should be provided") - # raise error if neither the image nor dockerfile exist since we need one if not image and not dockerfile: raise ValueError("Either image or dockerfile should be provided") - self.lang = lang + self.lang: str = lang + self.client: docker.DockerClient = docker.from_env() + self.image: Union[Image, str] = image + self.dockerfile: Optional[str] = dockerfile + self.container = None + self.path = None + self.keep_template = keep_template + self.is_create_template: bool = False def open(self): - raise NotImplementedError + warning_str = ( + "Since the `keep_image` flag is set to True the image and container will not be removed after the session " + "ends and remains for future use." + ) + if self.dockerfile: + self.path = os.path.dirname(self.dockerfile) + self.image, _ = self.client.images.build( + path=self.path, + dockerfile=os.path.basename(self.dockerfile), + tag="sandbox", + ) + self.is_create_template = True + f_str = f"Built image {self.image.tags[-1]} from {self.dockerfile}" + f_str = f"{f_str}. {warning_str}" if self.keep_template else f_str + print(f_str) + + if isinstance(self.image, str): + if not image_exists(self.client, self.image): + self.image = self.client.images.pull(self.image) + self.is_create_template = True + f_str = f"Pulled image {self.image.tags[-1]}" + f_str = f"{f_str}. {warning_str}" if self.keep_template else f_str + print(f_str) + else: + self.image = self.client.images.get(self.image) + print(f"Using image {self.image.tags[-1]}") + + self.container = self.client.containers.create( + self.image, name="sandbox", detach=True + ) def close(self): - raise NotImplementedError + if self.container: + self.container.remove(force=True) + self.container = None + + if self.is_create_template and not self.keep_template: + 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") def run(self, code: str, libraries: List = []): raise NotImplementedError @@ -38,4 +99,9 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): - self.close() \ No newline at end of file + self.close() + + +if __name__ == "__main__": + with SandboxSession(dockerfile="tests/busybox.Dockerfile") as session: + session.run("print('Hello, World!')") diff --git a/llm_sandbox/utils.py b/llm_sandbox/utils.py index 88f00bd..ea6f50d 100644 --- a/llm_sandbox/utils.py +++ b/llm_sandbox/utils.py @@ -1,9 +1,7 @@ import docker import docker.errors -from typing import List from docker import DockerClient -from docker.models.images import Image def image_exists(client: DockerClient, image: str) -> bool: diff --git a/tests/busybox.Dockerfile b/tests/busybox.Dockerfile index f51439e..9a3adf6 100644 --- a/tests/busybox.Dockerfile +++ b/tests/busybox.Dockerfile @@ -1 +1 @@ -FROM busybox:latest \ No newline at end of file +FROM busybox:latest