content: #  Copyright (c) ZenML GmbH 2024. 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 for the BentoML container deployment service."""
  
  import os
  import sys
  from typing import Any, Dict, List, Optional, Union
  
  import bentoml
  import docker.errors as docker_errors
  from bentoml.client import Client
  
  from zenml.client import Client as ZenMLClient
  from zenml.constants import DEFAULT_LOCAL_SERVICE_IP_ADDRESS
  from zenml.enums import ServiceState
  from zenml.integrations.bentoml.constants import (
      BENTOML_DEFAULT_PORT,
      BENTOML_HEALTHCHECK_URL_PATH,
      BENTOML_PREDICTION_URL_PATH,
  )
  from zenml.logger import get_logger
  from zenml.models.v2.misc.service import ServiceType
  from zenml.services.container.container_service import (
      ContainerService,
      ContainerServiceConfig,
  )
  from zenml.services.container.container_service_endpoint import (
      ContainerServiceEndpoint,
      ContainerServiceEndpointConfig,
  )
  from zenml.services.service import BaseDeploymentService
  from zenml.services.service_endpoint import ServiceEndpointProtocol
  from zenml.services.service_monitor import (
      HTTPEndpointHealthMonitor,
      HTTPEndpointHealthMonitorConfig,
  )
  
  logger = get_logger(__name__)
  
  BENTOML_CONTAINER_DEPLOYMENT_SERVICE_NAME = "bentoml-container-deployment"
  
  
  class BentoMLContainerDeploymentConfig(ContainerServiceConfig):
      """BentoML container deployment configuration."""
  
      model_name: str
      model_uri: str
      bento_tag: str
      bento_uri: Optional[str] = None
      platform: Optional[str] = None
      image: str = ""
      image_tag: Optional[str] = None
      features: Optional[List[str]] = None
      file: Optional[str] = None
      apis: List[str] = []
      working_dir: Optional[str] = None
      workers: int = 1
      backlog: int = 2048
      host: Optional[str] = None
      port: Optional[int] = None
  
  
  class BentoMLContainerDeploymentEndpointConfig(ContainerServiceEndpointConfig):
      """BentoML container deployment service configuration.
  
      Attributes:
          prediction_url_path: URI subpath for prediction requests
      """
  
      prediction_url_path: str
  
  
  class BentoMLContainerDeploymentEndpoint(ContainerServiceEndpoint):
      """A service endpoint exposed by the BentoML container deployment service.
  
      Attributes:
          config: service endpoint configuration
      """
  
      config: BentoMLContainerDeploymentEndpointConfig
  
      @property
      def prediction_url(self) -> Optional[str]:
          """Gets the prediction URL for the endpoint.
  
          Returns:
              the prediction URL for the endpoint
          """
          uri = self.status.uri
          if not uri:
              return None
          return os.path.join(uri, self.config.prediction_url_path)
  
  
  class BentoMLContainerDeploymentService(
      ContainerService, BaseDeploymentService
  ):
      """BentoML container deployment service."""
  
      SERVICE_TYPE = ServiceType(
          name=BENTOML_CONTAINER_DEPLOYMENT_SERVICE_NAME,
          type="model-serving",
          flavor="bentoml",
          description="BentoML container prediction service",
          logo_url="https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_deployer/bentoml.png",
      )
  
      config: BentoMLContainerDeploymentConfig
      endpoint: BentoMLContainerDeploymentEndpoint
  
      def __init__(
          self,
          config: Union[BentoMLContainerDeploymentConfig, Dict[str, Any]],
          **attrs: Any,
      ) -> None:
          """Initialize the BentoML deployment service.
  
          Args:
              config: service configuration
              attrs: additional attributes to set on the service
          """
          # ensure that the endpoint is created before the service is initialized
          # TODO [ENG-700]: implement a service factory or builder for BentoML
          #   deployment services
          if (
              isinstance(config, BentoMLContainerDeploymentConfig)
              and "endpoint" not in attrs
          ):
              endpoint = BentoMLContainerDeploymentEndpoint(
                  config=BentoMLContainerDeploymentEndpointConfig(
                      protocol=ServiceEndpointProtocol.HTTP,
                      port=config.port or BENTOML_DEFAULT_PORT,
                      ip_address=config.host or DEFAULT_LOCAL_SERVICE_IP_ADDRESS,
                      prediction_url_path=BENTOML_PREDICTION_URL_PATH,
                  ),
                  monitor=HTTPEndpointHealthMonitor(
                      config=HTTPEndpointHealthMonitorConfig(
                          healthcheck_uri_path=BENTOML_HEALTHCHECK_URL_PATH,
                      )
                  ),
              )
              attrs["endpoint"] = endpoint
          super().__init__(config=config, **attrs)
  
      # override the is_running property to check if the bentoml container is running
      @property
      def is_running(self) -> bool:
          """Check if the service is currently running.
  
          This method will actively poll the external service to get its status
          and will return the result.
  
          Returns:
              True if the service is running and active (i.e. the endpoints are
              responsive, if any are configured), otherwise False.
          """
          self.update_status()
          return self.status.state == ServiceState.ACTIVE
  
      # override the container start method to use the root user
      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:
              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,
                  user="root",
                  labels={
                      "zenml-service-uuid": str(self.uuid),
                  },
                  working_dir="/home/bentoml/bento",
                  extra_hosts={"host.docker.internal": "host-gateway"},
              )
  
              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 _containerize_and_push_bento(self) -> None:
          """Containerize the bento and push it to the container registry.
  
          Raises:
              Exception: If the bento containerization fails.
          """
          zenml_client = ZenMLClient()
          container_registry = zenml_client.active_stack.container_registry
          # a tuple of config image and image tag
          if self.config.image and self.config.image_tag:
              image_tag = (self.config.image, self.config.image_tag)
          else:
              # if container registry is present in the stack, name the image
              # with the container registry uri, else name the image with the bento tag
              if container_registry:
                  image_name = (
                      f"{container_registry.config.uri}/{self.config.bento_tag}"
                  )
                  image_tag = (image_name,)  # type: ignore
                  self.config.image = image_name
              else:
                  # bentoml will use the bento tag as the name of the image
                  image_tag = (self.config.bento_tag,)  # type: ignore
                  self.config.image = self.config.bento_tag
          try:
              bentoml.container.build(
                  bento_tag=self.config.bento_tag,
                  backend="docker",  # hardcoding docker since container service only supports docker
                  image_tag=image_tag,
                  features=self.config.features,
                  file=self.config.file,
                  platform=self.config.platform,
              )
  
          except Exception as e:
              logger.error(f"Error containerizing the bento: {e}")
              raise e
  
          if container_registry:
              logger.info(
                  f"Pushing bento to container registry {container_registry.config.uri}"
              )
              # push the bento to the image registry
              container_registry.push_image(self.config.image)
          else:
              logger.warning(
                  "No container registry found in the active stack. "
                  "Please add a container registry to your stack to push "
                  "the bento to an image registry."
              )
  
      def provision(self) -> None:
          """Provision the service."""
          # containerize the bento
          self._containerize_and_push_bento()
          # run the container
          super().provision()
  
      def run(self) -> None:
          """Start the service.
  
          Raises:
              FileNotFoundError: If the bento file is not found.
              subprocess.CalledProcessError: If the bentoml serve command fails.
          """
          from bentoml._internal.service.loader import load
  
          logger.info("Starting BentoML container deployment service...")
  
          self.endpoint.prepare_for_start()
  
          if self.config.working_dir is None:
              if os.path.isdir(os.path.expanduser(self.config.bento_tag)):
                  self.config.working_dir = os.path.expanduser(
                      self.config.bento_tag
                  )
              else:
                  self.config.working_dir = "."
          if sys.path[0] != self.config.working_dir:
              sys.path.insert(0, self.config.working_dir)
  
          _ = load(bento_identifier=".", working_dir=self.config.working_dir)
          # run bentoml serve command inside the container
          # Use subprocess for better control and error handling
          import subprocess
  
          try:
              subprocess.run(["bentoml", "serve"], check=True)
          except subprocess.CalledProcessError as e:
              logger.error(f"Failed to start BentoML service: {e}")
              raise
          except FileNotFoundError:
              logger.error(
                  "BentoML command not found. Make sure it's installed and in the PATH."
              )
              raise
  
      @property
      def prediction_url(self) -> Optional[str]:
          """Get the URI where the http server is running.
  
          Returns:
              The URI where the http service can be accessed to get more information
              about the service and to make predictions.
          """
          if not self.is_running:
              return None
          return self.endpoint.prediction_url
  
      @property
      def prediction_apis_urls(self) -> Optional[List[str]]:
          """Get the URI where the prediction api services is answering requests.
  
          Returns:
              The URI where the prediction service apis can be contacted to process
              HTTP/REST inference requests, or None, if the service isn't running.
          """
          if not self.is_running:
              return None
  
          if self.config.apis:
              return [
                  f"{self.endpoint.prediction_url}/{api}"
                  for api in self.config.apis
              ]
          return None
  
      def predict(self, api_endpoint: str, data: Any) -> Any:
          """Make a prediction using the service.
  
          Args:
              data: data to make a prediction on
              api_endpoint: the api endpoint to make the prediction on
  
          Returns:
              The prediction result.
  
          Raises:
              Exception: if the service is not running
              ValueError: if the prediction endpoint is unknown.
          """
          if not self.is_running:
              raise Exception(
                  "BentoML prediction service is not running. "
                  "Please start the service before making predictions."
              )
          if self.endpoint.prediction_url is not None:
              client = Client.from_url(
                  self.endpoint.prediction_url.replace("http://", "").rstrip("/")
              )
              result = client.call(api_endpoint, data)
          else:
              raise ValueError("No endpoint known for prediction.")
          return result
  