truefoundry 0.2.10__py3-none-any.whl → 0.3.0__py3-none-any.whl

truefoundry/deploy/builder/__init__.py

@@ -0,0 +1,134 @@
+from typing import Dict, List, Optional, Union
+
+from truefoundry.deploy.auto_gen.models import (
+    DockerFileBuild,
+    PythonBuild,
+    TaskDockerFileBuild,
+    TaskPythonBuild,
+)
+from truefoundry.deploy.builder.builders import get_builder
+from truefoundry.deploy.builder.builders.tfy_notebook_buildpack.dockerfile_template import (
+    NotebookImageBuild,
+)
+from truefoundry.pydantic_v1 import BaseModel, Field
+
+
+class _BuildConfig(BaseModel):
+    __root__: Union[
+        DockerFileBuild,
+        PythonBuild,
+        NotebookImageBuild,
+        TaskPythonBuild,
+        TaskDockerFileBuild,
+    ] = Field(discriminator="type")
+
+
+def build(
+    build_configuration: Union[BaseModel, Dict],
+    tag: str,
+    extra_opts: Optional[List[str]] = None,
+):
+    build_configuration = _BuildConfig.parse_obj(build_configuration).__root__
+    if isinstance(build_configuration, TaskPythonBuild):
+        build_configuration_dict = build_configuration.dict()
+        build_configuration_dict.update({"type": "tfy-python-buildpack", "command": ""})
+        build_configuration = PythonBuild.parse_obj(build_configuration_dict)
+    if isinstance(build_configuration, TaskDockerFileBuild):
+        build_configuration_dict = build_configuration.dict()
+        build_configuration_dict.update({"type": "dockerfile"})
+        build_configuration = DockerFileBuild.parse_obj(build_configuration_dict)
+
+    builder = get_builder(build_configuration.type)
+    return builder(
+        build_configuration=build_configuration,
+        tag=tag,
+        extra_opts=extra_opts,
+    )
+
+
+if __name__ == "__main__":
+    # TODO: remove these and write tests for Build class to dockerfile content.
+    import os
+    from tempfile import TemporaryDirectory
+
+    with TemporaryDirectory() as local_dir:
+        dockerfile_path = os.path.join(local_dir, "Dockerfile.test")
+        with open(dockerfile_path, "w", encoding="utf8") as fp:
+            fp.write("from postgres:latest")
+
+        build_config = DockerFileBuild(
+            type="dockerfile",
+            build_context_path=local_dir,
+            dockerfile_path=dockerfile_path,
+        )
+
+        build(tag="docker-test", build_configuration=build_config)
+
+    with TemporaryDirectory() as local_dir:
+        requirements_path = os.path.join(local_dir, "requirements.txt")
+        with open(requirements_path, "w", encoding="utf8") as fp:
+            fp.write("requests")
+
+        python_path = os.path.join(local_dir, "main.py")
+        with open(python_path, "w", encoding="utf8") as fp:
+            fp.write("import requests; print(requests.__version__);")
+
+        build_config = PythonBuild(
+            type="tfy-python-buildpack",
+            build_context_path=local_dir,
+            command=["python", "main.py"],
+            python_version="3.7.13",
+        )
+        build(tag="python-test-exec", build_configuration=build_config)
+        build_config = PythonBuild(
+            type="tfy-python-buildpack",
+            build_context_path=local_dir,
+            command="python main.py",
+            python_version="3.8",
+        )
+        build(tag="python-test-shell", build_configuration=build_config)
+
+    with TemporaryDirectory() as local_dir:
+        requirements_path = os.path.join(local_dir, "requirements.txt")
+        with open(requirements_path, "w", encoding="utf8") as fp:
+            fp.write("numpy")
+
+        python_path = os.path.join(local_dir, "main.py")
+        with open(python_path, "w", encoding="utf8") as fp:
+            fp.write("import numpy; import requests; print(requests.__version__);")
+
+        build_config = PythonBuild(
+            type="tfy-python-buildpack",
+            build_context_path=local_dir,
+            command=["python", "main.py"],
+            python_version="3.7.13",
+            pip_packages=["requests"],
+        )
+        build(tag="python-pip-and-req-file", build_configuration=build_config)
+
+    with TemporaryDirectory() as local_dir:
+        python_path = os.path.join(local_dir, "main.py")
+        with open(python_path, "w", encoding="utf8") as fp:
+            fp.write("import requests; print(requests.__version__);")
+
+        build_config = PythonBuild(
+            type="tfy-python-buildpack",
+            build_context_path=local_dir,
+            command=["python", "main.py"],
+            python_version="3.7.13",
+            pip_packages=["requests>0.1"],
+        )
+        build(tag="python-only-pip", build_configuration=build_config)
+
+    with TemporaryDirectory() as local_dir:
+        python_path = os.path.join(local_dir, "main.py")
+        with open(python_path, "w", encoding="utf8") as fp:
+            fp.write("print(1)")
+
+        build_config = PythonBuild(
+            type="tfy-python-buildpack",
+            build_context_path=local_dir,
+            command=["python", "main.py"],
+            python_version="3.7.13",
+        )
+        build(tag="python-only-no-pip-no-req", build_configuration=build_config)

truefoundry/deploy/builder/builders/__init__.py

@@ -0,0 +1,22 @@
+from typing import Callable, Dict
+
+from truefoundry.deploy.builder.builders import (
+    dockerfile,
+    tfy_notebook_buildpack,
+    tfy_python_buildpack,
+)
+
+BUILD_REGISTRY: Dict[str, Callable] = {
+    "dockerfile": dockerfile.build,
+    "tfy-python-buildpack": tfy_python_buildpack.build,
+    "tfy-notebook-buildpack": tfy_notebook_buildpack.build,
+}
+
+__all__ = ["get_builder"]
+
+
+def get_builder(build_configuration_type: str) -> Callable:
+    if build_configuration_type not in BUILD_REGISTRY:
+        raise NotImplementedError(f"Builder for {build_configuration_type} not found")
+
+    return BUILD_REGISTRY[build_configuration_type]

truefoundry/deploy/builder/builders/dockerfile.py

@@ -0,0 +1,57 @@
+import os
+from typing import Dict, List, Optional
+
+from truefoundry.deploy.auto_gen.models import DockerFileBuild
+from truefoundry.deploy.builder.docker_service import build_docker_image
+from truefoundry.logger import logger
+
+__all__ = ["build"]
+
+
+def _get_expanded_and_absolute_path(path: str):
+    return os.path.abspath(os.path.expanduser(path))
+
+
+def _build_docker_image(
+    tag: str,
+    path: str = ".",
+    file: Optional[str] = None,
+    build_args: Optional[Dict[str, str]] = None,
+    extra_opts: Optional[List[str]] = None,
+):
+    path = _get_expanded_and_absolute_path(path)
+
+    if file:
+        file = _get_expanded_and_absolute_path(file)
+
+    build_docker_image(
+        path=path,
+        tag=tag,
+        # TODO: can we pick target platform(s) picked from cluster
+        platform="linux/amd64",
+        dockerfile=file,
+        build_args=build_args,
+        extra_opts=extra_opts,
+    )
+
+
+def build(
+    tag: str,
+    build_configuration: DockerFileBuild,
+    extra_opts: Optional[List[str]] = None,
+):
+    dockerfile_path = _get_expanded_and_absolute_path(
+        build_configuration.dockerfile_path
+    )
+    with open(dockerfile_path) as f:
+        dockerfile_content = f.read()
+        logger.info("Dockerfile content:-")
+        logger.info(dockerfile_content)
+
+    _build_docker_image(
+        tag=tag,
+        path=build_configuration.build_context_path,
+        file=build_configuration.dockerfile_path,
+        build_args=build_configuration.build_args,
+        extra_opts=extra_opts,
+    )

truefoundry/deploy/builder/builders/tfy_notebook_buildpack/__init__.py

@@ -0,0 +1,46 @@
+import os
+from tempfile import TemporaryDirectory
+from typing import List, Optional
+
+from truefoundry.deploy.auto_gen.models import DockerFileBuild
+from truefoundry.deploy.builder.builders import dockerfile
+from truefoundry.deploy.builder.builders.tfy_notebook_buildpack.dockerfile_template import (
+    NotebookImageBuild,
+    generate_dockerfile_content,
+)
+
+__all__ = ["generate_dockerfile_content", "build"]
+
+
+def _convert_to_dockerfile_build_config(
+    build_configuration: NotebookImageBuild, local_dir: str
+) -> DockerFileBuild:
+    dockerfile_content = generate_dockerfile_content(
+        build_configuration=build_configuration,
+        local_dir=local_dir,
+    )
+    dockerfile_path = os.path.join(local_dir, "Dockerfile")
+    with open(dockerfile_path, "w", encoding="utf8") as fp:
+        fp.write(dockerfile_content)
+    return DockerFileBuild(
+        type="dockerfile",
+        dockerfile_path=dockerfile_path,
+        build_context_path=local_dir,
+    )
+
+
+def build(
+    tag: str,
+    build_configuration: NotebookImageBuild,
+    extra_opts: Optional[List[str]] = None,
+):
+    with TemporaryDirectory() as local_dir:
+        docker_build_configuration = _convert_to_dockerfile_build_config(
+            build_configuration,
+            local_dir=local_dir,
+        )
+        dockerfile.build(
+            tag=tag,
+            build_configuration=docker_build_configuration,
+            extra_opts=extra_opts,
+        )

truefoundry/deploy/builder/builders/tfy_notebook_buildpack/dockerfile_template.py

@@ -0,0 +1,66 @@
+import hashlib
+import os
+from typing import Literal, Optional
+
+from mako.template import Template
+
+from truefoundry.pydantic_v1 import BaseModel
+
+
+class NotebookImageBuild(BaseModel):
+    type: Literal["tfy-notebook-buildpack"] = "tfy-notebook-buildpack"
+    base_image_uri: str
+    build_script: Optional[str] = None
+
+
+DOCKERFILE_TEMPLATE = Template(
+    """
+FROM ${base_image_uri}
+USER root
+
+% if build_script_docker_commands is not None:
+${build_script_docker_commands}
+% endif
+
+USER $NB_UID
+"""
+)
+
+
+def generate_build_script_docker_commands(
+    build_script: Optional[str], local_dir: str
+) -> Optional[str]:
+    if not build_script:
+        return None
+    build_script_path = None
+    if build_script:
+        # we add build script's hash to the file name to ensure docker cache invalidation
+        script_hash = hashlib.sha256(build_script.encode("utf-8")).hexdigest()
+        build_script_path = os.path.join(local_dir, f"build-script-{script_hash}.sh")
+        with open(build_script_path, "w") as fp:
+            fp.write(build_script)
+        build_script_path = os.path.relpath(build_script_path, local_dir)
+    run_build_script_command = f"""\
+COPY {build_script_path} /tmp/user-build-script.sh
+RUN mkdir -p /var/log/ && DEBIAN_FRONTEND=noninteractive bash -ex /tmp/user-build-script.sh 2>&1 | tee /var/log/user-build-script-output.log
+"""
+    return run_build_script_command
+
+
+def generate_dockerfile_content(
+    build_configuration: NotebookImageBuild, local_dir: str
+) -> str:
+    build_script_docker_commands = generate_build_script_docker_commands(
+        build_script=build_configuration.build_script,
+        local_dir=local_dir,
+    )
+
+    template_args = {
+        "base_image_uri": build_configuration.base_image_uri,
+        "build_script_docker_commands": build_script_docker_commands,
+    }
+
+    template = DOCKERFILE_TEMPLATE
+
+    dockerfile_content = template.render(**template_args)
+    return dockerfile_content

truefoundry/deploy/builder/builders/tfy_python_buildpack/__init__.py

@@ -0,0 +1,44 @@
+import os
+from tempfile import TemporaryDirectory
+from typing import List, Optional
+
+from truefoundry.deploy.auto_gen.models import DockerFileBuild, PythonBuild
+from truefoundry.deploy.builder.builders import dockerfile
+from truefoundry.deploy.builder.builders.tfy_python_buildpack.dockerfile_template import (
+    generate_dockerfile_content,
+)
+
+__all__ = ["generate_dockerfile_content", "build"]
+
+
+def _convert_to_dockerfile_build_config(
+    build_configuration: PythonBuild,
+    dockerfile_path: str,
+) -> DockerFileBuild:
+    dockerfile_content = generate_dockerfile_content(
+        build_configuration=build_configuration
+    )
+    with open(dockerfile_path, "w", encoding="utf8") as fp:
+        fp.write(dockerfile_content)
+
+    return DockerFileBuild(
+        type="dockerfile",
+        dockerfile_path=dockerfile_path,
+        build_context_path=build_configuration.build_context_path,
+    )
+
+
+def build(
+    tag: str,
+    build_configuration: PythonBuild,
+    extra_opts: Optional[List[str]] = None,
+):
+    with TemporaryDirectory() as local_dir:
+        docker_build_configuration = _convert_to_dockerfile_build_config(
+            build_configuration, dockerfile_path=os.path.join(local_dir, "Dockerfile")
+        )
+        dockerfile.build(
+            tag=tag,
+            build_configuration=docker_build_configuration,
+            extra_opts=extra_opts,
+        )

truefoundry/deploy/builder/builders/tfy_python_buildpack/dockerfile_template.py

@@ -0,0 +1,158 @@
+import os
+from typing import Dict, List, Optional
+
+from mako.template import Template
+
+from truefoundry.deploy.auto_gen.models import PythonBuild
+from truefoundry.deploy.v2.lib.patched_models import CUDAVersion
+
+# TODO (chiragjn): Switch to a non-root user inside the container
+
+_POST_PYTHON_INSTALL_TEMPLATE = """
+% if apt_install_command is not None:
+RUN ${apt_install_command}
+% endif
+
+% if requirements_path is not None:
+COPY ${requirements_path} ${requirements_destination_path}
+% endif
+
+% if pip_install_command is not None:
+RUN ${pip_install_command}
+% endif
+
+COPY . /app
+WORKDIR /app
+"""
+
+DOCKERFILE_TEMPLATE = Template(
+    """
+FROM --platform=linux/amd64 python:${python_version}
+ENV PATH=/virtualenvs/venv/bin:$PATH
+RUN apt update && \
+    DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends git && \
+    python -m venv /virtualenvs/venv/ && \
+    rm -rf /var/lib/apt/lists/*
+"""
+    + _POST_PYTHON_INSTALL_TEMPLATE
+)
+
+CUDA_DOCKERFILE_TEMPLATE = Template(
+    """
+FROM --platform=linux/amd64 nvidia/cuda:${nvidia_cuda_image_tag}
+ENV PATH=/virtualenvs/venv/bin:$PATH
+RUN echo "deb https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu $(cat /etc/os-release | grep UBUNTU_CODENAME | cut -d = -f 2) main" >> /etc/apt/sources.list && \
+    echo "deb-src https://ppa.launchpadcontent.net/deadsnakes/ppa/ubuntu $(cat /etc/os-release | grep UBUNTU_CODENAME | cut -d = -f 2) main" >> /etc/apt/sources.list && \
+    apt-key adv --keyserver keyserver.ubuntu.com --recv-keys F23C5A6CF475977595C89F51BA6932366A755776 && \
+    apt update && \
+    DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends git python${python_version}-dev python${python_version}-venv && \
+    python${python_version} -m venv /virtualenvs/venv/ && \
+    rm -rf /var/lib/apt/lists/*
+"""
+    + _POST_PYTHON_INSTALL_TEMPLATE
+)
+
+CUDA_VERSION_TO_IMAGE_TAG: Dict[str, str] = {
+    CUDAVersion.CUDA_11_0_CUDNN8.value: "11.0.3-cudnn8-runtime-ubuntu20.04",
+    CUDAVersion.CUDA_11_1_CUDNN8.value: "11.1.1-cudnn8-runtime-ubuntu20.04",
+    CUDAVersion.CUDA_11_2_CUDNN8.value: "11.2.2-cudnn8-runtime-ubuntu20.04",
+    CUDAVersion.CUDA_11_3_CUDNN8.value: "11.3.1-cudnn8-runtime-ubuntu20.04",
+    CUDAVersion.CUDA_11_4_CUDNN8.value: "11.4.3-cudnn8-runtime-ubuntu20.04",
+    CUDAVersion.CUDA_11_5_CUDNN8.value: "11.5.2-cudnn8-runtime-ubuntu20.04",
+    CUDAVersion.CUDA_11_6_CUDNN8.value: "11.6.2-cudnn8-runtime-ubuntu20.04",
+    CUDAVersion.CUDA_11_7_CUDNN8.value: "11.7.1-cudnn8-runtime-ubuntu22.04",
+    CUDAVersion.CUDA_11_8_CUDNN8.value: "11.8.0-cudnn8-runtime-ubuntu22.04",
+    CUDAVersion.CUDA_12_0_CUDNN8.value: "12.0.1-cudnn8-runtime-ubuntu22.04",
+    CUDAVersion.CUDA_12_1_CUDNN8.value: "12.1.1-cudnn8-runtime-ubuntu22.04",
+    CUDAVersion.CUDA_12_2_CUDNN8.value: "12.2.2-cudnn8-runtime-ubuntu22.04",
+}
+
+
+def resolve_requirements_txt_path(build_configuration: PythonBuild) -> Optional[str]:
+    if build_configuration.requirements_path:
+        return build_configuration.requirements_path
+
+    # TODO: what if there is a requirements.txt but user does not wants us to use it.
+    possible_requirements_txt_path = os.path.join(
+        build_configuration.build_context_path, "requirements.txt"
+    )
+
+    if os.path.isfile(possible_requirements_txt_path):
+        return os.path.relpath(
+            possible_requirements_txt_path, start=build_configuration.build_context_path
+        )
+
+    return None
+
+
+def generate_apt_install_command(apt_packages: Optional[List[str]]) -> Optional[str]:
+    packages_list = None
+    if apt_packages:
+        packages_list = " ".join(p.strip() for p in apt_packages if p.strip())
+    if not packages_list:
+        return None
+    apt_update_command = "apt update"
+    apt_install_command = f"DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends {packages_list}"
+    clear_apt_lists_command = "rm -rf /var/lib/apt/lists/*"
+    return " && ".join(
+        [apt_update_command, apt_install_command, clear_apt_lists_command]
+    )
+
+
+def generate_pip_install_command(
+    requirements_path: Optional[str], pip_packages: Optional[List[str]]
+) -> Optional[str]:
+    upgrade_pip_command = "python -m pip install -U pip setuptools wheel"
+    final_pip_install_command = None
+    pip_install_base_command = "python -m pip install --use-pep517 --no-cache-dir"
+    if requirements_path:
+        final_pip_install_command = f"{pip_install_base_command} -r {requirements_path}"
+
+    if pip_packages:
+        final_pip_install_command = (
+            final_pip_install_command or pip_install_base_command
+        )
+        final_pip_install_command += " " + " ".join(
+            f"'{package}'" for package in pip_packages
+        )
+
+    if not final_pip_install_command:
+        return None
+
+    return " && ".join([upgrade_pip_command, final_pip_install_command])
+
+
+def generate_dockerfile_content(
+    build_configuration: PythonBuild,
+) -> str:
+    # TODO (chiragjn): Handle recursive references to other requirements files e.g. `-r requirements-gpu.txt`
+    requirements_path = resolve_requirements_txt_path(build_configuration)
+    requirements_destination_path = (
+        "/tmp/requirements.txt" if requirements_path else None
+    )
+    pip_install_command = generate_pip_install_command(
+        requirements_path=requirements_destination_path,
+        pip_packages=build_configuration.pip_packages,
+    )
+    apt_install_command = generate_apt_install_command(
+        apt_packages=build_configuration.apt_packages
+    )
+
+    template_args = {
+        "python_version": build_configuration.python_version,
+        "apt_install_command": apt_install_command,
+        "requirements_path": requirements_path,
+        "requirements_destination_path": requirements_destination_path,
+        "pip_install_command": pip_install_command,
+    }
+
+    if build_configuration.cuda_version:
+        template = CUDA_DOCKERFILE_TEMPLATE
+        template_args["nvidia_cuda_image_tag"] = CUDA_VERSION_TO_IMAGE_TAG.get(
+            build_configuration.cuda_version, build_configuration.cuda_version
+        )
+    else:
+        template = DOCKERFILE_TEMPLATE
+
+    dockerfile_content = template.render(**template_args)
+    return dockerfile_content

truefoundry/deploy/builder/docker_service.py

@@ -0,0 +1,168 @@
+import os
+import subprocess
+from typing import Dict, List, Optional
+
+import docker
+from rich.console import Console
+
+from truefoundry.deploy.lib.clients.shell_client import Shell
+from truefoundry.logger import logger
+
+__all__ = [
+    "build_docker_image",
+    "push_docker_image",
+    "pull_docker_image",
+    "push_docker_image_with_latest_tag",
+    "env_has_docker",
+]
+
+
+def _get_build_args_string(build_args: Optional[Dict[str, str]] = None) -> str:
+    if not build_args:
+        return None
+    result = []
+    for param, value in build_args.items():
+        result.extend(["--build-arg", f"{param.strip()}={value}"])
+    return result
+
+
+def _get_docker_client():
+    try:
+        return docker.from_env()
+    except Exception as ex:
+        raise Exception("Could not connect to Docker") from ex
+
+
+def env_has_docker():
+    try:
+        _get_docker_client()
+        return True
+    except Exception:
+        return False
+
+
+# this is required since push does throw an error if it
+# fails - so we have to parse the response logs to catch the error
+# the other option is to run `docker login` as a subprocess, but it's
+# not recommended to provide password to subprocess
+def _catch_error_in_push(response: List[dict]):
+    for line in response:
+        if line.get("error") is not None:
+            raise Exception(
+                f'Failed to push to registry with message \'{line.get("error")}\''
+            )
+
+
+def _run_cmds(command: List[str]):
+    console = Console(color_system=None, markup=False, soft_wrap=True)
+    with subprocess.Popen(
+        command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+    ) as sp:
+        for line in sp.stdout:
+            console.print(line.decode("utf-8").strip())
+        sp.communicate()
+        if sp.returncode is not None and sp.returncode != 0:
+            raise Exception(f"Command: {command} failed")
+
+
+def build_docker_image(
+    path: str,
+    tag: str,
+    platform: str,
+    dockerfile: str,
+    extra_opts: Optional[List[str]] = None,
+    build_args: Optional[Dict[str, str]] = None,
+):
+    use_depot = bool(os.environ.get("USE_DEPOT"))
+    depot_project_id = os.environ.get("DEPOT_PROJECT_KEY")
+    logger.info("Starting docker build...")
+    if use_depot and depot_project_id:
+        try:
+            command = [
+                "depot",
+                "build",
+                "--project",
+                depot_project_id,
+                "-f",
+                dockerfile,
+                "-t",
+                tag,
+                path,
+            ]
+            final_build_args = _get_build_args_string(build_args=build_args)
+            if final_build_args:
+                command.extend(final_build_args)
+            command.append("--push")  # keep push at last
+            Shell().execute_shell_command(command=command)
+        except Exception as e:
+            raise Exception("Error while building Docker image using Depot") from e
+    else:
+        try:
+            # TODO (chiragjn): Maybe consider using client.images.build
+            build_args_list = []
+            if build_args:
+                for k, v in build_args.items():
+                    build_args_list += ["--build-arg", f"{k}={v}"]
+
+            docker_build_cmd = [
+                "docker",
+                "build",
+                "-t",
+                tag,
+                "-f",
+                dockerfile,
+                "--platform",
+                platform,
+            ]
+            docker_build_cmd += [path]
+            docker_build_cmd += build_args_list
+            docker_build_cmd += extra_opts if extra_opts else []
+            _run_cmds(docker_build_cmd)
+        except Exception as e:
+            raise Exception(f"Error while building Docker image: {e}") from e
+
+
+def push_docker_image(
+    image_uri: str,
+    docker_login_username: str,
+    docker_login_password: str,
+):
+    client = _get_docker_client()
+    auth_config = {"username": docker_login_username, "password": docker_login_password}
+    logger.info(f"Pushing {image_uri}")
+    response = client.images.push(
+        repository=image_uri, auth_config=auth_config, decode=True, stream=True
+    )
+    _catch_error_in_push(response=response)
+
+
+def push_docker_image_with_latest_tag(
+    image_uri: str,
+    docker_login_username: str,
+    docker_login_password: str,
+):
+    client = _get_docker_client()
+    auth_config = {"username": docker_login_username, "password": docker_login_password}
+    repository_without_tag, _ = image_uri.rsplit(":", 1)
+    image = client.images.get(image_uri)
+    image.tag(repository=repository_without_tag, tag="latest")
+
+    logger.info(f"Pushing {repository_without_tag}:latest")
+    response = client.images.push(
+        repository=repository_without_tag,
+        tag="latest",
+        auth_config=auth_config,
+        decode=True,
+        stream=True,
+    )
+    _catch_error_in_push(response=response)
+
+
+def pull_docker_image(
+    image_uri: str,
+    docker_login_username: str,
+    docker_login_password: str,
+):
+    auth_config = {"username": docker_login_username, "password": docker_login_password}
+    logger.info(f"Pulling cache image {image_uri}")
+    _get_docker_client().images.pull(repository=image_uri, auth_config=auth_config)