zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/integrations/bentoml/model_deployers/bentoml_model_deployer.py

@@ -24,12 +24,20 @@ from zenml.integrations.bentoml.flavors.bentoml_model_deployer_flavor import (
     BentoMLModelDeployerConfig,
     BentoMLModelDeployerFlavor,
 )
-from zenml.integrations.bentoml.services.bentoml_deployment import (
-    BentoMLDeploymentConfig,
-    BentoMLDeploymentService,
+from zenml.integrations.bentoml.services.bentoml_container_deployment import (
+    BENTOML_CONTAINER_DEPLOYMENT_SERVICE_NAME,
+    BentoMLContainerDeploymentConfig,
+    BentoMLContainerDeploymentService,
+)
+from zenml.integrations.bentoml.services.bentoml_local_deployment import (
+    BENTOML_LOCAL_DEPLOYMENT_SERVICE_NAME,
+    BentoMLLocalDeploymentConfig,
+    BentoMLLocalDeploymentService,
 )
 from zenml.logger import get_logger
 from zenml.model_deployers import BaseModelDeployer, BaseModelDeployerFlavor
+from zenml.services.container.container_service import ContainerServiceStatus
+from zenml.services.local.local_service import LocalDaemonServiceStatus
 from zenml.services.service import BaseService, ServiceConfig
 from zenml.utils.io_utils import create_dir_recursive_if_not_exists
 
@@ -100,8 +108,8 @@ class BentoMLModelDeployer(BaseModelDeployer):
         return self._service_path
 
     @staticmethod
-    def get_model_server_info(  # type: ignore[override]
-        service_instance: "BentoMLDeploymentService",
+    def get_model_server_info(
+        service_instance: BaseService,
     ) -> Dict[str, Optional[str]]:
         """Return implementation specific information on the model server.
 
@@ -110,26 +118,61 @@ class BentoMLModelDeployer(BaseModelDeployer):
 
         Returns:
             A dictionary containing the model server information.
+
+        Raises:
+            ValueError: If the service type is not supported.
         """
+        if (
+            service_instance.SERVICE_TYPE.name
+            == BENTOML_CONTAINER_DEPLOYMENT_SERVICE_NAME
+        ):
+            service_instance = cast(
+                BentoMLContainerDeploymentService, service_instance
+            )
+        elif (
+            service_instance.SERVICE_TYPE.name
+            == BENTOML_LOCAL_DEPLOYMENT_SERVICE_NAME
+        ):
+            service_instance = cast(
+                BentoMLLocalDeploymentService, service_instance
+            )
+        else:
+            raise ValueError(
+                f"Unsupported service type: {service_instance.SERVICE_TYPE.name}"
+            )
+
         predictions_apis_urls = ""
-        if service_instance.prediction_apis_urls is not None:
+        if service_instance.prediction_apis_urls is not None:  # type: ignore
             predictions_apis_urls = ", ".join(
                 [
                     api
-                    for api in service_instance.prediction_apis_urls
+                    for api in service_instance.prediction_apis_urls  # type: ignore
                     if api is not None
                 ]
             )
 
+        service_config = service_instance.config
+        assert isinstance(
+            service_config,
+            (BentoMLLocalDeploymentConfig, BentoMLContainerDeploymentConfig),
+        )
+
+        service_status = service_instance.status
+        assert isinstance(
+            service_status, (ContainerServiceStatus, LocalDaemonServiceStatus)
+        )
+
         return {
             "HEALTH_CHECK_URL": service_instance.get_healthcheck_url(),
             "PREDICTION_URL": service_instance.get_prediction_url(),
-            "BENTO_TAG": service_instance.config.bento,
-            "MODEL_NAME": service_instance.config.model_name,
-            "MODEL_URI": service_instance.config.model_uri,
-            "BENTO_URI": service_instance.config.bento_uri,
-            "SERVICE_PATH": service_instance.status.runtime_path,
-            "DAEMON_PID": str(service_instance.status.pid),
+            "BENTO_TAG": service_config.bento_tag,
+            "MODEL_NAME": service_config.model_name,
+            "MODEL_URI": service_config.model_uri,
+            "BENTO_URI": service_config.bento_uri,
+            "SERVICE_PATH": service_status.runtime_path,
+            "DAEMON_PID": str(service_status.pid)
+            if hasattr(service_status, "pid")
+            else None,
             "PREDICTION_APIS_URLS": predictions_apis_urls,
         }
 
@@ -179,7 +222,6 @@ class BentoMLModelDeployer(BaseModelDeployer):
             The ZenML BentoML deployment service object that can be used to
             interact with the BentoML model http server.
         """
-        config = cast(BentoMLDeploymentConfig, config)
         service = self._create_new_service(
             id=id, timeout=timeout, config=config
         )
@@ -190,11 +232,20 @@ class BentoMLModelDeployer(BaseModelDeployer):
         self,
         timeout: int,
         force: bool,
-        existing_service: BentoMLDeploymentService,
+        existing_service: BaseService,
     ) -> None:
         # stop the older service
         existing_service.stop(timeout=timeout, force=force)
 
+        # assert that the service is either a BentoMLLocalDeploymentService or a BentoMLContainerDeploymentService
+        if not isinstance(
+            existing_service,
+            (BentoMLLocalDeploymentService, BentoMLContainerDeploymentService),
+        ):
+            raise ValueError(
+                f"Unsupported service type: {type(existing_service)}"
+            )
+
         # delete the old configuration file
         if existing_service.status.runtime_path:
             shutil.rmtree(existing_service.status.runtime_path)
@@ -203,8 +254,8 @@ class BentoMLModelDeployer(BaseModelDeployer):
     # of workers etc.the step implementation will create a new config using
     # all values from the user and add values like pipeline name, model_uri
     def _create_new_service(
-        self, id: UUID, timeout: int, config: BentoMLDeploymentConfig
-    ) -> BentoMLDeploymentService:
+        self, id: UUID, timeout: int, config: ServiceConfig
+    ) -> BaseService:
         """Creates a new BentoMLDeploymentService.
 
         Args:
@@ -216,11 +267,27 @@ class BentoMLModelDeployer(BaseModelDeployer):
         Returns:
             The BentoMLDeploymentService object that can be used to interact
             with the BentoML model server.
+
+        Raises:
+            ValueError: If the service type is not supported.
         """
+        assert isinstance(
+            config,
+            (BentoMLLocalDeploymentConfig, BentoMLContainerDeploymentConfig),
+        )
         # set the root runtime path with the stack component's UUID
         config.root_runtime_path = self.local_path
         # create a new service for the new model
-        service = BentoMLDeploymentService(uuid=id, config=config)
+        # if the config is of type BentoMLLocalDeploymentConfig, create a
+        # BentoMLLocalDeploymentService, otherwise create a
+        # BentoMLContainerDeploymentService
+        service: BaseService
+        if isinstance(config, BentoMLLocalDeploymentConfig):
+            service = BentoMLLocalDeploymentService(uuid=id, config=config)
+        elif isinstance(config, BentoMLContainerDeploymentConfig):
+            service = BentoMLContainerDeploymentService(uuid=id, config=config)
+        else:
+            raise ValueError(f"Unsupported service type: {type(config)}")
         service.start(timeout=timeout)
 
         return service
@@ -274,7 +341,6 @@ class BentoMLModelDeployer(BaseModelDeployer):
             timeout: Timeout in seconds to wait for the service to stop.
             force: If True, force the service to stop.
         """
-        service = cast(BentoMLDeploymentService, service)
         self._clean_up_existing_service(
             existing_service=service, timeout=timeout, force=force
         )

zenml/integrations/bentoml/services/__init__.py

@@ -13,9 +13,20 @@
 #  permissions and limitations under the License.
 """Initialization for BentoML services."""
 
-from zenml.integrations.bentoml.services.bentoml_deployment import (  # noqa
-    BentoMLDeploymentConfig,
-    BentoMLDeploymentService,
+from zenml.integrations.bentoml.services.bentoml_local_deployment import (  # noqa
+    BentoMLLocalDeploymentConfig,
+    BentoMLLocalDeploymentService,
 )
+from zenml.integrations.bentoml.services.bentoml_container_deployment import (  # noqa
+    BentoMLContainerDeploymentConfig,
+    BentoMLContainerDeploymentService,
+)
+from zenml.integrations.bentoml.services.deployment_type import BentoMLDeploymentType
 
-__all__ = ["BentoMLDeploymentConfig", "BentoMLDeploymentService"]
+__all__ = [
+    "BentoMLLocalDeploymentConfig",
+    "BentoMLLocalDeploymentService",
+    "BentoMLContainerDeploymentConfig",
+    "BentoMLContainerDeploymentService",
+    "BentoMLDeploymentType",
+]

zenml/integrations/bentoml/services/bentoml_container_deployment.py

@@ -0,0 +1,399 @@
+#  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.integrations.bentoml.constants import (
+    BENTOML_DEFAULT_PORT,
+    BENTOML_HEALTHCHECK_URL_PATH,
+    BENTOML_PREDICTION_URL_PATH,
+)
+from zenml.logger import get_logger
+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,
+)
+from zenml.services.service_status import ServiceState
+from zenml.services.service_type import ServiceType
+
+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