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.
  """Implementation of a containerized ZenML service."""
  
  import os
  import pathlib
  import sys
  import tempfile
  import time
  from abc import abstractmethod
  from typing import Any, Dict, Generator, List, Optional, Tuple
  
  import docker.errors as docker_errors
  from docker.client import DockerClient
  from docker.models.containers import Container
  from pydantic import Field
  
  from zenml.constants import ENV_ZENML_CONFIG_PATH
  from zenml.enums import ServiceState
  from zenml.logger import get_logger
  from zenml.services.container.container_service_endpoint import (
      ContainerServiceEndpoint,
  )
  from zenml.services.service import BaseService, ServiceConfig
  from zenml.services.service_status import ServiceStatus
  from zenml.utils import docker_utils
  from zenml.utils.docker_utils import check_docker
  from zenml.utils.io_utils import (
      create_dir_recursive_if_not_exists,
      get_global_config_directory,
  )
  
  logger = get_logger(__name__)
  
  
  SERVICE_CONTAINER_PATH = "/service"
  SERVICE_LOG_FILE_NAME = "service.log"
  SERVICE_CONFIG_FILE_NAME = "service.json"
  SERVICE_CONTAINER_GLOBAL_CONFIG_DIR = "zenconfig"
  SERVICE_CONTAINER_GLOBAL_CONFIG_PATH = os.path.join(
      "/", SERVICE_CONTAINER_GLOBAL_CONFIG_DIR
  )
  DOCKER_ZENML_SERVER_DEFAULT_IMAGE = "zenmldocker/zenml"
  ENV_ZENML_SERVICE_CONTAINER = "ZENML_SERVICE_CONTAINER"
  
  
  class ContainerServiceConfig(ServiceConfig):
      """containerized service configuration.
  
      Attributes:
          root_runtime_path: the root path where the service stores its files.
          singleton: set to True to store the service files directly in the
              `root_runtime_path` directory instead of creating a subdirectory for
              each service instance. Only has effect if the `root_runtime_path` is
              also set.
          image: the container image to use for the service.
      """
  
      root_runtime_path: Optional[str] = None
      singleton: bool = False
      image: str = DOCKER_ZENML_SERVER_DEFAULT_IMAGE
  
  
  class ContainerServiceStatus(ServiceStatus):
      """containerized service status.
  
      Attributes:
          runtime_path: the path where the service files (e.g. the configuration
              file used to start the service daemon and the logfile) are located
      """
  
      runtime_path: Optional[str] = None
  
      @property
      def config_file(self) -> Optional[str]:
          """Get the path to the service configuration file.
  
          Returns:
              The path to the configuration file, or None, if the
              service has never been started before.
          """
          if not self.runtime_path:
              return None
          return os.path.join(self.runtime_path, SERVICE_CONFIG_FILE_NAME)
  
      @property
      def log_file(self) -> Optional[str]:
          """Get the path to the log file where the service output is/has been logged.
  
          Returns:
              The path to the log file, or None, if the service has never been
              started before.
          """
          if not self.runtime_path:
              return None
          return os.path.join(self.runtime_path, SERVICE_LOG_FILE_NAME)
  
  
  class ContainerService(BaseService):
      """A service represented by a containerized process.
  
      This class extends the base service class with functionality concerning
      the life-cycle management and tracking of external services implemented as
      docker containers.
  
      To define a containerized service, subclass this class and implement the
      `run` method. Upon `start`, the service will spawn a container that
      ends up calling the `run` method.
  
      For example,
  
      ```python
  
      from zenml.services import ServiceType, ContainerService, ContainerServiceConfig
      import time
  
      class SleepingServiceConfig(ContainerServiceConfig):
  
          wake_up_after: int
  
      class SleepingService(ContainerService):
  
          SERVICE_TYPE = ServiceType(
              name="sleeper",
              description="Sleeping container",
              type="container",
              flavor="sleeping",
          )
          config: SleepingServiceConfig
  
          def run(self) -> None:
              time.sleep(self.config.wake_up_after)
  
      service = SleepingService(config=SleepingServiceConfig(wake_up_after=10))
      service.start()
      ```
  
      NOTE: the `SleepingService` class and its parent module have to be
      discoverable as part of a ZenML `Integration`, otherwise the daemon will
      fail with the following error:
  
      ```
      TypeError: Cannot load service with unregistered service type:
      name='sleeper' type='container' flavor='sleeping' description='Sleeping container'
      ```
  
      Attributes:
          config: service configuration
          status: service status
          endpoint: optional service endpoint
      """
  
      config: ContainerServiceConfig = Field(
          default_factory=ContainerServiceConfig
      )
      status: ContainerServiceStatus = Field(
          default_factory=ContainerServiceStatus
      )
      # TODO [ENG-705]: allow multiple endpoints per service
      endpoint: Optional[ContainerServiceEndpoint] = None
  
      _docker_client: Optional[DockerClient] = None
  
      @property
      def docker_client(self) -> DockerClient:
          """Initialize and/or return the docker client.
  
          Returns:
              The docker client.
          """
          if self._docker_client is None:
              self._docker_client = (
                  docker_utils._try_get_docker_client_from_env()
              )
          return self._docker_client
  
      @property
      def container_id(self) -> str:
          """Get the ID of the docker container for a service.
  
          Returns:
              The ID of the docker container for the service.
          """
          return f"zenml-{str(self.uuid)}"
  
      def get_service_status_message(self) -> str:
          """Get a message about the current operational state of the service.
  
          Returns:
              A message providing information about the current operational
              state of the service.
          """
          msg = super().get_service_status_message()
          msg += f"  Container ID: `{self.container_id}`\n"
          if self.status.log_file:
              msg += (
                  f"For more information on the service status, please see the "
                  f"following log file: {self.status.log_file}\n"
              )
          return msg
  
      def check_status(self) -> Tuple[ServiceState, str]:
          """Check the the current operational state of the docker container.
  
          Returns:
              The operational state of the docker container and a message
              providing additional information about that state (e.g. a
              description of the error, if one is encountered).
          """
          # Check if Docker is available first
          if not check_docker():
              return (ServiceState.INACTIVE, "Docker daemon is not running")
  
          container: Optional[Container] = None
          try:
              container = self.docker_client.containers.get(self.container_id)
          except docker_errors.NotFound:
              # container doesn't exist yet or was removed
              pass
  
          if container is None:
              return ServiceState.INACTIVE, "Docker container is not present"
          elif container.status == "running":
              return ServiceState.ACTIVE, "Docker container is running"
          elif container.status == "exited":
              return (
                  ServiceState.ERROR,
                  "Docker container has exited.",
              )
          else:
              return (
                  ServiceState.INACTIVE,
                  f"Docker container is {container.status}",
              )
  
      def _setup_runtime_path(self) -> None:
          """Set up the runtime path for the service.
  
          This method sets up the runtime path for the service.
          """
          # reuse the config file and logfile location from a previous run,
          # if available
          if not self.status.runtime_path or not os.path.exists(
              self.status.runtime_path
          ):
              if self.config.root_runtime_path:
                  if self.config.singleton:
                      self.status.runtime_path = self.config.root_runtime_path
                  else:
                      self.status.runtime_path = os.path.join(
                          self.config.root_runtime_path,
                          str(self.uuid),
                      )
                  create_dir_recursive_if_not_exists(self.status.runtime_path)
              else:
                  self.status.runtime_path = tempfile.mkdtemp(
                      prefix="zenml-service-"
                  )
  
      def _get_container_cmd(self) -> Tuple[List[str], Dict[str, str]]:
          """Get the command to run the service container.
  
          The default implementation provided by this class is the following:
  
            * this ContainerService instance and its configuration
            are serialized as JSON and saved to a file
            * the entrypoint.py script is launched as a docker container
            and pointed to the serialized service file
            * the entrypoint script re-creates the ContainerService instance
            from the serialized configuration, then calls the `run`
            method that must be implemented by the subclass
  
          Subclasses that need a different command to launch the container
          should override this method.
  
          Returns:
              Command needed to launch the docker container and the environment
              variables to set, in the formats accepted by subprocess.Popen.
          """
          # to avoid circular imports, import here
          from zenml.services.container import entrypoint
  
          assert self.status.config_file is not None
          assert self.status.log_file is not None
  
          with open(self.status.config_file, "w") as f:
              f.write(self.model_dump_json(indent=4))
          pathlib.Path(self.status.log_file).touch()
  
          command = [
              "python",
              "-m",
              entrypoint.__name__,
              "--config-file",
              os.path.join(SERVICE_CONTAINER_PATH, SERVICE_CONFIG_FILE_NAME),
          ]
  
          command_env = {
              ENV_ZENML_SERVICE_CONTAINER: "true",
          }
          for k, v in os.environ.items():
              if k.startswith("ZENML_"):
                  command_env[k] = v
          # the global configuration is mounted into the container at a
          # different location
          command_env[ENV_ZENML_CONFIG_PATH] = (
              SERVICE_CONTAINER_GLOBAL_CONFIG_PATH
          )
  
          return command, command_env
  
      def _get_container_volumes(self) -> Dict[str, Dict[str, str]]:
          """Get the volumes to mount into the service container.
  
          The default implementation provided by this class mounts the
          following directories into the container:
  
            * the service runtime path
            * the global configuration directory
  
          Subclasses that need to mount additional volumes should override
          this method.
  
          Returns:
              A dictionary mapping host paths to dictionaries containing
              the mount options for each volume.
          """
          volumes: Dict[str, Dict[str, str]] = {}
  
          assert self.status.runtime_path is not None
  
          volumes[self.status.runtime_path] = {
              "bind": SERVICE_CONTAINER_PATH,
              "mode": "rw",
          }
  
          volumes[get_global_config_directory()] = {
              "bind": SERVICE_CONTAINER_GLOBAL_CONFIG_PATH,
              "mode": "rw",
          }
  
          return volumes
  
      @property
      def container(self) -> Optional[Container]:
          """Get the docker container for the service.
  
          Returns:
              The docker container for the service, or None if the container
              does not exist.
          """
          try:
              return self.docker_client.containers.get(self.container_id)
          except docker_errors.NotFound:
              # container doesn't exist yet or was removed
              return None
  
      def _start_container(self) -> None:
          """Start the service docker container associated with this service."""
          container = self.container
  
          if container:
              # the container exists, check if it is running
              if container.status == "running":
                  logger.debug(
                      "Container for service '%s' is already running",
                      self,
                  )
                  return
  
              # the container is stopped or in an error state, remove it
              logger.debug(
                  "Removing previous container for service '%s'",
                  self,
              )
              container.remove(force=True)
  
          logger.debug("Starting container for service '%s'...", self)
  
          try:
              self.docker_client.images.get(self.config.image)
          except docker_errors.ImageNotFound:
              logger.debug(
                  "Pulling container image '%s' for service '%s'...",
                  self.config.image,
                  self,
              )
              self.docker_client.images.pull(self.config.image)
  
          self._setup_runtime_path()
  
          ports: Dict[int, Optional[int]] = {}
          if self.endpoint:
              self.endpoint.prepare_for_start()
              if self.endpoint.status.port:
                  ports[self.endpoint.status.port] = self.endpoint.status.port
  
          command, env = self._get_container_cmd()
          volumes = self._get_container_volumes()
  
          try:
              uid_args: Dict[str, Any] = {}
              if sys.platform == "win32":
                  # File permissions are not checked on Windows. This if clause
                  # prevents mypy from complaining about unused 'type: ignore'
                  # statements
                  pass
              else:
                  # Run the container in the context of the local UID/GID
                  # to ensure that the local database can be shared
                  # with the container.
                  logger.debug(
                      "Setting UID and GID to local user/group in container."
                  )
                  uid_args = dict(
                      user=os.getuid(),
                      group_add=[os.getgid()],
                  )
  
              container = self.docker_client.containers.run(
                  name=self.container_id,
                  image=self.config.image,
                  entrypoint=command,
                  detach=True,
                  volumes=volumes,
                  environment=env,
                  remove=False,
                  auto_remove=False,
                  ports=ports,
                  labels={
                      "zenml-service-uuid": str(self.uuid),
                  },
                  working_dir=SERVICE_CONTAINER_PATH,
                  extra_hosts={"host.docker.internal": "host-gateway"},
                  **uid_args,
              )
  
              logger.debug(
                  "Docker container for service '%s' started with ID: %s",
                  self,
                  self.container_id,
              )
          except docker_errors.DockerException as e:
              logger.error(
                  "Docker container for service '%s' failed to start: %s",
                  self,
                  e,
              )
  
      def _stop_daemon(self, force: bool = False) -> None:
          """Stop the service docker container associated with this service.
  
          Args:
              force: if True, the service container will be forcefully stopped
          """
          container = self.container
          if not container:
              # service container is not running
              logger.debug(
                  "Docker container for service '%s' no longer running",
                  self,
              )
              return
  
          logger.debug("Stopping container for service '%s' ...", self)
          if force:
              container.kill()
              container.remove(force=True)
          else:
              container.stop()
              container.remove()
  
      def provision(self) -> None:
          """Provision the service."""
          self._start_container()
  
      def deprovision(self, force: bool = False) -> None:
          """Deprovision the service.
  
          Args:
              force: if True, the service container will be forcefully stopped
          """
          self._stop_daemon(force)
  
      def get_logs(
          self, follow: bool = False, tail: Optional[int] = None
      ) -> Generator[str, bool, None]:
          """Retrieve the service logs.
  
          Args:
              follow: if True, the logs will be streamed as they are written
              tail: only retrieve the last NUM lines of log output.
  
          Yields:
              A generator that can be accessed to get the service logs.
          """
          if not self.status.log_file or not os.path.exists(
              self.status.log_file
          ):
              return
  
          with open(self.status.log_file, "r") as f:
              if tail:
                  # TODO[ENG-864]: implement a more efficient tailing mechanism that
                  #   doesn't read the entire file
                  lines = f.readlines()[-tail:]
                  for line in lines:
                      yield line.rstrip("\n")
                  if not follow:
                      return
              line = ""
              while True:
                  partial_line = f.readline()
                  if partial_line:
                      line += partial_line
                      if line.endswith("\n"):
                          stop = yield line.rstrip("\n")
                          if stop:
                              break
                          line = ""
                  elif follow:
                      time.sleep(1)
                  else:
                      break
  
      @abstractmethod
      def run(self) -> None:
          """Run the containerized service logic associated with this service.
  
          Subclasses must implement this method to provide the containerized
          service functionality. This method will be executed in the context of
          the running container, not in the context of the process that calls the
          `start` method.
          """
  