<reponame>mldock/mldock
"""Local API methods"""
import os
import logging
from pathlib import Path
import requests
import docker

from mldock.config_managers.cli import CliConfigureManager
from mldock.platform_helpers import utils

logger = logging.getLogger("mldock")

# log formatting
def pretty_build_logs(line: dict):
    """format docker build logs"""
    if line is None:
        return None

    error = line.get("error", None)
    error_detail = line.get("errorDetail", None)

    if error is not None:
        logger.error("{}\n{}".format(error, error_detail))

    stream = line.get("stream", "")

    if ("Step" in stream) & (":" in stream):
        logger.info(" ==> {MESSAGE}".format(MESSAGE=stream))
    else:
        logger.debug(" ==> {MESSAGE}".format(MESSAGE=stream))

    return None


class DockerLLCManager:
    """Docker APIClient context manager"""

    client = None

    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def __enter__(self):
        self.client = docker.APIClient(**self.kwargs)
        return self.client

    def __exit__(self, exc_type, exc_value, exc_tb):
        logger.info(self.client.prune_builds())
        self.client.close()


class DockerManager:
    """Docker Client context manager"""

    client = None

    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def __enter__(self):
        self.client = docker.from_env(**self.kwargs)
        return self.client

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.client.close()


# API methods


def generate_mldock_mount_volume(
    host_prefix: str,
    container_prefix: str,
    working_dir: str,
    base_container_path: str,
    mode: str = "rw",
):
    """
    generate host and container paths for docker volume mounts.

    args:
        host_prefix (str): host prefix to use relative to working_dir
        container_prefix (str): container prefix to use relative to base_container_path
        working_dir (str): host working directory to build host volumes from
        base_container_path (str): container base path to build container volumes from
        mode (str): (default='rw') docker volume mount read/write mode
    """
    host_path = Path(working_dir, host_prefix).absolute().as_posix()
    container_path = Path(base_container_path, container_prefix).as_posix()
    return host_path, {"bind": container_path, "mode": mode}


def make_multiple_mldock_mount_volumes(
    volumes: dict, working_dir: str, base_container_path: str, mode: str = "rw"
):
    """
    Generate multiple docker volume mounts for a collection of volume mount mappings.

    args:
        volumes (dict): volume mount mappings
        working_dir (str): host working directory to build host volumes from
        base_container_path (str): container base path to build container volumes from
        mode (str): (default='rw') docker volume mount read/write mode
    """
    docker_volumes = {}
    for volume in volumes:

        host_mount, container_mount = generate_mldock_mount_volume(
            host_prefix=volume["host_prefix"],
            container_prefix=volume["container_prefix"],
            working_dir=working_dir,
            base_container_path=base_container_path,
            mode=mode,
        )

        docker_volumes.update({host_mount: container_mount})
    return docker_volumes


def train_model(working_dir, docker_tag, image_name, entrypoint, **kwargs):
    """
    Trains ML model(s) locally
    :param working_dir: [str], source root directory
    :param docker_tag: [str], the Docker tag for the image
    :param image_name: [str], The name of the Docker image
    """

    config_manager = CliConfigureManager()

    try:
        with DockerManager() as client:
            base_ml_path = kwargs.get("base_ml_path", "/opt/ml")

            container_volumes = make_multiple_mldock_mount_volumes(
                volumes=[
                    {"host_prefix": "config", "container_prefix": "input/config"},
                    {"host_prefix": "data", "container_prefix": "input/data"},
                    {"host_prefix": "model", "container_prefix": "model"},
                    {"host_prefix": "output", "container_prefix": "output"},
                ],
                working_dir=working_dir,
                base_container_path=base_ml_path,
                mode="rw",
            )

            if config_manager.local.get("auth_type", None) is not None:

                credentials_volume = utils.get_sdk_credentials_volume_mount(
                    config_manager.local.get("auth_type")
                )

                container_volumes.update(credentials_volume)

            container = client.containers.run(
                image="{IMAGE}:{TAG}".format(IMAGE=image_name, TAG=docker_tag),
                entrypoint=entrypoint,
                command=kwargs.get("cmd", "train"),
                environment=kwargs.get("env", {}),
                remove=True,
                tty=True,
                volumes=container_volumes,
                detach=True,
                stream=True,
            )
            logs = container.logs(follow=True).decode("utf-8")

            logger.info(logs)
    except (KeyboardInterrupt, SystemExit) as exception:
        logger.error(exception)
        container.kill()
        raise
    except (
        docker.errors.APIError,
        docker.errors.ContainerError,
        docker.errors.ImageNotFound,
    ) as exception:
        logger.error(exception)
        raise
    except Exception as exception:
        logger.error(exception)
        raise


def deploy_model(working_dir, docker_tag, image_name, entrypoint, **kwargs):
    """
    Deploys ML models(s) locally
    :param working_dir: [str], source root directory
    :param docker_tag: [str], the Docker tag for the image
    :param image_name: [str], The name of the Docker image
    """

    config_manager = CliConfigureManager()
    try:
        with DockerManager() as client:
            base_ml_path = kwargs.get("base_ml_path", "/opt/ml")

            container_volumes = make_multiple_mldock_mount_volumes(
                volumes=[
                    {"host_prefix": "config", "container_prefix": "input/config"},
                    {"host_prefix": "data", "container_prefix": "input/data"},
                    {"host_prefix": "model", "container_prefix": "model"},
                    {"host_prefix": "output", "container_prefix": "output"},
                ],
                working_dir=working_dir,
                base_container_path=base_ml_path,
                mode="rw",
            )

            if config_manager.local.get("auth_type", None) is not None:

                credentials_volume = utils.get_sdk_credentials_volume_mount(
                    config_manager.local.get("auth_type")
                )

                container_volumes.update(credentials_volume)

            container = client.containers.run(
                image="{IMAGE}:{TAG}".format(IMAGE=image_name, TAG=docker_tag),
                entrypoint=entrypoint,
                command=kwargs.get("cmd", "serve"),
                environment=kwargs.get("env", {}),
                ports=kwargs.get("port", {8080: 8080}),
                remove=True,
                tty=True,
                volumes=container_volumes,
                auto_remove=True,
                detach=True,
                stream=True,
            )

            if kwargs.get("verbose", False):
                logs = container.logs(follow=True).decode("utf-8")

                logger.info(logs)

    except requests.exceptions.ConnectionError as exception:
        logger.info(
            "\nWorker socket connection timed out. "
            "\n\nThis can happen when connection pool is overloaded. "
            "We suggest setting environment variable: "
            "DOCKER_CLIENT_TIMEOUT=120 to give the client more time."
        )

    except (KeyboardInterrupt, SystemExit) as exception:
        logger.error(exception)
        container.kill()
        raise
    except (
        docker.errors.APIError,
        docker.errors.ContainerError,
        docker.errors.ImageNotFound,
    ) as exception:
        logger.error(exception)
        raise
    except Exception as exception:
        logger.error(exception)
        raise


def docker_build(
    image_name: str,
    dockerfile_path: str,
    module_path: str,
    requirements_file_path: str,
    **kwargs
):
    """Runs the build executable using docker python sdk.

    Args:
        script_path (str): relative path to script when run on root
        base_path (str):
        image_name (str):
        dockerfile_path (str):
        module_path (str):
        target_dir_name (str):
        requirements_file_path (str):
        docker_tag (str): the docker tag to build
    """

    try:

        with DockerLLCManager(
            max_pool_size=20,
            base_url=os.environ.get("DOCKER_HOST", "unix:///var/run/docker.sock"),
        ) as client:

            full_image_name = "{IMAGE}:{TAG}".format(
                IMAGE=image_name, TAG=kwargs.get("docker_tag", "latest")
            )
            logs = client.build(
                tag=full_image_name,
                path=".",
                dockerfile=Path(dockerfile_path).as_posix(),
                buildargs={
                    "module_path": Path(module_path).as_posix(),
                    "target_dir_name": Path(
                        kwargs.get("target_dir_name", "src")
                    ).as_posix(),
                    "requirements_file_path": Path(requirements_file_path).as_posix(),
                    "container_platform": kwargs.get("container_platform", None),
                },
                nocache=kwargs.get("no_cache", False),
                rm=True,
                decode=True,
            )
            for line in logs:
                yield line

    except (KeyboardInterrupt, SystemExit) as exception:
        logger.error(exception)
        raise
    except (
        docker.errors.APIError,
        docker.errors.ContainerError,
        docker.errors.ImageNotFound,
    ) as exception:
        logger.error(exception)
        raise
    except (docker.errors.DockerException) as exception:
        if "NewConnectionError" in str(exception):
            raise Exception(
                (
                    "Unable to connect to docker daemon. "
                    "Set environment variable: DOCKER_HOST"
                )
            ) from exception
        raise
    except Exception as exception:
        logger.error(exception)
        raise


def search_mldock_containers(**kwargs):
    """Performs a docker ps with filters. Incorporating the label to find only MLDOCK Containers"""
    mldock_filters = {}
    client = docker.from_env()
    ancestor = kwargs.get("ancestor", None)
    if ancestor is not None:
        mldock_filters.update({"ancestor": ancestor})

    label = kwargs.get("label", ["MLDOCK__IS_MLDOCK_CONTAINER=true"])
    mldock_filters.update({"label": label})
    containers = client.containers.list(filters=mldock_filters)

    return containers

"""Local API methods"""
import os
import logging
from pathlib import Path
import requests
import docker

from mldock.config_managers.cli import CliConfigureManager
from mldock.platform_helpers import utils

logger = logging.getLogger("mldock")

# log formatting
def pretty_build_logs(line: dict):
    """format docker build logs"""
    if line is None:
        return None

    error = line.get("error", None)
    error_detail = line.get("errorDetail", None)

    if error is not None:
        logger.error("{}\n{}".format(error, error_detail))

    stream = line.get("stream", "")

    if ("Step" in stream) & (":" in stream):
        logger.info(" ==> {MESSAGE}".format(MESSAGE=stream))
    else:
        logger.debug(" ==> {MESSAGE}".format(MESSAGE=stream))

    return None


class DockerLLCManager:
    """Docker APIClient context manager"""

    client = None

    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def __enter__(self):
        self.client = docker.APIClient(**self.kwargs)
        return self.client

    def __exit__(self, exc_type, exc_value, exc_tb):
        logger.info(self.client.prune_builds())
        self.client.close()


class DockerManager:
    """Docker Client context manager"""

    client = None

    def __init__(self, **kwargs):
        self.kwargs = kwargs

    def __enter__(self):
        self.client = docker.from_env(**self.kwargs)
        return self.client

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.client.close()


# API methods


def generate_mldock_mount_volume(
    host_prefix: str,
    container_prefix: str,
    working_dir: str,
    base_container_path: str,
    mode: str = "rw",
):
    """
    generate host and container paths for docker volume mounts.

    args:
        host_prefix (str): host prefix to use relative to working_dir
        container_prefix (str): container prefix to use relative to base_container_path
        working_dir (str): host working directory to build host volumes from
        base_container_path (str): container base path to build container volumes from
        mode (str): (default='rw') docker volume mount read/write mode
    """
    host_path = Path(working_dir, host_prefix).absolute().as_posix()
    container_path = Path(base_container_path, container_prefix).as_posix()
    return host_path, {"bind": container_path, "mode": mode}


def make_multiple_mldock_mount_volumes(
    volumes: dict, working_dir: str, base_container_path: str, mode: str = "rw"
):
    """
    Generate multiple docker volume mounts for a collection of volume mount mappings.

    args:
        volumes (dict): volume mount mappings
        working_dir (str): host working directory to build host volumes from
        base_container_path (str): container base path to build container volumes from
        mode (str): (default='rw') docker volume mount read/write mode
    """
    docker_volumes = {}
    for volume in volumes:

        host_mount, container_mount = generate_mldock_mount_volume(
            host_prefix=volume["host_prefix"],
            container_prefix=volume["container_prefix"],
            working_dir=working_dir,
            base_container_path=base_container_path,
            mode=mode,
        )

        docker_volumes.update({host_mount: container_mount})
    return docker_volumes


def train_model(working_dir, docker_tag, image_name, entrypoint, **kwargs):
    """
    Trains ML model(s) locally
    :param working_dir: [str], source root directory
    :param docker_tag: [str], the Docker tag for the image
    :param image_name: [str], The name of the Docker image
    """

    config_manager = CliConfigureManager()

    try:
        with DockerManager() as client:
            base_ml_path = kwargs.get("base_ml_path", "/opt/ml")

            container_volumes = make_multiple_mldock_mount_volumes(
                volumes=[
                    {"host_prefix": "config", "container_prefix": "input/config"},
                    {"host_prefix": "data", "container_prefix": "input/data"},
                    {"host_prefix": "model", "container_prefix": "model"},
                    {"host_prefix": "output", "container_prefix": "output"},
                ],
                working_dir=working_dir,
                base_container_path=base_ml_path,
                mode="rw",
            )

            if config_manager.local.get("auth_type", None) is not None:

                credentials_volume = utils.get_sdk_credentials_volume_mount(
                    config_manager.local.get("auth_type")
                )

                container_volumes.update(credentials_volume)

            container = client.containers.run(
                image="{IMAGE}:{TAG}".format(IMAGE=image_name, TAG=docker_tag),
                entrypoint=entrypoint,
                command=kwargs.get("cmd", "train"),
                environment=kwargs.get("env", {}),
                remove=True,
                tty=True,
                volumes=container_volumes,
                detach=True,
                stream=True,
            )
            logs = container.logs(follow=True).decode("utf-8")

            logger.info(logs)
    except (KeyboardInterrupt, SystemExit) as exception:
        logger.error(exception)
        container.kill()
        raise
    except (
        docker.errors.APIError,
        docker.errors.ContainerError,
        docker.errors.ImageNotFound,
    ) as exception:
        logger.error(exception)
        raise
    except Exception as exception:
        logger.error(exception)
        raise


def deploy_model(working_dir, docker_tag, image_name, entrypoint, **kwargs):
    """
    Deploys ML models(s) locally
    :param working_dir: [str], source root directory
    :param docker_tag: [str], the Docker tag for the image
    :param image_name: [str], The name of the Docker image
    """

    config_manager = CliConfigureManager()
    try:
        with DockerManager() as client:
            base_ml_path = kwargs.get("base_ml_path", "/opt/ml")

            container_volumes = make_multiple_mldock_mount_volumes(
                volumes=[
                    {"host_prefix": "config", "container_prefix": "input/config"},
                    {"host_prefix": "data", "container_prefix": "input/data"},
                    {"host_prefix": "model", "container_prefix": "model"},
                    {"host_prefix": "output", "container_prefix": "output"},
                ],
                working_dir=working_dir,
                base_container_path=base_ml_path,
                mode="rw",
            )

            if config_manager.local.get("auth_type", None) is not None:

                credentials_volume = utils.get_sdk_credentials_volume_mount(
                    config_manager.local.get("auth_type")
                )

                container_volumes.update(credentials_volume)

            container = client.containers.run(
                image="{IMAGE}:{TAG}".format(IMAGE=image_name, TAG=docker_tag),
                entrypoint=entrypoint,
                command=kwargs.get("cmd", "serve"),
                environment=kwargs.get("env", {}),
                ports=kwargs.get("port", {8080: 8080}),
                remove=True,
                tty=True,
                volumes=container_volumes,
                auto_remove=True,
                detach=True,
                stream=True,
            )

            if kwargs.get("verbose", False):
                logs = container.logs(follow=True).decode("utf-8")

                logger.info(logs)

    except requests.exceptions.ConnectionError as exception:
        logger.info(
            "\nWorker socket connection timed out. "
            "\n\nThis can happen when connection pool is overloaded. "
            "We suggest setting environment variable: "
            "DOCKER_CLIENT_TIMEOUT=120 to give the client more time."
        )

    except (KeyboardInterrupt, SystemExit) as exception:
        logger.error(exception)
        container.kill()
        raise
    except (
        docker.errors.APIError,
        docker.errors.ContainerError,
        docker.errors.ImageNotFound,
    ) as exception:
        logger.error(exception)
        raise
    except Exception as exception:
        logger.error(exception)
        raise


def docker_build(
    image_name: str,
    dockerfile_path: str,
    module_path: str,
    requirements_file_path: str,
    **kwargs
):
    """Runs the build executable using docker python sdk.

    Args:
        script_path (str): relative path to script when run on root
        base_path (str):
        image_name (str):
        dockerfile_path (str):
        module_path (str):
        target_dir_name (str):
        requirements_file_path (str):
        docker_tag (str): the docker tag to build
    """

    try:

        with DockerLLCManager(
            max_pool_size=20,
            base_url=os.environ.get("DOCKER_HOST", "unix:///var/run/docker.sock"),
        ) as client:

            full_image_name = "{IMAGE}:{TAG}".format(
                IMAGE=image_name, TAG=kwargs.get("docker_tag", "latest")
            )
            logs = client.build(
                tag=full_image_name,
                path=".",
                dockerfile=Path(dockerfile_path).as_posix(),
                buildargs={
                    "module_path": Path(module_path).as_posix(),
                    "target_dir_name": Path(
                        kwargs.get("target_dir_name", "src")
                    ).as_posix(),
                    "requirements_file_path": Path(requirements_file_path).as_posix(),
                    "container_platform": kwargs.get("container_platform", None),
                },
                nocache=kwargs.get("no_cache", False),
                rm=True,
                decode=True,
            )
            for line in logs:
                yield line

    except (KeyboardInterrupt, SystemExit) as exception:
        logger.error(exception)
        raise
    except (
        docker.errors.APIError,
        docker.errors.ContainerError,
        docker.errors.ImageNotFound,
    ) as exception:
        logger.error(exception)
        raise
    except (docker.errors.DockerException) as exception:
        if "NewConnectionError" in str(exception):
            raise Exception(
                (
                    "Unable to connect to docker daemon. "
                    "Set environment variable: DOCKER_HOST"
                )
            ) from exception
        raise
    except Exception as exception:
        logger.error(exception)
        raise


def search_mldock_containers(**kwargs):
    """Performs a docker ps with filters. Incorporating the label to find only MLDOCK Containers"""
    mldock_filters = {}
    client = docker.from_env()
    ancestor = kwargs.get("ancestor", None)
    if ancestor is not None:
        mldock_filters.update({"ancestor": ancestor})

    label = kwargs.get("label", ["MLDOCK__IS_MLDOCK_CONTAINER=true"])
    mldock_filters.update({"label": label})
    containers = client.containers.list(filters=mldock_filters)

    return containers
Local API methods # log formatting format docker build logs Docker APIClient context manager Docker Client context manager # API methods generate host and container paths for docker volume mounts.      args:         host_prefix (str): host prefix to use relative to working_dir         container_prefix (str): container prefix to use relative to base_container_path         working_dir (str): host working directory to build host volumes from         base_container_path (str): container base path to build container volumes from         mode (str): (default='rw') docker volume mount read/write mode Generate multiple docker volume mounts for a collection of volume mount mappings.      args:         volumes (dict): volume mount mappings         working_dir (str): host working directory to build host volumes from         base_container_path (str): container base path to build container volumes from         mode (str): (default='rw') docker volume mount read/write mode Trains ML model(s) locally     :param working_dir: [str], source root directory     :param docker_tag: [str], the Docker tag for the image     :param image_name: [str], The name of the Docker image Deploys ML models(s) locally     :param working_dir: [str], source root directory     :param docker_tag: [str], the Docker tag for the image     :param image_name: [str], The name of the Docker image Runs the build executable using docker python sdk.      Args:         script_path (str): relative path to script when run on root         base_path (str):         image_name (str):         dockerfile_path (str):         module_path (str):         target_dir_name (str):         requirements_file_path (str):         docker_tag (str): the docker tag to build Performs a docker ps with filters. Incorporating the label to find only MLDOCK Containers