Rename model version to a model

.github/workflows/ci.yml

@@ -56,5 +56,5 @@ jobs:
         with:
           stack-name: ${{ matrix.stack-name }}
           python-version: ${{ matrix.python-version }}
-          ref-zenml: ${{ inputs.ref-zenml || 'feature/OSS-2609-endpoint_artifacts_rename_new' }}
+          ref-zenml: ${{ inputs.ref-zenml || 'feature/OSSK-342-rename-model-version-to-a-model' }}
           ref-template: ${{ inputs.ref-template || github.ref }}

README.md

@@ -94,7 +94,7 @@ We will be going section by section diving into implementation details and shari
 Training pipeline is designed to create a new Model Control Plane version and promote it to inference stage upon successfully passing the quality assurance at the end of the pipeline. This ensures that we always infer only on quality-assured Model Control Plane version and provides a seamless integration of required artifacts of this Model Control Plane version later on inference runs.
 This is achieved by providing this configuration in `train_config.yaml` used to configure our pipeline:
 ```yaml
-model_version:
+model:
   name: your_product_name
 ```
 
@@ -143,13 +143,13 @@ After the steps are executed we need to collect results (one best model per each
 ```python
 from zenml import get_step_context
 
-model_version = get_step_context().model_version
+model = get_step_context().model
 
 best_model = None
 best_metric = -1
 # consume artifacts attached to current model version in Model Control Plane
 for step_name in step_names:
-    hp_output = model_version.get_data_artifact(
+    hp_output = model.get_data_artifact(
         step_name=step_name, name="hp_result"
     )
     model: ClassifierMixin = hp_output.load()
@@ -239,7 +239,7 @@ By doing so we ensure that the best-performing version will be used for inferenc
 The Deployment pipeline is designed to run with inference Model Control Plane version context. This ensures that we always infer only on quality-assured Model Control Plane version and provide seamless integration of required artifacts created during training of this Model Control Plane version.
 This is achieved by providing this configuration in `deployer_config.yaml` used to configure our pipeline:
 ```yaml
-model_version:
+model:
   name: your_product_name
   version: production
 ```
@@ -259,7 +259,7 @@ NOTE: In this template a prediction service is only created for local orchestrat
 The Batch Inference pipeline is designed to run with inference Model Control Plane version context. This ensures that we always infer only on quality-assured Model Control Plane version and provide seamless integration of required artifacts created during training of this Model Control Plane version.
 This is achieved by providing this configuration in `inference_config.yaml` used to configure our pipeline:
 ```yaml
-model_version:
+model:
   name: your_product_name
   version: production
 ```
@@ -324,18 +324,18 @@ NOTE: On non-local orchestrators a `model` artifact will be loaded into memory t
 def inference_predict(
     dataset_inf: pd.DataFrame,
 ) -> Annotated[pd.Series, "predictions"]:
-    model_version = get_step_context().model_version
+    model = get_step_context().model
 
     # get predictor
-    predictor_service: Optional[MLFlowDeploymentService] = model_version.get_endpoint_artifact(
+    predictor_service: Optional[MLFlowDeploymentService] = model.get_endpoint_artifact(
         "mlflow_deployment"
     ).load()
     if predictor_service is not None:
         # run prediction from service
         predictions = predictor_service.predict(request=dataset_inf)
     else:
         # run prediction from memory
-        predictor = model_version.get_model_artifact("model").load()
+        predictor = model.get_model_artifact("model").load()
         predictions = predictor.predict(dataset_inf)
 
     predictions = pd.Series(predictions, name="predicted")

template/configs/deployer_config.yaml

@@ -21,7 +21,7 @@ steps:
       notify_on_success: False
 
 # configuration of the Model Control Plane
-model_version:
+model:
   name: {{ product_name }}
   version: {{ target_environment }}
 

template/configs/inference_config.yaml

@@ -21,7 +21,7 @@ steps:
       notify_on_success: False
 
 # configuration of the Model Control Plane
-model_version:
+model:
   name: {{ product_name }}
   version: {{ target_environment }}
 

template/configs/train_config.yaml

@@ -31,7 +31,7 @@ steps:
       notify_on_success: False
 
 # configuration of the Model Control Plane
-model_version:
+model:
   name: {{ product_name }}
   license: {{ open_source_license }}
   description: {{ product_name }} E2E Batch Use Case

template/steps/deployment/deployment_deploy.py

@@ -42,12 +42,12 @@ def deployment_deploy() -> (
     """
     ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ###
     if Client().active_stack.orchestrator.flavor == "local":
-        model_version = get_step_context().model_version
+        model = get_step_context().model
 
         # deploy predictor service
         deployment_service = mlflow_model_registry_deployer_step.entrypoint(
-            registry_model_name=model_version.name,
-            registry_model_version=model_version.metadata["model_registry_version"],
+            registry_model_name=model.name,
+            registry_model_version=model.metadata["model_registry_version"],
             replace_existing=True,
         )
     else:

template/steps/inference/inference_predict.py

@@ -35,12 +35,12 @@ def inference_predict(
         The predictions as pandas series
     """
     ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ###
-    model_version = get_step_context().model_version
+    model = get_step_context().model
 
     # get predictor
     predictor_service: Optional[
         MLFlowDeploymentService
-    ] = model_version.load_artifact("mlflow_deployment")
+    ] = model.load_artifact("mlflow_deployment")
     if predictor_service is not None:
         # run prediction from service
         predictions = predictor_service.predict(request=dataset_inf)
@@ -50,7 +50,7 @@ def inference_predict(
             "as the orchestrator is not local."
         )
         # run prediction from memory
-        predictor = model_version.load_artifact("model")
+        predictor = model.load_artifact("model")
         predictions = predictor.predict(dataset_inf)
 
     predictions = pd.Series(predictions, name="predicted")

template/steps/promotion/{% if metric_compare_promotion %}compute_performance_metrics_on_current_data.py{% endif %}

@@ -6,7 +6,7 @@ from typing_extensions import Annotated
 import pandas as pd
 from sklearn.metrics import accuracy_score
 from zenml import step, get_step_context
-from zenml.model.model_version import ModelVersion
+from zenml import Model
 from zenml.logger import get_logger
 
 logger = get_logger(__name__)
@@ -42,8 +42,8 @@ def compute_performance_metrics_on_current_data(
     logger.info("Evaluating model metrics...")
 
     # Get model version numbers from Model Control Plane
-    latest_version = get_step_context().model_version
-    current_version = ModelVersion(name=latest_version.name, version=target_env)
+    latest_version = get_step_context().model
+    current_version = Model(name=latest_version.name, version=target_env)
 
     latest_version_number = latest_version.number
     current_version_number = current_version.number

template/steps/promotion/{% if metric_compare_promotion %}promote_with_metric_compare.py{% endif %}

@@ -1,7 +1,7 @@
 # {% include 'template/license_header' %}
 
 from zenml import get_step_context, step
-from zenml.model.model_version import ModelVersion
+from zenml import Model
 from zenml.logger import get_logger
 
 from utils import promote_in_model_registry
@@ -44,8 +44,8 @@ def promote_with_metric_compare(
     should_promote = True
 
     # Get model version numbers from Model Control Plane
-    latest_version = get_step_context().model_version
-    current_version = ModelVersion(name=latest_version.name, version=target_env)
+    latest_version = get_step_context().model
+    current_version = Model(name=latest_version.name, version=target_env)
 
     current_version_number = current_version.number
 
@@ -69,8 +69,8 @@ def promote_with_metric_compare(
 
     if should_promote:
         # Promote in Model Control Plane
-        model_version = get_step_context().model_version
-        model_version.set_stage(stage=target_env, force=True)
+        model = get_step_context().model
+        model.set_stage(stage=target_env, force=True)
         logger.info(f"Current model version was promoted to '{target_env}'.")
 
         # Promote in Model Registry

template/steps/promotion/{% if not metric_compare_promotion %}promote_latest_version.py{% endif %}

@@ -1,7 +1,7 @@
 # {% include 'template/license_header' %}
 
 from zenml import get_step_context, step
-from zenml.model.model_version import ModelVersion
+from zenml import Model
 from zenml.logger import get_logger
 
 from utils import promote_in_model_registry
@@ -22,13 +22,13 @@ def promote_latest_version(
 
     ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ###
     # Get model version numbers from Model Control Plane
-    latest_version = get_step_context().model_version
-    current_version = ModelVersion(name=latest_version.name, version=target_env)
+    latest_version = get_step_context().model
+    current_version = Model(name=latest_version.name, version=target_env)
     logger.info(f"Promoting latest model version `{latest_version}`")
 
     # Promote in Model Control Plane
-    model_version = get_step_context().model_version
-    model_version.set_stage(stage=target_env, force=True)
+    model = get_step_context().model
+    model.set_stage(stage=target_env, force=True)
     logger.info(f"Current model version was promoted to '{target_env}'.")
 
     # Promote in Model Registry

template/steps/training/model_trainer.py

@@ -84,8 +84,8 @@ def model_trainer(
     if model_registry:
         versions = model_registry.list_model_versions(name=name)
         if versions:
-            model_version = get_step_context().model_version
-            model_version.log_metadata({"model_registry_version": versions[-1].version})
+            model_ = get_step_context().model
+            model_.log_metadata({"model_registry_version": versions[-1].version})
     ### YOUR CODE ENDS HERE ###
 
     return model

template/steps/{% if hyperparameters_tuning %}hp_tuning{% endif %}/hp_tuning_select_best_model.py

@@ -25,17 +25,17 @@ def hp_tuning_select_best_model(
         The best possible model class and its' parameters.
     """
     ### ADD YOUR OWN CODE HERE - THIS IS JUST AN EXAMPLE ###
-    model_version = get_step_context().model_version
+    model = get_step_context().model
 
     best_model = None
     best_metric = -1
     # consume artifacts attached to current model version in Model Control Plane
     for step_name in step_names:
-        hp_output = model_version.get_data_artifact("hp_result")
-        model: ClassifierMixin = hp_output.load()
+        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)
         if best_model is None or best_metric < metric:
-            best_model = model
+            best_model = model_
     ### YOUR CODE ENDS HERE ###
     return best_model