diff --git a/examples/e2e/.copier-answers.yml b/examples/e2e/.copier-answers.yml index f8f1b62c353..70e3f86e323 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: 2023.11.24 +_commit: 2023.12.06 _src_path: gh:zenml-io/template-e2e-batch data_quality_checks: true email: '' diff --git a/examples/e2e/steps/deployment/deployment_deploy.py b/examples/e2e/steps/deployment/deployment_deploy.py index 7fa359acc10..b0b905b52c9 100644 --- a/examples/e2e/steps/deployment/deployment_deploy.py +++ b/examples/e2e/steps/deployment/deployment_deploy.py @@ -65,9 +65,9 @@ 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/examples/e2e/steps/promotion/promote_with_metric_compare.py b/examples/e2e/steps/promotion/promote_with_metric_compare.py index ad02eadad39..0bf1a931950 100644 --- a/examples/e2e/steps/promotion/promote_with_metric_compare.py +++ b/examples/e2e/steps/promotion/promote_with_metric_compare.py @@ -90,21 +90,17 @@ 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, @@ -113,11 +109,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/examples/e2e/steps/training/model_trainer.py b/examples/e2e/steps/training/model_trainer.py index 80824b05a94..2fde436abf0 100644 --- a/examples/e2e/steps/training/model_trainer.py +++ b/examples/e2e/steps/training/model_trainer.py @@ -20,7 +20,7 @@ from sklearn.base import ClassifierMixin from typing_extensions import Annotated -from zenml import ArtifactConfig, log_artifact_metadata, step +from zenml import ArtifactConfig, get_step_context, step from zenml.client import Client from zenml.integrations.mlflow.experiment_trackers import ( MLFlowExperimentTracker, @@ -103,9 +103,9 @@ 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 ###