zenml-nightly 0.70.0.dev20241129__py3-none-any.whl → 0.70.0.dev20241130__py3-none-any.whl

zenml/models/v2/core/run_metadata.py

@@ -13,16 +13,16 @@
 #  permissions and limitations under the License.
 """Models representing run metadata."""
 
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 from uuid import UUID
 
-from pydantic import Field
+from pydantic import Field, model_validator
 
-from zenml.enums import MetadataResourceTypes
 from zenml.metadata.metadata_types import MetadataType, MetadataTypeEnum
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedRequest,
 )
+from zenml.models.v2.misc.run_metadata import RunMetadataResource
 
 # ------------------ Request Model ------------------
 
@@ -30,14 +30,12 @@ from zenml.models.v2.base.scoped import (
 class RunMetadataRequest(WorkspaceScopedRequest):
     """Request model for run metadata."""
 
-    resource_id: UUID = Field(
-        title="The ID of the resource that this metadata belongs to.",
-    )
-    resource_type: MetadataResourceTypes = Field(
-        title="The type of the resource that this metadata belongs to.",
+    resources: List[RunMetadataResource] = Field(
+        title="The list of resources that this metadata belongs to."
     )
     stack_component_id: Optional[UUID] = Field(
-        title="The ID of the stack component that this metadata belongs to."
+        title="The ID of the stack component that this metadata belongs to.",
+        default=None,
     )
     values: Dict[str, "MetadataType"] = Field(
         title="The metadata to be created.",
@@ -45,3 +43,26 @@ class RunMetadataRequest(WorkspaceScopedRequest):
     types: Dict[str, "MetadataTypeEnum"] = Field(
         title="The types of the metadata to be created.",
     )
+    publisher_step_id: Optional[UUID] = Field(
+        title="The ID of the step execution that published this metadata.",
+        default=None,
+    )
+
+    @model_validator(mode="after")
+    def validate_values_keys(self) -> "RunMetadataRequest":
+        """Validates if the keys in the metadata are properly defined.
+
+        Returns:
+            self
+
+        Raises:
+            ValueError: if one of the key in the metadata contains `:`
+        """
+        invalid_keys = [key for key in self.values.keys() if ":" in key]
+        if invalid_keys:
+            raise ValueError(
+                "You can not use colons (`:`) in the key names when you "
+                "are creating metadata for your ZenML objects. Please change "
+                f"the following keys: {invalid_keys}"
+            )
+        return self

zenml/models/v2/core/step_run.py

@@ -594,6 +594,7 @@ class StepRunFilter(WorkspaceScopedFilter):
         from zenml.zen_stores.schemas import (
             ModelSchema,
             ModelVersionSchema,
+            RunMetadataResourceSchema,
             RunMetadataSchema,
             StepRunSchema,
         )
@@ -612,10 +613,11 @@ class StepRunFilter(WorkspaceScopedFilter):
 
             for key, value in self.run_metadata.items():
                 additional_filter = and_(
-                    RunMetadataSchema.resource_id == StepRunSchema.id,
-                    RunMetadataSchema.resource_type
+                    RunMetadataResourceSchema.resource_id == StepRunSchema.id,
+                    RunMetadataResourceSchema.resource_type
                     == MetadataResourceTypes.STEP_RUN,
-                    RunMetadataSchema.key == key,
+                    RunMetadataResourceSchema.run_metadata_id
+                    == RunMetadataSchema.id,
                     self.generate_custom_query_conditions_for_column(
                         value=value,
                         table=RunMetadataSchema,

zenml/models/v2/misc/run_metadata.py

@@ -0,0 +1,38 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Utility classes for modeling run metadata."""
+
+from datetime import datetime
+from uuid import UUID
+
+from pydantic import BaseModel, Field
+
+from zenml.enums import MetadataResourceTypes
+from zenml.metadata.metadata_types import MetadataType
+
+
+class RunMetadataResource(BaseModel):
+    """Utility class to help identify resources to tag metadata to."""
+
+    id: UUID = Field(title="The ID of the resource.")
+    type: MetadataResourceTypes = Field(title="The type of the resource.")
+
+
+class RunMetadataEntry(BaseModel):
+    """Utility class to sort/list run metadata entries."""
+
+    value: MetadataType = Field(title="The value for the run metadata entry")
+    created: datetime = Field(
+        title="The timestamp when this resource was created."
+    )

zenml/orchestrators/input_utils.py

@@ -20,7 +20,7 @@ from zenml.client import Client
 from zenml.config.step_configurations import Step
 from zenml.enums import ArtifactSaveType, StepRunInputArtifactType
 from zenml.exceptions import InputResolutionError
-from zenml.utils import pagination_utils
+from zenml.utils import pagination_utils, string_utils
 
 if TYPE_CHECKING:
     from zenml.models import PipelineRunResponse
@@ -53,7 +53,8 @@ def resolve_step_inputs(
     current_run_steps = {
         run_step.name: run_step
         for run_step in pagination_utils.depaginate(
-            Client().list_run_steps, pipeline_run_id=pipeline_run.id
+            Client().list_run_steps,
+            pipeline_run_id=pipeline_run.id,
         )
     }
 
@@ -66,11 +67,23 @@ def resolve_step_inputs(
                 f"No step `{input_.step_name}` found in current run."
             )
 
+        # Try to get the substitutions from the pipeline run first, as we
+        # already have a hydrated version of that. In the unlikely case
+        # that the pipeline run is outdated, we fetch it from the step
+        # run instead which will costs us one hydration call.
+        substitutions = (
+            pipeline_run.step_substitutions.get(step_run.name)
+            or step_run.config.substitutions
+        )
+        output_name = string_utils.format_name_template(
+            input_.output_name, substitutions=substitutions
+        )
+
         try:
-            outputs = step_run.outputs[input_.output_name]
+            outputs = step_run.outputs[output_name]
         except KeyError:
             raise InputResolutionError(
-                f"No step output `{input_.output_name}` found for step "
+                f"No step output `{output_name}` found for step "
                 f"`{input_.step_name}`."
             )
 
@@ -83,12 +96,12 @@ def resolve_step_inputs(
             # This should never happen, there can only be a single regular step
             # output for a name
             raise InputResolutionError(
-                f"Too many step outputs for output `{input_.output_name}` of "
+                f"Too many step outputs for output `{output_name}` of "
                 f"step `{input_.step_name}`."
             )
         elif len(step_outputs) == 0:
             raise InputResolutionError(
-                f"No step output `{input_.output_name}` found for step "
+                f"No step output `{output_name}` found for step "
                 f"`{input_.step_name}`."
             )
 

zenml/orchestrators/publish_utils.py

@@ -21,6 +21,7 @@ from zenml.enums import ExecutionStatus, MetadataResourceTypes
 from zenml.models import (
     PipelineRunResponse,
     PipelineRunUpdate,
+    RunMetadataResource,
     StepRunResponse,
     StepRunUpdate,
 )
@@ -129,8 +130,11 @@ def publish_pipeline_run_metadata(
     for stack_component_id, metadata in pipeline_run_metadata.items():
         client.create_run_metadata(
             metadata=metadata,
-            resource_id=pipeline_run_id,
-            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+            resources=[
+                RunMetadataResource(
+                    id=pipeline_run_id, type=MetadataResourceTypes.PIPELINE_RUN
+                )
+            ],
             stack_component_id=stack_component_id,
         )
 
@@ -150,7 +154,10 @@ def publish_step_run_metadata(
     for stack_component_id, metadata in step_run_metadata.items():
         client.create_run_metadata(
             metadata=metadata,
-            resource_id=step_run_id,
-            resource_type=MetadataResourceTypes.STEP_RUN,
+            resources=[
+                RunMetadataResource(
+                    id=step_run_id, type=MetadataResourceTypes.STEP_RUN
+                )
+            ],
             stack_component_id=stack_component_id,
         )

zenml/orchestrators/step_launcher.py

@@ -310,8 +310,8 @@ class StepLauncher:
             and a boolean indicating whether the run was created or reused.
         """
         start_time = datetime.utcnow()
-        run_name = orchestrator_utils.get_run_name(
-            run_name_template=self._deployment.run_name_template,
+        run_name = string_utils.format_name_template(
+            name_template=self._deployment.run_name_template,
             substitutions=self._deployment.pipeline_configuration._get_full_substitutions(
                 start_time
             ),

zenml/orchestrators/step_run_utils.py

@@ -309,6 +309,9 @@ def create_cached_step_runs(
         for invocation_id in cache_candidates:
             visited_invocations.add(invocation_id)
 
+            # Make sure the request factory has the most up to date pipeline
+            # run to avoid hydration calls
+            request_factory.pipeline_run = pipeline_run
             try:
                 step_run_request = request_factory.create_request(
                     invocation_id

zenml/orchestrators/step_runner.py

@@ -56,7 +56,7 @@ from zenml.steps.utils import (
     parse_return_type_annotations,
     resolve_type_annotation,
 )
-from zenml.utils import materializer_utils, source_utils
+from zenml.utils import materializer_utils, source_utils, string_utils
 from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
@@ -292,16 +292,15 @@ class StepRunner:
         """
         collections.append(output_annotations)
         for k, v in list(output_annotations.items()):
-            _evaluated_name = None
-            if v.artifact_config:
-                _evaluated_name = v.artifact_config._evaluated_name(
-                    step_run.config.substitutions
+            name = k
+            if v.artifact_config and v.artifact_config.name:
+                name = string_utils.format_name_template(
+                    v.artifact_config.name,
+                    substitutions=step_run.config.substitutions,
                 )
-            if _evaluated_name is None:
-                _evaluated_name = k
 
             for d in collections:
-                d[_evaluated_name] = d.pop(k)
+                d[name] = d.pop(k)
 
     def _load_step(self) -> "BaseStep":
         """Load the step instance.

zenml/orchestrators/utils.py

@@ -32,7 +32,6 @@ from zenml.constants import (
 from zenml.enums import AuthScheme, StackComponentType, StoreType
 from zenml.logger import get_logger
 from zenml.stack import StackComponent
-from zenml.utils.string_utils import format_name_template
 
 logger = get_logger(__name__)
 
@@ -196,29 +195,6 @@ def get_config_environment_vars(
     return environment_vars
 
 
-def get_run_name(run_name_template: str, substitutions: Dict[str, str]) -> str:
-    """Fill out the run name template to get a complete run name.
-
-    Args:
-        run_name_template: The run name template to fill out.
-        substitutions: The substitutions to use in the template.
-
-    Raises:
-        ValueError: If the run name is empty.
-
-    Returns:
-        The run name derived from the template.
-    """
-    run_name = format_name_template(
-        run_name_template, substitutions=substitutions
-    )
-
-    if run_name == "":
-        raise ValueError("Empty run names are not allowed.")
-
-    return run_name
-
-
 class register_artifact_store_filesystem:
     """Context manager for the artifact_store/filesystem_registry dependency.
 

zenml/pipelines/run_utils.py

@@ -23,9 +23,8 @@ from zenml.models import (
     StackResponse,
 )
 from zenml.orchestrators.publish_utils import publish_failed_pipeline_run
-from zenml.orchestrators.utils import get_run_name
 from zenml.stack import Flavor, Stack
-from zenml.utils import code_utils, notebook_utils, source_utils
+from zenml.utils import code_utils, notebook_utils, source_utils, string_utils
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
@@ -68,8 +67,8 @@ def create_placeholder_run(
         return None
     start_time = datetime.utcnow()
     run_request = PipelineRunRequest(
-        name=get_run_name(
-            run_name_template=deployment.run_name_template,
+        name=string_utils.format_name_template(
+            name_template=deployment.run_name_template,
             substitutions=deployment.pipeline_configuration._get_full_substitutions(
                 start_time
             ),

zenml/steps/utils.py

@@ -42,6 +42,7 @@ from zenml.enums import (
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
+from zenml.models import RunMetadataResource
 from zenml.steps.step_context import get_step_context
 from zenml.utils import settings_utils, source_code_utils, typing_utils
 
@@ -489,8 +490,11 @@ def log_step_metadata(
         step_run_id = pipeline_run.steps[step_name].id
     client.create_run_metadata(
         metadata=metadata,
-        resource_id=step_run_id,
-        resource_type=MetadataResourceTypes.STEP_RUN,
+        resources=[
+            RunMetadataResource(
+                id=step_run_id, type=MetadataResourceTypes.STEP_RUN
+            )
+        ],
     )