zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/bentoml/model_deployers/bentoml_model_deployer.py

@@ -12,27 +12,32 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Implementation of the BentoML Model Deployer."""
+
 import os
 import shutil
-from pathlib import Path
-from typing import ClassVar, Dict, List, Optional, Type, cast
+from typing import ClassVar, Dict, Optional, Type, cast
 from uuid import UUID
 
 from zenml.config.global_config import GlobalConfiguration
 from zenml.constants import DEFAULT_SERVICE_START_STOP_TIMEOUT
-from zenml.integrations.bentoml.constants import BENTOML_DEFAULT_PORT
 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 import ServiceRegistry
-from zenml.services.local.local_service import SERVICE_DAEMON_CONFIG_FILE_NAME
+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
 
@@ -43,9 +48,9 @@ class BentoMLModelDeployer(BaseModelDeployer):
     """BentoML model deployer stack component implementation."""
 
     NAME: ClassVar[str] = "BentoML"
-    FLAVOR: ClassVar[
-        Type[BaseModelDeployerFlavor]
-    ] = BentoMLModelDeployerFlavor
+    FLAVOR: ClassVar[Type[BaseModelDeployerFlavor]] = (
+        BentoMLModelDeployerFlavor
+    )
 
     _service_path: Optional[str] = None
 
@@ -103,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.
 
@@ -113,32 +118,68 @@ 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 {
-            "PREDICTION_URL": service_instance.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),
+            "HEALTH_CHECK_URL": service_instance.get_healthcheck_url(),
+            "PREDICTION_URL": service_instance.get_prediction_url(),
+            "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,
         }
 
-    def deploy_model(
+    def perform_deploy_model(
         self,
+        id: UUID,
         config: ServiceConfig,
-        replace: bool = False,
         timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT,
     ) -> BaseService:
         """Create a new BentoML deployment service or update an existing one.
@@ -170,10 +211,8 @@ class BentoMLModelDeployer(BaseModelDeployer):
         and the others are deleted.
 
         Args:
+            id: the UUID of the BentoML model deployer.
             config: the configuration of the model to be deployed with BentoML.
-            replace: set this flag to True to find and update an equivalent
-                BentoML deployment server with the new model instead of
-                creating and starting a new deployment server.
             timeout: the timeout in seconds to wait for the BentoML server
                 to be provisioned and successfully started or updated. If set
                 to 0, the method will return immediately after the BentoML
@@ -183,60 +222,30 @@ 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 = None
-
-        # if replace is True, remove all existing services
-        if replace is True:
-            existing_services = self.find_model_server(
-                pipeline_name=config.pipeline_name,
-                pipeline_step_name=config.pipeline_step_name,
-                model_name=config.model_name,
-            )
-
-            for existing_service in existing_services:
-                if service is None:
-                    # keep the most recently created service
-                    service = cast(BentoMLDeploymentService, existing_service)
-                try:
-                    # delete the older services and don't wait for them to
-                    # be deprovisioned
-                    self._clean_up_existing_service(
-                        existing_service=cast(
-                            BentoMLDeploymentService, existing_service
-                        ),
-                        timeout=timeout,
-                        force=True,
-                    )
-                except RuntimeError:
-                    # ignore errors encountered while stopping old services
-                    pass
-        if service:
-            logger.info(
-                f"Updating an existing BentoML deployment service: {service}"
-            )
-
-            # set the root runtime path with the stack component's UUID
-            config.root_runtime_path = self.local_path
-            service.stop(timeout=timeout, force=True)
-            service.update(config)
-            service.start(timeout=timeout)
-        else:
-            # create a new BentoMLDeploymentService instance
-            service = self._create_new_service(timeout, config)
-            logger.info(f"Created a new BentoML deployment service: {service}")
-
-        return cast(BaseService, service)
+        service = self._create_new_service(
+            id=id, timeout=timeout, config=config
+        )
+        logger.info(f"Created a new BentoML deployment service: {service}")
+        return service
 
     def _clean_up_existing_service(
         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)
@@ -245,213 +254,93 @@ 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, timeout: int, config: BentoMLDeploymentConfig
-    ) -> BentoMLDeploymentService:
+        self, id: UUID, timeout: int, config: ServiceConfig
+    ) -> BaseService:
         """Creates a new BentoMLDeploymentService.
 
         Args:
-            timeout: the timeout in seconds to wait for the BentoML http server
+            id: the ID of the BentoML deployment service to be created or updated.
+            timeout: the timeout in seconds to wait for the BentoML server
                 to be provisioned and successfully started or updated.
             config: the configuration of the model to be deployed with BentoML.
 
         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(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
 
-    def find_model_server(
-        self,
-        running: bool = False,
-        service_uuid: Optional[UUID] = None,
-        pipeline_name: Optional[str] = None,
-        run_name: Optional[str] = None,
-        pipeline_step_name: Optional[str] = None,
-        model_name: Optional[str] = None,
-        model_uri: Optional[str] = None,
-        model_type: Optional[str] = None,
-    ) -> List[BaseService]:
-        """Finds one or more model servers that match the given criteria.
-
-        Args:
-            running: If true, only running services will be returned.
-            service_uuid: The UUID of the service that was originally used
-                to deploy the model.
-            pipeline_name: Name of the pipeline that the deployed model was part
-                of.
-            run_name: ID of the pipeline run which the deployed model
-                was part of.
-            pipeline_step_name: The name of the pipeline model deployment step
-                that deployed the model.
-            model_name: Name of the deployed model.
-            model_uri: URI of the deployed model.
-            model_type: Type/format of the deployed model. Not used in this
-                BentoML case.
-
-        Returns:
-            One or more Service objects representing model servers that match
-            the input search criteria.
-
-        Raises:
-            TypeError: if any of the input arguments are of an invalid type.
-        """
-        services = []
-        config = BentoMLDeploymentConfig(
-            model_name=model_name or "",
-            bento="",
-            port=BENTOML_DEFAULT_PORT,
-            model_uri=model_uri or "",
-            working_dir="",
-            pipeline_name=pipeline_name or "",
-            pipeline_run_id=run_name or "",
-            run_name=run_name or "",
-            pipeline_step_name=pipeline_step_name or "",
-        )
-
-        # find all services that match the input criteria
-        for root, _, files in os.walk(self.local_path):
-            if service_uuid and Path(root).name != str(service_uuid):
-                continue
-            for file in files:
-                if file == SERVICE_DAEMON_CONFIG_FILE_NAME:
-                    service_config_path = os.path.join(root, file)
-                    logger.debug(
-                        "Loading service daemon configuration from %s",
-                        service_config_path,
-                    )
-                    existing_service_config = None
-                    with open(service_config_path, "r") as f:
-                        existing_service_config = f.read()
-                    existing_service = (
-                        ServiceRegistry().load_service_from_json(
-                            existing_service_config
-                        )
-                    )
-                    if not isinstance(
-                        existing_service, BentoMLDeploymentService
-                    ):
-                        raise TypeError(
-                            f"Expected service type BentoMLDeploymentService but got "
-                            f"{type(existing_service)} instead"
-                        )
-                    existing_service.update_status()
-                    if self._matches_search_criteria(existing_service, config):
-                        if not running or existing_service.is_running:
-                            services.append(
-                                cast(BaseService, existing_service)
-                            )
-
-        return services
-
-    def _matches_search_criteria(
-        self,
-        existing_service: BentoMLDeploymentService,
-        config: BentoMLDeploymentConfig,
-    ) -> bool:
-        """Returns true if a service matches the input criteria.
-
-        If any of the values in the input criteria are None, they are ignored.
-        This allows listing services just by common pipeline names or step
-        names, etc.
-
-        Args:
-            existing_service: The materialized Service instance derived from
-                the config of the older (existing) service
-            config: The BentoMlDeploymentConfig object passed to the
-                deploy_model function holding parameters of the new service
-                to be created.
-
-        Returns:
-            True if the service matches the input criteria.
-        """
-        existing_service_config = existing_service.config
-
-        # check if the existing service matches the input criteria
-        if (
-            (
-                not config.pipeline_name
-                or existing_service_config.pipeline_name
-                == config.pipeline_name
-            )
-            and (
-                not config.model_name
-                or existing_service_config.model_name == config.model_name
-            )
-            and (
-                not config.pipeline_step_name
-                or existing_service_config.pipeline_step_name
-                == config.pipeline_step_name
-            )
-            and (
-                not config.run_name
-                or existing_service_config.run_name == config.run_name
-            )
-        ):
-            return True
-
-        return False
-
-    def stop_model_server(
+    def perform_stop_model(
         self,
-        uuid: UUID,
+        service: BaseService,
         timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT,
         force: bool = False,
-    ) -> None:
+    ) -> BaseService:
         """Method to stop a model server.
 
         Args:
-            uuid: UUID of the model server to stop.
+            service: The service to stop.
             timeout: Timeout in seconds to wait for the service to stop.
             force: If True, force the service to stop.
-        """
-        # get list of all services
-        existing_services = self.find_model_server(service_uuid=uuid)
 
-        # if the service exists, stop it
-        if existing_services:
-            existing_services[0].stop(timeout=timeout, force=force)
+        Returns:
+            The stopped service.
+        """
+        service.stop(timeout=timeout, force=force)
+        return service
 
-    def start_model_server(
-        self, uuid: UUID, timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT
-    ) -> None:
+    def perform_start_model(
+        self,
+        service: BaseService,
+        timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT,
+    ) -> BaseService:
         """Method to start a model server.
 
         Args:
-            uuid: UUID of the model server to start.
+            service: The service to start.
             timeout: Timeout in seconds to wait for the service to start.
-        """
-        # get list of all services
-        existing_services = self.find_model_server(service_uuid=uuid)
 
-        # if the service exists, start it
-        if existing_services:
-            existing_services[0].start(timeout=timeout)
+        Returns:
+            The started service.
+        """
+        service.start(timeout=timeout)
+        return service
 
-    def delete_model_server(
+    def perform_delete_model(
         self,
-        uuid: UUID,
+        service: BaseService,
         timeout: int = DEFAULT_SERVICE_START_STOP_TIMEOUT,
         force: bool = False,
     ) -> None:
         """Method to delete all configuration of a model server.
 
         Args:
-            uuid: UUID of the model server to delete.
+            service: The service to delete.
             timeout: Timeout in seconds to wait for the service to stop.
             force: If True, force the service to stop.
         """
-        # get list of all services
-        existing_services = self.find_model_server(service_uuid=uuid)
-
-        # if the service exists, clean it up
-        if existing_services:
-            service = cast(BentoMLDeploymentService, existing_services[0])
-            self._clean_up_existing_service(
-                existing_service=service, timeout=timeout, force=force
-            )
+        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",
+]