Simplify Metadata handling (#3096)

* Initial commit, nuking all metadata responses and seeing what breaks

* Removed last remnant of LazyLoader

* Reintroducing the lazy loaders.

* Add LazyRunMetadataResponse to EntrypointFunctionDefinition

* Test for lazy loaders works now

* Fixed tests, reformatted

* Use updated template

* Auto-update of Starter template

* Updated more templates

* Fixed failing test

* Fixed step run schemas

* Auto-update of E2E template

* Auto-update of NLP template

* Fixed tests, removed additional .value access

* Further fixing

* Fixed linting issues

* Reformatted

* Linted, formatted and tested again

* Typing

* Maybe fix everything

* Apply some feedback

---------

Co-authored-by: Andrei Vishniakov <[email protected]>
Co-authored-by: GitHub Actions <[email protected]>
Co-authored-by: Michael Schuster <[email protected]>
Co-authored-by: Baris Can Durak <[email protected]>
Co-authored-by: Barış Can Durak <[email protected]>

Author: 5 people

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

@@ -46,7 +46,7 @@ jobs:
           python-version: ${{ inputs.python-version }}
           stack-name: local
           ref-zenml: ${{ github.ref }}
-          ref-template: 2024.10.10  # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
+          ref-template: 2024.10.21  # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
       - name: Clean-up
         run: |
           rm -rf ./local_checkout
@@ -118,7 +118,7 @@ jobs:
           python-version: ${{ inputs.python-version }}
           stack-name: local
           ref-zenml: ${{ github.ref }}
-          ref-template: 2024.09.23  # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
+          ref-template: 2024.10.21  # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
       - name: Clean-up
         run: |
           rm -rf ./local_checkout
@@ -189,7 +189,7 @@ jobs:
           python-version: ${{ inputs.python-version }}
           stack-name: local
           ref-zenml: ${{ github.ref }}
-          ref-template: 2024.09.24  # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
+          ref-template: 2024.10.21  # Make sure it is aligned with ZENML_PROJECT_TEMPLATES from src/zenml/cli/base.py
       - name: Clean-up
         run: |
           rm -rf ./local_checkout

docs/book/component-guide/orchestrators/kubeflow.md

@@ -198,7 +198,7 @@ Kubeflow comes with its own UI that you can use to find further details about yo
 from zenml.client import Client
 
 pipeline_run = Client().get_pipeline_run("<PIPELINE_RUN_NAME>")
-orchestrator_url = pipeline_run.run_metadata["orchestrator_url"].value
+orchestrator_url = pipeline_run.run_metadata["orchestrator_url"]
 ```
 
 #### Additional configuration

examples/e2e/.copier-answers.yml

@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2024.10.10
+_commit: 2024.10.21
 _src_path: gh:zenml-io/template-e2e-batch
 data_quality_checks: true
 email: [email protected]

examples/e2e/steps/deployment/deployment_deploy.py

@@ -67,7 +67,7 @@ def deployment_deploy() -> (
             registry_model_name=model.name,
             registry_model_version=model.run_metadata[
                 "model_registry_version"
-            ].value,
+            ],
             replace_existing=True,
         )
     else:

examples/e2e/steps/hp_tuning/hp_tuning_select_best_model.py

@@ -50,7 +50,7 @@ def hp_tuning_select_best_model(
         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)
+        metric = float(hp_output.run_metadata["metric"])
         if best_model is None or best_metric < metric:
             best_model = model_
     ### YOUR CODE ENDS HERE ###

examples/e2e/steps/promotion/promote_with_metric_compare.py

@@ -92,14 +92,14 @@ def promote_with_metric_compare(
         # Promote in Model Registry
         latest_version_model_registry_number = latest_version.run_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.run_metadata["model_registry_version"].value
+                current_version.run_metadata["model_registry_version"]
             )
         promote_in_model_registry(
             latest_version=latest_version_model_registry_number,
@@ -111,7 +111,7 @@ def promote_with_metric_compare(
     else:
         promoted_version = current_version.run_metadata[
             "model_registry_version"
-        ].value
+        ]
 
     logger.info(
         f"Current model version in `{target_env}` is `{promoted_version}` registered in Model Registry"

examples/e2e_nlp/.copier-answers.yml

@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2024.09.23
+_commit: 2024.10.21
 _src_path: gh:zenml-io/template-nlp
 accelerator: cpu
 cloud_of_choice: aws

examples/e2e_nlp/gradio/requirements.txt

@@ -9,4 +9,4 @@ pandas==1.5.3
 session_info==1.0.0
 scikit-learn==1.5.0
 transformers==4.28.1
-IPython==7.34.0
+IPython==8.10.0

examples/e2e_nlp/steps/promotion/promote_get_metrics.py

@@ -56,9 +56,7 @@ def promote_get_metrics() -> (
 
     # Get current model version metric in current run
     model = get_step_context().model
-    current_metrics = (
-        model.get_model_artifact("model").run_metadata["metrics"].value
-    )
+    current_metrics = model.get_model_artifact("model").run_metadata["metrics"]
     logger.info(f"Current model version metrics are {current_metrics}")
 
     # Get latest saved model version metric in target environment
@@ -72,11 +70,9 @@ def promote_get_metrics() -> (
     except KeyError:
         latest_version = None
     if latest_version:
-        latest_metrics = (
-            latest_version.get_model_artifact("model")
-            .run_metadata["metrics"]
-            .value
-        )
+        latest_metrics = latest_version.get_model_artifact(
+            "model"
+        ).run_metadata["metrics"]
         logger.info(f"Latest model version metrics are {latest_metrics}")
     else:
         logger.info("No currently promoted model version found.")

examples/mlops_starter/.copier-answers.yml

@@ -1,5 +1,5 @@
 # Changes here will be overwritten by Copier
-_commit: 2024.09.24
+_commit: 2024.10.21
 _src_path: gh:zenml-io/template-starter
 email: [email protected]
 full_name: ZenML GmbH

examples/mlops_starter/quickstart.ipynb

@@ -994,8 +994,8 @@
     "@pipeline\n",
     "def inference(preprocess_pipeline_id: UUID):\n",
     "    \"\"\"Model batch inference pipeline\"\"\"\n",
-    "    # random_state = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).metadata[\"random_state\"].value\n",
-    "    # target = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).run_metadata['target'].value\n",
+    "    # random_state = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).metadata[\"random_state\"]\n",
+    "    # target = client.get_artifact_version(name_id_or_prefix=preprocess_pipeline_id).run_metadata['target']\n",
     "    random_state = 42\n",
     "    target = \"target\"\n",
     "\n",

examples/mlops_starter/run.py

@@ -239,8 +239,8 @@ def main(
         #  to get the random state and target column
         random_state = preprocess_pipeline_artifact.run_metadata[
             "random_state"
-        ].value
-        target = preprocess_pipeline_artifact.run_metadata["target"].value
+        ]
+        target = preprocess_pipeline_artifact.run_metadata["target"]
         run_args_inference["random_state"] = random_state
         run_args_inference["target"] = target
 

examples/mlops_starter/steps/model_promoter.py

@@ -58,11 +58,9 @@ def model_promoter(accuracy: float, stage: str = "production") -> bool:
         try:
             stage_model = client.get_model_version(current_model.name, stage)
             # We compare their metrics
-            prod_accuracy = (
-                stage_model.get_artifact("sklearn_classifier")
-                .run_metadata["test_accuracy"]
-                .value
-            )
+            prod_accuracy = stage_model.get_artifact(
+                "sklearn_classifier"
+            ).run_metadata["test_accuracy"]
             if float(accuracy) > float(prod_accuracy):
                 # If current model has better metrics, we promote it
                 is_promoted = True

src/zenml/cli/base.py

@@ -79,15 +79,15 @@ def copier_github_url(self) -> str:
 ZENML_PROJECT_TEMPLATES = dict(
     e2e_batch=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-e2e-batch",
-        github_tag="2024.10.10",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.10.21",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     starter=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-starter",
-        github_tag="2024.09.24",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.10.21",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     nlp=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-nlp",
-        github_tag="2024.09.23",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.10.21",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     llm_finetuning=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-llm-finetuning",

src/zenml/cli/model.py

@@ -59,19 +59,14 @@ def _model_to_print(model: ModelResponse) -> Dict[str, Any]:
 def _model_version_to_print(
     model_version: ModelVersionResponse,
 ) -> Dict[str, Any]:
-    run_metadata = None
-    if model_version.run_metadata:
-        run_metadata = {
-            k: v.value for k, v in model_version.run_metadata.items()
-        }
     return {
         "id": model_version.id,
         "model": model_version.model.name,
         "name": model_version.name,
         "number": model_version.number,
         "description": model_version.description,
         "stage": model_version.stage,
-        "run_metadata": run_metadata,
+        "run_metadata": model_version.run_metadata,
         "tags": [t.name for t in model_version.tags],
         "data_artifacts_count": len(model_version.data_artifact_ids),
         "model_artifacts_count": len(model_version.model_artifact_ids),

src/zenml/client.py

@@ -136,9 +136,7 @@
     PipelineResponse,
     PipelineRunFilter,
     PipelineRunResponse,
-    RunMetadataFilter,
     RunMetadataRequest,
-    RunMetadataResponse,
     RunTemplateFilter,
     RunTemplateRequest,
     RunTemplateResponse,
@@ -4417,7 +4415,7 @@ def create_run_metadata(
         resource_id: UUID,
         resource_type: MetadataResourceTypes,
         stack_component_id: Optional[UUID] = None,
-    ) -> List[RunMetadataResponse]:
+    ) -> None:
         """Create run metadata.
 
         Args:
@@ -4430,7 +4428,7 @@ def create_run_metadata(
                 the metadata.
 
         Returns:
-            The created metadata, as string to model dictionary.
+            None
         """
         from zenml.metadata.metadata_types import get_metadata_type
 
@@ -4465,74 +4463,8 @@ def create_run_metadata(
             values=values,
             types=types,
         )
-        return self.zen_store.create_run_metadata(run_metadata)
-
-    def list_run_metadata(
-        self,
-        sort_by: str = "created",
-        page: int = PAGINATION_STARTING_PAGE,
-        size: int = PAGE_SIZE_DEFAULT,
-        logical_operator: LogicalOperators = LogicalOperators.AND,
-        id: Optional[Union[UUID, str]] = None,
-        created: Optional[Union[datetime, str]] = None,
-        updated: Optional[Union[datetime, str]] = None,
-        workspace_id: Optional[UUID] = None,
-        user_id: Optional[UUID] = None,
-        resource_id: Optional[UUID] = None,
-        resource_type: Optional[MetadataResourceTypes] = None,
-        stack_component_id: Optional[UUID] = None,
-        key: Optional[str] = None,
-        value: Optional["MetadataType"] = None,
-        type: Optional[str] = None,
-        hydrate: bool = False,
-    ) -> Page[RunMetadataResponse]:
-        """List run metadata.
-
-        Args:
-            sort_by: The field to sort the results by.
-            page: The page number to return.
-            size: The number of results to return per page.
-            logical_operator: The logical operator to use for filtering.
-            id: The ID of the metadata.
-            created: The creation time of the metadata.
-            updated: The last update time of the metadata.
-            workspace_id: The ID of the workspace the metadata belongs to.
-            user_id: The ID of the user that created the metadata.
-            resource_id: The ID of the resource the metadata belongs to.
-            resource_type: The type of the resource the metadata belongs to.
-            stack_component_id: The ID of the stack component that produced
-                the metadata.
-            key: The key of the metadata.
-            value: The value of the metadata.
-            type: The type of the metadata.
-            hydrate: Flag deciding whether to hydrate the output model(s)
-                by including metadata fields in the response.
-
-        Returns:
-            The run metadata.
-        """
-        metadata_filter_model = RunMetadataFilter(
-            sort_by=sort_by,
-            page=page,
-            size=size,
-            logical_operator=logical_operator,
-            id=id,
-            created=created,
-            updated=updated,
-            workspace_id=workspace_id,
-            user_id=user_id,
-            resource_id=resource_id,
-            resource_type=resource_type,
-            stack_component_id=stack_component_id,
-            key=key,
-            value=value,
-            type=type,
-        )
-        metadata_filter_model.set_scope_workspace(self.active_workspace.id)
-        return self.zen_store.list_run_metadata(
-            metadata_filter_model,
-            hydrate=hydrate,
-        )
+        self.zen_store.create_run_metadata(run_metadata)
+        return None
 
     # -------------------------------- Secrets ---------------------------------
 

src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py

@@ -566,7 +566,7 @@ def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
 
         # Fetch the status of the _PipelineExecution
         if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
-            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID].value
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID]
         elif run.orchestrator_run_id is not None:
             run_id = run.orchestrator_run_id
         else:

src/zenml/integrations/azure/orchestrators/azureml_orchestrator.py

@@ -482,7 +482,7 @@ def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
 
         # Fetch the status of the PipelineJob
         if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
-            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID].value
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID]
         elif run.orchestrator_run_id is not None:
             run_id = run.orchestrator_run_id
         else:

src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -835,7 +835,7 @@ def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
 
         # Fetch the status of the PipelineJob
         if METADATA_ORCHESTRATOR_RUN_ID in run.run_metadata:
-            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID].value
+            run_id = run.run_metadata[METADATA_ORCHESTRATOR_RUN_ID]
         elif run.orchestrator_run_id is not None:
             run_id = run.orchestrator_run_id
         else:

src/zenml/metadata/lazy_load.py

@@ -13,10 +13,25 @@
 #  permissions and limitations under the License.
 """Run Metadata Lazy Loader definition."""
 
-from typing import TYPE_CHECKING, Optional
+from typing import Optional
 
-if TYPE_CHECKING:
-    from zenml.models import RunMetadataResponse
+from pydantic import BaseModel
+
+from zenml.metadata.metadata_types import MetadataType
+
+
+class LazyRunMetadataResponse(BaseModel):
+    """Lazy run metadata response.
+
+    Used if the run metadata is accessed from the model in
+    a pipeline context available only during pipeline compilation.
+    """
+
+    lazy_load_artifact_name: Optional[str] = None
+    lazy_load_artifact_version: Optional[str] = None
+    lazy_load_metadata_name: Optional[str] = None
+    lazy_load_model_name: str
+    lazy_load_model_version: Optional[str] = None
 
 
 class RunMetadataLazyGetter:
@@ -47,7 +62,7 @@ def __init__(
         self._lazy_load_artifact_name = _lazy_load_artifact_name
         self._lazy_load_artifact_version = _lazy_load_artifact_version
 
-    def __getitem__(self, key: str) -> "RunMetadataResponse":
+    def __getitem__(self, key: str) -> MetadataType:
         """Get the metadata for the given key.
 
         Args:
@@ -56,9 +71,7 @@ def __getitem__(self, key: str) -> "RunMetadataResponse":
         Returns:
             The metadata lazy loader wrapper for the given key.
         """
-        from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
-
-        return LazyRunMetadataResponse(
+        return LazyRunMetadataResponse(  # type: ignore[return-value]
             lazy_load_model_name=self._lazy_load_model_name,
             lazy_load_model_version=self._lazy_load_model_version,
             lazy_load_artifact_name=self._lazy_load_artifact_name,