zenml-nightly 0.70.0.dev20241125__py3-none-any.whl → 0.71.0.dev20241220__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

@@ -14,11 +14,11 @@
 """Utils for strings."""
 
 import base64
-import datetime
 import functools
 import random
 import string
-from typing import Any, Callable, Dict, TypeVar, cast
+from datetime import datetime
+from typing import Any, Callable, Dict, Optional, TypeVar, cast
 
 from pydantic import BaseModel
 
@@ -147,30 +147,55 @@ def validate_name(model: BaseModel) -> None:
 
 def format_name_template(
     name_template: str,
-    **kwargs: 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.
-        **kwargs: The arguments to replace 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 substitutions.
+        ValueError: If the formatted name is empty.
     """
-    kwargs["date"] = kwargs.get(
-        "date",
-        datetime.datetime.now(datetime.timezone.utc).strftime("%Y_%m_%d"),
-    )
-    kwargs["time"] = kwargs.get(
-        "time",
-        datetime.datetime.now(datetime.timezone.utc).strftime("%H_%M_%S_%f"),
-    )
-    return name_template.format(**kwargs)
+    substitutions = substitutions or {}
+
+    if ("date" not in substitutions and "{date}" in name_template) or (
+        "time" not in substitutions and "{time}" in name_template
+    ):
+        from zenml import get_step_context
+
+        try:
+            pr = get_step_context().pipeline_run
+            start_time = pr.start_time
+            substitutions.update(pr.config.substitutions)
+        except RuntimeError:
+            start_time = None
+
+        if start_time is None:
+            start_time = datetime.utcnow()
+        substitutions.setdefault("date", start_time.strftime("%Y_%m_%d"))
+        substitutions.setdefault("time", start_time.strftime("%H_%M_%S_%f"))
+
+    try:
+        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:

zenml/utils/visualization_utils.py

@@ -13,9 +13,10 @@
 #  permissions and limitations under the License.
 """Utility functions for dashboard visualizations."""
 
+import json
 from typing import TYPE_CHECKING, Optional
 
-from IPython.core.display import HTML, Image, Markdown, display
+from IPython.core.display import HTML, JSON, Image, Markdown, display
 
 from zenml.artifacts.utils import load_artifact_visualization
 from zenml.enums import VisualizationType
@@ -63,6 +64,8 @@ def visualize_artifact(
             assert isinstance(visualization.value, str)
             table = format_csv_visualization_as_html(visualization.value)
             display(HTML(table))
+        elif visualization.type == VisualizationType.JSON:
+            display(JSON(json.loads(visualization.value)))
         else:
             display(visualization.value)
 

zenml/zen_server/auth.py

@@ -413,10 +413,7 @@ def authenticate_credentials(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
-            if pipeline_run_status in [
-                ExecutionStatus.FAILED,
-                ExecutionStatus.COMPLETED,
-            ]:
+            if pipeline_run_status.is_finished:
                 error = (
                     f"The execution of pipeline run "
                     f"{decoded_token.pipeline_run_id} has already concluded and "
@@ -461,10 +458,7 @@ def authenticate_credentials(
                 logger.error(error)
                 raise CredentialsNotValid(error)
 
-            if step_run_status in [
-                ExecutionStatus.FAILED,
-                ExecutionStatus.COMPLETED,
-            ]:
+            if step_run_status.is_finished:
                 error = (
                     f"The execution of step run "
                     f"{decoded_token.step_run_id} has already concluded and "
@@ -821,7 +815,10 @@ def generate_access_token(
         response: The FastAPI response object.
         device: The device used for authentication.
         api_key: The service account API key used for authentication.
-        expires_in: The number of seconds until the token expires.
+        expires_in: The number of seconds until the token expires. If not set,
+            the default value is determined automatically based on the server
+            configuration and type of token. If set to 0, the token will not
+            expire.
         schedule_id: The ID of the schedule to scope the token to.
         pipeline_run_id: The ID of the pipeline run to scope the token to.
         step_run_id: The ID of the step run to scope the token to.
@@ -835,7 +832,9 @@ def generate_access_token(
     # according to the values configured in the server config. Device tokens are
     # handled separately from regular user tokens.
     expires: Optional[datetime] = None
-    if expires_in:
+    if expires_in == 0:
+        expires_in = None
+    elif expires_in is not None:
         expires = datetime.utcnow() + timedelta(seconds=expires_in)
     elif device:
         # If a device was used for authentication, the token will expire

zenml/zen_server/cloud_utils.py

@@ -170,7 +170,9 @@ class ZenMLCloudConnection:
             token = self._fetch_auth_token()
             self._session.headers.update({"Authorization": "Bearer " + token})
 
-            retries = Retry(total=5, backoff_factor=0.1)
+            retries = Retry(
+                total=5, backoff_factor=0.1, status_forcelist=[502, 504]
+            )
             self._session.mount(
                 "https://",
                 HTTPAdapter(

zenml/zen_server/dashboard/assets/{404-NVXKFp-x.js → 404-Cqu3EDCm.js}

@@ -1 +1 @@
-import{j as e}from"./@radix-DeK6qiuw.js";import{f as s,r as t}from"./index-CCOPpudF.js";import{E as r}from"./EmptyState-BzdlCwp3.js";import{S as a}from"./help-Cc9bBIJH.js";import{L as o}from"./@react-router-B3Z5rLr2.js";import"./@tanstack-DT5WLu9C.js";import"./@reactflow-CK0KJUen.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};
+import{j as e}from"./@radix-DeK6qiuw.js";import{f as s,r as t}from"./index-FO-p0GU7.js";import{E as r}from"./EmptyState-BzdlCwp3.js";import{S as a}from"./help-Cc9bBIJH.js";import{L as o}from"./@react-router-B3Z5rLr2.js";import"./@tanstack-DT5WLu9C.js";import"./@reactflow-D2Y7BWwz.js";function d(){return e.jsx("div",{className:"flex min-h-screen w-full flex-col",children:e.jsx(r,{icon:e.jsx(a,{className:"h-[120px] w-[120px] fill-neutral-300"}),children:e.jsxs("div",{className:"text-center",children:[e.jsx("h1",{className:"mb-2 text-display-xs font-semibold",children:"We can't find the page you are looking for"}),e.jsx("p",{className:"text-lg text-theme-text-secondary",children:"You can try typing a different URL or we can bring you back to your Homepage."}),e.jsx("div",{className:"mt-5 flex justify-center",children:e.jsx(s,{size:"md",asChild:!0,children:e.jsx(o,{className:"w-min self-center whitespace-nowrap",to:t.home,children:e.jsx("span",{className:"px-0.5",children:"Go to Home"})})})})]})})})}export{d as default};