content: #  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
  #
  #  Licensed under the Apache License, Version 2.0 (the "License");
  #  you may not use this file except in compliance with the License.
  #  You may obtain a copy of the License at:
  #
  #       https://www.apache.org/licenses/LICENSE-2.0
  #
  #  Unless required by applicable law or agreed to in writing, software
  #  distributed under the License is distributed on an "AS IS" BASIS,
  #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  #  or implied. See the License for the specific language governing
  #  permissions and limitations under the License.
  """Local Docker image builder implementation."""
  
  import shutil
  import subprocess
  import tempfile
  from typing import TYPE_CHECKING, Optional, Type, cast
  
  from pydantic import Field
  
  from zenml.image_builders import (
      BaseImageBuilder,
      BaseImageBuilderConfig,
      BaseImageBuilderFlavor,
  )
  from zenml.utils import docker_utils
  
  if TYPE_CHECKING:
      from zenml.config.docker_settings import DockerBuildOptions
      from zenml.container_registries import BaseContainerRegistry
      from zenml.image_builders import BuildContext
  
  
  class LocalImageBuilderConfig(BaseImageBuilderConfig):
      """Local image builder configuration."""
  
      use_subprocess_call: bool = Field(
          default=False,
          description="Whether to use a subprocess `docker build` call to "
          "build the image instead of using the Docker python SDK.",
      )
  
  
  class LocalImageBuilder(BaseImageBuilder):
      """Local image builder implementation."""
  
      @property
      def config(self) -> LocalImageBuilderConfig:
          """The stack component configuration.
  
          Returns:
              The configuration.
          """
          return cast(LocalImageBuilderConfig, self._config)
  
      @property
      def is_building_locally(self) -> bool:
          """Whether the image builder builds the images on the client machine.
  
          Returns:
              True if the image builder builds locally, False otherwise.
          """
          return True
  
      @staticmethod
      def _check_prerequisites() -> None:
          """Checks that all prerequisites are installed.
  
          Raises:
              RuntimeError: If any of the prerequisites are not installed or
                  running.
          """
          if not shutil.which("docker"):
              raise RuntimeError(
                  "`docker` is required to run the local image builder."
              )
  
          if not docker_utils.check_docker():
              # For 3., this is not supported by the python docker library
              # https://github.com/docker/docker-py/issues/3146
              raise RuntimeError(
                  "Unable to connect to the Docker daemon. There are three "
                  "common causes for this:\n"
                  "1) The Docker daemon isn't running.\n"
                  "2) The Docker client isn't configured correctly. The client "
                  "loads its configuration from the following file: "
                  "$HOME/.docker/config.json. If your configuration file is in a "
                  "different location, set the `DOCKER_CONFIG` environment "
                  "variable to the directory that contains your `config.json` "
                  "file.\n"
                  "3) If your Docker CLI is working fine but you ran into this "
                  "issue, you might be using a non-default Docker context which "
                  "is not supported by the Docker python library. To verify "
                  "this, run `docker context ls` and check which context has a "
                  "`*` next to it. If this is not the `default` context, copy "
                  "the `DOCKER ENDPOINT` value of that context and set the "
                  "`DOCKER_HOST` environment variable to that value."
              )
  
      def build(
          self,
          image_name: str,
          build_context: "BuildContext",
          docker_build_options: Optional["DockerBuildOptions"] = None,
          container_registry: Optional["BaseContainerRegistry"] = None,
      ) -> str:
          """Builds and optionally pushes an image using the local Docker client.
  
          Args:
              image_name: Name of the image to build and push.
              build_context: The build context to use for the image.
              docker_build_options: Docker build options.
              container_registry: Optional container registry to push to.
  
          Returns:
              The Docker image repo digest.
          """
          self._check_prerequisites()
  
          if self.config.use_subprocess_call:
              self._build_with_subprocess_call(
                  image_name=image_name,
                  build_context=build_context,
                  docker_build_options=docker_build_options,
              )
          else:
              self._build_with_python_sdk(
                  image_name=image_name,
                  build_context=build_context,
                  docker_build_options=docker_build_options,
                  container_registry=container_registry,
              )
  
          if container_registry:
              return container_registry.push_image(image_name)
          else:
              return image_name
  
      def _build_with_python_sdk(
          self,
          image_name: str,
          build_context: "BuildContext",
          docker_build_options: Optional["DockerBuildOptions"] = None,
          container_registry: Optional["BaseContainerRegistry"] = None,
      ) -> None:
          """Builds an image using the Python Docker SDK.
  
          Args:
              image_name: Name of the image to build and push.
              build_context: The build context to use for the image.
              docker_build_options: Docker build options.
              container_registry: Optional container registry.
          """
          if container_registry:
              # Use the container registry's docker client, which may be
              # authenticated to access additional registries
              docker_client = container_registry.docker_client
          else:
              docker_client = docker_utils._try_get_docker_client_from_env()
  
          build_options = (
              docker_build_options.to_docker_python_sdk_options()
              if docker_build_options
              else {}
          )
  
          with tempfile.TemporaryFile(mode="w+b") as f:
              build_context.write_archive(f)
  
              # We use the client api directly here, so we can stream the logs
              output_stream = docker_client.images.client.api.build(
                  fileobj=f,
                  custom_context=True,
                  tag=image_name,
                  **build_options,
              )
          docker_utils._process_stream(output_stream)
  
      def _build_with_subprocess_call(
          self,
          image_name: str,
          build_context: "BuildContext",
          docker_build_options: Optional["DockerBuildOptions"] = None,
      ) -> None:
          """Builds an image using a subprocess `docker build` call.
  
          Args:
              image_name: Name of the image to build and push.
              build_context: The build context to use for the image.
              docker_build_options: Docker build options.
  
          Raises:
              RuntimeError: If the subprocess call fails.
          """
          with tempfile.TemporaryFile(mode="w+b") as f:
              build_context.write_archive(f)
  
              # The `-` signals that the build context will be passed as a tar
              # file to stdin
              command = ["docker", "build", "-", "-t", image_name]
  
              if docker_build_options:
                  command.extend(docker_build_options.to_docker_cli_options())
  
              process = subprocess.Popen(command, stdin=f)
  
              result = process.wait()
              if result != 0:
                  raise RuntimeError(
                      f"Failed to build image {image_name}. Please check the "
                      "logs above for more information."
                  )
  
  
  class LocalImageBuilderFlavor(BaseImageBuilderFlavor):
      """Local image builder flavor."""
  
      @property
      def name(self) -> str:
          """The flavor name.
  
          Returns:
              The flavor name.
          """
          return "local"
  
      @property
      def docs_url(self) -> Optional[str]:
          """A url to point at docs explaining this flavor.
  
          Returns:
              A flavor docs url.
          """
          return self.generate_default_docs_url()
  
      @property
      def sdk_docs_url(self) -> Optional[str]:
          """A url to point at docs explaining this flavor.
  
          Returns:
              A flavor docs url.
          """
          return self.generate_default_sdk_docs_url()
  
      @property
      def logo_url(self) -> str:
          """A url to represent the flavor in the dashboard.
  
          Returns:
              The flavor logo.
          """
          return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/image_builder/local.svg"
  
      @property
      def config_class(self) -> Type[LocalImageBuilderConfig]:
          """Config class.
  
          Returns:
              The config class.
          """
          return LocalImageBuilderConfig
  
      @property
      def implementation_class(self) -> Type[LocalImageBuilder]:
          """Implementation class.
  
          Returns:
              The implementation class.
          """
          return LocalImageBuilder
  