mlflow--mlflow/mlflow/projects/docker.py

import logging
import os
import posixpath
import shutil
import subprocess
import tempfile
import urllib.parse
import urllib.request

import docker

from mlflow import tracking
from mlflow.environment_variables import MLFLOW_TRACKING_URI
from mlflow.exceptions import ExecutionException
from mlflow.projects.utils import MLFLOW_DOCKER_WORKDIR_PATH
from mlflow.utils import file_utils, process
from mlflow.utils.databricks_utils import get_databricks_env_vars
from mlflow.utils.file_utils import _handle_readonly_on_windows
from mlflow.utils.git_utils import get_git_commit
from mlflow.utils.mlflow_tags import MLFLOW_DOCKER_IMAGE_ID, MLFLOW_DOCKER_IMAGE_URI

_logger = logging.getLogger(__name__)

_GENERATED_DOCKERFILE_NAME = "Dockerfile.mlflow-autogenerated"
_MLFLOW_DOCKER_TRACKING_DIR_PATH = "/mlflow/tmp/mlruns"
_PROJECT_TAR_ARCHIVE_NAME = "mlflow-project-docker-build-context"


def validate_docker_installation():
    """
    Verify if Docker is installed and running on host machine.
    """
    if shutil.which("docker") is None:
        raise ExecutionException(
            "Could not find Docker executable. "
            "Ensure Docker is installed as per the instructions "
            "at https://docs.docker.com/install/overview/."
        )

    cmd = ["docker", "info"]
    prc = process._exec_cmd(
        cmd,
        throw_on_error=False,
        capture_output=False,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
    )
    if prc.returncode != 0:
        joined_cmd = " ".join(cmd)
        raise ExecutionException(
            f"Ran `{joined_cmd}` to ensure docker daemon is running but it failed "
            f"with the following output:\n{prc.stdout}"
        )


def validate_docker_env(project):
    if not project.name:
        raise ExecutionException(
            "Project name in MLProject must be specified when using docker for image tagging."
        )
    if not project.docker_env.get("image"):
        raise ExecutionException(
            "Project with docker environment must specify the docker image "
            "to use via an 'image' field under the 'docker_env' field."
        )


def build_docker_image(work_dir, repository_uri, base_image, run_id, build_image, docker_auth):
    """
    Build a docker image containing the project in `work_dir`, using the base image.
    """
    image_uri = _get_docker_image_uri(repository_uri=repository_uri, work_dir=work_dir)
    client = docker.from_env()
    if docker_auth is not None:
        client.login(**docker_auth)

    if not build_image:
        if not client.images.list(name=base_image):
            _logger.info(f"Pulling {base_image}")
            image = client.images.pull(base_image)
        else:
            _logger.info(f"{base_image} already exists")
            image = client.images.get(base_image)
        image_uri = base_image
    else:
        dockerfile = (
            f"FROM {base_image}\n COPY {_PROJECT_TAR_ARCHIVE_NAME}/ {MLFLOW_DOCKER_WORKDIR_PATH}\n"
            f" WORKDIR {MLFLOW_DOCKER_WORKDIR_PATH}\n"
        )
        build_ctx_path = _create_docker_build_ctx(work_dir, dockerfile)
        with open(build_ctx_path, "rb") as docker_build_ctx:
            _logger.info("=== Building docker image %s ===", image_uri)
            image, _ = client.images.build(
                tag=image_uri,
                forcerm=True,
                dockerfile=posixpath.join(_PROJECT_TAR_ARCHIVE_NAME, _GENERATED_DOCKERFILE_NAME),
                fileobj=docker_build_ctx,
                custom_context=True,
                encoding="gzip",
            )
        try:
            os.remove(build_ctx_path)
        except Exception:
            _logger.info("Temporary docker context file %s was not deleted.", build_ctx_path)
    tracking.MlflowClient().set_tag(run_id, MLFLOW_DOCKER_IMAGE_URI, image_uri)
    tracking.MlflowClient().set_tag(run_id, MLFLOW_DOCKER_IMAGE_ID, image.id)
    return image


def _get_docker_image_uri(repository_uri, work_dir):
    """
    Args:
        repository_uri: The URI of the Docker repository with which to tag the image. The
            repository URI is used as the prefix of the image URI.
        work_dir: Path to the working directory in which to search for a git commit hash
    """
    repository_uri = repository_uri or "docker-project"
    # Optionally include first 7 digits of git SHA in tag name, if available.
    git_commit = get_git_commit(work_dir)
    version_string = ":" + git_commit[:7] if git_commit else ""
    return repository_uri + version_string


def _create_docker_build_ctx(work_dir, dockerfile_contents):
    """
    Creates build context tarfile containing Dockerfile and project code, returning path to tarfile
    """
    directory = tempfile.mkdtemp()
    try:
        dst_path = os.path.join(directory, "mlflow-project-contents")
        shutil.copytree(src=work_dir, dst=dst_path)
        with open(os.path.join(dst_path, _GENERATED_DOCKERFILE_NAME), "w") as handle:
            handle.write(dockerfile_contents)
        _, result_path = tempfile.mkstemp()
        file_utils.make_tarfile(
            output_filename=result_path, source_dir=dst_path, archive_name=_PROJECT_TAR_ARCHIVE_NAME
        )
    finally:
        shutil.rmtree(directory, onerror=_handle_readonly_on_windows)
    return result_path


def get_docker_tracking_cmd_and_envs(tracking_uri):
    cmds = []
    env_vars = {}

    local_path, container_tracking_uri = _get_local_uri_or_none(tracking_uri)
    if local_path is not None:
        cmds = ["-v", f"{local_path}:{_MLFLOW_DOCKER_TRACKING_DIR_PATH}"]
        env_vars[MLFLOW_TRACKING_URI.name] = container_tracking_uri
    env_vars.update(get_databricks_env_vars(tracking_uri))
    return cmds, env_vars


def _get_local_uri_or_none(uri):
    if uri == "databricks":
        return None, None
    parsed_uri = urllib.parse.urlparse(uri)
    if not parsed_uri.netloc and parsed_uri.scheme in ("", "file", "sqlite"):
        path = urllib.request.url2pathname(parsed_uri.path)
        if parsed_uri.scheme == "sqlite":
            uri = file_utils.path_to_local_sqlite_uri(_MLFLOW_DOCKER_TRACKING_DIR_PATH)
        else:
            uri = file_utils.path_to_local_file_uri(_MLFLOW_DOCKER_TRACKING_DIR_PATH)
        return path, uri
    else:
        return None, None