From 2fb8a7d8b18b169418f669e3866ce848fe3c2a7d Mon Sep 17 00:00:00 2001 From: Safoine El khabich <34200873+safoinme@users.noreply.github.com> Date: Mon, 3 Jun 2024 12:58:48 +0100 Subject: [PATCH 01/23] initial commit on vertex ai deployer and model registry --- src/zenml/integrations/gcp/__init__.py | 4 + .../integrations/gcp/flavors/__init__.py | 6 + .../flavors/vertex_model_deployer_flavor.py | 149 ++++++++++ .../gcp/model_deployers/__init__.py | 20 ++ .../model_deployers/vertex_model_deployer.py | 242 ++++++++++++++++ .../integrations/gcp/services/__init__.py | 19 ++ .../gcp/services/vertex_deployment.py | 263 ++++++++++++++++++ 7 files changed, 703 insertions(+) create mode 100644 src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py create mode 100644 src/zenml/integrations/gcp/model_deployers/__init__.py create mode 100644 src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py create mode 100644 src/zenml/integrations/gcp/services/__init__.py create mode 100644 src/zenml/integrations/gcp/services/vertex_deployment.py diff --git a/src/zenml/integrations/gcp/__init__.py b/src/zenml/integrations/gcp/__init__.py index 0c955c8eeec..5d5da217977 100644 --- a/src/zenml/integrations/gcp/__init__.py +++ b/src/zenml/integrations/gcp/__init__.py @@ -33,6 +33,10 @@ GCP_VERTEX_ORCHESTRATOR_FLAVOR = "vertex" GCP_VERTEX_STEP_OPERATOR_FLAVOR = "vertex" +# Model deployer constants +VERTEX_MODEL_DEPLOYER_FLAVOR = "vertex" +VERTEX_SERVICE_ARTIFACT = "vertex_deployment_service" + # Service connector constants GCP_CONNECTOR_TYPE = "gcp" GCP_RESOURCE_TYPE = "gcp-generic" diff --git a/src/zenml/integrations/gcp/flavors/__init__.py b/src/zenml/integrations/gcp/flavors/__init__.py index 73bb6259aa5..1328ec75b2d 100644 --- a/src/zenml/integrations/gcp/flavors/__init__.py +++ b/src/zenml/integrations/gcp/flavors/__init__.py @@ -29,6 +29,10 @@ VertexStepOperatorConfig, VertexStepOperatorFlavor, ) +from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( + VertexModelDeployerConfig, + VertexModelDeployerFlavor, +) __all__ = [ "GCPArtifactStoreFlavor", @@ -39,4 +43,6 @@ "VertexOrchestratorConfig", "VertexStepOperatorFlavor", "VertexStepOperatorConfig", + "VertexModelDeployerFlavor", + "VertexModelDeployerConfig", ] diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py new file mode 100644 index 00000000000..cb798cc19c8 --- /dev/null +++ b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py @@ -0,0 +1,149 @@ +# 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. +"""Vertex AI model deployer flavor.""" + +from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type + +from pydantic import BaseModel + +from zenml.integrations.gcp import VERTEX_MODEL_DEPLOYER_FLAVOR +from zenml.model_deployers.base_model_deployer import ( + BaseModelDeployerConfig, + BaseModelDeployerFlavor, +) +from zenml.utils.secret_utils import SecretField + +if TYPE_CHECKING: + from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( + VertexModelDeployer, + ) + + +class VertexBaseConfig(BaseModel): + """Vertex AI Inference Endpoint configuration.""" + + location: Optional[str] = None + version: Optional[str] = None + serving_container_image_uri: Optional[str] = None + artifact_uri: Optional[str] = None + model_id: Optional[str] = None + is_default_version: Optional[bool] = None + serving_container_command: Optional[Sequence[str]] = None, + serving_container_args: Optional[Sequence[str]] = None, + serving_container_environment_variables: Optional[ + Dict[str, str] + ] = None, + serving_container_ports: Optional[Sequence[int]] = None, + serving_container_grpc_ports: Optional[Sequence[int]] = None, + deployed_model_display_name: Optional[str] = None + traffic_percentage: Optional[int] = 0 + traffic_split: Optional[Dict[str, int]] = None + machine_type: Optional[str] = None + accelerator_type: Optional[str] = None + accelerator_count: Optional[int] = None + min_replica_count: Optional[int] = None + max_replica_count: Optional[int] = None + service_account: Optional[str] = None + metadata: Optional[Dict[str, str]] = None + network: Optional[str] = None + encryption_spec_key_name: Optional[str] = None + sync=True, + deploy_request_timeout: Optional[int] = None + autoscaling_target_cpu_utilization: Optional[float] = None + autoscaling_target_accelerator_duty_cycle: Optional[float] = None + enable_access_logging: Optional[bool] = None + disable_container_logging: Optional[bool] = None + + + + + +class VertexModelDeployerConfig( + BaseModelDeployerConfig, VertexBaseConfig +): + """Configuration for the Vertex AI model deployer. + + Attributes: + model_name: The name of the model. + project_id: The project ID. + location: The location of the model. + version: The version of the model. + """ + + # The namespace to list endpoints for. Set to `"*"` to list all endpoints + # from all namespaces (i.e. personal namespace and all orgs the user belongs to). + model_name: str + + + +class VertexModelDeployerFlavor(BaseModelDeployerFlavor): + """Vertex AI Endpoint model deployer flavor.""" + + @property + def name(self) -> str: + """Name of the flavor. + + Returns: + The name of the flavor. + """ + return VERTEX_MODEL_DEPLOYER_FLAVOR + + @property + def docs_url(self) -> Optional[str]: + """A url to point at docs explaining this flavor. + + Returns: + A flavor docs url. + """ + return self.generate_default_docs_url() + + @property + def sdk_docs_url(self) -> Optional[str]: + """A url to point at SDK docs explaining this flavor. + + Returns: + A flavor SDK docs url. + """ + return self.generate_default_sdk_docs_url() + + @property + def logo_url(self) -> str: + """A url to represent the flavor in the dashboard. + + Returns: + The flavor logo. + """ + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_registry/vertexai.png" + + @property + def config_class(self) -> Type[VertexModelDeployerConfig]: + """Returns `VertexModelDeployerConfig` config class. + + Returns: + The config class. + """ + return VertexModelDeployerConfig + + @property + def implementation_class(self) -> Type["VertexModelDeployer"]: + """Implementation class for this flavor. + + Returns: + The implementation class. + """ + from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( + VertexModelDeployer, + ) + + return VertexModelDeployer diff --git a/src/zenml/integrations/gcp/model_deployers/__init__.py b/src/zenml/integrations/gcp/model_deployers/__init__.py new file mode 100644 index 00000000000..99ee319f891 --- /dev/null +++ b/src/zenml/integrations/gcp/model_deployers/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) ZenML GmbH 2023. 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: +# +# http://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 Vertex AI model deployers.""" + +from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( # noqa + VertexMdelDeployer, +) + +__all__ = ["VertexMdelDeployer"] diff --git a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py new file mode 100644 index 00000000000..35a20890a0e --- /dev/null +++ b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py @@ -0,0 +1,242 @@ +# 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 Vertex AI 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.gcp import VERTEX_SERVICE_ARTIFACT +from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( + VertexModelDeployerConfig, + VertexModelDeployerFlavor, +) +from zenml.integrations.gcp.google_credentials_mixin import GoogleCredentialsMixin +from zenml.integrations.gcp.services.vertex_deployment import ( + VertexDeploymentService, + VertexServiceConfig, +) +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 VertexModelDeployer(BaseModelDeployer, GoogleCredentialsMixin): + """Vertex implementation of the BaseModelDeployer.""" + + @property + def config(self) -> VertexModelDeployerConfig: + """Config class for the Vertex AI Model deployer settings class. + + Returns: + The configuration. + """ + return cast(VertexModelDeployerConfig, 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 Vertex AI 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: VertexServiceConfig + ) -> VertexDeploymentService: + """Creates a new VertexDeploymentService. + + Args: + id: the UUID of the model to be deployed with Vertex AI model deployer. + timeout: the timeout in seconds to wait for the Vertex AI inference endpoint + to be provisioned and successfully started or updated. + config: the configuration of the model to be deployed with Vertex AI model deployer. + + Returns: + The VertexServiceConfig object that can be used to interact + with the Vertex AI inference endpoint. + """ + # create a new service for the new model + service = VertexDeploymentService(uuid=id, config=config) + + logger.info( + f"Creating an artifact {VERTEX_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: VertexDeploymentService, + ) -> None: + """Stop existing services. + + Args: + timeout: the timeout in seconds to wait for the Vertex AI + deployment to be stopped. + force: if True, force the service to stop + existing_service: Existing Vertex AI 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 Vertex AI 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 Vertex AI. + config: the configuration of the model to be deployed with Vertex AI. + timeout: the timeout in seconds to wait for the Vertex AI endpoint + to be provisioned and successfully started or updated. If set + to 0, the method will return immediately after the Vertex AI + server is provisioned, without waiting for it to fully start. + + Returns: + The ZenML Vertex AI deployment service object that can be used to + interact with the remote Vertex AI inference endpoint server. + """ + with track_handler(AnalyticsEvent.MODEL_DEPLOYED) as analytics_handler: + config = cast(VertexServiceConfig, config) + # create a new VertexDeploymentService instance + service = self._create_new_service( + id=id, timeout=timeout, config=config + ) + logger.info( + f"Creating a new Vertex AI 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(VertexDeploymentService, service) + self._clean_up_existing_service( + existing_service=service, timeout=timeout, force=force + ) + + @staticmethod + def get_model_server_info( # type: ignore[override] + service_instance: "VertexDeploymentService", + ) -> Dict[str, Optional[str]]: + """Return implementation specific information that might be relevant to the user. + + Args: + service_instance: Instance of a VertexDeploymentService + + Returns: + Model server information. + """ + return { + "PREDICTION_URL": service_instance.get_prediction_url(), + "HEALTH_CHECK_URL": service_instance.get_healthcheck_url(), + } diff --git a/src/zenml/integrations/gcp/services/__init__.py b/src/zenml/integrations/gcp/services/__init__.py new file mode 100644 index 00000000000..b9f858b5302 --- /dev/null +++ b/src/zenml/integrations/gcp/services/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) ZenML GmbH 2023. 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: +# +# http://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 MLflow Service.""" + +from zenml.integrations.mlflow.services.mlflow_deployment import ( # noqa + MLFlowDeploymentConfig, + MLFlowDeploymentService, +) diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py new file mode 100644 index 00000000000..98df9d28e46 --- /dev/null +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -0,0 +1,263 @@ +# 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 Vertex AI Deployment service.""" + +from typing import TYPE_CHECKING, Any, Generator, Optional, Tuple, cast + +from pydantic import Field + +from google.cloud import aiplatform + +from zenml.client import Client +from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( + VertexBaseConfig, +) +from zenml.logger import get_logger +from zenml.services import ServiceState, ServiceStatus, ServiceType +from zenml.services.service import BaseDeploymentService, ServiceConfig + +if TYPE_CHECKING: + from google.auth.credentials import Credentials + +logger = get_logger(__name__) + +POLLING_TIMEOUT = 1200 +UUID_SLICE_LENGTH: int = 8 + + +class VertexServiceConfig(VertexBaseConfig, ServiceConfig): + """Vertex AI service configurations.""" + + +class VertexServiceStatus(ServiceStatus): + """Vertex AI service status.""" + + +class VertexDeploymentService(BaseDeploymentService): + """Vertex AI model deployment service. + + Attributes: + SERVICE_TYPE: a service type descriptor with information describing + the Vertex AI deployment service class + config: service configuration + """ + + SERVICE_TYPE = ServiceType( + name="vertex-deployment", + type="model-serving", + flavor="vertex", + description="Vertex AI inference endpoint prediction service", + ) + config: VertexServiceConfig + status: VertexServiceStatus = Field( + default_factory=lambda: VertexServiceStatus() + ) + + def __init__(self, config: VertexServiceConfig, credentials: Tuple["Credentials", str], **attrs: Any): + """Initialize the Vertex AI deployment service. + + Args: + config: service configuration + attrs: additional attributes to set on the service + """ + super().__init__(config=config, **attrs) + self._config = config + self._project, self._credentials = credentials # Store credentials as a private attribute + + @property + def config(self) -> VertexServiceConfig: + """Returns the config of the deployment service. + + Returns: + The config of the deployment service. + """ + return cast(VertexServiceConfig, self._config) + + def get_token(self) -> str: + """Get the Vertex AI token. + + Raises: + ValueError: If token not found. + + Returns: + Vertex AI 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.gcp.model_deployers.vertex_model_deployer import ( + VertexModelDeployer, + ) + + model_deployer = client.active_stack.model_deployer + if not isinstance(model_deployer, VertexModelDeployer): + raise ValueError( + "VertexModelDeployer 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 vertex_model(self) -> aiplatform.Model: + """Get the deployed Vertex AI inference endpoint. + + Returns: + Vertex AI inference endpoint. + """ + return aiplatform.Model(f"projects/{self.__project}/locations/{self.config.location}/models/{self.config.model_id}") + + @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 + + def provision(self) -> None: + """Provision or update remote Vertex AI deployment instance. + + Raises: + Exception: If any unexpected error while creating inference endpoint. + """ + try: + # Attempt to create and wait for the inference endpoint + vertex_endpoint = self.vertex_model.deploy( + deployed_model_display_name=self.config.deployed_model_display_name, + traffic_percentage=self.config.traffic_percentage, + traffic_split=self.config.traffic_split, + machine_type=self.config.machine_type, + min_replica_count=self.config.min_replica_count, + max_replica_count=self.config.max_replica_count, + accelerator_type=self.config.accelerator_type, + accelerator_count=self.config.accelerator_count, + service_account=self.config.service_account, + metadata=self.config.metadata, + deploy_request_timeout=self.config.deploy_request_timeout, + autoscaling_target_cpu_utilization=self.config.autoscaling_target_cpu_utilization, + autoscaling_target_accelerator_duty_cycle=self.config.autoscaling_target_accelerator_duty_cycle, + enable_access_logging=self.config.enable_access_logging, + disable_container_logging=self.config.disable_container_logging, + encryption_spec_key_name=self.config.encryption_spec_key_name, + deploy_request_timeout=self.config.deploy_request_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 Vertex AI inference endpoint: {e}" + ) + + # Check if the endpoint URL is available after provisioning + if hf_endpoint.url: + logger.info( + f"Vertex AI inference endpoint successfully deployed and available. Endpoint URL: {hf_endpoint.url}" + ) + else: + logger.error( + "Failed to start Vertex AI inference endpoint service: No URL available, please check the Vertex AI console for more details." + ) + + def check_status(self) -> Tuple[ServiceState, str]: + """Check the the current operational state of the Vertex AI deployment. + + Returns: + The operational state of the Vertex AI deployment and a message + providing additional information about that state (e.g. a + description of the error, if one is encountered). + """ + pass + + def deprovision(self, force: bool = False) -> None: + """Deprovision the remote Vertex AI deployment instance. + + Args: + force: if True, the remote deployment instance will be + forcefully deprovisioned. + """ + try: + self.vertex_model.undeploy() + except HfHubHTTPError: + logger.error( + "Vertex AI 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( + "Vertex AI 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( + "Vertex AI 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 Vertex AI Inference Endpoint. + + Returns: + A unique name for the Vertex AI Inference Endpoint. + """ + return ( + f"{self.config.service_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}" + ) From c03f2a01a0a50a536195c4d0386936a1a08dcbce Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 6 Jun 2024 08:06:03 +0100 Subject: [PATCH 02/23] vertex model --- .../model_registries/vertex_model_registry.py | 314 ++++++++++++++++++ 1 file changed, 314 insertions(+) create mode 100644 src/zenml/integrations/gcp/model_registries/vertex_model_registry.py diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py new file mode 100644 index 00000000000..48403965ca5 --- /dev/null +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -0,0 +1,314 @@ +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple, cast + +from google.cloud import aiplatform +from google.cloud.aiplatform import Model, ModelRegistry, ModelVersion +from google.cloud.aiplatform.exceptions import NotFound + +from zenml.enums import StackComponentType +from zenml.stack import Flavor, StackComponent +from zenml.stack.stack_component import StackComponentConfig +from zenml.model_registries.base_model_registry import ( + BaseModelRegistry, + ModelRegistryModelMetadata, + ModelVersionStage, + RegisteredModel, + RegistryModelVersion, +) +from zenml.stack.stack_validator import StackValidator +from zenml.logger import get_logger + +logger = get_logger(__name__) + +class VertexAIModelRegistry(BaseModelRegistry): + """Register models using Vertex AI.""" + + def __init__(self): + super().__init__() + aiplatform.init() # Initialize the Vertex AI SDK + + @property + def config(self) -> StackComponentConfig: + """Returns the config of the model registries.""" + return cast(StackComponentConfig, self._config) + + def register_model( + self, + name: str, + description: Optional[str] = None, + metadata: Optional[Dict[str, str]] = None, + ) -> RegisteredModel: + """Register a model to the Vertex AI model registry.""" + try: + model = Model( + display_name=name, + description=description, + labels=metadata + ) + model.upload() + return RegisteredModel(name=name, description=description, metadata=metadata) + except Exception as e: + raise RuntimeError(f"Failed to register model: {str(e)}") + + def delete_model( + self, + name: str, + ) -> None: + """Delete a model from the Vertex AI model registry.""" + try: + model = Model(model_name=name) + model.delete() + except NotFound: + raise KeyError(f"Model with name {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to delete model: {str(e)}") + + def update_model( + self, + name: str, + description: Optional[str] = None, + metadata: Optional[Dict[str, str]] = None, + remove_metadata: Optional[List[str]] = None, + ) -> RegisteredModel: + """Update a model in the Vertex AI model registry.""" + try: + model = Model(model_name=name) + if description: + model.update(description=description) + if metadata: + for key, value in metadata.items(): + model.labels[key] = value + if remove_metadata: + for key in remove_metadata: + if key in model.labels: + del model.labels[key] + model.update() + return self.get_model(name) + except NotFound: + raise KeyError(f"Model with name {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to update model: {str(e)}") + + def get_model(self, name: str) -> RegisteredModel: + """Get a model from the Vertex AI model registry.""" + try: + model = Model(model_name=name) + model_resource = model.gca_resource + return RegisteredModel( + name=model_resource.display_name, + description=model_resource.description, + metadata=model_resource.labels + ) + except NotFound: + raise KeyError(f"Model with name {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to get model: {str(e)}") + + def list_models( + self, + name: Optional[str] = None, + metadata: Optional[Dict[str, str]] = None, + ) -> List[RegisteredModel]: + """List models in the Vertex AI model registry.""" + filter_expression = "" + if name: + filter_expression += f"display_name={name}" + if metadata: + for key, value in metadata.items(): + filter_expression += f"labels.{key}={value} " + try: + models = Model.list(filter=filter_expression) + return [ + RegisteredModel( + name=model.display_name, + description=model.description, + metadata=model.labels + ) + for model in models + ] + except Exception as e: + raise RuntimeError(f"Failed to list models: {str(e)}") + + def register_model_version( + self, + name: str, + version: Optional[str] = None, + model_source_uri: Optional[str] = None, + description: Optional[str] = None, + metadata: Optional[ModelRegistryModelMetadata] = None, + **kwargs: Any, + ) -> RegistryModelVersion: + """Register a model version to the Vertex AI model registry.""" + try: + model = Model(model_name=name) + version_info = model.upload_version( + display_name=version, + description=description, + artifact_uri=model_source_uri, + labels=metadata.dict() if metadata else None + ) + return RegistryModelVersion( + version=version_info.version_id, + model_source_uri=model_source_uri, + model_format="Custom", + registered_model=self.get_model(name), + description=description, + created_at=version_info.create_time, + last_updated_at=version_info.update_time, + stage=ModelVersionStage.NONE, + metadata=metadata + ) + except NotFound: + raise KeyError(f"Model with name {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to register model version: {str(e)}") + + def delete_model_version( + self, + name: str, + version: str, + ) -> None: + """Delete a model version from the Vertex AI model registry.""" + try: + model = Model(model_name=name) + version_info = ModelVersion(model_name=f"{name}@{version}") + version_info.delete() + except NotFound: + raise KeyError(f"Model version {version} of model {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to delete model version: {str(e)}") + + def update_model_version( + self, + name: str, + version: str, + description: Optional[str] = None, + metadata: Optional[ModelRegistryModelMetadata] = None, + remove_metadata: Optional[List[str]] = None, + stage: Optional[ModelVersionStage] = None, + ) -> RegistryModelVersion: + """Update a model version in the Vertex AI model registry.""" + try: + model_version = ModelVersion(model_name=f"{name}@{version}") + if description: + model_version.update(description=description) + if metadata: + for key, value in metadata.dict().items(): + model_version.labels[key] = value + if remove_metadata: + for key in remove_metadata: + if key in model_version.labels: + del model_version.labels[key] + model_version.update() + if stage: + # Handle stage update if needed + pass + return self.get_model_version(name, version) + except NotFound: + raise KeyError(f"Model version {version} of model {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to update model version: {str(e)}") + + def get_model_version( + self, name: str, version: str + ) -> RegistryModelVersion: + """Get a model version from the Vertex AI model registry.""" + try: + model_version = ModelVersion(model_name=f"{name}@{version}") + return RegistryModelVersion( + version=model_version.version_id, + model_source_uri=model_version.gca_resource.artifact_uri, + model_format="Custom", + registered_model=self.get_model(name), + description=model_version.description, + created_at=model_version.create_time, + last_updated_at=model_version.update_time, + stage=ModelVersionStage.NONE, + metadata=ModelRegistryModelMetadata(**model_version.labels) + ) + except NotFound: + raise KeyError(f"Model version {version} of model {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to get model version: {str(e)}") + + def list_model_versions( + self, + name: Optional[str] = None, + model_source_uri: Optional[str] = None, + metadata: Optional[ModelRegistryModelMetadata] = None, + stage: Optional[ModelVersionStage] = None, + count: Optional[int] = None, + created_after: Optional[datetime] = None, + created_before: Optional[datetime] = None, + order_by_date: Optional[str] = None, + **kwargs: Any, + ) -> List[RegistryModelVersion]: + """List model versions from the Vertex AI model registry.""" + filter_expression = "" + if name: + filter_expression += f"display_name={name}" + if metadata: + for key, value in metadata.dict().items(): + filter_expression += f"labels.{key}={value} " + try: + model = Model(model_name=name) + versions = model.list_versions(filter=filter_expression) + return [ + RegistryModelVersion( + version=v.version_id, + model_source_uri=v.artifact_uri, + model_format="Custom", + registered_model=self.get_model(name), + description=v.description, + created_at=v.create_time, + + + last_updated_at=v.update_time, + stage=ModelVersionStage.NONE, + metadata=ModelRegistryModelMetadata(**v.labels) + ) + for v in versions + ] + except Exception as e: + raise RuntimeError(f"Failed to list model versions: {str(e)}") + + def load_model_version( + self, + name: str, + version: str, + **kwargs: Any, + ) -> Any: + """Load a model version from the Vertex AI model registry.""" + try: + model_version = ModelVersion(model_name=f"{name}@{version}") + return model_version + except NotFound: + raise KeyError(f"Model version {version} of model {name} does not exist.") + except Exception as e: + raise RuntimeError(f"Failed to load model version: {str(e)}") + + def get_model_uri_artifact_store( + self, + model_version: RegistryModelVersion, + ) -> str: + """Get the model URI artifact store.""" + return model_version.model_source_uri + + +class VertexAIModelRegistryFlavor(Flavor): + """Base class for all ZenML model registry flavors.""" + + @property + def type(self) -> StackComponentType: + """Type of the flavor.""" + return StackComponentType.MODEL_REGISTRY + + @property + def config_class(self) -> Type[StackComponentConfig]: + """Config class for this flavor.""" + return StackComponentConfig + + @property + def implementation_class(self) -> Type[StackComponent]: + """Returns the implementation class for this flavor.""" + return VertexAIModelRegistry From 4eeeb277cd2d76cfbdc178234e37cf22ce8acf75 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Mon, 15 Jul 2024 13:52:27 +0100 Subject: [PATCH 03/23] vertex deployer --- .../integrations/gcp/model_deployers/vertex_model_deployer.py | 4 +++- src/zenml/integrations/gcp/services/vertex_deployment.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py index 35a20890a0e..f16167c25a8 100644 --- a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py +++ b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py @@ -24,7 +24,9 @@ VertexModelDeployerConfig, VertexModelDeployerFlavor, ) -from zenml.integrations.gcp.google_credentials_mixin import GoogleCredentialsMixin +from zenml.integrations.gcp.google_credentials_mixin import ( + GoogleCredentialsMixin, +) from zenml.integrations.gcp.services.vertex_deployment import ( VertexDeploymentService, VertexServiceConfig, diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py index 98df9d28e46..07fd8ed2260 100644 --- a/src/zenml/integrations/gcp/services/vertex_deployment.py +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -170,7 +170,7 @@ def provision(self) -> None: ) # Check if the endpoint URL is available after provisioning - if hf_endpoint.url: + if vertex_endpoint. logger.info( f"Vertex AI inference endpoint successfully deployed and available. Endpoint URL: {hf_endpoint.url}" ) @@ -260,4 +260,4 @@ def _generate_an_endpoint_name(self) -> str: """ return ( f"{self.config.service_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}" - ) + ) \ No newline at end of file From 7c0ca3f15eda0ae8db3324f251eb1569b19a6fdf Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 18 Sep 2024 13:26:07 +0100 Subject: [PATCH 04/23] vertex registry code --- src/zenml/integrations/gcp/__init__.py | 5 + .../integrations/gcp/flavors/__init__.py | 6 + .../flavors/vertex_model_deployer_flavor.py | 32 +- .../flavors/vertex_model_registry_flavor.py | 130 +++++++ .../gcp/model_deployers/__init__.py | 4 +- .../model_deployers/vertex_model_deployer.py | 174 ++++++---- .../gcp/model_registries/__init__.py | 20 ++ .../model_registries/vertex_model_registry.py | 239 +++++++------ .../integrations/gcp/services/__init__.py | 8 +- .../gcp/services/vertex_deployment.py | 322 +++++++++++------- 10 files changed, 600 insertions(+), 340 deletions(-) create mode 100644 src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py create mode 100644 src/zenml/integrations/gcp/model_registries/__init__.py diff --git a/src/zenml/integrations/gcp/__init__.py b/src/zenml/integrations/gcp/__init__.py index 43bf09edda8..3c9de9a9348 100644 --- a/src/zenml/integrations/gcp/__init__.py +++ b/src/zenml/integrations/gcp/__init__.py @@ -34,6 +34,7 @@ GCP_VERTEX_STEP_OPERATOR_FLAVOR = "vertex" # Model deployer constants +VERTEX_MODEL_REGISTRY_FLAVOR = "vertex" VERTEX_MODEL_DEPLOYER_FLAVOR = "vertex" VERTEX_SERVICE_ARTIFACT = "vertex_deployment_service" @@ -76,6 +77,8 @@ def flavors(cls) -> List[Type[Flavor]]: GCPImageBuilderFlavor, VertexOrchestratorFlavor, VertexStepOperatorFlavor, + VertexModelDeployerFlavor, + VertexAIModelRegistryFlavor, ) return [ @@ -83,6 +86,8 @@ def flavors(cls) -> List[Type[Flavor]]: GCPImageBuilderFlavor, VertexOrchestratorFlavor, VertexStepOperatorFlavor, + VertexAIModelRegistryFlavor, + VertexModelDeployerFlavor, ] diff --git a/src/zenml/integrations/gcp/flavors/__init__.py b/src/zenml/integrations/gcp/flavors/__init__.py index 1328ec75b2d..cecf637cefd 100644 --- a/src/zenml/integrations/gcp/flavors/__init__.py +++ b/src/zenml/integrations/gcp/flavors/__init__.py @@ -33,6 +33,10 @@ VertexModelDeployerConfig, VertexModelDeployerFlavor, ) +from zenml.integrations.gcp.flavors.vertex_model_registry_flavor import ( + VertexAIModelRegistryConfig, + VertexAIModelRegistryFlavor, +) __all__ = [ "GCPArtifactStoreFlavor", @@ -45,4 +49,6 @@ "VertexStepOperatorConfig", "VertexModelDeployerFlavor", "VertexModelDeployerConfig", + "VertexAIModelRegistryFlavor", + "VertexAIModelRegistryConfig", ] diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py index cb798cc19c8..85d4bd52485 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Vertex AI model deployer flavor.""" -from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type +from typing import TYPE_CHECKING, Dict, Optional, Sequence, Type from pydantic import BaseModel @@ -22,7 +22,6 @@ BaseModelDeployerConfig, BaseModelDeployerFlavor, ) -from zenml.utils.secret_utils import SecretField if TYPE_CHECKING: from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( @@ -39,13 +38,11 @@ class VertexBaseConfig(BaseModel): artifact_uri: Optional[str] = None model_id: Optional[str] = None is_default_version: Optional[bool] = None - serving_container_command: Optional[Sequence[str]] = None, - serving_container_args: Optional[Sequence[str]] = None, - serving_container_environment_variables: Optional[ - Dict[str, str] - ] = None, - serving_container_ports: Optional[Sequence[int]] = None, - serving_container_grpc_ports: Optional[Sequence[int]] = None, + serving_container_command: Optional[Sequence[str]] = None + serving_container_args: Optional[Sequence[str]] = None + serving_container_environment_variables: Optional[Dict[str, str]] = None + serving_container_ports: Optional[Sequence[int]] = None + serving_container_grpc_ports: Optional[Sequence[int]] = None deployed_model_display_name: Optional[str] = None traffic_percentage: Optional[int] = 0 traffic_split: Optional[Dict[str, int]] = None @@ -58,33 +55,26 @@ class VertexBaseConfig(BaseModel): metadata: Optional[Dict[str, str]] = None network: Optional[str] = None encryption_spec_key_name: Optional[str] = None - sync=True, + sync: Optional[bool] = True deploy_request_timeout: Optional[int] = None autoscaling_target_cpu_utilization: Optional[float] = None autoscaling_target_accelerator_duty_cycle: Optional[float] = None enable_access_logging: Optional[bool] = None disable_container_logging: Optional[bool] = None - - - -class VertexModelDeployerConfig( - BaseModelDeployerConfig, VertexBaseConfig -): +class VertexModelDeployerConfig(BaseModelDeployerConfig, VertexBaseConfig): """Configuration for the Vertex AI model deployer. Attributes: - model_name: The name of the model. project_id: The project ID. location: The location of the model. - version: The version of the model. """ # The namespace to list endpoints for. Set to `"*"` to list all endpoints # from all namespaces (i.e. personal namespace and all orgs the user belongs to). - model_name: str - + project_id: str + location: Optional[str] = None class VertexModelDeployerFlavor(BaseModelDeployerFlavor): @@ -124,7 +114,7 @@ def logo_url(self) -> str: Returns: The flavor logo. """ - return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/model_registry/vertexai.png" + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/artifact_store/gcp.png" @property def config_class(self) -> Type[VertexModelDeployerConfig]: diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py new file mode 100644 index 00000000000..1c13e95a95e --- /dev/null +++ b/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py @@ -0,0 +1,130 @@ +# Copyright (c) ZenML GmbH 2023. 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. +"""VertexAI model registry flavor.""" + +from typing import TYPE_CHECKING, Optional, Type + +from zenml.config.base_settings import BaseSettings +from zenml.integrations.gcp import ( + GCP_RESOURCE_TYPE, + VERTEX_MODEL_REGISTRY_FLAVOR +) +from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( + VertexBaseConfig, +) +from zenml.integrations.gcp.google_credentials_mixin import ( + GoogleCredentialsConfigMixin, +) +from zenml.models import ServiceConnectorRequirements +from zenml.model_registries.base_model_registry import ( + BaseModelRegistryConfig, + BaseModelRegistryFlavor, +) + +if TYPE_CHECKING: + from zenml.integrations.gcp.model_registries import ( + VertexAIModelRegistry, + ) + +class VertexAIModelRegistrySettings(BaseSettings): + """Settings for the VertexAI model registry.""" + + location: str + + +class VertexAIModelRegistryConfig( + BaseModelRegistryConfig, + GoogleCredentialsConfigMixin, + VertexAIModelRegistrySettings +): + """Configuration for the VertexAI model registry.""" + + +class VertexAIModelRegistryFlavor(BaseModelRegistryFlavor): + """Model registry flavor for VertexAI models.""" + + @property + def name(self) -> str: + """Name of the flavor. + + Returns: + The name of the flavor. + """ + return VERTEX_MODEL_REGISTRY_FLAVOR + + @property + def service_connector_requirements( + self, + ) -> Optional[ServiceConnectorRequirements]: + """Service connector resource requirements for service connectors. + + Specifies resource requirements that are used to filter the available + service connector types that are compatible with this flavor. + + Returns: + Requirements for compatible service connectors, if a service + connector is required for this flavor. + """ + return ServiceConnectorRequirements( + resource_type=GCP_RESOURCE_TYPE, + ) + + @property + def docs_url(self) -> Optional[str]: + """A url to point at docs explaining this flavor. + + Returns: + A flavor docs url. + """ + return self.generate_default_docs_url() + + @property + def sdk_docs_url(self) -> Optional[str]: + """A url to point at SDK docs explaining this flavor. + + Returns: + A flavor SDK docs url. + """ + return self.generate_default_sdk_docs_url() + + @property + def logo_url(self) -> str: + """A url to represent the flavor in the dashboard. + + Returns: + The flavor logo. + """ + return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/artifact_store/gcp.png" + + @property + def config_class(self) -> Type[VertexAIModelRegistryConfig]: + """Returns `VertexAIModelRegistryConfig` config class. + + Returns: + The config class. + """ + return VertexAIModelRegistryConfig + + @property + def implementation_class(self) -> Type["VertexAIModelRegistry"]: + """Implementation class for this flavor. + + Returns: + The implementation class. + """ + from zenml.integrations.gcp.model_registries import ( + VertexAIModelRegistry, + ) + + return VertexAIModelRegistry diff --git a/src/zenml/integrations/gcp/model_deployers/__init__.py b/src/zenml/integrations/gcp/model_deployers/__init__.py index 99ee319f891..203f57c096f 100644 --- a/src/zenml/integrations/gcp/model_deployers/__init__.py +++ b/src/zenml/integrations/gcp/model_deployers/__init__.py @@ -14,7 +14,7 @@ """Initialization of the Vertex AI model deployers.""" from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( # noqa - VertexMdelDeployer, + VertexModelDeployer, ) -__all__ = ["VertexMdelDeployer"] +__all__ = ["VertexModelDeployer"] diff --git a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py index f16167c25a8..17ad388588d 100644 --- a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py +++ b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py @@ -13,7 +13,7 @@ # permissions and limitations under the License. """Implementation of the Vertex AI Model Deployer.""" -from typing import ClassVar, Dict, Optional, Tuple, Type, cast +from typing import ClassVar, Dict, List, Optional, Tuple, Type, cast from uuid import UUID from zenml.analytics.enums import AnalyticsEvent @@ -24,9 +24,6 @@ VertexModelDeployerConfig, VertexModelDeployerFlavor, ) -from zenml.integrations.gcp.google_credentials_mixin import ( - GoogleCredentialsMixin, -) from zenml.integrations.gcp.services.vertex_deployment import ( VertexDeploymentService, VertexServiceConfig, @@ -43,12 +40,18 @@ logger = get_logger(__name__) -class VertexModelDeployer(BaseModelDeployer, GoogleCredentialsMixin): - """Vertex implementation of the BaseModelDeployer.""" + +class VertexModelDeployer(BaseModelDeployer): + """Vertex AI endpoint model deployer.""" + + NAME: ClassVar[str] = "Vertex AI" + FLAVOR: ClassVar[Type["BaseModelDeployerFlavor"]] = ( + VertexModelDeployerFlavor + ) @property def config(self) -> VertexModelDeployerConfig: - """Config class for the Vertex AI Model deployer settings class. + """Returns the `VertexModelDeployerConfig` config. Returns: The configuration. @@ -60,14 +63,13 @@ def validator(self) -> Optional[StackValidator]: """Validates the stack. Returns: - A validator that checks that the stack contains a remote artifact - store. + A validator that checks that the stack contains required GCP components. """ - def _validate_if_secret_or_token_is_present( + def _validate_gcp_stack( stack: "Stack", ) -> Tuple[bool, str]: - """Check if secret or token is present in the stack. + """Check if GCP components are properly configured in the stack. Args: stack: The stack to validate. @@ -76,33 +78,34 @@ def _validate_if_secret_or_token_is_present( 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 Vertex AI model deployer requires either a secret name" - " or a token to be present in the stack." - ) + if not self.config.project_id or not self.config.location: + return False, ( + "The Vertex AI model deployer requires a GCP project and " + "location to be specified in the configuration." + ) + return True, "Stack is valid for Vertex AI model deployment." return StackValidator( - custom_validation_function=_validate_if_secret_or_token_is_present, + custom_validation_function=_validate_gcp_stack, ) - def _create_new_service( - self, id: UUID, timeout: int, config: VertexServiceConfig + def _create_deployment_service( + self, id: UUID, timeout: int, config: VertexModelDeployerConfig ) -> VertexDeploymentService: - """Creates a new VertexDeploymentService. + """Creates a new DatabricksDeploymentService. Args: - id: the UUID of the model to be deployed with Vertex AI model deployer. - timeout: the timeout in seconds to wait for the Vertex AI inference endpoint + id: the UUID of the model to be deployed with Databricks model deployer. + timeout: the timeout in seconds to wait for the Databricks inference endpoint to be provisioned and successfully started or updated. - config: the configuration of the model to be deployed with Vertex AI model deployer. + config: the configuration of the model to be deployed with Databricks model deployer. Returns: - The VertexServiceConfig object that can be used to interact - with the Vertex AI inference endpoint. + The VertexModelDeployerConfig object that can be used to interact + with the Databricks inference endpoint. """ # create a new service for the new model service = VertexDeploymentService(uuid=id, config=config) - logger.info( f"Creating an artifact {VERTEX_SERVICE_ARTIFACT} with service instance attached as metadata." " If there's an active pipeline and/or model this artifact will be associated with it." @@ -110,66 +113,48 @@ def _create_new_service( service.start(timeout=timeout) return service - def _clean_up_existing_service( - self, - timeout: int, - force: bool, - existing_service: VertexDeploymentService, - ) -> None: - """Stop existing services. - - Args: - timeout: the timeout in seconds to wait for the Vertex AI - deployment to be stopped. - force: if True, force the service to stop - existing_service: Existing Vertex AI 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 Vertex AI deployment service or update an existing one. - - This should serve the supplied model and deployment configuration. + """Deploy a model to Vertex AI. Args: - id: the UUID of the model to be deployed with Vertex AI. - config: the configuration of the model to be deployed with Vertex AI. - timeout: the timeout in seconds to wait for the Vertex AI endpoint - to be provisioned and successfully started or updated. If set - to 0, the method will return immediately after the Vertex AI - server is provisioned, without waiting for it to fully start. + id: the UUID of the service to be created. + config: the configuration of the model to be deployed. + timeout: the timeout for the deployment operation. Returns: - The ZenML Vertex AI deployment service object that can be used to - interact with the remote Vertex AI inference endpoint server. + The ZenML Vertex AI deployment service object. """ with track_handler(AnalyticsEvent.MODEL_DEPLOYED) as analytics_handler: config = cast(VertexServiceConfig, config) - # create a new VertexDeploymentService instance - service = self._create_new_service( - id=id, timeout=timeout, config=config - ) + service = self._create_deployment_service(id=id, config=config) logger.info( - f"Creating a new Vertex AI inference endpoint service: {service}" + f"Creating a new Vertex AI deployment service: {service}" ) - # Add telemetry with metadata that gets the stack metadata and - # differentiates between pure model and custom code deployments - stack = Client().active_stack + service.start(timeout=timeout) + + client = Client() + 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, + "store_type": client.zen_store.type.value, **stack_metadata, } + # Create a service artifact + client.create_artifact( + name=VERTEX_SERVICE_ARTIFACT, + artifact_store_id=client.active_stack.artifact_store.id, + producer=service, + ) + return service def perform_stop_model( @@ -178,7 +163,7 @@ def perform_stop_model( timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT, force: bool = False, ) -> BaseService: - """Method to stop a model server. + """Stop a Vertex AI deployment service. Args: service: The service to stop. @@ -196,7 +181,7 @@ def perform_start_model( service: BaseService, timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT, ) -> BaseService: - """Method to start a model server. + """Start a Vertex AI deployment service. Args: service: The service to start. @@ -214,7 +199,7 @@ def perform_delete_model( timeout: int = DEFAULT_DEPLOYMENT_START_STOP_TIMEOUT, force: bool = False, ) -> None: - """Method to delete all configuration of a model server. + """Delete a Vertex AI deployment service. Args: service: The service to delete. @@ -222,23 +207,66 @@ def perform_delete_model( force: If True, force the service to stop. """ service = cast(VertexDeploymentService, service) - self._clean_up_existing_service( - existing_service=service, timeout=timeout, force=force - ) + service.stop(timeout=timeout, force=force) + service.delete() @staticmethod - def get_model_server_info( # type: ignore[override] + def get_model_server_info( service_instance: "VertexDeploymentService", ) -> Dict[str, Optional[str]]: - """Return implementation specific information that might be relevant to the user. + """Get information about the deployed model server. Args: - service_instance: Instance of a VertexDeploymentService + service_instance: The VertexDeploymentService instance. Returns: - Model server information. + A dictionary containing information about the model server. """ return { - "PREDICTION_URL": service_instance.get_prediction_url(), + "PREDICTION_URL": service_instance.prediction_url, "HEALTH_CHECK_URL": service_instance.get_healthcheck_url(), } + + def find_model_server( + self, + running: Optional[bool] = None, + 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_version: Optional[str] = None, + ) -> List[BaseService]: + """Find deployed model servers in Vertex AI. + + Args: + running: Filter by running status. + service_uuid: Filter by service UUID. + pipeline_name: Filter by pipeline name. + run_name: Filter by run name. + pipeline_step_name: Filter by pipeline step name. + model_name: Filter by model name. + model_uri: Filter by model URI. + model_version: Filter by model version. + + Returns: + A list of services matching the given criteria. + """ + client = Client() + services = client.list_services( + service_type=VertexDeploymentService.SERVICE_TYPE, + running=running, + service_uuid=service_uuid, + pipeline_name=pipeline_name, + run_name=run_name, + pipeline_step_name=pipeline_step_name, + model_name=model_name, + model_uri=model_uri, + model_version=model_version, + ) + + return [ + VertexDeploymentService.from_model(service_model) + for service_model in services + ] diff --git a/src/zenml/integrations/gcp/model_registries/__init__.py b/src/zenml/integrations/gcp/model_registries/__init__.py new file mode 100644 index 00000000000..38622ef0da3 --- /dev/null +++ b/src/zenml/integrations/gcp/model_registries/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) ZenML GmbH 2023. 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: +# +# http://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 Vertex AI model deployers.""" + +from zenml.integrations.gcp.model_registries.vertex_model_registry import ( + VertexAIModelRegistry +) + +__all__ = ["VertexAIModelRegistry"] diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py index 48403965ca5..f3fd31d84d7 100644 --- a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -1,13 +1,30 @@ +# Copyright (c) ZenML GmbH 2023. 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: +# +# http://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. +"""Vertex AI model registry integration for ZenML.""" + from datetime import datetime -from typing import Any, Dict, List, Optional, Tuple, cast +from typing import Any, Dict, List, Optional, cast from google.cloud import aiplatform -from google.cloud.aiplatform import Model, ModelRegistry, ModelVersion -from google.cloud.aiplatform.exceptions import NotFound -from zenml.enums import StackComponentType -from zenml.stack import Flavor, StackComponent -from zenml.stack.stack_component import StackComponentConfig +from zenml.integrations.gcp.flavors.vertex_model_registry_flavor import ( + VertexAIModelRegistryConfig, +) +from zenml.integrations.gcp.google_credentials_mixin import ( + GoogleCredentialsMixin, +) +from zenml.logger import get_logger from zenml.model_registries.base_model_registry import ( BaseModelRegistry, ModelRegistryModelMetadata, @@ -15,22 +32,29 @@ RegisteredModel, RegistryModelVersion, ) -from zenml.stack.stack_validator import StackValidator -from zenml.logger import get_logger +from zenml.stack.stack_component import StackComponentConfig logger = get_logger(__name__) -class VertexAIModelRegistry(BaseModelRegistry): - """Register models using Vertex AI.""" - def __init__(self): - super().__init__() - aiplatform.init() # Initialize the Vertex AI SDK +class VertexAIModelRegistry(BaseModelRegistry, GoogleCredentialsMixin): + """Register models using Vertex AI.""" @property - def config(self) -> StackComponentConfig: - """Returns the config of the model registries.""" - return cast(StackComponentConfig, self._config) + def config(self) -> VertexAIModelRegistryConfig: + """Returns the config of the model registry. + + Returns: + The configuration. + """ + return cast(VertexAIModelRegistryConfig, self._config) + + def setup_aiplatform(self) -> None: + """Setup the Vertex AI platform.""" + credentials, project_id = self._get_authentication() + aiplatform.init( + project=project_id, location=self.config.location, credentials=credentials + ) def register_model( self, @@ -40,13 +64,16 @@ def register_model( ) -> RegisteredModel: """Register a model to the Vertex AI model registry.""" try: - model = Model( + model = aiplatform.Model.upload( display_name=name, description=description, - labels=metadata + labels=metadata, + serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest", # Placeholder + ) + breakpoint() + return RegisteredModel( + name=name, description=description, metadata=metadata ) - model.upload() - return RegisteredModel(name=name, description=description, metadata=metadata) except Exception as e: raise RuntimeError(f"Failed to register model: {str(e)}") @@ -56,10 +83,8 @@ def delete_model( ) -> None: """Delete a model from the Vertex AI model registry.""" try: - model = Model(model_name=name) + model = aiplatform.Model(model_name=name) model.delete() - except NotFound: - raise KeyError(f"Model with name {name} does not exist.") except Exception as e: raise RuntimeError(f"Failed to delete model: {str(e)}") @@ -72,35 +97,28 @@ def update_model( ) -> RegisteredModel: """Update a model in the Vertex AI model registry.""" try: - model = Model(model_name=name) + model = aiplatform.Model(model_name=name) if description: - model.update(description=description) + model.description = description if metadata: - for key, value in metadata.items(): - model.labels[key] = value + model.labels.update(metadata) if remove_metadata: for key in remove_metadata: - if key in model.labels: - del model.labels[key] + model.labels.pop(key, None) model.update() return self.get_model(name) - except NotFound: - raise KeyError(f"Model with name {name} does not exist.") except Exception as e: raise RuntimeError(f"Failed to update model: {str(e)}") def get_model(self, name: str) -> RegisteredModel: """Get a model from the Vertex AI model registry.""" try: - model = Model(model_name=name) - model_resource = model.gca_resource + model = aiplatform.Model(model_name=name) return RegisteredModel( - name=model_resource.display_name, - description=model_resource.description, - metadata=model_resource.labels + name=model.name, + description=model.description, + metadata=model.labels, ) - except NotFound: - raise KeyError(f"Model with name {name} does not exist.") except Exception as e: raise RuntimeError(f"Failed to get model: {str(e)}") @@ -110,19 +128,22 @@ def list_models( metadata: Optional[Dict[str, str]] = None, ) -> List[RegisteredModel]: """List models in the Vertex AI model registry.""" - filter_expression = "" + filter_expr = [] if name: - filter_expression += f"display_name={name}" + filter_expr.append(f"display_name={name}") if metadata: for key, value in metadata.items(): - filter_expression += f"labels.{key}={value} " + filter_expr.append(f"labels.{key}={value}") + + filter_str = " AND ".join(filter_expr) if filter_expr else None + try: - models = Model.list(filter=filter_expression) + models = aiplatform.Model.list(filter=filter_str) return [ RegisteredModel( name=model.display_name, description=model.description, - metadata=model.labels + metadata=model.labels, ) for model in models ] @@ -139,27 +160,32 @@ def register_model_version( **kwargs: Any, ) -> RegistryModelVersion: """Register a model version to the Vertex AI model registry.""" + metadata_dict = metadata.model_dump() if metadata else {} + serving_container_image_uri = metadata_dict.get( + "serving_container_image_uri", None + ) + is_default_version = metadata_dict.get("is_default_version", False) + self.setup_aiplatform() try: - model = Model(model_name=name) - version_info = model.upload_version( - display_name=version, - description=description, + version_info = aiplatform.Model.upload( artifact_uri=model_source_uri, - labels=metadata.dict() if metadata else None + display_name=f"{name}_{version}", + serving_container_image_uri="europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-3:latest", + description=description, + is_default_version=is_default_version, + labels=metadata_dict, ) return RegistryModelVersion( version=version_info.version_id, model_source_uri=model_source_uri, - model_format="Custom", - registered_model=self.get_model(name), + model_format="Custom", # Vertex AI doesn't provide this info directly + registered_model=self.get_model(version_info.name), description=description, created_at=version_info.create_time, last_updated_at=version_info.update_time, - stage=ModelVersionStage.NONE, - metadata=metadata + stage=ModelVersionStage.NONE, # Vertex AI doesn't have built-in stages + metadata=metadata, ) - except NotFound: - raise KeyError(f"Model with name {name} does not exist.") except Exception as e: raise RuntimeError(f"Failed to register model version: {str(e)}") @@ -170,11 +196,10 @@ def delete_model_version( ) -> None: """Delete a model version from the Vertex AI model registry.""" try: - model = Model(model_name=name) - version_info = ModelVersion(model_name=f"{name}@{version}") - version_info.delete() - except NotFound: - raise KeyError(f"Model version {version} of model {name} does not exist.") + model_version = aiplatform.ModelVersion( + model_name=f"{name}@{version}" + ) + model_version.delete() except Exception as e: raise RuntimeError(f"Failed to delete model version: {str(e)}") @@ -189,23 +214,19 @@ def update_model_version( ) -> RegistryModelVersion: """Update a model version in the Vertex AI model registry.""" try: - model_version = ModelVersion(model_name=f"{name}@{version}") + model_version = aiplatform.ModelVersion( + model_name=f"{name}@{version}" + ) if description: - model_version.update(description=description) + model_version.description = description if metadata: - for key, value in metadata.dict().items(): - model_version.labels[key] = value + model_version.labels.update(metadata.dict()) if remove_metadata: for key in remove_metadata: - if key in model_version.labels: - del model_version.labels[key] + model_version.labels.pop(key, None) model_version.update() - if stage: - # Handle stage update if needed - pass + # Note: Vertex AI doesn't have built-in stages, so we ignore the 'stage' parameter return self.get_model_version(name, version) - except NotFound: - raise KeyError(f"Model version {version} of model {name} does not exist.") except Exception as e: raise RuntimeError(f"Failed to update model version: {str(e)}") @@ -214,20 +235,20 @@ def get_model_version( ) -> RegistryModelVersion: """Get a model version from the Vertex AI model registry.""" try: - model_version = ModelVersion(model_name=f"{name}@{version}") + model_version = aiplatform.ModelVersion( + model_name=f"{name}@{version}" + ) return RegistryModelVersion( version=model_version.version_id, - model_source_uri=model_version.gca_resource.artifact_uri, - model_format="Custom", + model_source_uri=model_version.artifact_uri, + model_format="Custom", # Vertex AI doesn't provide this info directly registered_model=self.get_model(name), description=model_version.description, created_at=model_version.create_time, last_updated_at=model_version.update_time, - stage=ModelVersionStage.NONE, - metadata=ModelRegistryModelMetadata(**model_version.labels) + stage=ModelVersionStage.NONE, # Vertex AI doesn't have built-in stages + metadata=ModelRegistryModelMetadata(**model_version.labels), ) - except NotFound: - raise KeyError(f"Model version {version} of model {name} does not exist.") except Exception as e: raise RuntimeError(f"Failed to get model version: {str(e)}") @@ -244,31 +265,48 @@ def list_model_versions( **kwargs: Any, ) -> List[RegistryModelVersion]: """List model versions from the Vertex AI model registry.""" - filter_expression = "" + filter_expr = [] if name: - filter_expression += f"display_name={name}" + filter_expr.append(f"display_name={name}") if metadata: for key, value in metadata.dict().items(): - filter_expression += f"labels.{key}={value} " + filter_expr.append(f"labels.{key}={value}") + if created_after: + filter_expr.append(f"create_time>{created_after.isoformat()}") + if created_before: + filter_expr.append(f"create_time<{created_before.isoformat()}") + + filter_str = " AND ".join(filter_expr) if filter_expr else None + try: - model = Model(model_name=name) - versions = model.list_versions(filter=filter_expression) - return [ + model = aiplatform.Model(model_name=name) + versions = model.list_versions(filter=filter_str) + + results = [ RegistryModelVersion( version=v.version_id, model_source_uri=v.artifact_uri, - model_format="Custom", + model_format="Custom", # Vertex AI doesn't provide this info directly registered_model=self.get_model(name), description=v.description, created_at=v.create_time, - - last_updated_at=v.update_time, - stage=ModelVersionStage.NONE, - metadata=ModelRegistryModelMetadata(**v.labels) + stage=ModelVersionStage.NONE, # Vertex AI doesn't have built-in stages + metadata=ModelRegistryModelMetadata(**v.labels), ) for v in versions ] + + if order_by_date: + results.sort( + key=lambda x: x.created_at, + reverse=(order_by_date.lower() == "desc"), + ) + + if count: + results = results[:count] + + return results except Exception as e: raise RuntimeError(f"Failed to list model versions: {str(e)}") @@ -280,10 +318,10 @@ def load_model_version( ) -> Any: """Load a model version from the Vertex AI model registry.""" try: - model_version = ModelVersion(model_name=f"{name}@{version}") + model_version = aiplatform.ModelVersion( + model_name=f"{name}@{version}" + ) return model_version - except NotFound: - raise KeyError(f"Model version {version} of model {name} does not exist.") except Exception as e: raise RuntimeError(f"Failed to load model version: {str(e)}") @@ -293,22 +331,3 @@ def get_model_uri_artifact_store( ) -> str: """Get the model URI artifact store.""" return model_version.model_source_uri - - -class VertexAIModelRegistryFlavor(Flavor): - """Base class for all ZenML model registry flavors.""" - - @property - def type(self) -> StackComponentType: - """Type of the flavor.""" - return StackComponentType.MODEL_REGISTRY - - @property - def config_class(self) -> Type[StackComponentConfig]: - """Config class for this flavor.""" - return StackComponentConfig - - @property - def implementation_class(self) -> Type[StackComponent]: - """Returns the implementation class for this flavor.""" - return VertexAIModelRegistry diff --git a/src/zenml/integrations/gcp/services/__init__.py b/src/zenml/integrations/gcp/services/__init__.py index b9f858b5302..a1b89b40ea7 100644 --- a/src/zenml/integrations/gcp/services/__init__.py +++ b/src/zenml/integrations/gcp/services/__init__.py @@ -13,7 +13,9 @@ # permissions and limitations under the License. """Initialization of the MLflow Service.""" -from zenml.integrations.mlflow.services.mlflow_deployment import ( # noqa - MLFlowDeploymentConfig, - MLFlowDeploymentService, +from zenml.integrations.gcp.services.vertex_deployment import ( # noqa + VertexServiceConfig, + VertexDeploymentService, ) + +__all__ = ["VertexServiceConfig", "VertexDeploymentService"] \ No newline at end of file diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py index 07fd8ed2260..93c69512173 100644 --- a/src/zenml/integrations/gcp/services/vertex_deployment.py +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -13,36 +13,77 @@ # permissions and limitations under the License. """Implementation of the Vertex AI Deployment service.""" -from typing import TYPE_CHECKING, Any, Generator, Optional, Tuple, cast - -from pydantic import Field +import re +from typing import Any, Dict, Generator, List, Optional, Tuple +from google.api_core import exceptions from google.cloud import aiplatform +from pydantic import BaseModel, Field -from zenml.client import Client -from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( - VertexBaseConfig, -) +from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import VertexBaseConfig from zenml.logger import get_logger from zenml.services import ServiceState, ServiceStatus, ServiceType from zenml.services.service import BaseDeploymentService, ServiceConfig -if TYPE_CHECKING: - from google.auth.credentials import Credentials - logger = get_logger(__name__) POLLING_TIMEOUT = 1200 UUID_SLICE_LENGTH: int = 8 +def sanitize_labels(labels: Dict[str, str]) -> None: + """Update the label values to be valid Kubernetes labels. + + See: + https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set -class VertexServiceConfig(VertexBaseConfig, ServiceConfig): + Args: + labels: the labels to sanitize. + """ + for key, value in labels.items(): + # Kubernetes labels must be alphanumeric, no longer than + # 63 characters, and must begin and end with an alphanumeric + # character ([a-z0-9A-Z]) + labels[key] = re.sub(r"[^0-9a-zA-Z-_\.]+", "_", value)[:63].strip( + "-_." + ) + +class VertexAIDeploymentConfig(VertexBaseConfig, ServiceConfig): """Vertex AI service configurations.""" + def get_vertex_deployment_labels(self) -> Dict[str, str]: + """Generate labels for the VertexAI deployment from the service configuration. + + These labels are attached to the VertexAI deployment resource + and may be used as label selectors in lookup operations. + + Returns: + The labels for the VertexAI deployment. + """ + labels = {} + if self.pipeline_name: + labels["zenml_pipeline_name"] = self.pipeline_name + if self.pipeline_step_name: + labels["zenml_pipeline_step_name"] = self.pipeline_step_name + if self.model_name: + labels["zenml_model_name"] = self.model_name + if self.model_uri: + labels["zenml_model_uri"] = self.model_uri + sanitize_labels(labels) + return labels + + +class VertexPredictionServiceEndpoint(BaseModel): + """Vertex AI Prediction Service Endpoint.""" + + endpoint_name: str + endpoint_url: Optional[str] = None + class VertexServiceStatus(ServiceStatus): """Vertex AI service status.""" + endpoint: Optional[VertexPredictionServiceEndpoint] = None + class VertexDeploymentService(BaseDeploymentService): """Vertex AI model deployment service. @@ -59,12 +100,12 @@ class VertexDeploymentService(BaseDeploymentService): flavor="vertex", description="Vertex AI inference endpoint prediction service", ) - config: VertexServiceConfig + config: VertexAIDeploymentConfig status: VertexServiceStatus = Field( default_factory=lambda: VertexServiceStatus() ) - def __init__(self, config: VertexServiceConfig, credentials: Tuple["Credentials", str], **attrs: Any): + def __init__(self, config: VertexAIDeploymentConfig, **attrs: Any): """Initialize the Vertex AI deployment service. Args: @@ -72,55 +113,7 @@ def __init__(self, config: VertexServiceConfig, credentials: Tuple["Credentials" attrs: additional attributes to set on the service """ super().__init__(config=config, **attrs) - self._config = config - self._project, self._credentials = credentials # Store credentials as a private attribute - - @property - def config(self) -> VertexServiceConfig: - """Returns the config of the deployment service. - - Returns: - The config of the deployment service. - """ - return cast(VertexServiceConfig, self._config) - - def get_token(self) -> str: - """Get the Vertex AI token. - - Raises: - ValueError: If token not found. - - Returns: - Vertex AI 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.gcp.model_deployers.vertex_model_deployer import ( - VertexModelDeployer, - ) - - model_deployer = client.active_stack.model_deployer - if not isinstance(model_deployer, VertexModelDeployer): - raise ValueError( - "VertexModelDeployer 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 vertex_model(self) -> aiplatform.Model: - """Get the deployed Vertex AI inference endpoint. - - Returns: - Vertex AI inference endpoint. - """ - return aiplatform.Model(f"projects/{self.__project}/locations/{self.config.location}/models/{self.config.model_id}") + aiplatform.init(project=config.project, location=config.location) @property def prediction_url(self) -> Optional[str]: @@ -130,64 +123,67 @@ def prediction_url(self) -> Optional[str]: 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 + return ( + self.status.endpoint.endpoint_url if self.status.endpoint else None + ) - def provision(self) -> None: - """Provision or update remote Vertex AI deployment instance. + def get_endpoints(self) -> List[aiplatform.Endpoint]: + """Get all endpoints for the current project and location.""" + return aiplatform.Endpoint.list() - Raises: - Exception: If any unexpected error while creating inference endpoint. + def _generate_endpoint_name(self) -> str: + """Generate a unique name for the Vertex AI Inference Endpoint. + + Returns: + A unique name for the Vertex AI Inference Endpoint. """ + return f"{self.config.model_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}" + + def provision(self) -> None: + """Provision or update remote Vertex AI deployment instance.""" try: - # Attempt to create and wait for the inference endpoint - vertex_endpoint = self.vertex_model.deploy( - deployed_model_display_name=self.config.deployed_model_display_name, - traffic_percentage=self.config.traffic_percentage, - traffic_split=self.config.traffic_split, + model = aiplatform.Model( + model_name=self.config.model_name, + version=self.config.model_version, + ) + + endpoint = aiplatform.Endpoint.create( + display_name=self._generate_endpoint_name() + ) + + deployment = endpoint.deploy( + model=model, machine_type=self.config.machine_type, min_replica_count=self.config.min_replica_count, max_replica_count=self.config.max_replica_count, accelerator_type=self.config.accelerator_type, accelerator_count=self.config.accelerator_count, service_account=self.config.service_account, - metadata=self.config.metadata, - deploy_request_timeout=self.config.deploy_request_timeout, - autoscaling_target_cpu_utilization=self.config.autoscaling_target_cpu_utilization, - autoscaling_target_accelerator_duty_cycle=self.config.autoscaling_target_accelerator_duty_cycle, - enable_access_logging=self.config.enable_access_logging, - disable_container_logging=self.config.disable_container_logging, + network=self.config.network, encryption_spec_key_name=self.config.encryption_spec_key_name, - deploy_request_timeout=self.config.deploy_request_timeout, + explanation_metadata=self.config.explanation_metadata, + explanation_parameters=self.config.explanation_parameters, + sync=True, ) - 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 Vertex AI inference endpoint: {e}" + self.status.endpoint = VertexPredictionServiceEndpoint( + endpoint_name=endpoint.resource_name, + endpoint_url=endpoint.resource_name, ) + self.status.update_state(ServiceState.ACTIVE) - # Check if the endpoint URL is available after provisioning - if vertex_endpoint. logger.info( - f"Vertex AI inference endpoint successfully deployed and available. Endpoint URL: {hf_endpoint.url}" + f"Vertex AI inference endpoint successfully deployed. " + f"Endpoint: {endpoint.resource_name}" ) - else: - logger.error( - "Failed to start Vertex AI inference endpoint service: No URL available, please check the Vertex AI console for more details." - ) - - def check_status(self) -> Tuple[ServiceState, str]: - """Check the the current operational state of the Vertex AI deployment. - Returns: - The operational state of the Vertex AI deployment and a message - providing additional information about that state (e.g. a - description of the error, if one is encountered). - """ - pass + except Exception as e: + self.status.update_state( + new_state=ServiceState.ERROR, error=str(e) + ) + raise RuntimeError( + f"An error occurred while provisioning the Vertex AI inference endpoint: {e}" + ) def deprovision(self, force: bool = False) -> None: """Deprovision the remote Vertex AI deployment instance. @@ -196,43 +192,97 @@ def deprovision(self, force: bool = False) -> None: force: if True, the remote deployment instance will be forcefully deprovisioned. """ + if self.status.endpoint: + try: + endpoint = aiplatform.Endpoint( + endpoint_name=self.status.endpoint.endpoint_name + ) + endpoint.undeploy_all() + endpoint.delete(force=force) + self.status.endpoint = None + self.status.update_state(ServiceState.INACTIVE) + logger.info( + f"Vertex AI Inference Endpoint {self.status.endpoint.endpoint_name} has been deprovisioned." + ) + except exceptions.NotFound: + logger.warning( + f"Vertex AI Inference Endpoint {self.status.endpoint.endpoint_name} not found. It may have been already deleted." + ) + except Exception as e: + raise RuntimeError( + f"Failed to deprovision Vertex AI Inference Endpoint: {e}" + ) + + def check_status(self) -> Tuple[ServiceState, str]: + """Check the current operational state of the Vertex AI deployment. + + Returns: + The operational state of the Vertex AI deployment and a message + providing additional information about that state. + """ + if not self.status.endpoint: + return ServiceState.INACTIVE, "Endpoint not provisioned" + try: - self.vertex_model.undeploy() - except HfHubHTTPError: - logger.error( - "Vertex AI Inference Endpoint is deleted or cannot be found." + endpoint = aiplatform.Endpoint( + endpoint_name=self.status.endpoint.endpoint_name ) + deployments = endpoint.list_deployments() + + if not deployments: + return ServiceState.INACTIVE, "No active deployments" + + # Check the state of all deployments + for deployment in deployments: + if deployment.state == "ACTIVE": + return ServiceState.ACTIVE, "Deployment is active" + elif deployment.state == "DEPLOYING": + return ( + ServiceState.PENDING_STARTUP, + "Deployment is in progress", + ) + elif deployment.state in ["FAILED", "DELETING"]: + return ( + ServiceState.ERROR, + f"Deployment is in {deployment.state} state", + ) + + return ServiceState.INACTIVE, "No active deployments found" + + except exceptions.NotFound: + return ServiceState.INACTIVE, "Endpoint not found" + except Exception as e: + return ServiceState.ERROR, f"Error checking status: {str(e)}" - def predict(self, data: "Any", max_new_tokens: int) -> "Any": + def predict(self, instances: List[Any]) -> List[Any]: """Make a prediction using the service. Args: - data: input data - max_new_tokens: Number of new tokens to generate + instances: List of instances to predict. Returns: - The prediction result. + The prediction results. Raises: - Exception: if the service is not running - NotImplementedError: if task is not supported. + Exception: if the service is not running or prediction fails. """ if not self.is_running: raise Exception( "Vertex AI 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." + + if not self.status.endpoint: + raise Exception("Endpoint information is missing.") + + try: + endpoint = aiplatform.Endpoint( + endpoint_name=self.status.endpoint.endpoint_name ) - return result + response = endpoint.predict(instances=instances) + return response.predictions + except Exception as e: + raise RuntimeError(f"Prediction failed: {str(e)}") def get_logs( self, follow: bool = False, tail: Optional[int] = None @@ -247,17 +297,27 @@ def get_logs( A generator that can be accessed to get the service logs. """ logger.info( - "Vertex AI Endpoints provides access to the logs of " - "your Endpoints through the UI in the “Logs” tab of your Endpoint" + "Vertex AI Endpoints provides access to the logs through " + "Cloud Logging. Please check the Google Cloud Console for detailed logs." ) - return # type: ignore + yield "Logs are available in Google Cloud Console." - def _generate_an_endpoint_name(self) -> str: - """Generate a unique name for the Vertex AI Inference Endpoint. + @property + def is_running(self) -> bool: + """Check if the service is running. Returns: - A unique name for the Vertex AI Inference Endpoint. + True if the service is in the ACTIVE state, False otherwise. """ - return ( - f"{self.config.service_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}" - ) \ No newline at end of file + state, _ = self.check_status() + return state == ServiceState.ACTIVE + + def start(self) -> None: + """Start the Vertex AI deployment service.""" + if not self.is_running: + self.provision() + + def stop(self) -> None: + """Stop the Vertex AI deployment service.""" + if self.is_running: + self.deprovision() From 6769b6c98afb7fdb7bd9e945439ef77b1ceb2eec Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Wed, 18 Sep 2024 13:26:22 +0100 Subject: [PATCH 05/23] format --- .../flavors/vertex_model_registry_flavor.py | 20 +++++++++---------- .../model_registries/vertex_model_registry.py | 9 +++++---- .../gcp/services/vertex_deployment.py | 8 ++++++-- .../materializers/cloudpickle_materializer.py | 2 +- 4 files changed, 21 insertions(+), 18 deletions(-) diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py index 1c13e95a95e..22adc0f6a5d 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py @@ -18,35 +18,33 @@ from zenml.config.base_settings import BaseSettings from zenml.integrations.gcp import ( GCP_RESOURCE_TYPE, - VERTEX_MODEL_REGISTRY_FLAVOR -) -from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( - VertexBaseConfig, + VERTEX_MODEL_REGISTRY_FLAVOR, ) from zenml.integrations.gcp.google_credentials_mixin import ( GoogleCredentialsConfigMixin, ) -from zenml.models import ServiceConnectorRequirements from zenml.model_registries.base_model_registry import ( BaseModelRegistryConfig, BaseModelRegistryFlavor, ) +from zenml.models import ServiceConnectorRequirements if TYPE_CHECKING: from zenml.integrations.gcp.model_registries import ( VertexAIModelRegistry, ) + class VertexAIModelRegistrySettings(BaseSettings): """Settings for the VertexAI model registry.""" - + location: str - + class VertexAIModelRegistryConfig( - BaseModelRegistryConfig, - GoogleCredentialsConfigMixin, - VertexAIModelRegistrySettings + BaseModelRegistryConfig, + GoogleCredentialsConfigMixin, + VertexAIModelRegistrySettings, ): """Configuration for the VertexAI model registry.""" @@ -79,7 +77,7 @@ def service_connector_requirements( return ServiceConnectorRequirements( resource_type=GCP_RESOURCE_TYPE, ) - + @property def docs_url(self) -> Optional[str]: """A url to point at docs explaining this flavor. diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py index f3fd31d84d7..c46bb3956c2 100644 --- a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -32,7 +32,6 @@ RegisteredModel, RegistryModelVersion, ) -from zenml.stack.stack_component import StackComponentConfig logger = get_logger(__name__) @@ -43,17 +42,19 @@ class VertexAIModelRegistry(BaseModelRegistry, GoogleCredentialsMixin): @property def config(self) -> VertexAIModelRegistryConfig: """Returns the config of the model registry. - + Returns: The configuration. """ return cast(VertexAIModelRegistryConfig, self._config) - + def setup_aiplatform(self) -> None: """Setup the Vertex AI platform.""" credentials, project_id = self._get_authentication() aiplatform.init( - project=project_id, location=self.config.location, credentials=credentials + project=project_id, + location=self.config.location, + credentials=credentials, ) def register_model( diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py index 93c69512173..ccebdc09e75 100644 --- a/src/zenml/integrations/gcp/services/vertex_deployment.py +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -20,7 +20,9 @@ from google.cloud import aiplatform from pydantic import BaseModel, Field -from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import VertexBaseConfig +from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( + VertexBaseConfig, +) from zenml.logger import get_logger from zenml.services import ServiceState, ServiceStatus, ServiceType from zenml.services.service import BaseDeploymentService, ServiceConfig @@ -30,6 +32,7 @@ POLLING_TIMEOUT = 1200 UUID_SLICE_LENGTH: int = 8 + def sanitize_labels(labels: Dict[str, str]) -> None: """Update the label values to be valid Kubernetes labels. @@ -46,7 +49,8 @@ def sanitize_labels(labels: Dict[str, str]) -> None: labels[key] = re.sub(r"[^0-9a-zA-Z-_\.]+", "_", value)[:63].strip( "-_." ) - + + class VertexAIDeploymentConfig(VertexBaseConfig, ServiceConfig): """Vertex AI service configurations.""" diff --git a/src/zenml/materializers/cloudpickle_materializer.py b/src/zenml/materializers/cloudpickle_materializer.py index 399ca7f2336..a6813cb4191 100644 --- a/src/zenml/materializers/cloudpickle_materializer.py +++ b/src/zenml/materializers/cloudpickle_materializer.py @@ -29,7 +29,7 @@ logger = get_logger(__name__) -DEFAULT_FILENAME = "artifact.pkl" +DEFAULT_FILENAME = "model.pkl" DEFAULT_PYTHON_VERSION_FILENAME = "python_version.txt" From 9a03f34522da92382f7048cf213fb56102d2d20e Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Fri, 20 Sep 2024 14:37:04 +0100 Subject: [PATCH 06/23] Refactor model registration and add URI parameter --- .../promotion/promote_with_metric_compare.py | 14 +++++++ examples/e2e/steps/training/model_trainer.py | 40 +++++++++---------- 2 files changed, 33 insertions(+), 21 deletions(-) diff --git a/examples/e2e/steps/promotion/promote_with_metric_compare.py b/examples/e2e/steps/promotion/promote_with_metric_compare.py index d68409d2d54..fcd12935905 100644 --- a/examples/e2e/steps/promotion/promote_with_metric_compare.py +++ b/examples/e2e/steps/promotion/promote_with_metric_compare.py @@ -18,6 +18,7 @@ from utils import promote_in_model_registry from zenml import Model, get_step_context, step +from zenml.client import Client from zenml.logger import get_logger logger = get_logger(__name__) @@ -29,6 +30,7 @@ def promote_with_metric_compare( current_metric: float, mlflow_model_name: str, target_env: str, + uri: str, ) -> None: """Try to promote trained model. @@ -57,6 +59,18 @@ def promote_with_metric_compare( ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ### should_promote = True + model_registry = Client().active_stack.model_registry + + # Register model version + model_version = model_registry.register_model_version( + name=mlflow_model_name, + version="1", + model_source_uri=uri, + description="test_register_model_version", + ) + + breakpoint() + # Get model version numbers from Model Control Plane latest_version = get_step_context().model current_version = Model(name=latest_version.name, version=target_env) diff --git a/examples/e2e/steps/training/model_trainer.py b/examples/e2e/steps/training/model_trainer.py index 87a695f5695..43e8c3f4402 100644 --- a/examples/e2e/steps/training/model_trainer.py +++ b/examples/e2e/steps/training/model_trainer.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # +from typing import Tuple import mlflow import pandas as pd @@ -25,9 +26,10 @@ from zenml.integrations.mlflow.experiment_trackers import ( MLFlowExperimentTracker, ) -from zenml.integrations.mlflow.steps.mlflow_registry import ( - mlflow_register_model_step, -) + +# from zenml.integrations.mlflow.steps.mlflow_registry import ( +# mlflow_register_model_step, +# ) from zenml.logger import get_logger logger = get_logger(__name__) @@ -49,8 +51,11 @@ def model_trainer( model: ClassifierMixin, target: str, name: str, -) -> Annotated[ - ClassifierMixin, ArtifactConfig(name="model", is_model_artifact=True) +) -> Tuple[ + Annotated[ + ClassifierMixin, ArtifactConfig(name="model", is_model_artifact=True) + ], + Annotated[str, "uri"], ]: """Configure and train a model on the training dataset. @@ -82,6 +87,9 @@ def model_trainer( Returns: The trained model artifact. """ + step_context = get_step_context() + # Get the URI where the output will be saved. + uri = step_context.get_output_artifact_uri(output_name="model") ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ### # Initialize the model with the hyperparameters indicated in the step @@ -94,19 +102,9 @@ def model_trainer( ) # register mlflow model - mlflow_register_model_step.entrypoint( - model, - name=name, - ) - # keep track of mlflow version for future use - model_registry = Client().active_stack.model_registry - if model_registry: - version = model_registry.get_latest_model_version( - name=name, stage=None - ) - if version: - model_ = get_step_context().model - model_.log_metadata({"model_registry_version": version.version}) - ### YOUR CODE ENDS HERE ### - - return model + # mlflow_register_model_step.entrypoint( + # model, + # name=name, + # ) + + return model, uri From afc5c2b220c39902e567ea309abae90132de459f Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 21 Sep 2024 13:51:48 +0100 Subject: [PATCH 07/23] Refactor model registration and add URI parameter --- .../model_registries/vertex_model_registry.py | 81 ++++++------------- .../model_registries/base_model_registry.py | 4 +- 2 files changed, 28 insertions(+), 57 deletions(-) diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py index c46bb3956c2..97582074842 100644 --- a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -64,19 +64,9 @@ def register_model( metadata: Optional[Dict[str, str]] = None, ) -> RegisteredModel: """Register a model to the Vertex AI model registry.""" - try: - model = aiplatform.Model.upload( - display_name=name, - description=description, - labels=metadata, - serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-0:latest", # Placeholder - ) - breakpoint() - return RegisteredModel( - name=name, description=description, metadata=metadata - ) - except Exception as e: - raise RuntimeError(f"Failed to register model: {str(e)}") + raise NotImplementedError( + "Vertex AI does not support registering models, you can only register model versions, skipping model registration..." + ) def delete_model( self, @@ -97,19 +87,9 @@ def update_model( remove_metadata: Optional[List[str]] = None, ) -> RegisteredModel: """Update a model in the Vertex AI model registry.""" - try: - model = aiplatform.Model(model_name=name) - if description: - model.description = description - if metadata: - model.labels.update(metadata) - if remove_metadata: - for key in remove_metadata: - model.labels.pop(key, None) - model.update() - return self.get_model(name) - except Exception as e: - raise RuntimeError(f"Failed to update model: {str(e)}") + raise NotImplementedError( + "Vertex AI does not support updating models, you can only update model versions, skipping model registration..." + ) def get_model(self, name: str) -> RegisteredModel: """Get a model from the Vertex AI model registry.""" @@ -129,17 +109,14 @@ def list_models( metadata: Optional[Dict[str, str]] = None, ) -> List[RegisteredModel]: """List models in the Vertex AI model registry.""" - filter_expr = [] + filter_expr = 'labels.managed_by="ZenML"' if name: - filter_expr.append(f"display_name={name}") + filter_expr = filter_expr + f' AND display_name="{name}"' if metadata: for key, value in metadata.items(): - filter_expr.append(f"labels.{key}={value}") - - filter_str = " AND ".join(filter_expr) if filter_expr else None - + filter_expr = filter_expr + f' AND labels.{key}="{value}"' try: - models = aiplatform.Model.list(filter=filter_str) + models = aiplatform.Model.list(filter=filter_expr) return [ RegisteredModel( name=model.display_name, @@ -163,7 +140,9 @@ def register_model_version( """Register a model version to the Vertex AI model registry.""" metadata_dict = metadata.model_dump() if metadata else {} serving_container_image_uri = metadata_dict.get( - "serving_container_image_uri", None + "serving_container_image_uri", + None + or "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-3:latest", ) is_default_version = metadata_dict.get("is_default_version", False) self.setup_aiplatform() @@ -171,7 +150,7 @@ def register_model_version( version_info = aiplatform.Model.upload( artifact_uri=model_source_uri, display_name=f"{name}_{version}", - serving_container_image_uri="europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-3:latest", + serving_container_image_uri=serving_container_image_uri, description=description, is_default_version=is_default_version, labels=metadata_dict, @@ -215,18 +194,16 @@ def update_model_version( ) -> RegistryModelVersion: """Update a model version in the Vertex AI model registry.""" try: - model_version = aiplatform.ModelVersion( - model_name=f"{name}@{version}" - ) - if description: - model_version.description = description + model_version = aiplatform.Model(model_name=f"{name}@{version}") + labels = model_version.labels if metadata: - model_version.labels.update(metadata.dict()) + metadata_dict = metadata.model_dump() if metadata else {} + for key, value in metadata_dict.items(): + labels[key] = value if remove_metadata: for key in remove_metadata: - model_version.labels.pop(key, None) - model_version.update() - # Note: Vertex AI doesn't have built-in stages, so we ignore the 'stage' parameter + labels.pop(key, None) + model_version.update(description=description, labels=labels) return self.get_model_version(name, version) except Exception as e: raise RuntimeError(f"Failed to update model version: {str(e)}") @@ -236,14 +213,12 @@ def get_model_version( ) -> RegistryModelVersion: """Get a model version from the Vertex AI model registry.""" try: - model_version = aiplatform.ModelVersion( - model_name=f"{name}@{version}" - ) + model_version = aiplatform.Model(model_name=f"{name}@{version}") return RegistryModelVersion( version=model_version.version_id, model_source_uri=model_version.artifact_uri, model_format="Custom", # Vertex AI doesn't provide this info directly - registered_model=self.get_model(name), + registered_model=self.get_model(model_version.name), description=model_version.description, created_at=model_version.create_time, last_updated_at=model_version.update_time, @@ -288,7 +263,7 @@ def list_model_versions( version=v.version_id, model_source_uri=v.artifact_uri, model_format="Custom", # Vertex AI doesn't provide this info directly - registered_model=self.get_model(name), + registered_model=self.get_model(v.name), description=v.description, created_at=v.create_time, last_updated_at=v.update_time, @@ -297,13 +272,7 @@ def list_model_versions( ) for v in versions ] - - if order_by_date: - results.sort( - key=lambda x: x.created_at, - reverse=(order_by_date.lower() == "desc"), - ) - + if count: results = results[:count] diff --git a/src/zenml/model_registries/base_model_registry.py b/src/zenml/model_registries/base_model_registry.py index 578d97d396c..727632eaa15 100644 --- a/src/zenml/model_registries/base_model_registry.py +++ b/src/zenml/model_registries/base_model_registry.py @@ -20,6 +20,7 @@ from pydantic import BaseModel, ConfigDict +from zenml import __version__ from zenml.enums import StackComponentType from zenml.stack import Flavor, StackComponent from zenml.stack.stack_component import StackComponentConfig @@ -62,7 +63,8 @@ class ModelRegistryModelMetadata(BaseModel): model and its development process. """ - zenml_version: Optional[str] = None + managed_by: str = "ZenML" + zenml_version: str = __version__ zenml_run_name: Optional[str] = None zenml_pipeline_name: Optional[str] = None zenml_pipeline_uuid: Optional[str] = None From 2dc0d2d49ec043b71bd8a3f8065e94c90b12eea5 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Sat, 21 Sep 2024 14:06:18 +0100 Subject: [PATCH 08/23] Refactor model registration and remove unnecessary code --- .../integrations/gcp/model_registries/vertex_model_registry.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py index 97582074842..873f22b54e6 100644 --- a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -272,7 +272,7 @@ def list_model_versions( ) for v in versions ] - + if count: results = results[:count] From 54b6748f2a642585854a8d29b6dde193f6ba262f Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Fri, 25 Oct 2024 15:37:50 +0100 Subject: [PATCH 09/23] Refactor GCP service and flavor classes for Vertex AI deployment --- .../flavors/vertex_model_deployer_flavor.py | 46 ++++-- .../model_deployers/vertex_model_deployer.py | 132 +++++++++--------- .../model_registries/vertex_model_registry.py | 10 +- .../integrations/gcp/services/__init__.py | 4 +- .../gcp/services/vertex_deployment.py | 33 +++-- 5 files changed, 128 insertions(+), 97 deletions(-) diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py index 85d4bd52485..f225847a5df 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py @@ -17,11 +17,20 @@ from pydantic import BaseModel -from zenml.integrations.gcp import VERTEX_MODEL_DEPLOYER_FLAVOR +from zenml.integrations.gcp import ( + GCP_RESOURCE_TYPE, + VERTEX_MODEL_DEPLOYER_FLAVOR, +) +from zenml.integrations.gcp.google_credentials_mixin import ( + GoogleCredentialsConfigMixin, +) from zenml.model_deployers.base_model_deployer import ( BaseModelDeployerConfig, BaseModelDeployerFlavor, ) +from zenml.models.v2.misc.service_connector_type import ( + ServiceConnectorRequirements, +) if TYPE_CHECKING: from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( @@ -61,20 +70,14 @@ class VertexBaseConfig(BaseModel): autoscaling_target_accelerator_duty_cycle: Optional[float] = None enable_access_logging: Optional[bool] = None disable_container_logging: Optional[bool] = None + explanation_metadata: Optional[Dict[str, str]] = None + explanation_parameters: Optional[Dict[str, str]] = None -class VertexModelDeployerConfig(BaseModelDeployerConfig, VertexBaseConfig): - """Configuration for the Vertex AI model deployer. - - Attributes: - project_id: The project ID. - location: The location of the model. - """ - - # The namespace to list endpoints for. Set to `"*"` to list all endpoints - # from all namespaces (i.e. personal namespace and all orgs the user belongs to). - project_id: str - location: Optional[str] = None +class VertexModelDeployerConfig( + BaseModelDeployerConfig, VertexBaseConfig, GoogleCredentialsConfigMixin +): + """Configuration for the Vertex AI model deployer.""" class VertexModelDeployerFlavor(BaseModelDeployerFlavor): @@ -89,6 +92,23 @@ def name(self) -> str: """ return VERTEX_MODEL_DEPLOYER_FLAVOR + @property + def service_connector_requirements( + self, + ) -> Optional[ServiceConnectorRequirements]: + """Service connector resource requirements for service connectors. + + Specifies resource requirements that are used to filter the available + service connector types that are compatible with this flavor. + + Returns: + Requirements for compatible service connectors, if a service + connector is required for this flavor. + """ + return ServiceConnectorRequirements( + resource_type=GCP_RESOURCE_TYPE, + ) + @property def docs_url(self) -> Optional[str]: """A url to point at docs explaining this flavor. diff --git a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py index 17ad388588d..472c3666eeb 100644 --- a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py +++ b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py @@ -13,20 +13,31 @@ # permissions and limitations under the License. """Implementation of the Vertex AI Model Deployer.""" -from typing import ClassVar, Dict, List, Optional, Tuple, Type, cast +from typing import ClassVar, Dict, Optional, Tuple, Type, cast from uuid import UUID +from google.cloud import aiplatform + from zenml.analytics.enums import AnalyticsEvent from zenml.analytics.utils import track_handler from zenml.client import Client -from zenml.integrations.gcp import VERTEX_SERVICE_ARTIFACT +from zenml.enums import StackComponentType +from zenml.integrations.gcp import ( + VERTEX_SERVICE_ARTIFACT, +) from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( VertexModelDeployerConfig, VertexModelDeployerFlavor, ) +from zenml.integrations.gcp.google_credentials_mixin import ( + GoogleCredentialsMixin, +) +from zenml.integrations.gcp.model_registries.vertex_model_registry import ( + VertexAIModelRegistry, +) from zenml.integrations.gcp.services.vertex_deployment import ( + VertexAIDeploymentConfig, VertexDeploymentService, - VertexServiceConfig, ) from zenml.logger import get_logger from zenml.model_deployers import BaseModelDeployer @@ -41,7 +52,7 @@ logger = get_logger(__name__) -class VertexModelDeployer(BaseModelDeployer): +class VertexModelDeployer(BaseModelDeployer, GoogleCredentialsMixin): """Vertex AI endpoint model deployer.""" NAME: ClassVar[str] = "Vertex AI" @@ -58,35 +69,72 @@ def config(self) -> VertexModelDeployerConfig: """ return cast(VertexModelDeployerConfig, self._config) + def setup_aiplatform(self) -> None: + """Setup the Vertex AI platform.""" + credentials, project_id = self._get_authentication() + aiplatform.init( + project=project_id, + location=self.config.location, + credentials=credentials, + ) + @property def validator(self) -> Optional[StackValidator]: - """Validates the stack. + """Validates that the stack contains a model registry. + + Also validates that the artifact store is not local. Returns: - A validator that checks that the stack contains required GCP components. + A StackValidator instance. """ - def _validate_gcp_stack( - stack: "Stack", - ) -> Tuple[bool, str]: - """Check if GCP components are properly configured in the stack. + def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: + """Validates that all the stack components are not local. 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. + A tuple of (is_valid, error_message). """ - if not self.config.project_id or not self.config.location: + # Validate that the container registry is not local. + model_registry = stack.model_registry + if not model_registry and isinstance( + model_registry, VertexAIModelRegistry + ): return False, ( - "The Vertex AI model deployer requires a GCP project and " - "location to be specified in the configuration." + "The Vertex AI model deployer requires a Vertex AI model " + "registry to be present in the stack. Please add a Vertex AI " + "model registry to the stack." ) - return True, "Stack is valid for Vertex AI model deployment." + + # Validate that the rest of the components are not local. + for stack_comp in stack.components.values(): + # For Forward compatibility a list of components is returned, + # but only the first item is relevant for now + # TODO: [server] make sure the ComponentModel actually has + # a local_path property or implement similar check + local_path = stack_comp.local_path + if not local_path: + continue + return False, ( + f"The '{stack_comp.name}' {stack_comp.type.value} is a " + f"local stack component. The Vertex AI Pipelines " + f"orchestrator requires that all the components in the " + f"stack used to execute the pipeline have to be not local, " + f"because there is no way for Vertex to connect to your " + f"local machine. You should use a flavor of " + f"{stack_comp.type.value} other than '" + f"{stack_comp.flavor}'." + ) + + return True, "" return StackValidator( - custom_validation_function=_validate_gcp_stack, + required_components={ + StackComponentType.MODEL_REGISTRY, + }, + custom_validation_function=_validate_stack_requirements, ) def _create_deployment_service( @@ -130,8 +178,10 @@ def perform_deploy_model( The ZenML Vertex AI deployment service object. """ with track_handler(AnalyticsEvent.MODEL_DEPLOYED) as analytics_handler: - config = cast(VertexServiceConfig, config) - service = self._create_deployment_service(id=id, config=config) + config = cast(VertexAIDeploymentConfig, config) + service = self._create_deployment_service( + id=id, config=config, timeout=timeout + ) logger.info( f"Creating a new Vertex AI deployment service: {service}" ) @@ -226,47 +276,3 @@ def get_model_server_info( "PREDICTION_URL": service_instance.prediction_url, "HEALTH_CHECK_URL": service_instance.get_healthcheck_url(), } - - def find_model_server( - self, - running: Optional[bool] = None, - 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_version: Optional[str] = None, - ) -> List[BaseService]: - """Find deployed model servers in Vertex AI. - - Args: - running: Filter by running status. - service_uuid: Filter by service UUID. - pipeline_name: Filter by pipeline name. - run_name: Filter by run name. - pipeline_step_name: Filter by pipeline step name. - model_name: Filter by model name. - model_uri: Filter by model URI. - model_version: Filter by model version. - - Returns: - A list of services matching the given criteria. - """ - client = Client() - services = client.list_services( - service_type=VertexDeploymentService.SERVICE_TYPE, - running=running, - service_uuid=service_uuid, - pipeline_name=pipeline_name, - run_name=run_name, - pipeline_step_name=pipeline_step_name, - model_name=model_name, - model_uri=model_uri, - model_version=model_version, - ) - - return [ - VertexDeploymentService.from_model(service_model) - for service_model in services - ] diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py index 873f22b54e6..7f0332f3631 100644 --- a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -109,7 +109,8 @@ def list_models( metadata: Optional[Dict[str, str]] = None, ) -> List[RegisteredModel]: """List models in the Vertex AI model registry.""" - filter_expr = 'labels.managed_by="ZenML"' + self.setup_aiplatform() + filter_expr = 'labels.managed_by="zenml"' if name: filter_expr = filter_expr + f' AND display_name="{name}"' if metadata: @@ -138,6 +139,7 @@ def register_model_version( **kwargs: Any, ) -> RegistryModelVersion: """Register a model version to the Vertex AI model registry.""" + self.setup_aiplatform() metadata_dict = metadata.model_dump() if metadata else {} serving_container_image_uri = metadata_dict.get( "serving_container_image_uri", @@ -145,7 +147,7 @@ def register_model_version( or "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-3:latest", ) is_default_version = metadata_dict.get("is_default_version", False) - self.setup_aiplatform() + metadata_dict["managed_by"] = "zenml" try: version_info = aiplatform.Model.upload( artifact_uri=model_source_uri, @@ -175,6 +177,7 @@ def delete_model_version( version: str, ) -> None: """Delete a model version from the Vertex AI model registry.""" + self.setup_aiplatform() try: model_version = aiplatform.ModelVersion( model_name=f"{name}@{version}" @@ -193,6 +196,7 @@ def update_model_version( stage: Optional[ModelVersionStage] = None, ) -> RegistryModelVersion: """Update a model version in the Vertex AI model registry.""" + self.setup_aiplatform() try: model_version = aiplatform.Model(model_name=f"{name}@{version}") labels = model_version.labels @@ -212,6 +216,7 @@ def get_model_version( self, name: str, version: str ) -> RegistryModelVersion: """Get a model version from the Vertex AI model registry.""" + self.setup_aiplatform() try: model_version = aiplatform.Model(model_name=f"{name}@{version}") return RegistryModelVersion( @@ -241,6 +246,7 @@ def list_model_versions( **kwargs: Any, ) -> List[RegistryModelVersion]: """List model versions from the Vertex AI model registry.""" + self.setup_aiplatform() filter_expr = [] if name: filter_expr.append(f"display_name={name}") diff --git a/src/zenml/integrations/gcp/services/__init__.py b/src/zenml/integrations/gcp/services/__init__.py index a1b89b40ea7..be8c6508a37 100644 --- a/src/zenml/integrations/gcp/services/__init__.py +++ b/src/zenml/integrations/gcp/services/__init__.py @@ -14,8 +14,8 @@ """Initialization of the MLflow Service.""" from zenml.integrations.gcp.services.vertex_deployment import ( # noqa - VertexServiceConfig, + VertexAIDeploymentConfig, VertexDeploymentService, ) -__all__ = ["VertexServiceConfig", "VertexDeploymentService"] \ No newline at end of file +__all__ = ["VertexAIDeploymentConfig", "VertexDeploymentService"] \ No newline at end of file diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py index ccebdc09e75..72069a1f181 100644 --- a/src/zenml/integrations/gcp/services/vertex_deployment.py +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -20,6 +20,7 @@ from google.cloud import aiplatform from pydantic import BaseModel, Field +from zenml.client import Client from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( VertexBaseConfig, ) @@ -70,8 +71,6 @@ def get_vertex_deployment_labels(self) -> Dict[str, str]: labels["zenml_pipeline_step_name"] = self.pipeline_step_name if self.model_name: labels["zenml_model_name"] = self.model_name - if self.model_uri: - labels["zenml_model_uri"] = self.model_uri sanitize_labels(labels) return labels @@ -117,7 +116,6 @@ def __init__(self, config: VertexAIDeploymentConfig, **attrs: Any): attrs: additional attributes to set on the service """ super().__init__(config=config, **attrs) - aiplatform.init(project=config.project, location=config.location) @property def prediction_url(self) -> Optional[str]: @@ -145,17 +143,28 @@ def _generate_endpoint_name(self) -> str: def provision(self) -> None: """Provision or update remote Vertex AI deployment instance.""" + from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( + VertexModelDeployer, + ) + + zenml_client = Client() + model_deployer = zenml_client.active_stack.model_deployer + if isinstance(model_deployer, VertexModelDeployer): + model_deployer.setup_aiplatform() + else: + raise ValueError("Model deployer is not VertexModelDeployer") try: + breakpoint() model = aiplatform.Model( model_name=self.config.model_name, version=self.config.model_version, ) - + breakpoint() endpoint = aiplatform.Endpoint.create( display_name=self._generate_endpoint_name() ) - - deployment = endpoint.deploy( + breakpoint() + endpoint.deploy( model=model, machine_type=self.config.machine_type, min_replica_count=self.config.min_replica_count, @@ -169,7 +178,7 @@ def provision(self) -> None: explanation_parameters=self.config.explanation_parameters, sync=True, ) - + breakpoint() self.status.endpoint = VertexPredictionServiceEndpoint( endpoint_name=endpoint.resource_name, endpoint_url=endpoint.resource_name, @@ -315,13 +324,3 @@ def is_running(self) -> bool: """ state, _ = self.check_status() return state == ServiceState.ACTIVE - - def start(self) -> None: - """Start the Vertex AI deployment service.""" - if not self.is_running: - self.provision() - - def stop(self) -> None: - """Stop the Vertex AI deployment service.""" - if self.is_running: - self.deprovision() From a9804493162723fde5db745be6ecd395a39df9fa Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 31 Oct 2024 12:01:54 +0100 Subject: [PATCH 10/23] Refactor Vertex AI model registry and deployer configurations - Update model source URI retrieval in VertexAIModelRegistry. - Enhance BaseModelDeployer to check and start inactive services. - Set default replica counts to 1 and sync to False in VertexBaseConfig. - Rename and update documentation for deployment service creation in VertexModelDeployer. --- .../flavors/vertex_model_deployer_flavor.py | 7 +- .../model_deployers/vertex_model_deployer.py | 24 +- .../model_registries/vertex_model_registry.py | 2 +- .../gcp/services/vertex_deployment.py | 373 ++++++++++-------- .../model_deployers/base_model_deployer.py | 7 + 5 files changed, 229 insertions(+), 184 deletions(-) diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py index f225847a5df..5c8041c94e5 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py @@ -58,13 +58,13 @@ class VertexBaseConfig(BaseModel): machine_type: Optional[str] = None accelerator_type: Optional[str] = None accelerator_count: Optional[int] = None - min_replica_count: Optional[int] = None - max_replica_count: Optional[int] = None + min_replica_count: Optional[int] = 1 + max_replica_count: Optional[int] = 1 service_account: Optional[str] = None metadata: Optional[Dict[str, str]] = None network: Optional[str] = None encryption_spec_key_name: Optional[str] = None - sync: Optional[bool] = True + sync: Optional[bool] = False deploy_request_timeout: Optional[int] = None autoscaling_target_cpu_utilization: Optional[float] = None autoscaling_target_accelerator_duty_cycle: Optional[float] = None @@ -72,6 +72,7 @@ class VertexBaseConfig(BaseModel): disable_container_logging: Optional[bool] = None explanation_metadata: Optional[Dict[str, str]] = None explanation_parameters: Optional[Dict[str, str]] = None + existing_endpoint: Optional[str] = None class VertexModelDeployerConfig( diff --git a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py index 472c3666eeb..82ab4aca520 100644 --- a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py +++ b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py @@ -138,19 +138,19 @@ def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: ) def _create_deployment_service( - self, id: UUID, timeout: int, config: VertexModelDeployerConfig + self, id: UUID, timeout: int, config: VertexAIDeploymentConfig ) -> VertexDeploymentService: - """Creates a new DatabricksDeploymentService. + """Creates a new VertexAIDeploymentService. Args: - id: the UUID of the model to be deployed with Databricks model deployer. - timeout: the timeout in seconds to wait for the Databricks inference endpoint + id: the UUID of the model to be deployed with Vertex model deployer. + timeout: the timeout in seconds to wait for the Vertex inference endpoint to be provisioned and successfully started or updated. - config: the configuration of the model to be deployed with Databricks model deployer. + config: the configuration of the model to be deployed with Vertex model deployer. Returns: The VertexModelDeployerConfig object that can be used to interact - with the Databricks inference endpoint. + with the Vertex inference endpoint. """ # create a new service for the new model service = VertexDeploymentService(uuid=id, config=config) @@ -197,14 +197,6 @@ def perform_deploy_model( "store_type": client.zen_store.type.value, **stack_metadata, } - - # Create a service artifact - client.create_artifact( - name=VERTEX_SERVICE_ARTIFACT, - artifact_store_id=client.active_stack.artifact_store.id, - producer=service, - ) - return service def perform_stop_model( @@ -258,10 +250,10 @@ def perform_delete_model( """ service = cast(VertexDeploymentService, service) service.stop(timeout=timeout, force=force) - service.delete() + service.stop() @staticmethod - def get_model_server_info( + def get_model_server_info( # type: ignore[override] service_instance: "VertexDeploymentService", ) -> Dict[str, Optional[str]]: """Get information about the deployed model server. diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py index 7f0332f3631..66b7465bbaa 100644 --- a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -159,7 +159,7 @@ def register_model_version( ) return RegistryModelVersion( version=version_info.version_id, - model_source_uri=model_source_uri, + model_source_uri=version_info.resource_name, model_format="Custom", # Vertex AI doesn't provide this info directly registered_model=self.get_model(version_info.name), description=description, diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py index 72069a1f181..2353f0d8d47 100644 --- a/src/zenml/integrations/gcp/services/vertex_deployment.py +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -1,16 +1,3 @@ -# 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 Vertex AI Deployment service.""" import re @@ -30,7 +17,10 @@ logger = get_logger(__name__) -POLLING_TIMEOUT = 1200 +# Increase timeout for long-running operations +POLLING_TIMEOUT = ( + 1800 # Increased from 1200 to allow for longer deployment times +) UUID_SLICE_LENGTH: int = 8 @@ -52,26 +42,45 @@ def sanitize_labels(labels: Dict[str, str]) -> None: ) +def sanitize_vertex_label(value: str) -> str: + """Sanitize a label value to comply with Vertex AI requirements. + + Args: + value: The label value to sanitize + + Returns: + Sanitized label value + """ + # Convert to lowercase + value = value.lower() + # Replace any character that's not lowercase letter, number, dash or underscore + value = re.sub(r"[^a-z0-9\-_]", "-", value) + # Ensure it starts with a letter/number by prepending 'x' if needed + if not value[0].isalnum(): + value = f"x{value}" + # Truncate to 63 chars to stay under limit + return value[:63] + + class VertexAIDeploymentConfig(VertexBaseConfig, ServiceConfig): """Vertex AI service configurations.""" def get_vertex_deployment_labels(self) -> Dict[str, str]: - """Generate labels for the VertexAI deployment from the service configuration. - - These labels are attached to the VertexAI deployment resource - and may be used as label selectors in lookup operations. + """Generate labels for the VertexAI deployment from the service configuration.""" + labels = { + "managed-by": "zenml", # Changed from managed_by to managed-by + } - Returns: - The labels for the VertexAI deployment. - """ - labels = {} if self.pipeline_name: - labels["zenml_pipeline_name"] = self.pipeline_name + labels["pipeline-name"] = sanitize_vertex_label(self.pipeline_name) if self.pipeline_step_name: - labels["zenml_pipeline_step_name"] = self.pipeline_step_name + labels["step-name"] = sanitize_vertex_label( + self.pipeline_step_name + ) if self.model_name: - labels["zenml_model_name"] = self.model_name - sanitize_labels(labels) + labels["model-name"] = sanitize_vertex_label(self.model_name) + if self.service_name: + labels["service-name"] = sanitize_vertex_label(self.service_name) return labels @@ -80,6 +89,9 @@ class VertexPredictionServiceEndpoint(BaseModel): endpoint_name: str endpoint_url: Optional[str] = None + deployed_model_id: Optional[str] = ( + None # Added to track specific model deployment + ) class VertexServiceStatus(ServiceStatus): @@ -89,13 +101,7 @@ class VertexServiceStatus(ServiceStatus): class VertexDeploymentService(BaseDeploymentService): - """Vertex AI model deployment service. - - Attributes: - SERVICE_TYPE: a service type descriptor with information describing - the Vertex AI deployment service class - config: service configuration - """ + """Vertex AI model deployment service.""" SERVICE_TYPE = ServiceType( name="vertex-deployment", @@ -109,158 +115,202 @@ class VertexDeploymentService(BaseDeploymentService): ) def __init__(self, config: VertexAIDeploymentConfig, **attrs: Any): - """Initialize the Vertex AI deployment service. - - Args: - config: service configuration - attrs: additional attributes to set on the service - """ + """Initialize the Vertex AI deployment service.""" super().__init__(config=config, **attrs) + # Initialize aiplatform with project and location + from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( + VertexModelDeployer, + ) + + zenml_client = Client() + model_deployer = zenml_client.active_stack.model_deployer + if not isinstance(model_deployer, VertexModelDeployer): + raise ValueError("Model deployer is not VertexModelDeployer") + + model_deployer.setup_aiplatform() + @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.status.endpoint.endpoint_url if self.status.endpoint else None - ) + """The prediction URI exposed by the prediction service.""" + if not self.status.endpoint or not self.status.endpoint.endpoint_url: + return None + + # Construct proper prediction URL + return f"https://{self.config.location}-aiplatform.googleapis.com/v1/{self.status.endpoint.endpoint_url}" def get_endpoints(self) -> List[aiplatform.Endpoint]: """Get all endpoints for the current project and location.""" - return aiplatform.Endpoint.list() + try: + # Use proper filtering and pagination + return list( + aiplatform.Endpoint.list( + filter='labels.managed_by="zenml"', + location=self.config.location, + ) + ) + except Exception as e: + logger.error(f"Failed to list endpoints: {e}") + return [] def _generate_endpoint_name(self) -> str: - """Generate a unique name for the Vertex AI Inference Endpoint. - - Returns: - A unique name for the Vertex AI Inference Endpoint. - """ - return f"{self.config.model_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}" + """Generate a unique name for the Vertex AI Inference Endpoint.""" + # Make name more descriptive and conformant to Vertex AI naming rules + sanitized_model_name = re.sub( + r"[^a-zA-Z0-9-]", "-", self.config.model_name.lower() + ) + return f"{sanitized_model_name}-{str(self.uuid)[:UUID_SLICE_LENGTH]}" def provision(self) -> None: """Provision or update remote Vertex AI deployment instance.""" - from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( - VertexModelDeployer, - ) - - zenml_client = Client() - model_deployer = zenml_client.active_stack.model_deployer - if isinstance(model_deployer, VertexModelDeployer): - model_deployer.setup_aiplatform() - else: - raise ValueError("Model deployer is not VertexModelDeployer") try: - breakpoint() + if self.config.existing_endpoint: + # Use the existing endpoint + endpoint = aiplatform.Endpoint( + endpoint_name=self.config.existing_endpoint, + location=self.config.location, + ) + logger.info( + f"Using existing Vertex AI inference endpoint: {endpoint.resource_name}" + ) + else: + # Create the endpoint + endpoint_name = self._generate_endpoint_name() + endpoint = aiplatform.Endpoint.create( + display_name=endpoint_name, + location=self.config.location, + encryption_spec_key_name=self.config.encryption_spec_key_name, + labels=self.config.get_vertex_deployment_labels(), + ) + logger.info( + f"Vertex AI inference endpoint created: {endpoint.resource_name}" + ) + + # Then get the model model = aiplatform.Model( - model_name=self.config.model_name, - version=self.config.model_version, + model_name=self.config.model_id, + location=self.config.location, ) - breakpoint() - endpoint = aiplatform.Endpoint.create( - display_name=self._generate_endpoint_name() + logger.info( + f"Found existing model to deploy: {model.resource_name} to the endpoint." ) - breakpoint() + + # Deploy the model to the endpoint endpoint.deploy( model=model, + deployed_model_display_name=f"{endpoint_name}-deployment", machine_type=self.config.machine_type, min_replica_count=self.config.min_replica_count, max_replica_count=self.config.max_replica_count, accelerator_type=self.config.accelerator_type, accelerator_count=self.config.accelerator_count, service_account=self.config.service_account, - network=self.config.network, - encryption_spec_key_name=self.config.encryption_spec_key_name, explanation_metadata=self.config.explanation_metadata, explanation_parameters=self.config.explanation_parameters, - sync=True, + sync=self.config.sync, + ) + logger.info( + f"Model {model.resource_name} successfully deployed to endpoint {endpoint.resource_name}" ) - breakpoint() + + # Store both endpoint and deployment information self.status.endpoint = VertexPredictionServiceEndpoint( endpoint_name=endpoint.resource_name, endpoint_url=endpoint.resource_name, + deployed_model_id=model.resource_name, ) - self.status.update_state(ServiceState.ACTIVE) + self.status.update_state(ServiceState.PENDING_STARTUP) logger.info( - f"Vertex AI inference endpoint successfully deployed. " - f"Endpoint: {endpoint.resource_name}" + f"Vertex AI inference endpoint successfully deployed. Pending startup" + f"Endpoint: {endpoint.resource_name}, " ) except Exception as e: self.status.update_state( - new_state=ServiceState.ERROR, error=str(e) + new_state=ServiceState.ERROR, + error=f"Deployment failed: {str(e)}", ) raise RuntimeError( f"An error occurred while provisioning the Vertex AI inference endpoint: {e}" ) def deprovision(self, force: bool = False) -> None: - """Deprovision the remote Vertex AI deployment instance. - - Args: - force: if True, the remote deployment instance will be - forcefully deprovisioned. - """ - if self.status.endpoint: - try: - endpoint = aiplatform.Endpoint( - endpoint_name=self.status.endpoint.endpoint_name - ) - endpoint.undeploy_all() - endpoint.delete(force=force) - self.status.endpoint = None - self.status.update_state(ServiceState.INACTIVE) - logger.info( - f"Vertex AI Inference Endpoint {self.status.endpoint.endpoint_name} has been deprovisioned." - ) - except exceptions.NotFound: - logger.warning( - f"Vertex AI Inference Endpoint {self.status.endpoint.endpoint_name} not found. It may have been already deleted." - ) - except Exception as e: - raise RuntimeError( - f"Failed to deprovision Vertex AI Inference Endpoint: {e}" - ) + """Deprovision the remote Vertex AI deployment instance.""" + if not self.status.endpoint: + return - def check_status(self) -> Tuple[ServiceState, str]: - """Check the current operational state of the Vertex AI deployment. + try: + endpoint = aiplatform.Endpoint( + endpoint_name=self.status.endpoint.endpoint_name, + location=self.config.location, + ) + + # First undeploy the specific model if we have its ID + if self.status.endpoint.deployed_model_id: + try: + endpoint.undeploy( + deployed_model_id=self.status.endpoint.deployed_model_id, + sync=self.config.sync, + ) + except exceptions.NotFound: + logger.warning("Deployed model already undeployed") + + # Then delete the endpoint + endpoint.delete(force=force, sync=self.config.sync) + + self.status.endpoint = None + self.status.update_state(ServiceState.INACTIVE) + + logger.info("Vertex AI Inference Endpoint has been deprovisioned.") + + except exceptions.NotFound: + logger.warning( + "Vertex AI Inference Endpoint not found. It may have been already deleted." + ) + self.status.update_state(ServiceState.INACTIVE) + except Exception as e: + error_msg = ( + f"Failed to deprovision Vertex AI Inference Endpoint: {e}" + ) + logger.error(error_msg) + if not force: + raise RuntimeError(error_msg) - Returns: - The operational state of the Vertex AI deployment and a message - providing additional information about that state. - """ + def check_status(self) -> Tuple[ServiceState, str]: + """Check the current operational state of the Vertex AI deployment.""" if not self.status.endpoint: return ServiceState.INACTIVE, "Endpoint not provisioned" - try: + logger.info( + f"Checking status of Vertex AI Inference Endpoint: {self.status.endpoint.endpoint_name}" + ) endpoint = aiplatform.Endpoint( - endpoint_name=self.status.endpoint.endpoint_name + endpoint_name=self.status.endpoint.endpoint_name, + location=self.config.location, ) - deployments = endpoint.list_deployments() - - if not deployments: - return ServiceState.INACTIVE, "No active deployments" - - # Check the state of all deployments - for deployment in deployments: - if deployment.state == "ACTIVE": - return ServiceState.ACTIVE, "Deployment is active" - elif deployment.state == "DEPLOYING": - return ( - ServiceState.PENDING_STARTUP, - "Deployment is in progress", - ) - elif deployment.state in ["FAILED", "DELETING"]: - return ( - ServiceState.ERROR, - f"Deployment is in {deployment.state} state", + + # Get detailed deployment status + deployment = None + if self.status.endpoint.deployed_model_id: + deployments = [ + d + for d in endpoint.list_models() + if d.model == self.status.endpoint.deployed_model_id + ] + if deployments: + deployment = deployments[0] + logger.info( + f"Model {self.status.endpoint.deployed_model_id} was deployed to the endpoint" ) - return ServiceState.INACTIVE, "No active deployments found" + if not deployment: + logger.warning( + "No matching deployment found, endpoint may be inactive or failed to deploy" + ) + return ServiceState.INACTIVE, "No matching deployment found" + + return ServiceState.ACTIVE, "Deployment is ready" except exceptions.NotFound: return ServiceState.INACTIVE, "Endpoint not found" @@ -268,17 +318,7 @@ def check_status(self) -> Tuple[ServiceState, str]: return ServiceState.ERROR, f"Error checking status: {str(e)}" def predict(self, instances: List[Any]) -> List[Any]: - """Make a prediction using the service. - - Args: - instances: List of instances to predict. - - Returns: - The prediction results. - - Raises: - Exception: if the service is not running or prediction fails. - """ + """Make a prediction using the service.""" if not self.is_running: raise Exception( "Vertex AI endpoint inference service is not running. " @@ -290,37 +330,42 @@ def predict(self, instances: List[Any]) -> List[Any]: try: endpoint = aiplatform.Endpoint( - endpoint_name=self.status.endpoint.endpoint_name + endpoint_name=self.status.endpoint.endpoint_name, + location=self.config.location, ) - response = endpoint.predict(instances=instances) - return response.predictions + + # Add proper prediction parameters and handle sync/async + predictions = endpoint.predict( + instances=instances, + deployed_model_id=self.status.endpoint.deployed_model_id.split( + "/" + )[-1], + timeout=30, # Add reasonable timeout + ) + + if not predictions: + raise RuntimeError("No predictions returned") + except Exception as e: + logger.error(f"Prediction failed: {e}") raise RuntimeError(f"Prediction failed: {str(e)}") + return [predictions] + 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. - """ + """Retrieve the service logs.""" + # Note: Could be enhanced to actually fetch logs from Cloud Logging logger.info( "Vertex AI Endpoints provides access to the logs through " - "Cloud Logging. Please check the Google Cloud Console for detailed logs." + "Cloud Logging. Please check the Google Cloud Console for detailed logs. " + f"Location: {self.config.location}" ) yield "Logs are available in Google Cloud Console." @property def is_running(self) -> bool: - """Check if the service is running. - - Returns: - True if the service is in the ACTIVE state, False otherwise. - """ - state, _ = self.check_status() - return state == ServiceState.ACTIVE + """Check if the service is running.""" + self.update_status() + return self.status.state == ServiceState.ACTIVE diff --git a/src/zenml/model_deployers/base_model_deployer.py b/src/zenml/model_deployers/base_model_deployer.py index 40a65128f26..814e4f28175 100644 --- a/src/zenml/model_deployers/base_model_deployer.py +++ b/src/zenml/model_deployers/base_model_deployer.py @@ -32,6 +32,7 @@ from zenml.logger import get_logger from zenml.services import BaseService, ServiceConfig from zenml.services.service import BaseDeploymentService +from zenml.services.service_status import ServiceState from zenml.services.service_type import ServiceType from zenml.stack import StackComponent from zenml.stack.flavor import Flavor @@ -180,6 +181,12 @@ def deploy_model( logger.info( f"Existing model server found for {config.name or config.model_name} with the exact same configuration. Returning the existing service named {services[0].config.service_name}." ) + status, _ = services[0].check_status() + if status != ServiceState.ACTIVE: + logger.info( + f"Service found for {config.name or config.model_name} is not active. Starting the service." + ) + services[0].start(timeout=timeout) return services[0] else: # Find existing model server From 0a13214bbc5e094a820b18907d6edba508a4c48b Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 31 Oct 2024 13:44:54 +0100 Subject: [PATCH 11/23] Refactor model deployer configurations and add VertexAI model deployer --- .../component-guide/model-deployers/vertex.md | 179 ++++++++++++++++++ .../model-registries/vertex.md | 150 +++++++++++++++ docs/book/toc.md | 2 + 3 files changed, 331 insertions(+) create mode 100644 docs/book/component-guide/model-deployers/vertex.md create mode 100644 docs/book/component-guide/model-registries/vertex.md diff --git a/docs/book/component-guide/model-deployers/vertex.md b/docs/book/component-guide/model-deployers/vertex.md new file mode 100644 index 00000000000..812271da6e9 --- /dev/null +++ b/docs/book/component-guide/model-deployers/vertex.md @@ -0,0 +1,179 @@ +# Vertex AI Model Deployer + +[Vertex AI](https://cloud.google.com/vertex-ai) provides managed infrastructure for deploying machine learning models at scale. The Vertex AI Model Deployer in ZenML allows you to deploy models to Vertex AI endpoints, providing a scalable and managed solution for model serving. + +## When to use it? + +You should use the Vertex AI Model Deployer when: + +* You're already using Google Cloud Platform (GCP) and want to leverage its native ML infrastructure +* You need enterprise-grade model serving capabilities with autoscaling +* You want a fully managed solution for hosting ML models +* You need to handle high-throughput prediction requests +* You want to deploy models with GPU acceleration +* You need to monitor and track your model deployments + +This is particularly useful in the following scenarios: +* Deploying models to production with high availability requirements +* Serving models that need GPU acceleration +* Handling varying prediction workloads with autoscaling +* Integrating model serving with other GCP services + +{% hint style="warning" %} +The Vertex AI Model Deployer requires a Vertex AI Model Registry to be present in your stack. Make sure you have configured both components properly. +{% endhint %} + +## How to deploy it? + +The Vertex AI Model Deployer is provided by the GCP ZenML integration. First, install the integration: + +```shell +zenml integration install gcp -y +``` + +### Authentication and Service Connector Configuration + +The Vertex AI Model Deployer requires proper GCP authentication. The recommended way to configure this is using the ZenML Service Connector functionality: + +```shell +# Register the service connector with a service account key +zenml service-connector register vertex_deployer_connector \ + --type gcp \ + --auth-method=service-account \ + --project_id= \ + --service_account_json=@vertex-deployer-sa.json \ + --resource-type gcp-generic + +# Register the model deployer +zenml model-deployer register vertex_deployer \ + --flavor=vertex \ + --location=us-central1 + +# Connect the model deployer to the service connector +zenml model-deployer connect vertex_deployer --connector vertex_deployer_connector +``` + +{% hint style="info" %} +The service account needs the following permissions: +- `Vertex AI User` role for deploying models +- `Vertex AI Service Agent` role for managing model endpoints +{% endhint %} + +## How to use it + +### Deploy a model in a pipeline + +Here's an example of how to use the Vertex AI Model Deployer in a ZenML pipeline: + +```python +from typing_extensions import Annotated +from zenml import ArtifactConfig, get_step_context, step +from zenml.client import Client +from zenml.integrations.gcp.services.vertex_deployment import ( + VertexAIDeploymentConfig, + VertexDeploymentService, +) + +@step(enable_cache=False) +def model_deployer( + model_registry_uri: str, +) -> Annotated[ + VertexDeploymentService, + ArtifactConfig(name="vertex_deployment", is_deployment_artifact=True) +]: + """Model deployer step.""" + zenml_client = Client() + current_model = get_step_context().model + model_deployer = zenml_client.active_stack.model_deployer + + # Configure the deployment + vertex_deployment_config = VertexAIDeploymentConfig( + location="europe-west1", + name="zenml-vertex-quickstart", + model_name=current_model.name, + description="Vertex AI model deployment example", + model_id=model_registry_uri, + machine_type="n1-standard-4", # Optional: specify machine type + min_replica_count=1, # Optional: minimum number of replicas + max_replica_count=3, # Optional: maximum number of replicas + ) + + # Deploy the model + service = model_deployer.deploy_model( + config=vertex_deployment_config, + service_type=VertexDeploymentService.SERVICE_TYPE, + ) + + return service +``` + +### Configuration Options + +The Vertex AI Model Deployer accepts a rich set of configuration options through `VertexAIDeploymentConfig`: + +* Basic Configuration: + * `location`: GCP region for deployment (e.g., "us-central1") + * `name`: Name for the deployment endpoint + * `model_name`: Name of the model being deployed + * `model_id`: Model ID from the Vertex AI Model Registry + +* Infrastructure Configuration: + * `machine_type`: Type of machine to use (e.g., "n1-standard-4") + * `accelerator_type`: GPU accelerator type if needed + * `accelerator_count`: Number of GPUs per replica + * `min_replica_count`: Minimum number of serving replicas + * `max_replica_count`: Maximum number of serving replicas + +* Advanced Configuration: + * `service_account`: Custom service account for the deployment + * `network`: VPC network configuration + * `encryption_spec_key_name`: Customer-managed encryption key + * `enable_access_logging`: Enable detailed access logging + * `explanation_metadata`: Model explanation configuration + * `autoscaling_target_cpu_utilization`: Target CPU utilization for autoscaling + +### Running Predictions + +Once a model is deployed, you can run predictions using the service: + +```python +from zenml.integrations.gcp.model_deployers import VertexModelDeployer +from zenml.services import ServiceState + +# Get the deployed service +model_deployer = VertexModelDeployer.get_active_model_deployer() +services = model_deployer.find_model_server( + pipeline_name="deployment_pipeline", + pipeline_step_name="model_deployer", + model_name="my_model", +) + +if services: + service = services[0] + if service.is_running: + # Run prediction + prediction = service.predict( + instances=[{"feature1": 1.0, "feature2": 2.0}] + ) + print(f"Prediction: {prediction}") +``` + +### Limitations and Considerations + +1. **Stack Requirements**: + - Requires a Vertex AI Model Registry in the stack + - All stack components must be non-local + +2. **Authentication**: + - Requires proper GCP credentials with Vertex AI permissions + - Best practice is to use service connectors for authentication + +3. **Costs**: + - Vertex AI endpoints incur costs based on machine type and uptime + - Consider using autoscaling to optimize costs + +4. **Region Availability**: + - Service availability depends on Vertex AI regional availability + - Model and endpoint must be in the same region + +Check out the [SDK docs](https://sdkdocs.zenml.io) for more detailed information about the implementation. \ No newline at end of file diff --git a/docs/book/component-guide/model-registries/vertex.md b/docs/book/component-guide/model-registries/vertex.md new file mode 100644 index 00000000000..41a29ffb11d --- /dev/null +++ b/docs/book/component-guide/model-registries/vertex.md @@ -0,0 +1,150 @@ +# Vertex AI Model Registry + +[Vertex AI](https://cloud.google.com/vertex-ai) is Google Cloud's unified ML platform that helps you build, deploy, and scale ML models. The Vertex AI Model Registry is a centralized repository for managing your ML models throughout their lifecycle. ZenML's Vertex AI Model Registry integration allows you to register, version, and manage your models using Vertex AI's infrastructure. + +## When would you want to use it? + +You should consider using the Vertex AI Model Registry when: + +* You're already using Google Cloud Platform (GCP) and want to leverage its native ML infrastructure +* You need enterprise-grade model management capabilities with fine-grained access control +* You want to track model lineage and metadata in a centralized location +* You're building ML pipelines that need to integrate with other Vertex AI services +* You need to manage model deployment across different GCP environments + +This is particularly useful in the following scenarios: + +* Building production ML pipelines that need to integrate with GCP services +* Managing multiple versions of models across development and production environments +* Tracking model artifacts and metadata in a centralized location +* Deploying models to Vertex AI endpoints for serving + +{% hint style="warning" %} +Important: The Vertex AI Model Registry implementation only supports the model version interface, not the model interface. This means you cannot register, delete, or update models directly - you can only work with model versions. Operations like `register_model()`, `delete_model()`, and `update_model()` are not supported. +{% endhint %} + +## How do you deploy it? + +The Vertex AI Model Registry flavor is provided by the GCP ZenML integration. First, install the integration: + +```shell +zenml integration install gcp -y +``` + +### Authentication and Service Connector Configuration + +The Vertex AI Model Registry requires proper GCP authentication. The recommended way to configure this is using the ZenML Service Connector functionality. You have several options for authentication: + +1. Using a GCP Service Connector with a dedicated service account (Recommended): +```shell +# Register the service connector with a service account key +zenml service-connector register vertex_registry_connector \ + --type gcp \ + --auth-method=service-account \ + --project_id= \ + --service_account_json=@vertex-registry-sa.json \ + --resource-type gcp-generic + +# Register the model registry +zenml model-registry register vertex_registry \ + --flavor=vertex \ + --location=us-central1 + +# Connect the model registry to the service connector +zenml model-registry connect vertex_registry --connector vertex_registry_connector +``` + +2. Using local gcloud credentials: +```shell +# Register the model registry using local gcloud auth +zenml model-registry register vertex_registry \ + --flavor=vertex \ + --location=us-central1 +``` + +{% hint style="info" %} +The service account used needs the following permissions: +- `Vertex AI User` role for creating and managing model versions +- `Storage Object Viewer` role if accessing models stored in Google Cloud Storage +{% endhint %} + +## How do you use it? + +### Register models inside a pipeline + +Here's an example of how to use the Vertex AI Model Registry in your ZenML pipeline using the provided model registration step: + +```python +from typing_extensions import Annotated +from zenml import ArtifactConfig, get_step_context, step +from zenml.client import Client +from zenml.logger import get_logger + +logger = get_logger(__name__) + +@step(enable_cache=False) +def model_register() -> Annotated[str, ArtifactConfig(name="model_registry_uri")]: + """Model registration step.""" + # Get the current model from the context + current_model = get_step_context().model + + client = Client() + model_registry = client.active_stack.model_registry + model_version = model_registry.register_model_version( + name=current_model.name, + version=str(current_model.version), + model_source_uri=current_model.get_model_artifact("sklearn_classifier").uri, + description="ZenML model registered after promotion", + ) + logger.info( + f"Model version {model_version.version} registered in Model Registry" + ) + + return model_version.model_source_uri +``` + +### Configuration Options + +The Vertex AI Model Registry accepts the following configuration options: + +* `location`: The GCP region where the model registry will be created (e.g., "us-central1") +* `project_id`: (Optional) The GCP project ID. If not specified, will use the default project +* `credentials`: (Optional) GCP credentials configuration + +### Working with Model Versions + +Since the Vertex AI Model Registry only supports version-level operations, here's how to work with model versions: + +```shell +# List all model versions +zenml model-registry models list-versions + +# Get details of a specific model version +zenml model-registry models get-version -v + +# Delete a model version +zenml model-registry models delete-version -v +``` + +### Key Differences from MLflow Model Registry + +Unlike the MLflow Model Registry, the Vertex AI implementation has some important differences: + +1. **Version-Only Interface**: Vertex AI only supports model version operations. You cannot register, delete, or update models directly - only their versions. +2. **Authentication**: Uses GCP service connectors for authentication, similar to other Vertex AI services in ZenML. +3. **Staging Levels**: Vertex AI doesn't have built-in staging levels (like Production, Staging, etc.) - these are handled through metadata. +4. **Default Container Images**: Vertex AI requires a serving container image URI, which defaults to the scikit-learn prediction container if not specified. +5. **Managed Service**: As a fully managed service, you don't need to worry about infrastructure management, but you need valid GCP credentials. + +### Limitations + +Based on the implementation, there are some limitations to be aware of: + +1. The `register_model()`, `update_model()`, and `delete_model()` methods are not implemented as Vertex AI only supports registering model versions +2. Model stage transitions (Production, Staging, etc.) are not natively supported +3. Models must have a serving container image URI specified or will use the default scikit-learn image +4. All registered models are automatically labeled with `managed_by="zenml"` for tracking purposes + +Check out the [SDK docs](https://sdkdocs.zenml.io/latest/integration\_code\_docs/integrations-gcp/#zenml.integrations.gcp.model\_registry) to see more about the interface and implementation. + +
ZenML Scarf
\ No newline at end of file diff --git a/docs/book/toc.md b/docs/book/toc.md index 28b24626f03..e42a1b1516c 100644 --- a/docs/book/toc.md +++ b/docs/book/toc.md @@ -260,6 +260,7 @@ * [Develop a custom experiment tracker](component-guide/experiment-trackers/custom.md) * [Model Deployers](component-guide/model-deployers/model-deployers.md) * [MLflow](component-guide/model-deployers/mlflow.md) + * [VertexAI](component-guide/model-deployers/vertex.md) * [Seldon](component-guide/model-deployers/seldon.md) * [BentoML](component-guide/model-deployers/bentoml.md) * [Hugging Face](component-guide/model-deployers/huggingface.md) @@ -289,6 +290,7 @@ * [Develop a Custom Annotator](component-guide/annotators/custom.md) * [Model Registries](component-guide/model-registries/model-registries.md) * [MLflow Model Registry](component-guide/model-registries/mlflow.md) + * [VertexAI](component-guide/model-registries/vertex.md) * [Develop a Custom Model Registry](component-guide/model-registries/custom.md) * [Feature Stores](component-guide/feature-stores/feature-stores.md) * [Feast](component-guide/feature-stores/feast.md) From 53da68dd6fcb41099db959a98139cb4ef7ceb15d Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 31 Oct 2024 13:49:36 +0100 Subject: [PATCH 12/23] Refactor model deployer configurations and add VertexAI model deployer --- .../promotion/promote_with_metric_compare.py | 14 ------- examples/e2e/steps/training/model_trainer.py | 40 ++++++++++--------- 2 files changed, 21 insertions(+), 33 deletions(-) diff --git a/examples/e2e/steps/promotion/promote_with_metric_compare.py b/examples/e2e/steps/promotion/promote_with_metric_compare.py index 7affe356f44..038d219d32d 100644 --- a/examples/e2e/steps/promotion/promote_with_metric_compare.py +++ b/examples/e2e/steps/promotion/promote_with_metric_compare.py @@ -18,7 +18,6 @@ from utils import promote_in_model_registry from zenml import Model, get_step_context, step -from zenml.client import Client from zenml.logger import get_logger logger = get_logger(__name__) @@ -30,7 +29,6 @@ def promote_with_metric_compare( current_metric: float, mlflow_model_name: str, target_env: str, - uri: str, ) -> None: """Try to promote trained model. @@ -59,18 +57,6 @@ def promote_with_metric_compare( ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ### should_promote = True - model_registry = Client().active_stack.model_registry - - # Register model version - model_version = model_registry.register_model_version( - name=mlflow_model_name, - version="1", - model_source_uri=uri, - description="test_register_model_version", - ) - - breakpoint() - # Get model version numbers from Model Control Plane latest_version = get_step_context().model current_version = Model(name=latest_version.name, version=target_env) diff --git a/examples/e2e/steps/training/model_trainer.py b/examples/e2e/steps/training/model_trainer.py index 43e8c3f4402..87a695f5695 100644 --- a/examples/e2e/steps/training/model_trainer.py +++ b/examples/e2e/steps/training/model_trainer.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # -from typing import Tuple import mlflow import pandas as pd @@ -26,10 +25,9 @@ from zenml.integrations.mlflow.experiment_trackers import ( MLFlowExperimentTracker, ) - -# from zenml.integrations.mlflow.steps.mlflow_registry import ( -# mlflow_register_model_step, -# ) +from zenml.integrations.mlflow.steps.mlflow_registry import ( + mlflow_register_model_step, +) from zenml.logger import get_logger logger = get_logger(__name__) @@ -51,11 +49,8 @@ def model_trainer( model: ClassifierMixin, target: str, name: str, -) -> Tuple[ - Annotated[ - ClassifierMixin, ArtifactConfig(name="model", is_model_artifact=True) - ], - Annotated[str, "uri"], +) -> Annotated[ + ClassifierMixin, ArtifactConfig(name="model", is_model_artifact=True) ]: """Configure and train a model on the training dataset. @@ -87,9 +82,6 @@ def model_trainer( Returns: The trained model artifact. """ - step_context = get_step_context() - # Get the URI where the output will be saved. - uri = step_context.get_output_artifact_uri(output_name="model") ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ### # Initialize the model with the hyperparameters indicated in the step @@ -102,9 +94,19 @@ def model_trainer( ) # register mlflow model - # mlflow_register_model_step.entrypoint( - # model, - # name=name, - # ) - - return model, uri + mlflow_register_model_step.entrypoint( + model, + name=name, + ) + # keep track of mlflow version for future use + model_registry = Client().active_stack.model_registry + if model_registry: + version = model_registry.get_latest_model_version( + name=name, stage=None + ) + if version: + model_ = get_step_context().model + model_.log_metadata({"model_registry_version": version.version}) + ### YOUR CODE ENDS HERE ### + + return model From ce2019d32bed9c5c6dd70fab79397c61ea595765 Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Thu, 7 Nov 2024 14:45:21 +0100 Subject: [PATCH 13/23] Rename VertexAI model registry classes and update documentation for consistency --- .../component-guide/model-deployers/vertex.md | 6 +-- src/zenml/integrations/gcp/__init__.py | 4 +- .../integrations/gcp/flavors/__init__.py | 4 +- .../flavors/vertex_model_deployer_flavor.py | 3 +- .../flavors/vertex_model_registry_flavor.py | 2 +- .../model_deployers/vertex_model_deployer.py | 17 ++------ .../gcp/model_registries/__init__.py | 2 +- .../model_registries/vertex_model_registry.py | 1 - .../integrations/gcp/services/__init__.py | 6 +-- .../gcp/services/vertex_deployment.py | 43 ++++++++----------- .../model_registries/base_model_registry.py | 11 ++++- 11 files changed, 46 insertions(+), 53 deletions(-) diff --git a/docs/book/component-guide/model-deployers/vertex.md b/docs/book/component-guide/model-deployers/vertex.md index 812271da6e9..98453df6744 100644 --- a/docs/book/component-guide/model-deployers/vertex.md +++ b/docs/book/component-guide/model-deployers/vertex.md @@ -70,7 +70,7 @@ from typing_extensions import Annotated from zenml import ArtifactConfig, get_step_context, step from zenml.client import Client from zenml.integrations.gcp.services.vertex_deployment import ( - VertexAIDeploymentConfig, + VertexDeploymentConfig, VertexDeploymentService, ) @@ -87,7 +87,7 @@ def model_deployer( model_deployer = zenml_client.active_stack.model_deployer # Configure the deployment - vertex_deployment_config = VertexAIDeploymentConfig( + vertex_deployment_config = VertexDeploymentConfig( location="europe-west1", name="zenml-vertex-quickstart", model_name=current_model.name, @@ -109,7 +109,7 @@ def model_deployer( ### Configuration Options -The Vertex AI Model Deployer accepts a rich set of configuration options through `VertexAIDeploymentConfig`: +The Vertex AI Model Deployer accepts a rich set of configuration options through `VertexDeploymentConfig`: * Basic Configuration: * `location`: GCP region for deployment (e.g., "us-central1") diff --git a/src/zenml/integrations/gcp/__init__.py b/src/zenml/integrations/gcp/__init__.py index 3c9de9a9348..e25d9441427 100644 --- a/src/zenml/integrations/gcp/__init__.py +++ b/src/zenml/integrations/gcp/__init__.py @@ -78,7 +78,7 @@ def flavors(cls) -> List[Type[Flavor]]: VertexOrchestratorFlavor, VertexStepOperatorFlavor, VertexModelDeployerFlavor, - VertexAIModelRegistryFlavor, + VertexModelRegistryFlavor, ) return [ @@ -86,7 +86,7 @@ def flavors(cls) -> List[Type[Flavor]]: GCPImageBuilderFlavor, VertexOrchestratorFlavor, VertexStepOperatorFlavor, - VertexAIModelRegistryFlavor, + VertexModelRegistryFlavor, VertexModelDeployerFlavor, ] diff --git a/src/zenml/integrations/gcp/flavors/__init__.py b/src/zenml/integrations/gcp/flavors/__init__.py index cecf637cefd..25067703d55 100644 --- a/src/zenml/integrations/gcp/flavors/__init__.py +++ b/src/zenml/integrations/gcp/flavors/__init__.py @@ -35,7 +35,7 @@ ) from zenml.integrations.gcp.flavors.vertex_model_registry_flavor import ( VertexAIModelRegistryConfig, - VertexAIModelRegistryFlavor, + VertexModelRegistryFlavor, ) __all__ = [ @@ -49,6 +49,6 @@ "VertexStepOperatorConfig", "VertexModelDeployerFlavor", "VertexModelDeployerConfig", - "VertexAIModelRegistryFlavor", + "VertexModelRegistryFlavor", "VertexAIModelRegistryConfig", ] diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py index 5c8041c94e5..1b526cf0f2f 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_model_deployer_flavor.py @@ -73,6 +73,7 @@ class VertexBaseConfig(BaseModel): explanation_metadata: Optional[Dict[str, str]] = None explanation_parameters: Optional[Dict[str, str]] = None existing_endpoint: Optional[str] = None + labels: Optional[Dict[str, str]] = None class VertexModelDeployerConfig( @@ -82,7 +83,7 @@ class VertexModelDeployerConfig( class VertexModelDeployerFlavor(BaseModelDeployerFlavor): - """Vertex AI Endpoint model deployer flavor.""" + """Vertex AI model deployer flavor.""" @property def name(self) -> str: diff --git a/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py b/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py index 22adc0f6a5d..e16cf548685 100644 --- a/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py +++ b/src/zenml/integrations/gcp/flavors/vertex_model_registry_flavor.py @@ -49,7 +49,7 @@ class VertexAIModelRegistryConfig( """Configuration for the VertexAI model registry.""" -class VertexAIModelRegistryFlavor(BaseModelRegistryFlavor): +class VertexModelRegistryFlavor(BaseModelRegistryFlavor): """Model registry flavor for VertexAI models.""" @property diff --git a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py index 82ab4aca520..b152839356b 100644 --- a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py +++ b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py @@ -36,7 +36,7 @@ VertexAIModelRegistry, ) from zenml.integrations.gcp.services.vertex_deployment import ( - VertexAIDeploymentConfig, + VertexDeploymentConfig, VertexDeploymentService, ) from zenml.logger import get_logger @@ -97,11 +97,8 @@ def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: Returns: A tuple of (is_valid, error_message). """ - # Validate that the container registry is not local. model_registry = stack.model_registry - if not model_registry and isinstance( - model_registry, VertexAIModelRegistry - ): + if not isinstance(model_registry, VertexAIModelRegistry): return False, ( "The Vertex AI model deployer requires a Vertex AI model " "registry to be present in the stack. Please add a Vertex AI " @@ -110,10 +107,6 @@ def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: # Validate that the rest of the components are not local. for stack_comp in stack.components.values(): - # For Forward compatibility a list of components is returned, - # but only the first item is relevant for now - # TODO: [server] make sure the ComponentModel actually has - # a local_path property or implement similar check local_path = stack_comp.local_path if not local_path: continue @@ -138,7 +131,7 @@ def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: ) def _create_deployment_service( - self, id: UUID, timeout: int, config: VertexAIDeploymentConfig + self, id: UUID, timeout: int, config: VertexDeploymentConfig ) -> VertexDeploymentService: """Creates a new VertexAIDeploymentService. @@ -178,14 +171,13 @@ def perform_deploy_model( The ZenML Vertex AI deployment service object. """ with track_handler(AnalyticsEvent.MODEL_DEPLOYED) as analytics_handler: - config = cast(VertexAIDeploymentConfig, config) + config = cast(VertexDeploymentConfig, config) service = self._create_deployment_service( id=id, config=config, timeout=timeout ) logger.info( f"Creating a new Vertex AI deployment service: {service}" ) - service.start(timeout=timeout) client = Client() stack = client.active_stack @@ -250,7 +242,6 @@ def perform_delete_model( """ service = cast(VertexDeploymentService, service) service.stop(timeout=timeout, force=force) - service.stop() @staticmethod def get_model_server_info( # type: ignore[override] diff --git a/src/zenml/integrations/gcp/model_registries/__init__.py b/src/zenml/integrations/gcp/model_registries/__init__.py index 38622ef0da3..672c7c19619 100644 --- a/src/zenml/integrations/gcp/model_registries/__init__.py +++ b/src/zenml/integrations/gcp/model_registries/__init__.py @@ -11,7 +11,7 @@ # 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 Vertex AI model deployers.""" +"""Initialization of the Vertex AI model registry.""" from zenml.integrations.gcp.model_registries.vertex_model_registry import ( VertexAIModelRegistry diff --git a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py index 66b7465bbaa..fcc57001867 100644 --- a/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py +++ b/src/zenml/integrations/gcp/model_registries/vertex_model_registry.py @@ -147,7 +147,6 @@ def register_model_version( or "europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-3:latest", ) is_default_version = metadata_dict.get("is_default_version", False) - metadata_dict["managed_by"] = "zenml" try: version_info = aiplatform.Model.upload( artifact_uri=model_source_uri, diff --git a/src/zenml/integrations/gcp/services/__init__.py b/src/zenml/integrations/gcp/services/__init__.py index be8c6508a37..392a48e9694 100644 --- a/src/zenml/integrations/gcp/services/__init__.py +++ b/src/zenml/integrations/gcp/services/__init__.py @@ -11,11 +11,11 @@ # 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 MLflow Service.""" +"""Initialization of the Vertex Service.""" from zenml.integrations.gcp.services.vertex_deployment import ( # noqa - VertexAIDeploymentConfig, + VertexDeploymentConfig, VertexDeploymentService, ) -__all__ = ["VertexAIDeploymentConfig", "VertexDeploymentService"] \ No newline at end of file +__all__ = ["VertexDeploymentConfig", "VertexDeploymentService"] \ No newline at end of file diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py index 2353f0d8d47..32e8a9722f7 100644 --- a/src/zenml/integrations/gcp/services/vertex_deployment.py +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -1,3 +1,16 @@ +# Copyright (c) ZenML GmbH 2023. 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: +# +# http://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 Vertex AI Deployment service.""" import re @@ -24,24 +37,6 @@ UUID_SLICE_LENGTH: int = 8 -def sanitize_labels(labels: Dict[str, str]) -> None: - """Update the label values to be valid Kubernetes labels. - - See: - https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set - - Args: - labels: the labels to sanitize. - """ - for key, value in labels.items(): - # Kubernetes labels must be alphanumeric, no longer than - # 63 characters, and must begin and end with an alphanumeric - # character ([a-z0-9A-Z]) - labels[key] = re.sub(r"[^0-9a-zA-Z-_\.]+", "_", value)[:63].strip( - "-_." - ) - - def sanitize_vertex_label(value: str) -> str: """Sanitize a label value to comply with Vertex AI requirements. @@ -62,15 +57,13 @@ def sanitize_vertex_label(value: str) -> str: return value[:63] -class VertexAIDeploymentConfig(VertexBaseConfig, ServiceConfig): +class VertexDeploymentConfig(VertexBaseConfig, ServiceConfig): """Vertex AI service configurations.""" def get_vertex_deployment_labels(self) -> Dict[str, str]: """Generate labels for the VertexAI deployment from the service configuration.""" - labels = { - "managed-by": "zenml", # Changed from managed_by to managed-by - } - + labels = self.labels or {} + labels["managed_by"] = "zenml" if self.pipeline_name: labels["pipeline-name"] = sanitize_vertex_label(self.pipeline_name) if self.pipeline_step_name: @@ -109,12 +102,12 @@ class VertexDeploymentService(BaseDeploymentService): flavor="vertex", description="Vertex AI inference endpoint prediction service", ) - config: VertexAIDeploymentConfig + config: VertexDeploymentConfig status: VertexServiceStatus = Field( default_factory=lambda: VertexServiceStatus() ) - def __init__(self, config: VertexAIDeploymentConfig, **attrs: Any): + def __init__(self, config: VertexDeploymentConfig, **attrs: Any): """Initialize the Vertex AI deployment service.""" super().__init__(config=config, **attrs) diff --git a/src/zenml/model_registries/base_model_registry.py b/src/zenml/model_registries/base_model_registry.py index 727632eaa15..ffab13f974a 100644 --- a/src/zenml/model_registries/base_model_registry.py +++ b/src/zenml/model_registries/base_model_registry.py @@ -63,7 +63,7 @@ class ModelRegistryModelMetadata(BaseModel): model and its development process. """ - managed_by: str = "ZenML" + _managed_by: str = "zenml" zenml_version: str = __version__ zenml_run_name: Optional[str] = None zenml_pipeline_name: Optional[str] = None @@ -72,6 +72,15 @@ class ModelRegistryModelMetadata(BaseModel): zenml_step_name: Optional[str] = None zenml_workspace: Optional[str] = None + @property + def managed_by(self) -> str: + """Returns the managed_by attribute. + + Returns: + The managed_by attribute. + """ + return self._managed_by + @property def custom_attributes(self) -> Dict[str, str]: """Returns a dictionary of custom attributes. From 14f299837d25b3af60fbc5a305c987bc97050e69 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 13:47:34 +0000 Subject: [PATCH 14/23] Auto-update of LLM Finetuning template --- examples/llm_finetuning/.copier-answers.yml | 2 +- examples/llm_finetuning/steps/log_metadata.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/llm_finetuning/.copier-answers.yml b/examples/llm_finetuning/.copier-answers.yml index 09cb600a5fa..386863f54e8 100644 --- a/examples/llm_finetuning/.copier-answers.yml +++ b/examples/llm_finetuning/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.09.24 +_commit: 2024.08.29-1-g7af7693 _src_path: gh:zenml-io/template-llm-finetuning bf16: true cuda_version: cuda11.8 diff --git a/examples/llm_finetuning/steps/log_metadata.py b/examples/llm_finetuning/steps/log_metadata.py index 14371b78b6e..645f98cc8ea 100644 --- a/examples/llm_finetuning/steps/log_metadata.py +++ b/examples/llm_finetuning/steps/log_metadata.py @@ -34,7 +34,7 @@ def log_metadata_from_step_artifact( context = get_step_context() metadata_dict: Dict[str, Any] = ( - context.pipeline_run.steps[step_name].outputs[artifact_name].load() + context.pipeline_run.steps[step_name].outputs[artifact_name][0].load() ) metadata = {artifact_name: metadata_dict} From 0b30a61ddbb2b5a0c703cfc1a7f40e8a12a7bb03 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 13:47:55 +0000 Subject: [PATCH 15/23] Auto-update of Starter template --- examples/mlops_starter/.copier-answers.yml | 2 +- examples/mlops_starter/quickstart.ipynb | 4 ++-- examples/mlops_starter/run.py | 4 ++-- examples/mlops_starter/steps/model_promoter.py | 8 +++----- 4 files changed, 8 insertions(+), 10 deletions(-) diff --git a/examples/mlops_starter/.copier-answers.yml b/examples/mlops_starter/.copier-answers.yml index 8b1fb8187ed..e17f27ee551 100644 --- a/examples/mlops_starter/.copier-answers.yml +++ b/examples/mlops_starter/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.09.24 +_commit: 2024.10.21 _src_path: gh:zenml-io/template-starter email: info@zenml.io full_name: ZenML GmbH diff --git a/examples/mlops_starter/quickstart.ipynb b/examples/mlops_starter/quickstart.ipynb index df8c010b5ea..6fba7a0e8cc 100644 --- a/examples/mlops_starter/quickstart.ipynb +++ b/examples/mlops_starter/quickstart.ipynb @@ -994,8 +994,8 @@ "@pipeline\n", "def inference(preprocess_pipeline_id: UUID):\n", " \"\"\"Model batch inference pipeline\"\"\"\n", - " # random_state = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).metadata[\"random_state\"].value\n", - " # target = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).run_metadata['target'].value\n", + " # random_state = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).metadata[\"random_state\"]\n", + " # target = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).run_metadata['target']\n", " random_state = 42\n", " target = \"target\"\n", "\n", diff --git a/examples/mlops_starter/run.py b/examples/mlops_starter/run.py index d7b1a7f11b2..16a352588d6 100644 --- a/examples/mlops_starter/run.py +++ b/examples/mlops_starter/run.py @@ -239,8 +239,8 @@ def main( # to get the random state and target column random_state = preprocess_pipeline_artifact.run_metadata[ "random_state" - ].value - target = preprocess_pipeline_artifact.run_metadata["target"].value + ] + target = preprocess_pipeline_artifact.run_metadata["target"] run_args_inference["random_state"] = random_state run_args_inference["target"] = target diff --git a/examples/mlops_starter/steps/model_promoter.py b/examples/mlops_starter/steps/model_promoter.py index 52040638496..43d43ceac1f 100644 --- a/examples/mlops_starter/steps/model_promoter.py +++ b/examples/mlops_starter/steps/model_promoter.py @@ -58,11 +58,9 @@ def model_promoter(accuracy: float, stage: str = "production") -> bool: try: stage_model = client.get_model_version(current_model.name, stage) # We compare their metrics - prod_accuracy = ( - stage_model.get_artifact("sklearn_classifier") - .run_metadata["test_accuracy"] - .value - ) + prod_accuracy = stage_model.get_artifact( + "sklearn_classifier" + ).run_metadata["test_accuracy"] if float(accuracy) > float(prod_accuracy): # If current model has better metrics, we promote it is_promoted = True From 83dfe315977eab7a895f9afddc6f6360e7f83641 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 13:54:11 +0000 Subject: [PATCH 16/23] Auto-update of E2E template --- examples/e2e/.copier-answers.yml | 2 +- examples/e2e/steps/deployment/deployment_deploy.py | 2 +- examples/e2e/steps/hp_tuning/hp_tuning_select_best_model.py | 2 +- examples/e2e/steps/promotion/promote_with_metric_compare.py | 6 +++--- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/e2e/.copier-answers.yml b/examples/e2e/.copier-answers.yml index 74cc33d8594..b008b2c1e99 100644 --- a/examples/e2e/.copier-answers.yml +++ b/examples/e2e/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.10.10 +_commit: 2024.10.21 _src_path: gh:zenml-io/template-e2e-batch data_quality_checks: true email: info@zenml.io diff --git a/examples/e2e/steps/deployment/deployment_deploy.py b/examples/e2e/steps/deployment/deployment_deploy.py index 3fb0d879f3f..dad351e45be 100644 --- a/examples/e2e/steps/deployment/deployment_deploy.py +++ b/examples/e2e/steps/deployment/deployment_deploy.py @@ -67,7 +67,7 @@ def deployment_deploy() -> ( registry_model_name=model.name, registry_model_version=model.run_metadata[ "model_registry_version" - ].value, + ], replace_existing=True, ) else: diff --git a/examples/e2e/steps/hp_tuning/hp_tuning_select_best_model.py b/examples/e2e/steps/hp_tuning/hp_tuning_select_best_model.py index 7d5a6bc33ea..65e524ecd98 100644 --- a/examples/e2e/steps/hp_tuning/hp_tuning_select_best_model.py +++ b/examples/e2e/steps/hp_tuning/hp_tuning_select_best_model.py @@ -50,7 +50,7 @@ def hp_tuning_select_best_model( hp_output = model.get_data_artifact("hp_result") model_: ClassifierMixin = hp_output.load() # fetch metadata we attached earlier - metric = float(hp_output.run_metadata["metric"].value) + metric = float(hp_output.run_metadata["metric"]) if best_model is None or best_metric < metric: best_model = model_ ### YOUR CODE ENDS HERE ### diff --git a/examples/e2e/steps/promotion/promote_with_metric_compare.py b/examples/e2e/steps/promotion/promote_with_metric_compare.py index 038d219d32d..6bc580f47ba 100644 --- a/examples/e2e/steps/promotion/promote_with_metric_compare.py +++ b/examples/e2e/steps/promotion/promote_with_metric_compare.py @@ -92,14 +92,14 @@ def promote_with_metric_compare( # Promote in Model Registry latest_version_model_registry_number = latest_version.run_metadata[ "model_registry_version" - ].value + ] if current_version_number is None: current_version_model_registry_number = ( latest_version_model_registry_number ) else: current_version_model_registry_number = ( - current_version.run_metadata["model_registry_version"].value + current_version.run_metadata["model_registry_version"] ) promote_in_model_registry( latest_version=latest_version_model_registry_number, @@ -111,7 +111,7 @@ def promote_with_metric_compare( else: promoted_version = current_version.run_metadata[ "model_registry_version" - ].value + ] logger.info( f"Current model version in `{target_env}` is `{promoted_version}` registered in Model Registry" From 7888717a96b3c5cec491872e0e1595ba658be61f Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 13:57:27 +0000 Subject: [PATCH 17/23] Auto-update of NLP template --- examples/e2e_nlp/.copier-answers.yml | 2 +- examples/e2e_nlp/gradio/requirements.txt | 2 +- .../e2e_nlp/steps/promotion/promote_get_metrics.py | 12 ++++-------- 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/examples/e2e_nlp/.copier-answers.yml b/examples/e2e_nlp/.copier-answers.yml index 3ca2ba198fe..e509aae2760 100644 --- a/examples/e2e_nlp/.copier-answers.yml +++ b/examples/e2e_nlp/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.09.23 +_commit: 2024.10.21 _src_path: gh:zenml-io/template-nlp accelerator: cpu cloud_of_choice: aws diff --git a/examples/e2e_nlp/gradio/requirements.txt b/examples/e2e_nlp/gradio/requirements.txt index 1bddfdfb85b..b53f1df9e62 100644 --- a/examples/e2e_nlp/gradio/requirements.txt +++ b/examples/e2e_nlp/gradio/requirements.txt @@ -9,4 +9,4 @@ pandas==1.5.3 session_info==1.0.0 scikit-learn==1.5.0 transformers==4.28.1 -IPython==7.34.0 \ No newline at end of file +IPython==8.10.0 \ No newline at end of file diff --git a/examples/e2e_nlp/steps/promotion/promote_get_metrics.py b/examples/e2e_nlp/steps/promotion/promote_get_metrics.py index 7f2951a5865..b24ac42245c 100644 --- a/examples/e2e_nlp/steps/promotion/promote_get_metrics.py +++ b/examples/e2e_nlp/steps/promotion/promote_get_metrics.py @@ -56,9 +56,7 @@ def promote_get_metrics() -> ( # Get current model version metric in current run model = get_step_context().model - current_metrics = ( - model.get_model_artifact("model").run_metadata["metrics"].value - ) + current_metrics = model.get_model_artifact("model").run_metadata["metrics"] logger.info(f"Current model version metrics are {current_metrics}") # Get latest saved model version metric in target environment @@ -72,11 +70,9 @@ def promote_get_metrics() -> ( except KeyError: latest_version = None if latest_version: - latest_metrics = ( - latest_version.get_model_artifact("model") - .run_metadata["metrics"] - .value - ) + latest_metrics = latest_version.get_model_artifact( + "model" + ).run_metadata["metrics"] logger.info(f"Latest model version metrics are {latest_metrics}") else: logger.info("No currently promoted model version found.") From 70cc4a97c62b296b816da04a36b497ca374f69ea Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 15:07:24 +0000 Subject: [PATCH 18/23] Auto-update of LLM Finetuning template --- examples/llm_finetuning/.copier-answers.yml | 2 +- examples/llm_finetuning/steps/log_metadata.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/llm_finetuning/.copier-answers.yml b/examples/llm_finetuning/.copier-answers.yml index dd85e236760..4004897928b 100644 --- a/examples/llm_finetuning/.copier-answers.yml +++ b/examples/llm_finetuning/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.09.24-1-g378145b +_commit: 2024.10.30 _src_path: gh:zenml-io/template-llm-finetuning bf16: true cuda_version: cuda11.8 diff --git a/examples/llm_finetuning/steps/log_metadata.py b/examples/llm_finetuning/steps/log_metadata.py index 645f98cc8ea..14371b78b6e 100644 --- a/examples/llm_finetuning/steps/log_metadata.py +++ b/examples/llm_finetuning/steps/log_metadata.py @@ -34,7 +34,7 @@ def log_metadata_from_step_artifact( context = get_step_context() metadata_dict: Dict[str, Any] = ( - context.pipeline_run.steps[step_name].outputs[artifact_name][0].load() + context.pipeline_run.steps[step_name].outputs[artifact_name].load() ) metadata = {artifact_name: metadata_dict} From fcdec6eab8913c24f6c93f94c696b2ea7995f75b Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 15:07:44 +0000 Subject: [PATCH 19/23] Auto-update of Starter template --- examples/mlops_starter/.copier-answers.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/mlops_starter/.copier-answers.yml b/examples/mlops_starter/.copier-answers.yml index 21ba51bc459..fd6b937c7c9 100644 --- a/examples/mlops_starter/.copier-answers.yml +++ b/examples/mlops_starter/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.09.24-3-g2c1a682 +_commit: 2024.10.30 _src_path: gh:zenml-io/template-starter email: info@zenml.io full_name: ZenML GmbH From 0108c0f2b96b9ff333ad7380ac94a5a02373622a Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 15:14:35 +0000 Subject: [PATCH 20/23] Auto-update of E2E template --- examples/e2e/.copier-answers.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/e2e/.copier-answers.yml b/examples/e2e/.copier-answers.yml index 04c970cb9a4..cd687be59df 100644 --- a/examples/e2e/.copier-answers.yml +++ b/examples/e2e/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.10.10-5-g6edd133 +_commit: 2024.10.30 _src_path: gh:zenml-io/template-e2e-batch data_quality_checks: true email: info@zenml.io From 0c33f82ed1295c0838676ecb3f023fe294c83058 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Thu, 7 Nov 2024 15:17:27 +0000 Subject: [PATCH 21/23] Auto-update of NLP template --- examples/e2e_nlp/.copier-answers.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/e2e_nlp/.copier-answers.yml b/examples/e2e_nlp/.copier-answers.yml index a78ae8bac68..e13858e7da1 100644 --- a/examples/e2e_nlp/.copier-answers.yml +++ b/examples/e2e_nlp/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.10.21-3-ge37d83a +_commit: 2024.10.30 _src_path: gh:zenml-io/template-nlp accelerator: cpu cloud_of_choice: aws From 012cd6e1f6ed8c314e7113bc760d57639550c13d Mon Sep 17 00:00:00 2001 From: Safoine El khabich Date: Tue, 12 Nov 2024 15:32:58 +0100 Subject: [PATCH 22/23] Update default filenames and improve backward compatibility for sklearn materializer --- .../component-guide/model-deployers/vertex.md | 2 +- .../model-registries/vertex.md | 12 +- src/zenml/cli/model_registry.py | 3 +- .../model_deployers/vertex_model_deployer.py | 10 +- .../gcp/services/vertex_deployment.py | 103 +++++++++++++++--- .../materializers/sklearn_materializer.py | 84 +++++++++++--- .../materializers/cloudpickle_materializer.py | 2 +- .../model_registries/base_model_registry.py | 3 +- src/zenml/services/service.py | 3 + 9 files changed, 176 insertions(+), 46 deletions(-) diff --git a/docs/book/component-guide/model-deployers/vertex.md b/docs/book/component-guide/model-deployers/vertex.md index 98453df6744..727f2d763fe 100644 --- a/docs/book/component-guide/model-deployers/vertex.md +++ b/docs/book/component-guide/model-deployers/vertex.md @@ -76,7 +76,6 @@ from zenml.integrations.gcp.services.vertex_deployment import ( @step(enable_cache=False) def model_deployer( - model_registry_uri: str, ) -> Annotated[ VertexDeploymentService, ArtifactConfig(name="vertex_deployment", is_deployment_artifact=True) @@ -84,6 +83,7 @@ def model_deployer( """Model deployer step.""" zenml_client = Client() current_model = get_step_context().model + model_registry_uri = current_model.get_model_artifact("THE_MODEL_ARTIFACT_NAME_GIVEN_IN_TRAINING_STEP").uri model_deployer = zenml_client.active_stack.model_deployer # Configure the deployment diff --git a/docs/book/component-guide/model-registries/vertex.md b/docs/book/component-guide/model-registries/vertex.md index 41a29ffb11d..eef9096ce62 100644 --- a/docs/book/component-guide/model-registries/vertex.md +++ b/docs/book/component-guide/model-registries/vertex.md @@ -21,6 +21,12 @@ This is particularly useful in the following scenarios: {% hint style="warning" %} Important: The Vertex AI Model Registry implementation only supports the model version interface, not the model interface. This means you cannot register, delete, or update models directly - you can only work with model versions. Operations like `register_model()`, `delete_model()`, and `update_model()` are not supported. + +Unlike platforms like MLflow where you first create a model container and then add versions to it, Vertex AI combines model registration and versioning into a single operation: + +- When you upload a model, it automatically creates both the model and its first version +- Each subsequent upload with the same display name creates a new version +- You cannot create an empty model container without a version {% endhint %} ## How do you deploy it? @@ -141,9 +147,9 @@ Unlike the MLflow Model Registry, the Vertex AI implementation has some importan Based on the implementation, there are some limitations to be aware of: 1. The `register_model()`, `update_model()`, and `delete_model()` methods are not implemented as Vertex AI only supports registering model versions -2. Model stage transitions (Production, Staging, etc.) are not natively supported -3. Models must have a serving container image URI specified or will use the default scikit-learn image -4. All registered models are automatically labeled with `managed_by="zenml"` for tracking purposes +3. It's preferable for the models to be given a serving container image URI specified to avoid using the default scikit-learn prediction container and to ensure compatibility with Vertex AI endpoints +when deploying models. +4. All registered models by the integration are automatically labeled with `managed_by="zenml"` for tracking purposes Check out the [SDK docs](https://sdkdocs.zenml.io/latest/integration\_code\_docs/integrations-gcp/#zenml.integrations.gcp.model\_registry) to see more about the interface and implementation. diff --git a/src/zenml/cli/model_registry.py b/src/zenml/cli/model_registry.py index 93c55391817..838132b9e9a 100644 --- a/src/zenml/cli/model_registry.py +++ b/src/zenml/cli/model_registry.py @@ -18,6 +18,7 @@ import click +from zenml import __version__ from zenml.cli import utils as cli_utils from zenml.cli.cli import TagGroup, cli from zenml.enums import StackComponentType @@ -643,7 +644,7 @@ def register_model_version( # Parse metadata metadata = dict(metadata) if metadata else {} registered_metadata = ModelRegistryModelMetadata(**dict(metadata)) - registered_metadata.zenml_version = zenml_version + registered_metadata.zenml_version = zenml_version or __version__ registered_metadata.zenml_run_name = zenml_run_name registered_metadata.zenml_pipeline_name = zenml_pipeline_name registered_metadata.zenml_step_name = zenml_step_name diff --git a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py index b152839356b..3b6d31820cc 100644 --- a/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py +++ b/src/zenml/integrations/gcp/model_deployers/vertex_model_deployer.py @@ -22,9 +22,6 @@ from zenml.analytics.utils import track_handler from zenml.client import Client from zenml.enums import StackComponentType -from zenml.integrations.gcp import ( - VERTEX_SERVICE_ARTIFACT, -) from zenml.integrations.gcp.flavors.vertex_model_deployer_flavor import ( VertexModelDeployerConfig, VertexModelDeployerFlavor, @@ -142,14 +139,15 @@ def _create_deployment_service( config: the configuration of the model to be deployed with Vertex model deployer. Returns: - The VertexModelDeployerConfig object that can be used to interact + The VertexDeploymentService object that can be used to interact with the Vertex inference endpoint. """ # create a new service for the new model service = VertexDeploymentService(uuid=id, config=config) logger.info( - f"Creating an artifact {VERTEX_SERVICE_ARTIFACT} with service instance attached as metadata." - " If there's an active pipeline and/or model this artifact will be associated with it." + "Creating an artifact %s with service instance attached as metadata.", + "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 diff --git a/src/zenml/integrations/gcp/services/vertex_deployment.py b/src/zenml/integrations/gcp/services/vertex_deployment.py index 32e8a9722f7..c8a4e02f0a5 100644 --- a/src/zenml/integrations/gcp/services/vertex_deployment.py +++ b/src/zenml/integrations/gcp/services/vertex_deployment.py @@ -14,10 +14,12 @@ """Implementation of the Vertex AI Deployment service.""" import re +import time from typing import Any, Dict, Generator, List, Optional, Tuple from google.api_core import exceptions from google.cloud import aiplatform +from google.cloud import logging as vertex_logging from pydantic import BaseModel, Field from zenml.client import Client @@ -46,6 +48,10 @@ def sanitize_vertex_label(value: str) -> str: Returns: Sanitized label value """ + # Handle empty string + if not value: + return "" + # Convert to lowercase value = value.lower() # Replace any character that's not lowercase letter, number, dash or underscore @@ -81,10 +87,8 @@ class VertexPredictionServiceEndpoint(BaseModel): """Vertex AI Prediction Service Endpoint.""" endpoint_name: str + deployed_model_id: str endpoint_url: Optional[str] = None - deployed_model_id: Optional[str] = ( - None # Added to track specific model deployment - ) class VertexServiceStatus(ServiceStatus): @@ -107,10 +111,8 @@ class VertexDeploymentService(BaseDeploymentService): default_factory=lambda: VertexServiceStatus() ) - def __init__(self, config: VertexDeploymentConfig, **attrs: Any): - """Initialize the Vertex AI deployment service.""" - super().__init__(config=config, **attrs) - + def _initialize_gcp_clients(self) -> None: + """Initialize GCP clients with consistent credentials.""" # Initialize aiplatform with project and location from zenml.integrations.gcp.model_deployers.vertex_model_deployer import ( VertexModelDeployer, @@ -119,9 +121,29 @@ def __init__(self, config: VertexDeploymentConfig, **attrs: Any): zenml_client = Client() model_deployer = zenml_client.active_stack.model_deployer if not isinstance(model_deployer, VertexModelDeployer): - raise ValueError("Model deployer is not VertexModelDeployer") + raise RuntimeError( + "Active model deployer must be Vertex AI Model Deployer" + ) - model_deployer.setup_aiplatform() + # get credentials from model deployer + credentials, project_id = model_deployer._get_authentication() + + # Initialize aiplatform + aiplatform.init( + project=project_id, + location=self.config.location, + credentials=credentials, + ) + + # Initialize logging client + self.logging_client = vertex_logging.Client( + project=project_id, credentials=credentials + ) + + def __init__(self, config: VertexDeploymentConfig, **attrs: Any): + """Initialize the Vertex AI deployment service.""" + super().__init__(config=config, **attrs) + self._initialize_gcp_clients() @property def prediction_url(self) -> Optional[str]: @@ -187,6 +209,10 @@ def provision(self) -> None: logger.info( f"Found existing model to deploy: {model.resource_name} to the endpoint." ) + if not model: + raise RuntimeError( + f"Model {self.config.model_id} not found in the project." + ) # Deploy the model to the endpoint endpoint.deploy( @@ -332,7 +358,9 @@ def predict(self, instances: List[Any]) -> List[Any]: instances=instances, deployed_model_id=self.status.endpoint.deployed_model_id.split( "/" - )[-1], + )[-1] + if self.status.endpoint.deployed_model_id + else None, timeout=30, # Add reasonable timeout ) @@ -348,14 +376,53 @@ def predict(self, instances: List[Any]) -> List[Any]: def get_logs( self, follow: bool = False, tail: Optional[int] = None ) -> Generator[str, bool, None]: - """Retrieve the service logs.""" - # Note: Could be enhanced to actually fetch logs from Cloud Logging - logger.info( - "Vertex AI Endpoints provides access to the logs through " - "Cloud Logging. Please check the Google Cloud Console for detailed logs. " - f"Location: {self.config.location}" - ) - yield "Logs are available in Google Cloud Console." + """Retrieve the service logs from Cloud Logging. + + Args: + follow: If True, continuously yield new logs + tail: Number of most recent logs to return + """ + if not self.status.endpoint: + yield "No endpoint deployed yet" + return + + try: + # Create filter for Vertex AI endpoint logs + endpoint_id = self.status.endpoint.endpoint_name.split("/")[-1] + filter_str = ( + f'resource.type="aiplatform.googleapis.com/Endpoint" ' + f'resource.labels.endpoint_id="{endpoint_id}" ' + f'resource.labels.location="{self.config.location}"' + ) + + # Set time range for logs + if tail: + filter_str += f" limit {tail}" + + # Get log iterator + iterator = self.logging_client.list_entries( + filter_=filter_str, order_by=vertex_logging.DESCENDING + ) + + # Yield historical logs + for entry in iterator: + yield f"[{entry.timestamp}] {entry.severity}: {entry.payload.get('message', '')}" + + # If following logs, continue to stream new entries + if follow: + while True: + time.sleep(2) # Poll every 2 seconds + for entry in self.logging_client.list_entries( + filter_=filter_str, + order_by=vertex_logging.DESCENDING, + page_size=1, + ): + yield f"[{entry.timestamp}] {entry.severity}: {entry.payload.get('message', '')}" + + except Exception as e: + error_msg = f"Failed to retrieve logs: {str(e)}" + logger.error(error_msg) + yield error_msg @property def is_running(self) -> bool: diff --git a/src/zenml/integrations/sklearn/materializers/sklearn_materializer.py b/src/zenml/integrations/sklearn/materializers/sklearn_materializer.py index b11f7fe7080..d0b22d99e83 100644 --- a/src/zenml/integrations/sklearn/materializers/sklearn_materializer.py +++ b/src/zenml/integrations/sklearn/materializers/sklearn_materializer.py @@ -1,20 +1,9 @@ -# Copyright (c) ZenML GmbH 2021. 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 sklearn materializer.""" +import os from typing import Any, ClassVar, Tuple, Type +import cloudpickle from sklearn.base import ( BaseEstimator, BiclusterMixin, @@ -29,13 +18,20 @@ ) from zenml.enums import ArtifactType +from zenml.environment import Environment +from zenml.logger import get_logger from zenml.materializers.cloudpickle_materializer import ( + DEFAULT_FILENAME, CloudpickleMaterializer, ) +logger = get_logger(__name__) + +SKLEARN_MODEL_FILENAME = "model.pkl" + class SklearnMaterializer(CloudpickleMaterializer): - """Materializer to read data to and from sklearn.""" + """Materializer to read data to and from sklearn with backward compatibility.""" ASSOCIATED_TYPES: ClassVar[Tuple[Type[Any], ...]] = ( BaseEstimator, @@ -50,3 +46,63 @@ class SklearnMaterializer(CloudpickleMaterializer): TransformerMixin, ) ASSOCIATED_ARTIFACT_TYPE: ClassVar[ArtifactType] = ArtifactType.MODEL + + def load(self, data_type: Type[Any]) -> Any: + """Reads a sklearn model from pickle file with backward compatibility. + + Args: + data_type: The data type of the artifact. + + Returns: + The loaded sklearn model. + """ + # First try to load from model.pkl + model_filepath = os.path.join(self.uri, SKLEARN_MODEL_FILENAME) + artifact_filepath = os.path.join(self.uri, DEFAULT_FILENAME) + + # Check which file exists and load accordingly + if self.artifact_store.exists(model_filepath): + filepath = model_filepath + elif self.artifact_store.exists(artifact_filepath): + logger.info( + f"Loading from legacy filepath {artifact_filepath}. Future saves " + f"will use {model_filepath}" + ) + filepath = artifact_filepath + else: + raise FileNotFoundError( + f"Neither {model_filepath} nor {artifact_filepath} found in artifact store" + ) + + # validate python version before loading + source_python_version = self._load_python_version() + current_python_version = Environment().python_version() + if ( + source_python_version != "unknown" + and source_python_version != current_python_version + ): + logger.warning( + f"Your artifact was materialized under Python version " + f"'{source_python_version}' but you are currently using " + f"'{current_python_version}'. This might cause unexpected " + "behavior since pickle is not reproducible across Python " + "versions. Attempting to load anyway..." + ) + + # Load the model + with self.artifact_store.open(filepath, "rb") as fid: + return cloudpickle.load(fid) + + def save(self, data: Any) -> None: + """Saves a sklearn model to pickle file using the new filename. + + Args: + data: The sklearn model to save. + """ + # Save python version for validation on loading + self._save_python_version() + + # Save using the new filename + filepath = os.path.join(self.uri, SKLEARN_MODEL_FILENAME) + with self.artifact_store.open(filepath, "wb") as fid: + cloudpickle.dump(data, fid) diff --git a/src/zenml/materializers/cloudpickle_materializer.py b/src/zenml/materializers/cloudpickle_materializer.py index a6813cb4191..399ca7f2336 100644 --- a/src/zenml/materializers/cloudpickle_materializer.py +++ b/src/zenml/materializers/cloudpickle_materializer.py @@ -29,7 +29,7 @@ logger = get_logger(__name__) -DEFAULT_FILENAME = "model.pkl" +DEFAULT_FILENAME = "artifact.pkl" DEFAULT_PYTHON_VERSION_FILENAME = "python_version.txt" diff --git a/src/zenml/model_registries/base_model_registry.py b/src/zenml/model_registries/base_model_registry.py index ffab13f974a..3e3019dfb6f 100644 --- a/src/zenml/model_registries/base_model_registry.py +++ b/src/zenml/model_registries/base_model_registry.py @@ -20,7 +20,6 @@ from pydantic import BaseModel, ConfigDict -from zenml import __version__ from zenml.enums import StackComponentType from zenml.stack import Flavor, StackComponent from zenml.stack.stack_component import StackComponentConfig @@ -64,7 +63,7 @@ class ModelRegistryModelMetadata(BaseModel): """ _managed_by: str = "zenml" - zenml_version: str = __version__ + zenml_version: Optional[str] = None zenml_run_name: Optional[str] = None zenml_pipeline_name: Optional[str] = None zenml_pipeline_uuid: Optional[str] = None diff --git a/src/zenml/services/service.py b/src/zenml/services/service.py index 7b607aae611..0077a3a945f 100644 --- a/src/zenml/services/service.py +++ b/src/zenml/services/service.py @@ -35,6 +35,7 @@ from zenml.console import console from zenml.logger import get_logger +from zenml.model.model import Model from zenml.services.service_endpoint import BaseServiceEndpoint from zenml.services.service_monitor import HTTPEndpointHealthMonitor from zenml.services.service_status import ServiceState, ServiceStatus @@ -109,6 +110,7 @@ class ServiceConfig(BaseTypedModel): pipeline_name: name of the pipeline that spun up the service pipeline_step_name: name of the pipeline step that spun up the service run_name: name of the pipeline run that spun up the service. + zenml_model: the ZenML model object to be deployed. """ name: str = "" @@ -118,6 +120,7 @@ class ServiceConfig(BaseTypedModel): model_name: str = "" model_version: str = "" service_name: str = "" + zenml_model: Optional[Model] = None # TODO: In Pydantic v2, the `model_` is a protected namespaces for all # fields defined under base models. If not handled, this raises a warning. From ac2e69a9de424e17ae78bd42367d7b673a958c86 Mon Sep 17 00:00:00 2001 From: GitHub Actions Date: Tue, 12 Nov 2024 14:35:13 +0000 Subject: [PATCH 23/23] Auto-update of LLM Finetuning template --- examples/llm_finetuning/.copier-answers.yml | 2 +- examples/llm_finetuning/requirements.txt | 2 +- examples/llm_finetuning/steps/log_metadata.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/llm_finetuning/.copier-answers.yml b/examples/llm_finetuning/.copier-answers.yml index 4004897928b..2c547f98d61 100644 --- a/examples/llm_finetuning/.copier-answers.yml +++ b/examples/llm_finetuning/.copier-answers.yml @@ -1,5 +1,5 @@ # Changes here will be overwritten by Copier -_commit: 2024.10.30 +_commit: 2024.11.08 _src_path: gh:zenml-io/template-llm-finetuning bf16: true cuda_version: cuda11.8 diff --git a/examples/llm_finetuning/requirements.txt b/examples/llm_finetuning/requirements.txt index ac6d8625411..23d38eef649 100644 --- a/examples/llm_finetuning/requirements.txt +++ b/examples/llm_finetuning/requirements.txt @@ -1,6 +1,6 @@ zenml torch>=2.2.0 -datasets +datasets>=2.15 transformers peft bitsandbytes>=0.41.3 diff --git a/examples/llm_finetuning/steps/log_metadata.py b/examples/llm_finetuning/steps/log_metadata.py index 14371b78b6e..645f98cc8ea 100644 --- a/examples/llm_finetuning/steps/log_metadata.py +++ b/examples/llm_finetuning/steps/log_metadata.py @@ -34,7 +34,7 @@ def log_metadata_from_step_artifact( context = get_step_context() metadata_dict: Dict[str, Any] = ( - context.pipeline_run.steps[step_name].outputs[artifact_name].load() + context.pipeline_run.steps[step_name].outputs[artifact_name][0].load() ) metadata = {artifact_name: metadata_dict}