Refactor model deployment logic and update README with new configurations

Author: safoinme

vertex-registry-and-deployer/README.md

@@ -26,16 +26,16 @@ Once you have set up your Google Cloud credentials, we can create a stack and ru
 
 ```bash
 # Register the artifact store
-zenml artifact-store register gs_store -f gcp --path=gs://bucket-name
+zenml artifact-store register gs_store -f gcp --path=gs://zenml-vertex-test
 zenml artifact-store connect gs_store --connector gcp
 
 # Register the model registry
-zenml model-registry register vertex_registry --flavor=vertex --location=us-central1
+zenml model-registry register vertex_registry --flavor=vertex --location=europe-west1 
 zenml model-registry connect vertex_registry --connector gcp
 
 # Register Model Deployer
-zenml model-deployer register vertex_deployer --flavor=vertex --location=us-central1
-zenml model-deployer connect vertex_deployer --connector vertex_deployer_connector
+zenml model-deployer register vertex_deployer --flavor=vertex --location=europe-west1
+zenml model-deployer connect vertex_deployer --connector gcp
 
 # Register the stack
 zenml stack register vertex_stack --orchestrator default --artifact-store gs_store --model-registry vertex_registry --model-deployer vertex_deployer

vertex-registry-and-deployer/pipelines/inference.py

@@ -14,16 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from zenml import get_pipeline_context, pipeline
+from zenml import pipeline
 from zenml.logger import get_logger
 
 logger = get_logger(__name__)
 
 
 @pipeline
 def inference(random_state: int, target: str):
-    """
-    Model inference pipeline.
+    """Model inference pipeline.
 
     This is a pipeline that loads the inference data, processes it with
     the same preprocessing pipeline used in training, and runs inference
@@ -34,12 +33,12 @@ def inference(random_state: int, target: str):
         target: Name of target column in dataset.
     """
     # Get the production model artifact
-    model = get_pipeline_context().model.get_artifact("sklearn_classifier")
+    # model_artifact = get_pipeline_context().model.get_artifact("sklearn_classifier")
 
     # Get the preprocess pipeline artifact associated with this version
-    preprocess_pipeline = get_pipeline_context().model.get_artifact(
-        "preprocess_pipeline"
-    )
+    # preprocess_pipeline = get_pipeline_context().model.get_artifact(
+    #    "preprocess_pipeline"
+    # )
 
     # Link all the steps together by calling them and passing the output
     #  of one step as the input of the next step.

vertex-registry-and-deployer/pipelines/training.py

@@ -16,11 +16,10 @@
 #
 
 from typing import Optional
-from uuid import UUID
 
 from steps import model_deployer, model_promoter, model_register, model_trainer
+
 from zenml import pipeline
-from zenml.client import Client
 from zenml.logger import get_logger
 
 logger = get_logger(__name__)
@@ -47,9 +46,9 @@ def training(
     # Link all the steps together by calling them and passing the output
     # of one step as the input of the next step.
 
-    model, accuracy = model_trainer(target=target)
+    _, accuracy = model_trainer(target=target)
     is_promoted = model_promoter(accuracy=accuracy)
-    if is_promoted:
-        model_registry_uri = model_register()
-        model_deployer(model_registry_uri=model_registry_uri)
-    
+    model_registry_uri = model_register(is_promoted)
+    model_deployer(
+        model_registry_uri=model_registry_uri, is_promoted=is_promoted
+    )

vertex-registry-and-deployer/run.py

@@ -16,14 +16,14 @@
 #
 
 import os
-from typing import Optional
 
 import click
 import yaml
 from pipelines import (
     inference,
     training,
 )
+
 from zenml.client import Client
 from zenml.logger import get_logger
 

vertex-registry-and-deployer/steps/model_deployer.py

@@ -1,27 +1,9 @@
-# Apache Software License 2.0
-#
-# 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.
-#
-
-from typing import Any
-
 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,
 )
 from zenml.logger import get_logger
@@ -32,33 +14,46 @@
 @step(enable_cache=False)
 def model_deployer(
     model_registry_uri: str,
+    is_promoted: bool = False,
 ) -> Annotated[
-    VertexDeploymentService, ArtifactConfig(name="vertex_deployment", is_deployment_artifact=True)
+    VertexDeploymentService,
+    ArtifactConfig(name="vertex_deployment", is_deployment_artifact=True),
 ]:
     """Model deployer step.
-    
+
     Args:
         model_registry_uri: URI of the model in the model registry.
-    
+        is_promoted: Whether the model was promoted to production.
+
     Returns:
         The deployed model service.
     """
-    zenml_client = Client()
-    current_model = get_step_context().model
-    model_deployer = zenml_client.active_stack.model_deployer
-    vertex_deployment_config = VertexAIDeploymentConfig(
-        location="europe-west1",
-        name="zenml-vertex-quickstart",
-        model_name=current_model.name,
-        description="An example of deploying a model using the MLflow Model Deployer",
-        model_id=model_registry_uri,
-    )
-    service = model_deployer.deploy_model(
-        config=vertex_deployment_config,
-        service_type=VertexDeploymentService.SERVICE_TYPE,
-    )
+    if not is_promoted:
+        logger.info("Model not promoted, skipping deployment")
+        return None
+    else:
+        zenml_client = Client()
+        current_model = get_step_context().model
+        model_deployer = zenml_client.active_stack.model_deployer
+
+        # Create deployment configuration with proper model name and version
+        vertex_deployment_config = VertexDeploymentConfig(
+            location="europe-west1",
+            name=current_model.name,
+            display_name="zenml-vertex-quickstart",
+            model_name=model_registry_uri,  # This is the full resource name from registration
+            model_version=current_model.version,  # Specify the version explicitly
+            description="An example of deploying a model using the Vertex AI Model Deployer",
+            sync=True,  # Wait for deployment to complete
+            traffic_percentage=100,  # Route all traffic to this version
+        )
+
+        service = model_deployer.deploy_model(
+            config=vertex_deployment_config,
+            service_type=VertexDeploymentService.SERVICE_TYPE,
+        )
 
-    logger.info(
-        f"The deployed service info: {model_deployer.get_model_server_info(service)}"
-    )
-    return service
+        logger.info(
+            f"The deployed service info: {model_deployer.get_model_server_info(service)}"
+        )
+        return service

vertex-registry-and-deployer/steps/model_promoter.py

@@ -58,9 +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 accuracy > float(prod_accuracy):
                 # If current model has better metrics, we promote it
                 is_promoted = True

vertex-registry-and-deployer/steps/model_register.py

@@ -1,28 +1,77 @@
-# model_register.py
-
 from typing_extensions import Annotated
+
 from zenml import ArtifactConfig, get_step_context, step
 from zenml.client import Client
+from zenml.integrations.gcp.flavors.vertex_base_config import (
+    VertexAIContainerSpec,
+    VertexAIModelConfig,
+    VertexAIResourceSpec,
+)
 from zenml.logger import get_logger
+from zenml.model_registries.base_model_registry import (
+    ModelRegistryModelMetadata,
+)
 
 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
+def model_register(
+    is_promoted: bool = False,
+) -> Annotated[str, ArtifactConfig(name="model_registry_uri")]:
+    """Model registration step.
+
+    Registers a model version in the Vertex AI Model Registry with extended configuration
+    and returns the model's source URI. This configuration embeds details such as container,
+    resource, explanation settings etc. so that the deployment can reuse these pre-configured
+    settings from the registry.
+    """
+    if is_promoted:
+        # Get the current model from the step context
+        current_model = get_step_context().model
+
+        client = Client()
+        model_registry = client.active_stack.model_registry
+        # Create an extended model configuration using the base settings for Vertex AI
+        model_config = VertexAIModelConfig(
+            location="europe-west1",
+            container=VertexAIContainerSpec(
+                image_uri="europe-docker.pkg.dev/vertex-ai/prediction/sklearn-cpu.1-5:latest",
+                predict_route="predict",
+                health_route="health",
+                ports=[8080],
+            ),
+            resources=VertexAIResourceSpec(
+                machine_type="n1-standard-4",
+                min_replica_count=1,
+                max_replica_count=1,
+            ),
+            labels={"env": "production"},
+            description="Extended model configuration for Vertex AI",
+        )
+
+        # Register the model version with the extended configuration as metadata
+        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 version registered with extended configuration",
+            metadata=ModelRegistryModelMetadata(
+                zenml_pipeline_name=get_step_context().pipeline.name,
+                zenml_pipeline_run_uuid=str(
+                    get_step_context().pipeline_run.id
+                ),
+                zenml_step_name=get_step_context().step_run.name,
+            ),
+            config=model_config,
+        )
+        logger.info(
+            f"Model version {model_version.version} registered in Model Registry"
+        )
+
+        # Return the full resource name of the registered model
+        return model_version.registered_model.name
+    else:
+        return ""

vertex-registry-and-deployer/steps/model_trainer.py

@@ -1,4 +1,3 @@
-
 # Apache Software License 2.0
 #
 # Copyright (c) ZenML GmbH 2024. All rights reserved.
@@ -16,7 +15,7 @@
 # limitations under the License.
 #
 
-from typing import Optional, Tuple, Union
+from typing import Tuple, Union
 
 import pandas as pd
 from sklearn.base import ClassifierMixin
@@ -26,12 +25,16 @@
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import MinMaxScaler
 from typing_extensions import Annotated
-from zenml import ArtifactConfig, log_artifact_metadata, step
+
+from zenml import ArtifactConfig, step
+from zenml.enums import ArtifactType
 from zenml.logger import get_logger
+from zenml.utils.metadata_utils import log_metadata
 
 logger = get_logger(__name__)
 
-@step
+
+@step(enable_cache=False)
 def model_trainer(
     random_state: int = 42,
     test_size: float = 0.2,
@@ -41,7 +44,12 @@ def model_trainer(
     min_train_accuracy: float = 0.3,
     min_test_accuracy: float = 0.3,
 ) -> Tuple[
-    Annotated[ClassifierMixin, ArtifactConfig(name="sklearn_classifier", is_model_artifact=True)],
+    Annotated[
+        ClassifierMixin,
+        ArtifactConfig(
+            name="sklearn_classifier", artifact_type=ArtifactType.MODEL
+        ),
+    ],
     Annotated[float, ArtifactConfig(name="accuracy")],
 ]:
     # Load the dataset
@@ -99,21 +107,22 @@ def model_trainer(
         for message in messages:
             logger.warning(message)
 
-    log_artifact_metadata(
+    log_metadata(
         metadata={
             "train_accuracy": float(trn_acc),
             "test_accuracy": float(tst_acc),
         },
         artifact_name="sklearn_classifier",
+        infer_artifact=True,
     )
-
     return model, tst_acc
 
+
 class NADropper:
     """Support class to drop NA values in sklearn Pipeline."""
 
     def fit(self, *args, **kwargs):  # noqa: D102
         return self
 
     def transform(self, X: Union[pd.DataFrame, pd.Series]):  # noqa: D102
-        return X.dropna()
+        return X.dropna()