diff --git a/libs/langchain/langchain/utilities/docker_containers.py b/libs/langchain/langchain/utilities/docker_containers.py index 2efca4ccfd361..c86409d4167af 100644 --- a/libs/langchain/langchain/utilities/docker_containers.py +++ b/libs/langchain/langchain/utilities/docker_containers.py @@ -52,11 +52,16 @@ class DockerImage: def __init__(self, name: str): """Note that it does not pull the image from the internet. It only represents a tag so it must exist on your system. + It throws ValueError if docker image by that name does not exist locally. """ - self.name = name - # check if image exists - docker_client = get_docker_client() - if len(docker_client.images.list(name=name)) < 1: + splitted_name = name.split(":") + if len(splitted_name) == 1: + # by default, image has latest tag. + self.name = name + ":latest" + else: + self.name = name + + if not self.exists(name): raise ValueError( f"Invalid value: name={name} does not exist on your system." "Use DockerImage.from_tag() to pull it." @@ -65,6 +70,21 @@ def __init__(self, name: str): def __repr__(self) -> str: return f"DockerImage(name={self.name})" + @classmethod + def exists(cls, name: str) -> bool: + """Checks if the docker image exists""" + docker_client = get_docker_client() + return len(docker_client.images.list(name=name)) > 0 + + @classmethod + def remove(cls, name: str) -> None: + """WARNING: Removes image from the system, be cautious with this function. + It is irreversible operation!. + """ + if cls.exists(name): + docker_client = get_docker_client() + docker_client.images.remove(name) + @classmethod def from_tag( cls, @@ -78,10 +98,13 @@ def from_tag( Example: repository = "python" tag = "3.9-slim" """ docker_client = get_docker_client() + name = f"{repository}:{tag}" + if len(docker_client.images.list(name=name)) > 0: + return cls(name=name) docker_client.images.pull( repository=repository, tag=tag, auth_config=auth_config ) - return cls(name=f"{repository}:{tag}") + return cls(name=name) @classmethod def from_dockerfile( @@ -146,6 +169,13 @@ def __enter__(self) -> "DockerContainer": """Enters container context. It means that container is started and you can execute commands inside it. """ + self.unsafe_start() + return self + + def unsafe_start(self) -> None: + """Starts container without entering it. + Please prefer to use with DockerContainer statement. + """ assert self._container is None, "You cannot re-entry container" # tty=True is required to keep container alive self._container = self._client.containers.run( @@ -154,7 +184,6 @@ def __enter__(self) -> "DockerContainer": tty=True, **self._run_kwargs, ) - return self def __exit__( self, @@ -163,6 +192,7 @@ def __exit__( traceback: Optional[TracebackType], ) -> bool: """Cleanup container on exit.""" + assert self._container is not None, "You cannot exit unstarted container." if exc_type is not None: # re-throw exception. try to stop container and remove it try: @@ -171,10 +201,16 @@ def __exit__( print("Failed to stop and remove container to cleanup exception.", e) return False else: - self._cleanup() - self._container = None + self.unsafe_exit() return True + def unsafe_exit(self): + """Cleanup container on exit. Please prefer to use `with` statement.""" + if self._container is None: + return + self._cleanup() + self._container = None + def spawn_run( self, command: Union[str, List[str]], **kwargs: Any ) -> Tuple[int, bytes]: diff --git a/libs/langchain/tests/integration_tests/utilities/docker_test_data/Dockerfile b/libs/langchain/tests/integration_tests/utilities/docker_test_data/Dockerfile index fbc0d7f11be5f..7c19ef9a14310 100644 --- a/libs/langchain/tests/integration_tests/utilities/docker_test_data/Dockerfile +++ b/libs/langchain/tests/integration_tests/utilities/docker_test_data/Dockerfile @@ -1,7 +1,5 @@ # This is a test dockerfile that will be used to test the docker_containers. -FROM python:3.11-slim -RUN pip install cowsay -# This runs cowsay with moo so that we can test that the image works -# and additionally every other run command within container -# will execute but with moo prefix for the user input. -ENTRYPOINT ["python3", "-m", "cowsay", "moo"] \ No newline at end of file +FROM python:3.11-alpine +RUN pip install --no-cache-dir cowsay==6.0 +# This runs cowsay and it requires arguments like -t "hello world". +ENTRYPOINT ["python3", "-m", "cowsay"] \ No newline at end of file diff --git a/libs/langchain/tests/integration_tests/utilities/test_docker_containers.py b/libs/langchain/tests/integration_tests/utilities/test_docker_containers.py index 48e4c49ac3183..459cd9f2176d3 100644 --- a/libs/langchain/tests/integration_tests/utilities/test_docker_containers.py +++ b/libs/langchain/tests/integration_tests/utilities/test_docker_containers.py @@ -2,6 +2,7 @@ from typing import cast import pytest + from langchain.utilities.docker_containers import ( DockerContainer, DockerImage, @@ -31,14 +32,11 @@ def run_container_cowsay(image: DockerImage) -> None: # by ENTRYPOINT defined in dockerfile. try: container = DockerContainer(image) - ret_code, log = container.spawn_run("I like langchain!") + ret_code, log = container.spawn_run('-t "I like langchain!"') assert ret_code == 0 - assert ( - log.find(b"moo I like langchain") >= 0 - ), "Cowsay should say same words with moo" + assert log.find(b"I like langchain") >= 0, "Cowsay should say same words" finally: - docker_client = get_docker_client() - docker_client.images.remove(image.name) + DockerImage.remove(image) @pytest.mark.requires("docker")