Add production deployment pipeline and related steps for model deploy…

…ment

train_and_deploy/configs/deploy_production.yaml

@@ -0,0 +1,45 @@
+# Apache Software License 2.0
+# 
+# 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
+# 
+# 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.
+# 
+
+# environment configuration
+settings:
+  docker:
+    python_package_installer: uv
+    required_integrations:
+      - aws
+      - sklearn
+      - bentoml
+
+
+# configuration of steps  
+steps:
+  notify_on_success:
+    parameters:
+      notify_on_success: False
+
+# configuration of the Model Control Plane
+model:
+  name: gitguarden
+  version: staging
+
+# pipeline level extra configurations
+extra:
+  notify_on_failure: True
+
+
+parameters:
+  target_env: staging

train_and_deploy/pipelines/__init__.py

@@ -19,3 +19,4 @@
 from .batch_inference import gitguarden_batch_inference
 from .training import gitguarden_training
 from .local_deployment import gitguarden_local_deployment
+from .deploy_production import gitguarden_production_deployment

train_and_deploy/pipelines/deploy_production.py

@@ -0,0 +1,40 @@
+# Apache Software License 2.0
+#
+# 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
+#
+# 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 steps import dockerize_bento_model, notify_on_failure, notify_on_success, deploy_model_to_k8s
+
+from zenml import pipeline
+
+
+@pipeline(on_failure=notify_on_failure, enable_cache=False)
+def gitguarden_production_deployment(
+    target_env: str,
+):
+    """Model deployment pipeline.
+
+    This is a pipeline deploys trained model for future inference.
+    """
+    ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ###
+    # Link all the steps together by calling them and passing the output
+    # of one step as the input of the next step.
+    ########## Deployment stage ##########
+    # Get the production model artifact
+    bento_model_image = dockerize_bento_model(target_env=target_env)
+    deploy_model_to_k8s(bento_model_image)
+
+    notify_on_success(after=["deploy_model_to_k8s"])
+    ### YOUR CODE ENDS HERE ###

train_and_deploy/run.py

@@ -23,6 +23,7 @@
 from pipelines import (
     gitguarden_batch_inference,
     gitguarden_local_deployment,
+    gitguarden_production_deployment,
     gitguarden_training,
 )
 from zenml.logger import get_logger
@@ -133,6 +134,12 @@
     default=False,
     help="Whether to run the inference pipeline.",
 )
+@click.option(
+    "--production",
+    is_flag=True,
+    default=False,
+    help="Whether to run the production pipeline.",
+)
 def main(
     no_cache: bool = False,
     no_drop_na: bool = False,
@@ -145,6 +152,7 @@ def main(
     training: bool = True,
     deployment: bool = False,
     inference: bool = False,
+    production: bool = False,
 ):
     """Main entry point for the pipeline execution.
 
@@ -166,6 +174,7 @@ def main(
             thresholds are violated - the pipeline will fail. If `False` thresholds will
             not affect the pipeline.
         only_inference: If `True` only inference pipeline will be triggered.
+        production: If `True` only production pipeline will be triggered.
     """
     # Run a pipeline with the required parameters. This executes
     # all steps in the pipeline in the correct order using the orchestrator
@@ -225,6 +234,20 @@ def main(
         gitguarden_batch_inference.with_options(**pipeline_args)(
             **run_args_inference
         )
+    if production:
+        # Execute Production Pipeline
+        run_args_production = {}
+        pipeline_args["config_path"] = os.path.join(
+            os.path.dirname(os.path.realpath(__file__)),
+            "configs",
+            "deploy_production.yaml",
+        )
+        pipeline_args["run_name"] = (
+            f"gitguarden_production_run_{dt.now().strftime('%Y_%m_%d_%H_%M_%S')}"
+        )
+        gitguarden_production_deployment.with_options(**pipeline_args)(
+            **run_args_production
+        )
 
 
 if __name__ == "__main__":

train_and_deploy/service.py

@@ -1,14 +1,32 @@
+import bentoml
+import numpy as np
+from bentoml.validators import Shape
+from typing_extensions import Annotated
 
 
-import bentoml
-from bentoml.io import NumpyNdarray
+@bentoml.service
+class GitGuarden:
+    """
+    A simple service using a sklearn model
+    """
 
-gitguarden_runner = bentoml.sklearn.get("gitguarden").to_runner()
+    # Load in the class scope to declare the model as a dependency of the service
+    iris_model = bentoml.models.get("gitguarden:latest")
 
-svc = bentoml.Service(name="gitguarden_service", runners=[gitguarden_runner])
+    def __init__(self):
+        """
+        Initialize the service by loading the model from the model store
+        """
+        import joblib
 
-input_spec = NumpyNdarray(dtype="float", shape=(-1, 30))
+        self.model = joblib.load(self.iris_model.path_of("saved_model.pkl"))
 
-@svc.api(input=input_spec, output=NumpyNdarray())
-async def predict(input_arr):
-    return await gitguarden_runner.predict.async_run(input_arr)
+    @bentoml.api
+    def predict(
+        self,
+        input_series: Annotated[np.ndarray, Shape((-1, 30))],
+    ) -> np.ndarray:
+        """
+        Define API with preprocessing and model inference logic
+        """
+        return self.model.predict(input_series)

train_and_deploy/steps/__init__.py

@@ -31,4 +31,4 @@
     promote_with_metric_compare,
 )
 from .training import model_evaluator, model_trainer
-from .deployment import deployment_deploy, bento_builder
+from .deployment import deployment_deploy, bento_builder, dockerize_bento_model, deploy_model_to_k8s

train_and_deploy/steps/deployment/__init__.py

@@ -18,3 +18,5 @@
 
 from .deployment_deploy import deployment_deploy
 from .bento_builder import bento_builder
+from .dockerize_bento import dockerize_bento_model
+from .deploy_to_k8s import  deploy_model_to_k8s

train_and_deploy/steps/deployment/bento_builder.py

@@ -63,7 +63,7 @@ def bento_builder() -> (
         bento_model = bentoml.sklearn.save_model(model.name, model.load_artifact(name="model"))
         # Build the BentoML bundle
         bento = bentos.build(
-            service="service.py:svc",
+            service="service.py:GitGuarden",
             labels={
                 "zenml_version": zenml_version,
                 "model_name": model.name,

train_and_deploy/steps/deployment/deploy_to_k8s.py

@@ -0,0 +1,164 @@
+#  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.
+from pathlib import Path
+from typing import Dict, Optional
+
+import yaml
+from kubernetes import client, config
+from kubernetes.client.rest import ApiException
+from zenml import get_step_context, step
+from zenml.client import Client
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
+
+def apply_kubernetes_configuration(k8s_configs: list) -> None:
+    """Apply Kubernetes configurations using the K8s Python client.
+    
+    Args:
+        k8s_configs: List of Kubernetes configuration dictionaries
+    """
+    # Load Kubernetes configuration
+    try:
+        config.load_kube_config()
+    except:
+        config.load_incluster_config()  # For in-cluster deployment
+
+    # Initialize API clients
+    k8s_apps_v1 = client.AppsV1Api()
+    k8s_core_v1 = client.CoreV1Api()
+
+    for k8s_config in k8s_configs:
+        kind = k8s_config["kind"]
+        name = k8s_config["metadata"]["name"]
+        namespace = k8s_config["metadata"].get("namespace", "default")
+
+        try:
+            if kind == "Deployment":
+                # Check if deployment exists
+                try:
+                    k8s_apps_v1.read_namespaced_deployment(name, namespace)
+                    # Update existing deployment
+                    k8s_apps_v1.patch_namespaced_deployment(
+                        name=name,
+                        namespace=namespace,
+                        body=k8s_config
+                    )
+                    logger.info(f"Updated existing deployment: {name}")
+                except ApiException as e:
+                    if e.status == 404:
+                        # Create new deployment
+                        k8s_apps_v1.create_namespaced_deployment(
+                            namespace=namespace,
+                            body=k8s_config
+                        )
+                        logger.info(f"Created new deployment: {name}")
+                    else:
+                        raise e
+
+            elif kind == "Service":
+                # Check if service exists
+                try:
+                    k8s_core_v1.read_namespaced_service(name, namespace)
+                    # Update existing service
+                    k8s_core_v1.patch_namespaced_service(
+                        name=name,
+                        namespace=namespace,
+                        body=k8s_config
+                    )
+                    logger.info(f"Updated existing service: {name}")
+                except ApiException as e:
+                    if e.status == 404:
+                        # Create new service
+                        k8s_core_v1.create_namespaced_service(
+                            namespace=namespace,
+                            body=k8s_config
+                        )
+                        logger.info(f"Created new service: {name}")
+                    else:
+                        raise e
+
+        except ApiException as e:
+            logger.error(f"Error applying {kind} {name}: {e}")
+            raise e
+
+@step
+def deploy_model_to_k8s(
+    docker_image_tag: str,
+    namespace: str = "default"
+) -> Dict:
+    """Deploy a service to Kubernetes with the specified docker image and tag.
+    
+    Args:
+        docker_image: The full docker image name (e.g. "organization/image-name")
+        docker_image_tag: The tag to use for the docker image
+        namespace: Kubernetes namespace to deploy to (default: "default")
+        
+    Returns:
+        dict: Dictionary containing deployment information
+    """
+    # Get model name from context
+    model_name = get_step_context().model.name
+
+    # Read the K8s template
+    template_path = Path(__file__).parent / "k8s_template.yaml"
+    with open(template_path, "r") as f:
+        # Load all documents in the YAML file
+        k8s_configs = list(yaml.safe_load_all(f))
+
+    # Update both Service and Deployment configurations
+    for config in k8s_configs:
+        # Add namespace
+        config["metadata"]["namespace"] = namespace
+
+        # Update metadata labels and name
+        config["metadata"]["labels"]["app"] = model_name
+        config["metadata"]["name"] = model_name
+
+        if config["kind"] == "Service":
+            # Update service selector
+            config["spec"]["selector"]["app"] = model_name
+
+        elif config["kind"] == "Deployment":
+            # Update deployment selector and template
+            config["spec"]["selector"]["matchLabels"]["app"] = model_name
+            config["spec"]["template"]["metadata"]["labels"]["app"] = model_name
+
+            # Update the container image and name
+            containers = config["spec"]["template"]["spec"]["containers"]
+            for container in containers:
+                container["name"] = model_name
+                container["image"] = docker_image_tag
+
+    # Apply the configurations
+    try:
+        apply_kubernetes_configuration(k8s_configs)
+        deployment_status = "success"
+        logger.info(f"Successfully deployed model {model_name} with image: {docker_image_tag}")
+    except Exception as e:
+        deployment_status = "failed"
+        logger.error(f"Failed to deploy model {model_name}: {str(e)}")
+        raise e
+
+    # Return deployment information
+    deployment_info = {
+        "model_name": model_name,
+        "docker_image": docker_image_tag,
+        "namespace": namespace,
+        "status": deployment_status,
+        "service_port": 3000,
+        "configurations": k8s_configs
+    }
+
+    return deployment_info