ModelConfig is ModelVersion!

.github/workflows/update-templates-to-examples.yml

@@ -56,7 +56,7 @@ jobs:
           python-version: ${{ inputs.python-version }}
           stack-name: local
           ref-zenml: ${{ github.ref }}
-          ref-template: '2023.11.01' # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
+          ref-template: '2023.11.14' # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
 
       - name: Clean-up
         run: |
@@ -76,6 +76,7 @@ jobs:
 
       - name: Check-out fresh E2E template
         run: |
+          rm -rf examples/e2e
           mkdir -p examples/e2e
           printf '[email protected]' | zenml init --path examples/e2e --template e2e_batch --template-with-defaults
           bash scripts/format.sh

examples/e2e/.copier-answers.yml

@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2023.11.01
+_commit: 2023.11.14
 _src_path: gh:zenml-io/template-e2e-batch
 data_quality_checks: true
 email: ''

examples/e2e/configs/deployer_config.yaml

@@ -34,7 +34,7 @@ steps:
       notify_on_success: False
 
 # configuration of the Model Control Plane
-model_config:
+model_version:
   name: e2e_use_case
   version: staging
 

examples/e2e/configs/inference_config.yaml

@@ -34,7 +34,7 @@ steps:
       notify_on_success: False
 
 # configuration of the Model Control Plane
-model_config:
+model_version:
   name: e2e_use_case
   version: staging
 

examples/e2e/configs/train_config.yaml

@@ -44,7 +44,7 @@ steps:
       notify_on_success: False
 
 # configuration of the Model Control Plane
-model_config:
+model_version:
   name: e2e_use_case
   license: apache
   description: e2e_use_case E2E Batch Use Case
@@ -61,7 +61,6 @@ model_config:
   - sklearn
   - from template
   - ZenML delivered
-  create_new_model_version: true
 
 # pipeline level extra configurations
 extra:

examples/e2e/pipelines/training.py

@@ -101,7 +101,7 @@ def e2e_use_case_training(
             target=target,
         )
         after.append(step_name)
-    best_model = hp_tuning_select_best_model(after=after)
+    best_model = hp_tuning_select_best_model(step_names=after, after=after)
 
     ########## Training stage ##########
     model = model_trainer(

examples/e2e/steps/deployment/deployment_deploy.py

@@ -19,7 +19,6 @@
 from typing import Optional
 
 from typing_extensions import Annotated
-from utils import get_model_registry_version
 
 from zenml import get_step_context, step
 from zenml.client import Client
@@ -30,7 +29,7 @@
     mlflow_model_registry_deployer_step,
 )
 from zenml.logger import get_logger
-from zenml.model import DeploymentArtifactConfig
+from zenml.model import EndpointArtifactConfig
 
 logger = get_logger(__name__)
 
@@ -40,7 +39,7 @@ def deployment_deploy() -> (
     Annotated[
         Optional[MLFlowDeploymentService],
         "mlflow_deployment",
-        DeploymentArtifactConfig(),
+        EndpointArtifactConfig(),
     ]
 ):
     """Predictions step.
@@ -63,12 +62,14 @@ 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_config._get_model_version()
+        model_version = get_step_context().model_version
 
         # deploy predictor service
         deployment_service = mlflow_model_registry_deployer_step.entrypoint(
-            registry_model_name=model_version.model.name,
-            registry_model_version=get_model_registry_version(model_version),
+            registry_model_name=model_version.name,
+            registry_model_version=model_version.get_model_artifact("model")
+            .metadata["model_registry_version"]
+            .value,
             replace_existing=True,
         )
     else:

examples/e2e/steps/etl/inference_data_preprocessor.py

@@ -21,7 +21,7 @@
 from typing_extensions import Annotated
 
 from zenml import step
-from zenml.model import ArtifactConfig
+from zenml.model import DataArtifactConfig
 
 
 @step
@@ -32,7 +32,7 @@ def inference_data_preprocessor(
 ) -> Annotated[
     pd.DataFrame,
     "dataset_inf",
-    ArtifactConfig(overwrite=False, artifact_name="inference_dataset"),
+    DataArtifactConfig(overwrite=False, artifact_name="inference_dataset"),
 ]:
     """Data preprocessor step.
 

examples/e2e/steps/hp_tuning/hp_tuning_select_best_model.py

@@ -16,6 +16,8 @@
 #
 
 
+from typing import List
+
 from sklearn.base import ClassifierMixin
 from typing_extensions import Annotated
 
@@ -26,7 +28,9 @@
 
 
 @step
-def hp_tuning_select_best_model() -> Annotated[ClassifierMixin, "best_model"]:
+def hp_tuning_select_best_model(
+    step_names: List[str],
+) -> Annotated[ClassifierMixin, "best_model"]:
     """Find best model across all HP tuning attempts.
 
     This is an example of a model hyperparameter tuning step that loops
@@ -37,19 +41,19 @@ def hp_tuning_select_best_model() -> Annotated[ClassifierMixin, "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_config._get_model_version()
+    model_version = get_step_context().model_version
 
     best_model = None
     best_metric = -1
     # consume artifacts attached to current model version in Model Control Plane
-    for full_artifact_name in model_version.artifact_object_ids:
-        # if artifacts comes from one of HP tuning steps
-        if full_artifact_name.endswith("hp_result"):
-            hp_output = model_version.artifacts[full_artifact_name]["1"]
-            model: ClassifierMixin = hp_output.load()
-            # fetch metadata we attached earlier
-            metric = float(hp_output.metadata["metric"].value)
-            if best_model is None or best_metric < metric:
-                best_model = model
+    for step_name in step_names:
+        hp_output = model_version.get_data_artifact(
+            step_name=step_name, name="hp_result"
+        )
+        model: ClassifierMixin = hp_output.load()
+        # fetch metadata we attached earlier
+        metric = float(hp_output.metadata["metric"].value)
+        if best_model is None or best_metric < metric:
+            best_model = model
     ### YOUR CODE ENDS HERE ###
     return best_model

examples/e2e/steps/inference/inference_predict.py

@@ -26,15 +26,15 @@
     MLFlowDeploymentService,
 )
 from zenml.logger import get_logger
-from zenml.model import ArtifactConfig
+from zenml.model import DataArtifactConfig
 
 logger = get_logger(__name__)
 
 
 @step
 def inference_predict(
     dataset_inf: pd.DataFrame,
-) -> Annotated[pd.Series, "predictions", ArtifactConfig(overwrite=False)]:
+) -> Annotated[pd.Series, "predictions", DataArtifactConfig(overwrite=False)]:
     """Predictions step.
 
     This is an example of a predictions step that takes the data in and returns
@@ -54,12 +54,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_config._get_model_version()
+    model_version = get_step_context().model_version
 
     # get predictor
     predictor_service: Optional[
         MLFlowDeploymentService
-    ] = model_version.get_deployment("mlflow_deployment").load()
+    ] = model_version.get_endpoint_artifact("mlflow_deployment").load()
     if predictor_service is not None:
         # run prediction from service
         predictions = predictor_service.predict(request=dataset_inf)
@@ -69,7 +69,7 @@ def inference_predict(
             "as the orchestrator is not local."
         )
         # run prediction from memory
-        predictor = model_version.get_model_object("model").load()
+        predictor = model_version.get_model_artifact("model").load()
         predictions = predictor.predict(dataset_inf)
 
     predictions = pd.Series(predictions, name="predicted")

examples/e2e/steps/promotion/compute_performance_metrics_on_current_data.py

@@ -20,10 +20,10 @@
 import pandas as pd
 from sklearn.metrics import accuracy_score
 from typing_extensions import Annotated
-from utils import get_model_versions
 
-from zenml import step
+from zenml import get_step_context, step
 from zenml.logger import get_logger
+from zenml.model import ModelVersion
 
 logger = get_logger(__name__)
 
@@ -61,23 +61,33 @@ def compute_performance_metrics_on_current_data(
     logger.info("Evaluating model metrics...")
 
     # Get model version numbers from Model Control Plane
-    latest_version, current_version = get_model_versions(target_env)
+    latest_version = get_step_context().model_version
+    current_version = ModelVersion(
+        name=latest_version.name, version=target_env
+    )
 
-    # Get predictors
-    predictors = {
-        latest_version.number: latest_version.get_model_object("model").load(),
-        current_version.number: current_version.get_model_object(
-            "model"
-        ).load(),
-    }
+    latest_version_number = latest_version.number
+    current_version_number = current_version.number
+
+    if current_version_number is None:
+        current_version_number = -1
+        metrics = {latest_version_number: 1.0, current_version_number: 0.0}
+    else:
+        # Get predictors
+        predictors = {
+            latest_version_number: latest_version.get_model_artifact(
+                "model"
+            ).load(),
+            current_version_number: current_version.get_model_artifact(
+                "model"
+            ).load(),
+        }
 
-    if latest_version != current_version:
         metrics = {}
-        for version in [latest_version.number, current_version.number]:
+        for version in [latest_version_number, current_version_number]:
             # predict and evaluate
             predictions = predictors[version].predict(X)
             metrics[version] = accuracy_score(y, predictions)
-    else:
-        metrics = {latest_version.number: 1.0, current_version.number: 0.0}
+
     ### YOUR CODE ENDS HERE ###
-    return metrics[latest_version.number], metrics[current_version.number]
+    return metrics[latest_version_number], metrics[current_version_number]

examples/e2e/steps/promotion/promote_with_metric_compare.py

@@ -15,14 +15,11 @@
 # limitations under the License.
 #
 
-from utils import (
-    get_model_registry_version,
-    get_model_versions,
-    promote_in_model_registry,
-)
+from utils import promote_in_model_registry
 
 from zenml import get_step_context, step
 from zenml.logger import get_logger
+from zenml.model import ModelVersion
 
 logger = get_logger(__name__)
 
@@ -62,9 +59,14 @@ def promote_with_metric_compare(
     should_promote = True
 
     # Get model version numbers from Model Control Plane
-    latest_version, current_version = get_model_versions(target_env)
+    latest_version = get_step_context().model_version
+    current_version = ModelVersion(
+        name=latest_version.name, version=target_env
+    )
+
+    current_version_number = current_version.number
 
-    if latest_version.number == current_version.number:
+    if current_version_number is None:
         logger.info("No current model version found - promoting latest")
     else:
         logger.info(
@@ -81,21 +83,41 @@ def promote_with_metric_compare(
             )
             should_promote = False
 
-    promoted_version = get_model_registry_version(current_version)
     if should_promote:
         # Promote in Model Control Plane
-        model_version = get_step_context().model_config._get_model_version()
+        model_version = get_step_context().model_version
         model_version.set_stage(stage=target_env, force=True)
         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")
+            .metadata["model_registry_version"]
+            .value
+        )
+        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")
+                .metadata["model_registry_version"]
+                .value
+            )
         promote_in_model_registry(
-            latest_version=get_model_registry_version(latest_version),
-            current_version=get_model_registry_version(current_version),
+            latest_version=latest_version_model_registry_number,
+            current_version=current_version_model_registry_number,
             model_name=mlflow_model_name,
             target_env=target_env.capitalize(),
         )
-        promoted_version = get_model_registry_version(latest_version)
+        promoted_version = latest_version_model_registry_number
+    else:
+        promoted_version = (
+            current_version.get_model_artifact("model")
+            .metadata["model_registry_version"]
+            .value
+        )
 
     logger.info(
         f"Current model version in `{target_env}` is `{promoted_version}` registered in Model Registry"

examples/e2e/steps/training/model_trainer.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 #
 
-
 import mlflow
 import pandas as pd
 from sklearn.base import ClassifierMixin

examples/e2e/utils/__init__.py

@@ -17,5 +17,4 @@
 
 
 from .get_model_from_config import get_model_from_config
-from .model_versions import get_model_versions, get_model_registry_version
 from .promote_in_model_registry import promote_in_model_registry