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

zenml/utils/metadata_utils.py

@@ -13,28 +13,30 @@
 #  permissions and limitations under the License.
 """Utility functions to handle metadata for ZenML entities."""
 
-import contextlib
-from typing import Dict, Optional, Union, overload
+from typing import Dict, List, Optional, Union, overload
 from uuid import UUID
 
 from zenml.client import Client
-from zenml.enums import MetadataResourceTypes
+from zenml.enums import MetadataResourceTypes, ModelStages
 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
 
 logger = get_logger(__name__)
 
 
 @overload
-def log_metadata(metadata: Dict[str, MetadataType]) -> None: ...
+def log_metadata(
+    metadata: Dict[str, MetadataType],
+) -> None: ...
 
 
 @overload
 def log_metadata(
     *,
     metadata: Dict[str, MetadataType],
-    artifact_version_id: UUID,
+    step_id: UUID,
 ) -> None: ...
 
 
@@ -42,8 +44,8 @@ def log_metadata(
 def log_metadata(
     *,
     metadata: Dict[str, MetadataType],
-    artifact_name: str,
-    artifact_version: Optional[str] = None,
+    step_name: str,
+    run_id_name_or_prefix: Union[UUID, str],
 ) -> None: ...
 
 
@@ -51,7 +53,7 @@ def log_metadata(
 def log_metadata(
     *,
     metadata: Dict[str, MetadataType],
-    model_version_id: UUID,
+    run_id_name_or_prefix: Union[UUID, str],
 ) -> None: ...
 
 
@@ -59,8 +61,7 @@ def log_metadata(
 def log_metadata(
     *,
     metadata: Dict[str, MetadataType],
-    model_name: str,
-    model_version: str,
+    artifact_version_id: UUID,
 ) -> None: ...
 
 
@@ -68,7 +69,8 @@ def log_metadata(
 def log_metadata(
     *,
     metadata: Dict[str, MetadataType],
-    step_id: UUID,
+    artifact_name: str,
+    artifact_version: Optional[str] = None,
 ) -> None: ...
 
 
@@ -76,33 +78,53 @@ def log_metadata(
 def log_metadata(
     *,
     metadata: Dict[str, MetadataType],
-    run_id_name_or_prefix: Union[UUID, str],
+    infer_artifact: bool = False,
+    artifact_name: Optional[str] = None,
 ) -> None: ...
 
 
+# Model Metadata
 @overload
 def log_metadata(
     *,
     metadata: Dict[str, MetadataType],
-    step_name: str,
-    run_id_name_or_prefix: Union[UUID, str],
+    model_version_id: UUID,
 ) -> None: ...
 
 
+@overload
 def log_metadata(
+    *,
     metadata: Dict[str, MetadataType],
-    # Parameters to manually log metadata for steps and runs
+    model_name: str,
+    model_version: Union[ModelStages, int, str],
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    infer_model: bool = False,
+) -> None: ...
+
+
+def log_metadata(
+    metadata: Dict[str, MetadataType],
+    # Steps and runs
     step_id: Optional[UUID] = None,
     step_name: Optional[str] = None,
     run_id_name_or_prefix: Optional[Union[UUID, str]] = None,
-    # Parameters to manually log metadata for artifacts
+    # Artifacts
     artifact_version_id: Optional[UUID] = None,
     artifact_name: Optional[str] = None,
     artifact_version: Optional[str] = None,
-    # Parameters to manually log metadata for models
+    infer_artifact: bool = False,
+    # Models
     model_version_id: Optional[UUID] = None,
     model_name: Optional[str] = None,
-    model_version: Optional[str] = None,
+    model_version: Optional[Union[ModelStages, int, str]] = None,
+    infer_model: bool = False,
 ) -> None:
     """Logs metadata for various resource types in a generalized way.
 
@@ -114,9 +136,13 @@ def log_metadata(
         artifact_version_id: The ID of the artifact version
         artifact_name: The name of the artifact.
         artifact_version: The version of the artifact.
+        infer_artifact: Flag deciding whether the artifact version should be
+            inferred from the step context.
         model_version_id: The ID of the model version.
         model_name: The name of the model.
-        model_version: The version of the model
+        model_version: The version of the model.
+        infer_model: Flag deciding whether the model version should be
+            inferred from the step context.
 
     Raises:
         ValueError: If no identifiers are provided and the function is not
@@ -124,141 +150,147 @@ def log_metadata(
     """
     client = Client()
 
-    # If a step name is provided, we need a run_id_name_or_prefix and will log
-    # metadata for the steps pipeline and model accordingly.
-    if step_name is not None and run_id_name_or_prefix is not None:
-        run_model = client.get_pipeline_run(
-            name_id_or_prefix=run_id_name_or_prefix
-        )
-        step_model = run_model.steps[step_name]
+    resources: List[RunMetadataResource] = []
+    publisher_step_id = None
 
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=run_model.id,
-            resource_type=MetadataResourceTypes.PIPELINE_RUN,
-        )
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=step_model.id,
-            resource_type=MetadataResourceTypes.STEP_RUN,
-        )
-        if step_model.model_version:
-            client.create_run_metadata(
-                metadata=metadata,
-                resource_id=step_model.model_version.id,
-                resource_type=MetadataResourceTypes.MODEL_VERSION,
+    # Log metadata to a step by ID
+    if step_id is not None:
+        resources = [
+            RunMetadataResource(
+                id=step_id, type=MetadataResourceTypes.STEP_RUN
             )
+        ]
 
-    # If a step is identified by id, fetch it directly through the client,
-    # follow a similar procedure and log metadata for its pipeline and model
-    # as well.
-    elif step_id is not None:
-        step_model = client.get_run_step(step_run_id=step_id)
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=step_model.pipeline_run_id,
-            resource_type=MetadataResourceTypes.PIPELINE_RUN,
-        )
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=step_model.id,
-            resource_type=MetadataResourceTypes.STEP_RUN,
+    # Log metadata to a step by name and run ID
+    elif step_name is not None and run_id_name_or_prefix is not None:
+        step_model_id = (
+            client.get_pipeline_run(name_id_or_prefix=run_id_name_or_prefix)
+            .steps[step_name]
+            .id
         )
-        if step_model.model_version:
-            client.create_run_metadata(
-                metadata=metadata,
-                resource_id=step_model.model_version.id,
-                resource_type=MetadataResourceTypes.MODEL_VERSION,
+        resources = [
+            RunMetadataResource(
+                id=step_model_id, type=MetadataResourceTypes.STEP_RUN
             )
+        ]
 
-    # If a pipeline run id is identified, we need to log metadata to it and its
-    # model as well.
+    # Log metadata to a run by ID
     elif run_id_name_or_prefix is not None:
         run_model = client.get_pipeline_run(
             name_id_or_prefix=run_id_name_or_prefix
         )
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=run_model.id,
-            resource_type=MetadataResourceTypes.PIPELINE_RUN,
-        )
-        if run_model.model_version:
-            client.create_run_metadata(
-                metadata=metadata,
-                resource_id=run_model.model_version.id,
-                resource_type=MetadataResourceTypes.MODEL_VERSION,
+        resources = [
+            RunMetadataResource(
+                id=run_model.id, type=MetadataResourceTypes.PIPELINE_RUN
             )
+        ]
 
-    # If the user provides a model name and version, we use to model abstraction
-    # to fetch the model version and attach the corresponding metadata to it.
+    # Log metadata to a model version by name and version
     elif model_name is not None and model_version is not None:
-        from zenml import Model
-
-        mv = Model(name=model_name, version=model_version)
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=mv.id,
-            resource_type=MetadataResourceTypes.MODEL_VERSION,
+        model_version_model = client.get_model_version(
+            model_name_or_id=model_name,
+            model_version_name_or_number_or_id=model_version,
         )
+        resources = [
+            RunMetadataResource(
+                id=model_version_model.id,
+                type=MetadataResourceTypes.MODEL_VERSION,
+            )
+        ]
 
-    # If the user provides a model version id, we use the client to fetch it and
-    # attach the metadata to it.
+    # Log metadata to a model version by id
     elif model_version_id is not None:
-        model_version_id = client.get_model_version(
-            model_version_name_or_number_or_id=model_version_id
-        ).id
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=model_version_id,
-            resource_type=MetadataResourceTypes.MODEL_VERSION,
+        resources = [
+            RunMetadataResource(
+                id=model_version_id,
+                type=MetadataResourceTypes.MODEL_VERSION,
+            )
+        ]
+
+    # Log metadata to a model through the step context
+    elif infer_model is True:
+        try:
+            step_context = get_step_context()
+        except RuntimeError:
+            raise ValueError(
+                "If you are using the `infer_model` option, the function must "
+                "be called inside a step with configured `model` in decorator."
+                "Otherwise, you can provide a `model_version_id` or a "
+                "combination of `model_name` and `model_version`."
+            )
+
+        if step_context.model_version is None:
+            raise ValueError(
+                "The step context does not feature any model versions."
+            )
+
+        resources = [
+            RunMetadataResource(
+                id=step_context.model_version.id,
+                type=MetadataResourceTypes.MODEL_VERSION,
+            )
+        ]
+
+    # Log metadata to an artifact version by its name and version
+    elif artifact_name is not None and artifact_version is not None:
+        artifact_version_model = client.get_artifact_version(
+            name_id_or_prefix=artifact_name, version=artifact_version
         )
+        resources = [
+            RunMetadataResource(
+                id=artifact_version_model.id,
+                type=MetadataResourceTypes.ARTIFACT_VERSION,
+            )
+        ]
 
-    # If the user provides an artifact name, there are three possibilities. If
-    # an artifact version is also provided with the name, we use both to fetch
-    # the artifact version and use it to log the metadata. If no version is
-    # provided, if the function is called within a step we search the artifacts
-    # of the step if not we fetch the latest version and attach the metadata
-    # to the latest version.
-    elif artifact_name is not None:
-        if artifact_version:
-            artifact_version_model = client.get_artifact_version(
-                name_id_or_prefix=artifact_name, version=artifact_version
+    # Log metadata to an artifact version by its ID
+    elif artifact_version_id is not None:
+        resources = [
+            RunMetadataResource(
+                id=artifact_version_id,
+                type=MetadataResourceTypes.ARTIFACT_VERSION,
             )
-            client.create_run_metadata(
-                metadata=metadata,
-                resource_id=artifact_version_model.id,
-                resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+        ]
+
+    # Log metadata to an artifact version through the step context
+    elif infer_artifact is True:
+        try:
+            step_context = get_step_context()
+        except RuntimeError:
+            raise ValueError(
+                "When you are using the `infer_artifact` option when you call "
+                "`log_metadata`, it must be called inside a step with outputs."
+                "Otherwise, you can provide a `artifact_version_id` or a "
+                "combination of `artifact_name` and `artifact_version`."
             )
-        else:
-            step_context = None
-            with contextlib.suppress(RuntimeError):
-                step_context = get_step_context()
 
-            if step_context:
-                step_context.add_output_metadata(
-                    metadata=metadata, output_name=artifact_name
-                )
-            else:
-                artifact_version_model = client.get_artifact_version(
-                    name_id_or_prefix=artifact_name
+        step_output_names = list(step_context._outputs.keys())
+
+        if artifact_name is not None:
+            # If a name provided, ensure it is in the outputs
+            if artifact_name not in step_output_names:
+                raise ValueError(
+                    f"The provided artifact name`{artifact_name}` does not "
+                    f"exist in the step outputs: {step_output_names}."
                 )
-                client.create_run_metadata(
-                    metadata=metadata,
-                    resource_id=artifact_version_model.id,
-                    resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+        else:
+            # If no name provided, ensure there is only one output
+            if len(step_output_names) > 1:
+                raise ValueError(
+                    "There is more than one output. If you would like to use "
+                    "the `infer_artifact` option, you need to define an "
+                    "`artifact_name`."
                 )
 
-    # If the user directly provides an artifact_version_id, we use the client to
-    # fetch is and attach the metadata accordingly.
-    elif artifact_version_id is not None:
-        artifact_version_model = client.get_artifact_version(
-            name_id_or_prefix=artifact_version_id,
-        )
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=artifact_version_model.id,
-            resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+            if len(step_output_names) == 0:
+                raise ValueError("The step does not have any outputs.")
+
+            artifact_name = step_output_names[0]
+
+        step_context.add_output_metadata(
+            metadata=metadata, output_name=artifact_name
         )
+        return
 
     # If every additional value is None, that means we are calling it bare bones
     # and this call needs to happen during a step execution. We will use the
@@ -287,22 +319,14 @@ def log_metadata(
                 "of the step execution, please provide the required "
                 "identifiers."
             )
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=step_context.pipeline_run.id,
-            resource_type=MetadataResourceTypes.PIPELINE_RUN,
-        )
-        client.create_run_metadata(
-            metadata=metadata,
-            resource_id=step_context.step_run.id,
-            resource_type=MetadataResourceTypes.STEP_RUN,
-        )
-        if step_context.model_version:
-            client.create_run_metadata(
-                metadata=metadata,
-                resource_id=step_context.model_version.id,
-                resource_type=MetadataResourceTypes.MODEL_VERSION,
+
+        resources = [
+            RunMetadataResource(
+                id=step_context.step_run.id,
+                type=MetadataResourceTypes.STEP_RUN,
             )
+        ]
+        publisher_step_id = step_context.step_run.id
 
     else:
         raise ValueError(
@@ -310,26 +334,35 @@ def log_metadata(
             Unsupported way to call the `log_metadata`. Possible combinations "
             include:
             
-            # Inside a step
-            # Logs the metadata to the step, its run and possibly its model
+            # Automatic logging to a step (within a step)
             log_metadata(metadata={})
             
-            # Manually logging for a step
-            # Logs the metadata to the step, its run and possibly its model
+            # Manual logging to a step
             log_metadata(metadata={}, step_name=..., run_id_name_or_prefix=...)
             log_metadata(metadata={}, step_id=...)
             
-            # Manually logging for a run 
-            # Logs the metadata to the run, possibly its model
+            # Manual logging to a run
             log_metadata(metadata={}, run_id_name_or_prefix=...)
             
-            # Manually logging for a model
+            # Automatic logging to a model (within a step)
+            log_metadata(metadata={}, infer_model=True)
+            
+            # Manual logging to a model
             log_metadata(metadata={}, model_name=..., model_version=...)
             log_metadata(metadata={}, model_version_id=...)
             
-            # Manually logging for an artifact
-            log_metadata(metadata={}, artifact_name=...)  # inside a step 
+            # Automatic logging to an artifact (within a step)
+            log_metadata(metadata={}, infer_artifact=True)  # step with single output
+            log_metadata(metadata={}, artifact_name=..., infer_artifact=True)  # specific output of a step
+            
+            # Manual logging to an artifact
             log_metadata(metadata={}, artifact_name=..., artifact_version=...)
             log_metadata(metadata={}, artifact_version_id=...)
             """
         )
+
+    client.create_run_metadata(
+        metadata=metadata,
+        resources=resources,
+        publisher_step_id=publisher_step_id,
+    )

zenml/utils/string_utils.py

@@ -18,7 +18,7 @@ import functools
 import random
 import string
 from datetime import datetime
-from typing import Any, Callable, Dict, TypeVar, cast
+from typing import Any, Callable, Dict, Optional, TypeVar, cast
 
 from pydantic import BaseModel
 
@@ -147,24 +147,26 @@ def validate_name(model: BaseModel) -> None:
 
 def format_name_template(
     name_template: str,
-    substitutions: Dict[str, str],
+    substitutions: Optional[Dict[str, str]] = None,
 ) -> str:
     """Formats a name template with the given arguments.
 
-    By default, ZenML support Date and Time placeholders.
-    E.g. `my_run_{date}_{time}` will be formatted as `my_run_1970_01_01_00_00_00`.
-    Extra placeholders need to be explicitly passed in as kwargs.
+    Default substitutions for `{date}` and `{time}` placeholders will be used if
+    not included in the provided substitutions.
 
     Args:
         name_template: The name template to format.
-        substitutions: A dictionary of substitutions to use in the template.
+        substitutions: Substitutions to use in the template.
 
     Returns:
         The formatted name template.
 
     Raises:
-        KeyError: If a key in template is missing in the kwargs.
+        KeyError: If a key in template is missing in the substitutions.
+        ValueError: If the formatted name is empty.
     """
+    substitutions = substitutions or {}
+
     if ("date" not in substitutions and "{date}" in name_template) or (
         "time" not in substitutions and "{time}" in name_template
     ):
@@ -183,13 +185,18 @@ def format_name_template(
         substitutions.setdefault("time", start_time.strftime("%H_%M_%S_%f"))
 
     try:
-        return name_template.format(**substitutions)
+        formatted_name = name_template.format(**substitutions)
     except KeyError as e:
         raise KeyError(
             f"Could not format the name template `{name_template}`. "
             f"Missing key: {e}"
         )
 
+    if not formatted_name:
+        raise ValueError("Empty names are not allowed.")
+
+    return formatted_name
+
 
 def substitute_string(value: V, substitution_func: Callable[[str], str]) -> V:
     """Recursively substitute strings in objects.

zenml/zen_server/routers/workspaces_endpoints.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for workspaces."""
 
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 from uuid import UUID
 
 from fastapi import APIRouter, Depends, Security
@@ -102,6 +102,7 @@ from zenml.zen_server.rbac.endpoint_utils import (
 )
 from zenml.zen_server.rbac.models import Action, ResourceType
 from zenml.zen_server.rbac.utils import (
+    batch_verify_permissions_for_models,
     dehydrate_page,
     dehydrate_response_model,
     get_allowed_resource_ids,
@@ -998,24 +999,23 @@ def create_run_metadata(
             "is not supported."
         )
 
-    if run_metadata.resource_type == MetadataResourceTypes.PIPELINE_RUN:
-        run = zen_store().get_run(run_metadata.resource_id)
-        verify_permission_for_model(run, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.STEP_RUN:
-        step = zen_store().get_run_step(run_metadata.resource_id)
-        verify_permission_for_model(step, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.ARTIFACT_VERSION:
-        artifact_version = zen_store().get_artifact_version(
-            run_metadata.resource_id
-        )
-        verify_permission_for_model(artifact_version, action=Action.UPDATE)
-    elif run_metadata.resource_type == MetadataResourceTypes.MODEL_VERSION:
-        model_version = zen_store().get_model_version(run_metadata.resource_id)
-        verify_permission_for_model(model_version, action=Action.UPDATE)
-    else:
-        raise RuntimeError(
-            f"Unknown resource type: {run_metadata.resource_type}"
-        )
+    verify_models: List[Any] = []
+    for resource in run_metadata.resources:
+        if resource.type == MetadataResourceTypes.PIPELINE_RUN:
+            verify_models.append(zen_store().get_run(resource.id))
+        elif resource.type == MetadataResourceTypes.STEP_RUN:
+            verify_models.append(zen_store().get_run_step(resource.id))
+        elif resource.type == MetadataResourceTypes.ARTIFACT_VERSION:
+            verify_models.append(zen_store().get_artifact_version(resource.id))
+        elif resource.type == MetadataResourceTypes.MODEL_VERSION:
+            verify_models.append(zen_store().get_model_version(resource.id))
+        else:
+            raise RuntimeError(f"Unknown resource type: {resource.type}")
+
+    batch_verify_permissions_for_models(
+        models=verify_models,
+        action=Action.UPDATE,
+    )
 
     verify_permission(
         resource_type=ResourceType.RUN_METADATA, action=Action.CREATE

zenml/zen_stores/migrations/versions/b73bc71f1106_remove_component_spec_path.py

@@ -0,0 +1,36 @@
+"""Remove component spec path [b73bc71f1106].
+
+Revision ID: b73bc71f1106
+Revises: ec6307720f92
+Create Date: 2024-11-29 09:36:33.089945
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b73bc71f1106"
+down_revision = "ec6307720f92"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    """Upgrade database schema and/or data, creating a new revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("stack_component", schema=None) as batch_op:
+        batch_op.drop_column("component_spec_path")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade database schema and/or data back to the previous revision."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("stack_component", schema=None) as batch_op:
+        batch_op.add_column(
+            sa.Column("component_spec_path", sa.VARCHAR(), nullable=True)
+        )
+
+    # ### end Alembic commands ###