Moving tags to the body for artifacts and artifact versions (#2138)

* moving tags

* Auto-update of E2E template

---------

Co-authored-by: GitHub Actions <[email protected]>

Author: bcdurak

examples/e2e/.copier-answers.yml

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

examples/e2e/configs/train_config.yaml

@@ -32,13 +32,9 @@ steps:
   model_trainer:
     parameters:
       name: e2e_use_case
-  compute_performance_metrics_on_current_data:
-    parameters:
-      target_env: staging
   promote_with_metric_compare:
     parameters:
       mlflow_model_name: e2e_use_case
-      target_env: staging
   notify_on_success:
     parameters:
       notify_on_success: False
@@ -65,6 +61,9 @@ model_version:
 # pipeline level extra configurations
 extra:
   notify_on_failure: True
+# pipeline level parameters
+parameters:
+  target_env: staging
   # This set contains all the model configurations that you want 
   # to evaluate during hyperparameter tuning stage.
   model_search_space:

examples/e2e/pipelines/training.py

@@ -17,7 +17,7 @@
 
 
 import random
-from typing import List, Optional
+from typing import Any, Dict, List, Optional
 
 from steps import (
     compute_performance_metrics_on_current_data,
@@ -33,14 +33,16 @@
     train_data_splitter,
 )
 
-from zenml import get_pipeline_context, pipeline
+from zenml import pipeline
 from zenml.logger import get_logger
 
 logger = get_logger(__name__)
 
 
 @pipeline(on_failure=notify_on_failure)
 def e2e_use_case_training(
+    model_search_space: Dict[str, Any],
+    target_env: str,
     test_size: float = 0.2,
     drop_na: Optional[bool] = None,
     normalize: Optional[bool] = None,
@@ -57,6 +59,8 @@ def e2e_use_case_training(
     trains and evaluates a model.
 
     Args:
+        model_search_space: Search space for hyperparameter tuning
+        target_env: The environment to promote the model to
         test_size: Size of holdout set for training 0.0..1.0
         drop_na: If `True` NA values will be removed from dataset
         normalize: If `True` dataset will be normalized with MinMaxScaler
@@ -65,12 +69,10 @@ def e2e_use_case_training(
         min_test_accuracy: Threshold to stop execution if test set accuracy is lower
         fail_on_accuracy_quality_gates: If `True` and `min_train_accuracy` or `min_test_accuracy`
             are not met - execution will be interrupted early
-
     """
     ### 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.
-    pipeline_extra = get_pipeline_context().extra
     ########## ETL stage ##########
     raw_data, target, _ = data_loader(random_state=random.randint(0, 100))
     dataset_trn, dataset_tst = train_data_splitter(
@@ -87,9 +89,7 @@ def e2e_use_case_training(
     ########## Hyperparameter tuning stage ##########
     after = []
     search_steps_prefix = "hp_tuning_search_"
-    for config_name, model_search_configuration in pipeline_extra[
-        "model_search_space"
-    ].items():
+    for config_name, model_search_configuration in model_search_space.items():
         step_name = f"{search_steps_prefix}{config_name}"
         hp_tuning_single_search(
             id=step_name,
@@ -123,12 +123,15 @@ def e2e_use_case_training(
         latest_metric,
         current_metric,
     ) = compute_performance_metrics_on_current_data(
-        dataset_tst=dataset_tst, after=["model_evaluator"]
+        dataset_tst=dataset_tst,
+        target_env=target_env,
+        after=["model_evaluator"],
     )
 
     promote_with_metric_compare(
         latest_metric=latest_metric,
         current_metric=current_metric,
+        target_env=target_env,
     )
     last_step = "promote_with_metric_compare"
 

src/zenml/models/v2/core/artifact.py

@@ -63,6 +63,10 @@ class ArtifactUpdate(BaseModel):
 class ArtifactResponseBody(BaseResponseBody):
     """Response body for artifacts."""
 
+    tags: List[TagResponseModel] = Field(
+        title="Tags associated with the model",
+    )
+
 
 class ArtifactResponseMetadata(BaseResponseMetadata):
     """Response metadata for artifacts."""
@@ -71,9 +75,6 @@ class ArtifactResponseMetadata(BaseResponseMetadata):
         title="Whether the name is custom (True) or auto-generated (False).",
         default=False,
     )
-    tags: List[TagResponseModel] = Field(
-        title="Tags associated with the model",
-    )
 
 
 class ArtifactResponse(
@@ -98,22 +99,22 @@ def get_hydrated_version(self) -> "ArtifactResponse":
 
     # Body and metadata properties
     @property
-    def has_custom_name(self) -> bool:
-        """The `has_custom_name` property.
+    def tags(self) -> List[TagResponseModel]:
+        """The `tags` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().has_custom_name
+        return self.get_body().tags
 
     @property
-    def tags(self) -> List[TagResponseModel]:
-        """The `tags` property.
+    def has_custom_name(self) -> bool:
+        """The `has_custom_name` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().tags
+        return self.get_metadata().has_custom_name
 
     # Helper methods
     @property

src/zenml/models/v2/core/artifact_version.py

@@ -134,6 +134,9 @@ class ArtifactVersionResponseBody(WorkspaceScopedResponseBody):
     data_type: Source = Field(
         title="Data type of the artifact.",
     )
+    tags: List[TagResponseModel] = Field(
+        title="Tags associated with the model",
+    )
 
     _convert_source = convert_source_validator("materializer", "data_type")
 
@@ -149,9 +152,6 @@ class ArtifactVersionResponseMetadata(WorkspaceScopedResponseMetadata):
         title="ID of the step run that produced this artifact.",
         default=None,
     )
-    tags: List[TagResponseModel] = Field(
-        title="Tags associated with the model",
-    )
     visualizations: Optional[List["ArtifactVisualizationResponse"]] = Field(
         default=None, title="Visualizations of the artifact."
     )
@@ -215,31 +215,31 @@ def type(self) -> ArtifactType:
         return self.get_body().type
 
     @property
-    def artifact_store_id(self) -> Optional[UUID]:
-        """The `artifact_store_id` property.
+    def tags(self) -> List[TagResponseModel]:
+        """The `tags` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().artifact_store_id
+        return self.get_body().tags
 
     @property
-    def producer_step_run_id(self) -> Optional[UUID]:
-        """The `producer_step_run_id` property.
+    def artifact_store_id(self) -> Optional[UUID]:
+        """The `artifact_store_id` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().producer_step_run_id
+        return self.get_metadata().artifact_store_id
 
     @property
-    def tags(self) -> List[TagResponseModel]:
-        """The `tags` property.
+    def producer_step_run_id(self) -> Optional[UUID]:
+        """The `producer_step_run_id` property.
 
         Returns:
             the value of the property.
         """
-        return self.get_metadata().tags
+        return self.get_metadata().producer_step_run_id
 
     @property
     def visualizations(

src/zenml/zen_stores/schemas/artifact_schemas.py

@@ -113,14 +113,14 @@ def to_model(self, hydrate: bool = False) -> ArtifactResponse:
         body = ArtifactResponseBody(
             created=self.created,
             updated=self.updated,
+            tags=[t.tag.to_model() for t in self.tags],
         )
 
         # Create the metadata of the model
         metadata = None
         if hydrate:
             metadata = ArtifactResponseMetadata(
                 has_custom_name=self.has_custom_name,
-                tags=[t.tag.to_model() for t in self.tags],
             )
 
         return ArtifactResponse(
@@ -299,6 +299,7 @@ def to_model(self, hydrate: bool = False) -> ArtifactVersionResponse:
             data_type=data_type,
             created=self.created,
             updated=self.updated,
+            tags=[t.tag.to_model() for t in self.tags],
         )
 
         # Create the metadata of the model
@@ -318,7 +319,6 @@ def to_model(self, hydrate: bool = False) -> ArtifactVersionResponse:
                 producer_step_run_id=producer_step_run_id,
                 visualizations=[v.to_model() for v in self.visualizations],
                 run_metadata={m.key: m.to_model() for m in self.run_metadata},
-                tags=[t.tag.to_model() for t in self.tags],
             )
 
         return ArtifactVersionResponse(

tests/unit/conftest.py

@@ -484,10 +484,10 @@ def sample_artifact_model() -> ArtifactResponse:
         body=ArtifactResponseBody(
             created=datetime.now(),
             updated=datetime.now(),
+            tags=[],
         ),
         metadata=ArtifactResponseMetadata(
             has_custom_name=True,
-            tags=[],
         ),
     )
 
@@ -509,10 +509,10 @@ def sample_artifact_version_model(
             type=ArtifactType.DATA,
             materializer="sample_module.sample_materializer",
             data_type="sample_module.sample_data_type",
+            tags=[],
         ),
         metadata=ArtifactVersionResponseMetadata(
             workspace=sample_workspace_model,
-            tags=[],
         ),
     )