zenml-nightly 0.68.1.dev20241103__py3-none-any.whl → 0.70.0.dev20241115__py3-none-any.whl

zenml/steps/base_step.py

@@ -327,12 +327,12 @@ class BaseStep:
             The artifacts, external artifacts, model version artifacts/metadata and parameters for the step.
         """
         from zenml.artifacts.external_artifact import ExternalArtifact
+        from zenml.metadata.lazy_load import LazyRunMetadataResponse
         from zenml.model.lazy_load import ModelVersionDataLazyLoader
         from zenml.models.v2.core.artifact_version import (
             ArtifactVersionResponse,
             LazyArtifactVersionResponse,
         )
-        from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
 
         signature = inspect.signature(self.entrypoint, follow_wrapped=True)
 

zenml/steps/entrypoint_function_utils.py

@@ -32,6 +32,7 @@ from zenml.constants import ENFORCE_TYPE_ANNOTATIONS
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
+from zenml.metadata.lazy_load import LazyRunMetadataResponse
 from zenml.steps.utils import (
     OutputSignature,
     parse_return_type_annotations,
@@ -136,10 +137,7 @@ class EntrypointFunctionDefinition(NamedTuple):
             UnmaterializedArtifact,
         )
         from zenml.client_lazy_loader import ClientLazyLoader
-        from zenml.models import (
-            ArtifactVersionResponse,
-            RunMetadataResponse,
-        )
+        from zenml.models import ArtifactVersionResponse
 
         if key not in self.inputs:
             raise KeyError(
@@ -154,8 +152,8 @@ class EntrypointFunctionDefinition(NamedTuple):
                 StepArtifact,
                 ExternalArtifact,
                 ArtifactVersionResponse,
-                RunMetadataResponse,
                 ClientLazyLoader,
+                LazyRunMetadataResponse,
             ),
         ):
             # If we were to do any type validation for artifacts here, we

zenml/steps/step_context.py

@@ -26,6 +26,7 @@ from typing import (
 
 from zenml.exceptions import StepContextError
 from zenml.logger import get_logger
+from zenml.utils.callback_registry import CallbackRegistry
 from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
@@ -34,11 +35,12 @@ if TYPE_CHECKING:
     from zenml.metadata.metadata_types import MetadataType
     from zenml.model.model import Model
     from zenml.models import (
-        ArtifactVersionResponse,
         PipelineResponse,
         PipelineRunResponse,
         StepRunResponse,
     )
+    from zenml.models.v2.core.step_run import StepRunInputResponse
+
 
 logger = get_logger(__name__)
 
@@ -145,6 +147,7 @@ class StepContext(metaclass=SingletonMetaClass):
             )
             for key in output_materializers.keys()
         }
+        self._cleanup_registry = CallbackRegistry()
 
     @property
     def pipeline(self) -> "PipelineResponse":
@@ -189,7 +192,7 @@ class StepContext(metaclass=SingletonMetaClass):
         return self.model_version.to_model_class()
 
     @property
-    def inputs(self) -> Dict[str, "ArtifactVersionResponse"]:
+    def inputs(self) -> Dict[str, "StepRunInputResponse"]:
         """Returns the input artifacts of the current step.
 
         Returns:

zenml/steps/utils.py

@@ -26,7 +26,11 @@ from typing_extensions import Annotated
 
 from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.client import Client
-from zenml.enums import ExecutionStatus, MetadataResourceTypes
+from zenml.enums import (
+    ArtifactSaveType,
+    ExecutionStatus,
+    MetadataResourceTypes,
+)
 from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
@@ -438,6 +442,11 @@ def log_step_metadata(
             from within a step or if no pipeline name or ID is provided and
             the function is not called from within a step.
     """
+    logger.warning(
+        "The `log_step_metadata` function is deprecated and will soon be "
+        "removed. Please use `log_metadata` instead."
+    )
+
     step_context = None
     if not step_name:
         with contextlib.suppress(RuntimeError):
@@ -547,8 +556,10 @@ def run_as_single_step_pipeline(
     # 4. Load output artifacts
     step_run = next(iter(run.steps.values()))
     outputs = [
-        step_run.outputs[output_name].load()
+        artifact_version.load()
         for output_name in step_run.config.outputs.keys()
+        for artifact_version in step_run.outputs[output_name]
+        if artifact_version.save_type == ArtifactSaveType.STEP_OUTPUT
     ]
 
     if len(outputs) == 0:

zenml/utils/callback_registry.py

@@ -0,0 +1,71 @@
+#  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.
+"""Callback registry implementation."""
+
+from typing import Any, Callable, Dict, List, Tuple
+
+from typing_extensions import ParamSpec
+
+from zenml.logger import get_logger
+
+P = ParamSpec("P")
+
+logger = get_logger(__name__)
+
+
+class CallbackRegistry:
+    """Callback registry class."""
+
+    def __init__(self) -> None:
+        """Initializes the callback registry."""
+        self._callbacks: List[
+            Tuple[Callable[P, Any], Tuple[Any], Dict[str, Any]]
+        ] = []
+
+    def register_callback(
+        self, callback: Callable[P, Any], *args: P.args, **kwargs: P.kwargs
+    ) -> None:
+        """Register a callback.
+
+        Args:
+            callback: The callback to register.
+            *args: Arguments to call the callback with.
+            **kwargs: Keyword arguments to call the callback with.
+        """
+        self._callbacks.append((callback, args, kwargs))  # type: ignore[arg-type]
+
+    def reset(self) -> None:
+        """Reset the callbacks."""
+        self._callbacks = []
+
+    def execute_callbacks(self, raise_on_exception: bool) -> None:
+        """Execute all registered callbacks.
+
+        Args:
+            raise_on_exception: If True, exceptions raised during the execution
+                of the callbacks will be raised. If False, a warning with the
+                exception will be logged instead.
+
+        Raises:
+            Exception: Exceptions raised in any of the callbacks if
+                `raise_on_exception` is set to True.
+        """
+        for callback, args, kwargs in self._callbacks:
+            try:
+                callback(*args, **kwargs)
+            except Exception as e:
+                if raise_on_exception:
+                    raise e
+                else:
+                    logger.warning("Failed to run callback: %s", str(e))

zenml/utils/metadata_utils.py

@@ -0,0 +1,335 @@
+#  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 functions to handle metadata for ZenML entities."""
+
+import contextlib
+from typing import Dict, Optional, Union, overload
+from uuid import UUID
+
+from zenml.client import Client
+from zenml.enums import MetadataResourceTypes
+from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
+from zenml.steps.step_context import get_step_context
+
+logger = get_logger(__name__)
+
+
+@overload
+def log_metadata(metadata: Dict[str, MetadataType]) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    artifact_version_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    artifact_name: str,
+    artifact_version: Optional[str] = None,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    model_version_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    model_name: str,
+    model_version: str,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    step_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    run_id_name_or_prefix: Union[UUID, str],
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    step_name: str,
+    run_id_name_or_prefix: Union[UUID, str],
+) -> None: ...
+
+
+def log_metadata(
+    metadata: Dict[str, MetadataType],
+    # Parameters to manually log metadata for 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
+    artifact_version_id: Optional[UUID] = None,
+    artifact_name: Optional[str] = None,
+    artifact_version: Optional[str] = None,
+    # Parameters to manually log metadata for models
+    model_version_id: Optional[UUID] = None,
+    model_name: Optional[str] = None,
+    model_version: Optional[str] = None,
+) -> None:
+    """Logs metadata for various resource types in a generalized way.
+
+    Args:
+        metadata: The metadata to log.
+        step_id: The ID of the step.
+        step_name: The name of the step.
+        run_id_name_or_prefix: The id, name or prefix of the run
+        artifact_version_id: The ID of the artifact version
+        artifact_name: The name of the artifact.
+        artifact_version: The version of the artifact.
+        model_version_id: The ID of the model version.
+        model_name: The name of the model.
+        model_version: The version of the model
+
+    Raises:
+        ValueError: If no identifiers are provided and the function is not
+            called from within a step.
+    """
+    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]
+
+        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,
+            )
+
+    # 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,
+        )
+        if step_model.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=step_model.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    # If a pipeline run id is identified, we need to log metadata to it and its
+    # model as well.
+    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,
+            )
+
+    # 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.
+    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,
+        )
+
+    # If the user provides a model version id, we use the client to fetch it and
+    # attach the metadata to it.
+    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,
+        )
+
+    # 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
+            )
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=artifact_version_model.id,
+                resource_type=MetadataResourceTypes.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
+                )
+                client.create_run_metadata(
+                    metadata=metadata,
+                    resource_id=artifact_version_model.id,
+                    resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+                )
+
+    # 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 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
+    # step context to fetch the step, run and possibly the model version and
+    # attach the metadata accordingly.
+    elif all(
+        v is None
+        for v in [
+            step_id,
+            step_name,
+            run_id_name_or_prefix,
+            artifact_version_id,
+            artifact_name,
+            artifact_version,
+            model_version_id,
+            model_name,
+            model_version,
+        ]
+    ):
+        try:
+            step_context = get_step_context()
+        except RuntimeError:
+            raise ValueError(
+                "You are calling 'log_metadata()' outside of a step execution. "
+                "If you would like to add metadata to a ZenML entity outside "
+                "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,
+            )
+
+    else:
+        raise ValueError(
+            """
+            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
+            log_metadata(metadata={})
+            
+            # Manually logging for a step
+            # Logs the metadata to the step, its run and possibly its model
+            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
+            log_metadata(metadata={}, run_id_name_or_prefix=...)
+            
+            # Manually logging for 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 
+            log_metadata(metadata={}, artifact_name=..., artifact_version=...)
+            log_metadata(metadata={}, artifact_version_id=...)
+            """
+        )