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

zenml/integrations/huggingface/model_deployers/huggingface_model_deployer.py

@@ -0,0 +1,247 @@
+#  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 of the Hugging Face Model Deployer."""
+
+from typing import ClassVar, Dict, Optional, Tuple, Type, cast
+from uuid import UUID
+
+from zenml.analytics.enums import AnalyticsEvent
+from zenml.analytics.utils import track_handler
+from zenml.client import Client
+from zenml.integrations.huggingface import HUGGINGFACE_SERVICE_ARTIFACT
+from zenml.integrations.huggingface.flavors.huggingface_model_deployer_flavor import (
+    HuggingFaceModelDeployerConfig,
+    HuggingFaceModelDeployerFlavor,
+)
+from zenml.integrations.huggingface.services.huggingface_deployment import (
+    HuggingFaceDeploymentService,
+    HuggingFaceServiceConfig,
+)
+from zenml.logger import get_logger
+from zenml.model_deployers import BaseModelDeployer
+from zenml.model_deployers.base_model_deployer import (
+    DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+    BaseModelDeployerFlavor,
+)
+from zenml.services import BaseService, ServiceConfig
+from zenml.stack.stack import Stack
+from zenml.stack.stack_validator import StackValidator
+
+logger = get_logger(__name__)
+
+
+class HuggingFaceModelDeployer(BaseModelDeployer):
+    """Hugging Face endpoint model deployer."""
+
+    NAME: ClassVar[str] = "HuggingFace"
+    FLAVOR: ClassVar[Type[BaseModelDeployerFlavor]] = (
+        HuggingFaceModelDeployerFlavor
+    )
+
+    @property
+    def config(self) -> HuggingFaceModelDeployerConfig:
+        """Config class for the Hugging Face Model deployer settings class.
+
+        Returns:
+            The configuration.
+        """
+        return cast(HuggingFaceModelDeployerConfig, self._config)
+
+    @property
+    def validator(self) -> Optional[StackValidator]:
+        """Validates the stack.
+
+        Returns:
+            A validator that checks that the stack contains a remote artifact
+            store.
+        """
+
+        def _validate_if_secret_or_token_is_present(
+            stack: "Stack",
+        ) -> Tuple[bool, str]:
+            """Check if secret or token is present in the stack.
+
+            Args:
+                stack: The stack to validate.
+
+            Returns:
+                A tuple with a boolean indicating whether the stack is valid
+                and a message describing the validation result.
+            """
+            return bool(self.config.token or self.config.secret_name), (
+                "The Hugging Face model deployer requires either a secret name"
+                " or a token to be present in the stack."
+            )
+
+        return StackValidator(
+            custom_validation_function=_validate_if_secret_or_token_is_present,
+        )
+
+    def _create_new_service(
+        self, id: UUID, timeout: int, config: HuggingFaceServiceConfig
+    ) -> HuggingFaceDeploymentService:
+        """Creates a new Hugging FaceDeploymentService.
+
+        Args:
+            id: the UUID of the model to be deployed with Hugging Face model deployer.
+            timeout: the timeout in seconds to wait for the Hugging Face inference endpoint
+                to be provisioned and successfully started or updated.
+            config: the configuration of the model to be deployed with Hugging Face model deployer.
+
+        Returns:
+            The HuggingFaceServiceConfig object that can be used to interact
+            with the Hugging Face inference endpoint.
+        """
+        # create a new service for the new model
+        service = HuggingFaceDeploymentService(uuid=id, config=config)
+
+        logger.info(
+            f"Creating an artifact {HUGGINGFACE_SERVICE_ARTIFACT} with service instance attached as metadata."
+            " If there's an active pipeline and/or model this artifact will be associated with it."
+        )
+        service.start(timeout=timeout)
+        return service
+
+    def _clean_up_existing_service(
+        self,
+        timeout: int,
+        force: bool,
+        existing_service: HuggingFaceDeploymentService,
+    ) -> None:
+        """Stop existing services.
+
+        Args:
+            timeout: the timeout in seconds to wait for the Hugging Face
+                deployment to be stopped.
+            force: if True, force the service to stop
+            existing_service: Existing Hugging Face deployment service
+        """
+        # stop the older service
+        existing_service.stop(timeout=timeout, force=force)
+
+    def perform_deploy_model(
+        self,
+        id: UUID,
+        config: ServiceConfig,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+    ) -> BaseService:
+        """Create a new Hugging Face deployment service or update an existing one.
+
+        This should serve the supplied model and deployment configuration.
+
+        Args:
+            id: the UUID of the model to be deployed with Hugging Face.
+            config: the configuration of the model to be deployed with Hugging Face.
+            timeout: the timeout in seconds to wait for the Hugging Face endpoint
+                to be provisioned and successfully started or updated. If set
+                to 0, the method will return immediately after the Hugging Face
+                server is provisioned, without waiting for it to fully start.
+
+        Returns:
+            The ZenML Hugging Face deployment service object that can be used to
+            interact with the remote Hugging Face inference endpoint server.
+        """
+        with track_handler(AnalyticsEvent.MODEL_DEPLOYED) as analytics_handler:
+            config = cast(HuggingFaceServiceConfig, config)
+            # create a new HuggingFaceDeploymentService instance
+            service = self._create_new_service(
+                id=id, timeout=timeout, config=config
+            )
+            logger.info(
+                f"Creating a new Hugging Face inference endpoint service: {service}"
+            )
+            # Add telemetry with metadata that gets the stack metadata and
+            # differentiates between pure model and custom code deployments
+            stack = Client().active_stack
+            stack_metadata = {
+                component_type.value: component.flavor
+                for component_type, component in stack.components.items()
+            }
+            analytics_handler.metadata = {
+                "store_type": Client().zen_store.type.value,
+                **stack_metadata,
+            }
+
+        return service
+
+    def perform_stop_model(
+        self,
+        service: BaseService,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+        force: bool = False,
+    ) -> BaseService:
+        """Method to stop a model server.
+
+        Args:
+            service: The service to stop.
+            timeout: Timeout in seconds to wait for the service to stop.
+            force: If True, force the service to stop.
+
+        Returns:
+            The stopped service.
+        """
+        service.stop(timeout=timeout, force=force)
+        return service
+
+    def perform_start_model(
+        self,
+        service: BaseService,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+    ) -> BaseService:
+        """Method to start a model server.
+
+        Args:
+            service: The service to start.
+            timeout: Timeout in seconds to wait for the service to start.
+
+        Returns:
+            The started service.
+        """
+        service.start(timeout=timeout)
+        return service
+
+    def perform_delete_model(
+        self,
+        service: BaseService,
+        timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT,
+        force: bool = False,
+    ) -> None:
+        """Method to delete all configuration of a model server.
+
+        Args:
+            service: The service to delete.
+            timeout: Timeout in seconds to wait for the service to stop.
+            force: If True, force the service to stop.
+        """
+        service = cast(HuggingFaceDeploymentService, service)
+        self._clean_up_existing_service(
+            existing_service=service, timeout=timeout, force=force
+        )
+
+    @staticmethod
+    def get_model_server_info(  # type: ignore[override]
+        service_instance: "HuggingFaceDeploymentService",
+    ) -> Dict[str, Optional[str]]:
+        """Return implementation specific information that might be relevant to the user.
+
+        Args:
+            service_instance: Instance of a HuggingFaceDeploymentService
+
+        Returns:
+            Model server information.
+        """
+        return {
+            "PREDICTION_URL": service_instance.get_prediction_url(),
+            "HEALTH_CHECK_URL": service_instance.get_healthcheck_url(),
+        }

zenml/integrations/huggingface/services/__init__.py

@@ -0,0 +1,19 @@
+#  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.
+"""Initialization of the Hugging Face Service."""
+
+from zenml.integrations.huggingface.services.huggingface_deployment import (  # noqa
+    HuggingFaceDeploymentService,
+    HuggingFaceServiceConfig,
+)

zenml/integrations/huggingface/services/huggingface_deployment.py

@@ -0,0 +1,292 @@
+#  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 of the Hugging Face Deployment service."""
+
+from typing import Any, Generator, Optional, Tuple
+
+from huggingface_hub import (
+    InferenceClient,
+    InferenceEndpoint,
+    InferenceEndpointError,
+    InferenceEndpointStatus,
+    create_inference_endpoint,
+    get_inference_endpoint,
+)
+from huggingface_hub.utils import HfHubHTTPError
+from pydantic import Field
+
+from zenml.client import Client
+from zenml.integrations.huggingface.flavors.huggingface_model_deployer_flavor import (
+    HuggingFaceBaseConfig,
+)
+from zenml.logger import get_logger
+from zenml.services import ServiceState, ServiceStatus, ServiceType
+from zenml.services.service import BaseDeploymentService, ServiceConfig
+
+logger = get_logger(__name__)
+
+POLLING_TIMEOUT = 1200
+UUID_SLICE_LENGTH: int = 8
+
+
+class HuggingFaceServiceConfig(HuggingFaceBaseConfig, ServiceConfig):
+    """Hugging Face service configurations."""
+
+
+class HuggingFaceServiceStatus(ServiceStatus):
+    """Hugging Face service status."""
+
+
+class HuggingFaceDeploymentService(BaseDeploymentService):
+    """Hugging Face model deployment service.
+
+    Attributes:
+        SERVICE_TYPE: a service type descriptor with information describing
+            the Hugging Face deployment service class
+        config: service configuration
+    """
+
+    SERVICE_TYPE = ServiceType(
+        name="huggingface-deployment",
+        type="model-serving",
+        flavor="huggingface",
+        description="Hugging Face inference endpoint prediction service",
+    )
+    config: HuggingFaceServiceConfig
+    status: HuggingFaceServiceStatus = Field(
+        default_factory=lambda: HuggingFaceServiceStatus()
+    )
+
+    def __init__(self, config: HuggingFaceServiceConfig, **attrs: Any):
+        """Initialize the Hugging Face deployment service.
+
+        Args:
+            config: service configuration
+            attrs: additional attributes to set on the service
+        """
+        super().__init__(config=config, **attrs)
+
+    def get_token(self) -> str:
+        """Get the Hugging Face token.
+
+        Raises:
+            ValueError: If token not found.
+
+        Returns:
+            Hugging Face token.
+        """
+        client = Client()
+        token = None
+        if self.config.secret_name:
+            secret = client.get_secret(self.config.secret_name)
+            token = secret.secret_values["token"]
+        else:
+            from zenml.integrations.huggingface.model_deployers.huggingface_model_deployer import (
+                HuggingFaceModelDeployer,
+            )
+
+            model_deployer = client.active_stack.model_deployer
+            if not isinstance(model_deployer, HuggingFaceModelDeployer):
+                raise ValueError(
+                    "HuggingFaceModelDeployer is not active in the stack."
+                )
+            token = model_deployer.config.token or None
+        if not token:
+            raise ValueError("Token not found.")
+        return token
+
+    @property
+    def hf_endpoint(self) -> InferenceEndpoint:
+        """Get the deployed Hugging Face inference endpoint.
+
+        Returns:
+            Huggingface inference endpoint.
+        """
+        return get_inference_endpoint(
+            name=self._generate_an_endpoint_name(),
+            token=self.get_token(),
+            namespace=self.config.namespace,
+        )
+
+    @property
+    def prediction_url(self) -> Optional[str]:
+        """The prediction URI exposed by the prediction service.
+
+        Returns:
+            The prediction URI exposed by the prediction service, or None if
+            the service is not yet ready.
+        """
+        return self.hf_endpoint.url if self.is_running else None
+
+    @property
+    def inference_client(self) -> InferenceClient:
+        """Get the Hugging Face InferenceClient from Inference Endpoint.
+
+        Returns:
+            Hugging Face inference client.
+        """
+        return self.hf_endpoint.client
+
+    def provision(self) -> None:
+        """Provision or update remote Hugging Face deployment instance.
+
+        Raises:
+            Exception: If any unexpected error while creating inference endpoint.
+        """
+        try:
+            # Attempt to create and wait for the inference endpoint
+            hf_endpoint = create_inference_endpoint(
+                name=self._generate_an_endpoint_name(),
+                repository=self.config.repository,
+                framework=self.config.framework,
+                accelerator=self.config.accelerator,
+                instance_size=self.config.instance_size,
+                instance_type=self.config.instance_type,
+                region=self.config.region,
+                vendor=self.config.vendor,
+                account_id=self.config.account_id,
+                min_replica=self.config.min_replica,
+                max_replica=self.config.max_replica,
+                revision=self.config.revision,
+                task=self.config.task,
+                custom_image=self.config.custom_image,
+                type=self.config.endpoint_type,
+                token=self.get_token(),
+                namespace=self.config.namespace,
+            ).wait(timeout=POLLING_TIMEOUT)
+
+        except Exception as e:
+            self.status.update_state(
+                new_state=ServiceState.ERROR, error=str(e)
+            )
+            # Catch-all for any other unexpected errors
+            raise Exception(
+                f"An unexpected error occurred while provisioning the Hugging Face inference endpoint: {e}"
+            )
+
+        # Check if the endpoint URL is available after provisioning
+        if hf_endpoint.url:
+            logger.info(
+                f"Hugging Face inference endpoint successfully deployed and available. Endpoint URL: {hf_endpoint.url}"
+            )
+        else:
+            logger.error(
+                "Failed to start Hugging Face inference endpoint service: No URL available, please check the Hugging Face console for more details."
+            )
+
+    def check_status(self) -> Tuple[ServiceState, str]:
+        """Check the the current operational state of the Hugging Face deployment.
+
+        Returns:
+            The operational state of the Hugging Face deployment and a message
+            providing additional information about that state (e.g. a
+            description of the error, if one is encountered).
+        """
+        try:
+            status = self.hf_endpoint.status
+            if status == InferenceEndpointStatus.RUNNING:
+                return (ServiceState.ACTIVE, "")
+
+            elif status == InferenceEndpointStatus.SCALED_TO_ZERO:
+                return (
+                    ServiceState.SCALED_TO_ZERO,
+                    "Hugging Face Inference Endpoint is scaled to zero, but still running. It will be started on demand.",
+                )
+
+            elif status == InferenceEndpointStatus.FAILED:
+                return (
+                    ServiceState.ERROR,
+                    "Hugging Face Inference Endpoint deployment is inactive or not found",
+                )
+            elif status == InferenceEndpointStatus.PENDING:
+                return (ServiceState.PENDING_STARTUP, "")
+            return (ServiceState.PENDING_STARTUP, "")
+        except (InferenceEndpointError, HfHubHTTPError):
+            return (
+                ServiceState.INACTIVE,
+                "Hugging Face Inference Endpoint deployment is inactive or not found",
+            )
+
+    def deprovision(self, force: bool = False) -> None:
+        """Deprovision the remote Hugging Face deployment instance.
+
+        Args:
+            force: if True, the remote deployment instance will be
+                forcefully deprovisioned.
+        """
+        try:
+            self.hf_endpoint.delete()
+        except HfHubHTTPError:
+            logger.error(
+                "Hugging Face Inference Endpoint is deleted or cannot be found."
+            )
+
+    def predict(self, data: "Any", max_new_tokens: int) -> "Any":
+        """Make a prediction using the service.
+
+        Args:
+            data: input data
+            max_new_tokens: Number of new tokens to generate
+
+        Returns:
+            The prediction result.
+
+        Raises:
+            Exception: if the service is not running
+            NotImplementedError: if task is not supported.
+        """
+        if not self.is_running:
+            raise Exception(
+                "Hugging Face endpoint inference service is not running. "
+                "Please start the service before making predictions."
+            )
+        if self.prediction_url is not None:
+            if self.hf_endpoint.task == "text-generation":
+                result = self.inference_client.task_generation(
+                    data, max_new_tokens=max_new_tokens
+                )
+        else:
+            # TODO: Add support for all different supported tasks
+            raise NotImplementedError(
+                "Tasks other than text-generation is not implemented."
+            )
+        return result
+
+    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.
+
+        Returns:
+            A generator that can be accessed to get the service logs.
+        """
+        logger.info(
+            "Hugging Face Endpoints provides access to the logs of "
+            "your Endpoints through the UI in the “Logs” tab of your Endpoint"
+        )
+        return  # type: ignore
+
+    def _generate_an_endpoint_name(self) -> str:
+        """Generate a unique name for the Hugging Face Inference Endpoint.
+
+        Returns:
+            A unique name for the Hugging Face Inference Endpoint.
+        """
+        return (
+            f"{self.config.service_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}"
+        )

zenml/integrations/huggingface/steps/__init__.py

@@ -0,0 +1,21 @@
+#  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.
+"""Initialization for Hugging Face model deployer step."""
+
+from zenml.integrations.huggingface.steps.huggingface_deployer import (
+    huggingface_model_deployer_step,
+)
+from zenml.integrations.huggingface.steps.accelerate_runner import (
+    run_with_accelerate,
+)