From 30a41465708a9aedc3d3160e4815b0e8d75c70d7 Mon Sep 17 00:00:00 2001 From: Andrei Vishniakov <31008759+avishniakov@users.noreply.github.com> Date: Fri, 15 Dec 2023 14:51:31 +0100 Subject: [PATCH] Model version level metadata (#26) * model version metadata * update zenml ref --- .github/workflows/ci.yml | 3 +-- template/steps/deployment/deployment_deploy.py | 4 +--- ..._promotion %}promote_with_metric_compare.py{% endif %} | 6 +++--- ...mpare_promotion %}promote_latest_version.py{% endif %} | 4 ++-- template/steps/training/model_trainer.py | 8 +++----- 5 files changed, 10 insertions(+), 15 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1223b4e..4f5d8e9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -56,6 +56,5 @@ jobs: with: stack-name: ${{ matrix.stack-name }} python-version: ${{ matrix.python-version }} - ref-zenml: ${{ inputs.ref-zenml || 'feature/OSS-2529-passing-pipeline-parameters-as-yaml-and-document' }} + ref-zenml: ${{ inputs.ref-zenml || 'feature/OSS-2574-add-metadata-to-model-versions' }} ref-template: ${{ inputs.ref-template || github.ref }} - \ No newline at end of file diff --git a/template/steps/deployment/deployment_deploy.py b/template/steps/deployment/deployment_deploy.py index 9aa6227..d840b1b 100644 --- a/template/steps/deployment/deployment_deploy.py +++ b/template/steps/deployment/deployment_deploy.py @@ -47,9 +47,7 @@ def deployment_deploy() -> ( # deploy predictor service deployment_service = mlflow_model_registry_deployer_step.entrypoint( registry_model_name=model_version.name, - registry_model_version=model_version.get_model_artifact("model") - .run_metadata["model_registry_version"] - .value, + registry_model_version=model_version.metadata["model_registry_version"], replace_existing=True, ) else: diff --git a/template/steps/promotion/{% if metric_compare_promotion %}promote_with_metric_compare.py{% endif %} b/template/steps/promotion/{% if metric_compare_promotion %}promote_with_metric_compare.py{% endif %} index 712c420..70e702e 100644 --- a/template/steps/promotion/{% if metric_compare_promotion %}promote_with_metric_compare.py{% endif %} +++ b/template/steps/promotion/{% if metric_compare_promotion %}promote_with_metric_compare.py{% endif %} @@ -74,11 +74,11 @@ def promote_with_metric_compare( logger.info(f"Current model version was promoted to '{target_env}'.") # Promote in Model Registry - latest_version_model_registry_number = latest_version.get_model_artifact("model").run_metadata["model_registry_version"].value + latest_version_model_registry_number = latest_version.metadata["model_registry_version"] if current_version_number is None: current_version_model_registry_number = latest_version_model_registry_number else: - current_version_model_registry_number = current_version.get_model_artifact("model").run_metadata["model_registry_version"].value + current_version_model_registry_number = current_version.metadata["model_registry_version"] promote_in_model_registry( latest_version=latest_version_model_registry_number, current_version=current_version_model_registry_number, @@ -87,7 +87,7 @@ def promote_with_metric_compare( ) promoted_version = latest_version_model_registry_number else: - promoted_version = current_version.get_model_artifact("model").run_metadata["model_registry_version"].value + promoted_version = current_version.metadata["model_registry_version"] logger.info( f"Current model version in `{target_env}` is `{promoted_version}` registered in Model Registry" diff --git a/template/steps/promotion/{% if not metric_compare_promotion %}promote_latest_version.py{% endif %} b/template/steps/promotion/{% if not metric_compare_promotion %}promote_latest_version.py{% endif %} index 9b8273f..7bf4728 100644 --- a/template/steps/promotion/{% if not metric_compare_promotion %}promote_latest_version.py{% endif %} +++ b/template/steps/promotion/{% if not metric_compare_promotion %}promote_latest_version.py{% endif %} @@ -32,11 +32,11 @@ def promote_latest_version( logger.info(f"Current model version was promoted to '{target_env}'.") # Promote in Model Registry - latest_version_model_registry_number = latest_version.get_model_artifact("model").run_metadata["model_registry_version"].value + latest_version_model_registry_number = latest_version.metadata["model_registry_version"] if current_version.number is None: current_version_model_registry_number = latest_version_model_registry_number else: - current_version_model_registry_number = current_version.get_model_artifact("model").run_metadata["model_registry_version"].value + current_version_model_registry_number = current_version.metadata["model_registry_version"] promote_in_model_registry( latest_version=latest_version_model_registry_number, current_version=current_version_model_registry_number, diff --git a/template/steps/training/model_trainer.py b/template/steps/training/model_trainer.py index 39f7926..111f84e 100644 --- a/template/steps/training/model_trainer.py +++ b/template/steps/training/model_trainer.py @@ -5,7 +5,7 @@ import mlflow import pandas as pd from sklearn.base import ClassifierMixin -from zenml import ArtifactConfig, log_artifact_metadata, step +from zenml import ArtifactConfig, log_artifact_metadata, step, get_step_context from zenml.client import Client from zenml.integrations.mlflow.experiment_trackers import MLFlowExperimentTracker from zenml.integrations.mlflow.steps.mlflow_registry import mlflow_register_model_step @@ -84,10 +84,8 @@ def model_trainer( if model_registry: versions = model_registry.list_model_versions(name=name) if versions: - log_artifact_metadata( - metadata={"model_registry_version": versions[-1].version}, - artifact_name="model", - ) + model_version = get_step_context().model_version + model_version.log_metadata({"model_registry_version": versions[-1].version}) ### YOUR CODE ENDS HERE ### return model