zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/{new/steps → steps}/decorated_step.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Internal BaseStep subclass used by the step decorator."""
+
 from typing import Any
 
 from zenml.config.source import Source

zenml/steps/entrypoint_function_utils.py

@@ -12,6 +12,7 @@
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
 """Util functions for step and pipeline entrypoint functions."""
+
 import inspect
 from typing import (
     TYPE_CHECKING,
@@ -19,18 +20,19 @@ from typing import (
     Callable,
     Dict,
     NamedTuple,
-    Optional,
+    NoReturn,
     Sequence,
     Type,
     Union,
 )
 
-from pydantic import BaseConfig, ValidationError, create_model
+from pydantic import ConfigDict, ValidationError, create_model
 
 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,
@@ -40,45 +42,13 @@ from zenml.utils import yaml_utils
 
 if TYPE_CHECKING:
     from zenml.config.source import Source
-    from zenml.new.pipelines.pipeline import Pipeline
-    from zenml.steps import BaseStep
+    from zenml.pipelines.pipeline_definition import Pipeline
 
     MaterializerClassOrSource = Union[str, "Source", Type["BaseMaterializer"]]
 
 logger = get_logger(__name__)
 
 
-def get_step_entrypoint_signature(step: "BaseStep") -> inspect.Signature:
-    """Get the entrypoint signature of a step.
-
-    Args:
-        step: The step for which to get the entrypoint signature.
-
-    Returns:
-        The entrypoint function signature.
-    """
-    from zenml.steps import BaseParameters, StepContext
-
-    signature = inspect.signature(step.entrypoint, follow_wrapped=True)
-
-    def _is_param_of_class(annotation: Any, class_: Type[Any]) -> bool:
-        annotation = resolve_type_annotation(annotation)
-        return inspect.isclass(annotation) and issubclass(annotation, class_)
-
-    parameters = list(signature.parameters.values())
-
-    # Filter out deprecated args: step context and legacy parameters
-    parameters = [
-        param
-        for param in parameters
-        if not _is_param_of_class(param.annotation, class_=BaseParameters)
-        and not _is_param_of_class(param.annotation, class_=StepContext)
-    ]
-
-    signature = signature.replace(parameters=parameters)
-    return signature
-
-
 class StepArtifact:
     """Class to represent step output artifacts."""
 
@@ -102,6 +72,20 @@ class StepArtifact:
         self.annotation = annotation
         self.pipeline = pipeline
 
+    def __iter__(self) -> NoReturn:
+        """Raise a custom error if someone is trying to iterate this object.
+
+        Raises:
+            StepInterfaceError: If trying to iterate this object.
+        """
+        raise StepInterfaceError(
+            "Unable to unpack step artifact. This error is probably because "
+            "you're trying to unpack the return value of your step but the "
+            "step only returns a single artifact. For more information on how "
+            "to add type annotations to your step to indicate multiple "
+            "artifacts visit https://docs.zenml.io/how-to/build-pipelines/step-output-typing-and-annotation#type-annotations."
+        )
+
 
 def validate_reserved_arguments(
     signature: inspect.Signature, reserved_arguments: Sequence[str]
@@ -127,15 +111,10 @@ class EntrypointFunctionDefinition(NamedTuple):
         inputs: The entrypoint function inputs.
         outputs: The entrypoint function outputs. This dictionary maps output
             names to output annotations.
-        context: Optional parameter representing the `StepContext` input.
-        legacy_params: Optional parameter representing the `BaseParameters`
-            input.
     """
 
     inputs: Dict[str, inspect.Parameter]
     outputs: Dict[str, OutputSignature]
-    context: Optional[inspect.Parameter]
-    legacy_params: Optional[inspect.Parameter]
 
     def validate_input(self, key: str, value: Any) -> None:
         """Validates an input to the step entrypoint function.
@@ -157,10 +136,8 @@ class EntrypointFunctionDefinition(NamedTuple):
         from zenml.artifacts.unmaterialized_artifact import (
             UnmaterializedArtifact,
         )
-        from zenml.models import (
-            ArtifactVersionResponse,
-            RunMetadataResponse,
-        )
+        from zenml.client_lazy_loader import ClientLazyLoader
+        from zenml.models import ArtifactVersionResponse
 
         if key not in self.inputs:
             raise KeyError(
@@ -175,7 +152,8 @@ class EntrypointFunctionDefinition(NamedTuple):
                 StepArtifact,
                 ExternalArtifact,
                 ArtifactVersionResponse,
-                RunMetadataResponse,
+                ClientLazyLoader,
+                LazyRunMetadataResponse,
             ),
         ):
             # If we were to do any type validation for artifacts here, we
@@ -196,7 +174,7 @@ class EntrypointFunctionDefinition(NamedTuple):
                 f"'{key}' is not JSON serializable and can not be passed as "
                 "a parameter. This input can either be provided by the "
                 "output of another step or as an external artifact: "
-                "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/configure-steps-pipelines#pass-any-kind-of-data-to-your-steps"
+                "https://docs.zenml.io/user-guide/starter-guide/manage-artifacts#managing-artifacts-not-produced-by-zenml-pipelines"
             )
 
         try:
@@ -217,23 +195,24 @@ class EntrypointFunctionDefinition(NamedTuple):
             parameter: The function parameter for which the value was provided.
             value: The input value.
         """
-
-        class ModelConfig(BaseConfig):
-            arbitrary_types_allowed = False
+        # We allow passing None for optional annotations that would otherwise
+        # not be allowed as a parameter
+        config_dict = ConfigDict(arbitrary_types_allowed=value is None)
 
         # Create a pydantic model with just a single required field with the
         # type annotation of the parameter to verify the input type including
         # pydantics type coercion
         validation_model_class = create_model(
             "input_validation_model",
-            __config__=ModelConfig,
+            __config__=config_dict,
             value=(parameter.annotation, ...),
         )
         validation_model_class(value=value)
 
 
 def validate_entrypoint_function(
-    func: Callable[..., Any], reserved_arguments: Sequence[str] = ()
+    func: Callable[..., Any],
+    reserved_arguments: Sequence[str] = (),
 ) -> EntrypointFunctionDefinition:
     """Validates a step entrypoint function.
 
@@ -244,36 +223,20 @@ def validate_entrypoint_function(
     Raises:
         StepInterfaceError: If the entrypoint function has variable arguments
             or keyword arguments.
-        StepInterfaceError: If the entrypoint function has multiple
-            `BaseParameter` arguments.
-        StepInterfaceError: If the entrypoint function has multiple
-            `StepContext` arguments.
         RuntimeError: If type annotations should be enforced and a type
             annotation is missing.
 
     Returns:
         A validated definition of the entrypoint function.
     """
-    from zenml.steps import BaseParameters, StepContext
-
     signature = inspect.signature(func, follow_wrapped=True)
     validate_reserved_arguments(
         signature=signature, reserved_arguments=reserved_arguments
     )
 
     inputs = {}
-    context: Optional[inspect.Parameter] = None
-    legacy_params: Optional[inspect.Parameter] = None
 
     signature_parameters = list(signature.parameters.items())
-    if signature_parameters and signature_parameters[0][0] == "self":
-        # TODO: Once we get rid of the old step decorator, we can also remove
-        # the `BaseStepMeta` class which right now calls this function on an
-        # unbound instance method when using the class-based API. If we get rid
-        # of that, this check and removal of the `self` parameter is not
-        # necessary anymore
-        signature_parameters = signature_parameters[1:]
-
     for key, parameter in signature_parameters:
         if parameter.kind in {parameter.VAR_POSITIONAL, parameter.VAR_KEYWORD}:
             raise StepInterfaceError(
@@ -293,37 +256,14 @@ def validate_entrypoint_function(
             parameter = parameter.replace(annotation=Any)
 
         annotation = resolve_type_annotation(annotation)
-        if inspect.isclass(annotation) and issubclass(
-            annotation, BaseParameters
-        ):
-            if legacy_params is not None:
-                raise StepInterfaceError(
-                    f"Found multiple parameter arguments "
-                    f"('{legacy_params.name}' and '{key}') "
-                    f"for function {func.__name__}."
-                )
-            legacy_params = parameter
-
-        elif inspect.isclass(annotation) and issubclass(
-            annotation, StepContext
-        ):
-            if context is not None:
-                raise StepInterfaceError(
-                    f"Found multiple context arguments "
-                    f"('{context.name}' and '{key}') "
-                    f"for function {func.__name__}."
-                )
-            context = parameter
-        else:
-            inputs[key] = parameter
+        inputs[key] = parameter
 
     outputs = parse_return_type_annotations(
-        func=func, enforce_type_annotations=ENFORCE_TYPE_ANNOTATIONS
+        func=func,
+        enforce_type_annotations=ENFORCE_TYPE_ANNOTATIONS,
     )
 
     return EntrypointFunctionDefinition(
         inputs=inputs,
         outputs=outputs,
-        context=context,
-        legacy_params=legacy_params,
     )

zenml/{new/steps → steps}/step_context.py

@@ -17,28 +17,30 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
+    List,
     Mapping,
     Optional,
     Sequence,
     Type,
 )
 
-from zenml.exceptions import EntityExistsError, StepContextError
+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:
     from zenml.artifacts.artifact_config import ArtifactConfig
-    from zenml.config.step_run_info import StepRunInfo
     from zenml.materializers.base_materializer import BaseMaterializer
     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__)
 
@@ -94,8 +96,6 @@ class StepContext(metaclass=SingletonMetaClass):
         output_materializers: Mapping[str, Sequence[Type["BaseMaterializer"]]],
         output_artifact_uris: Mapping[str, str],
         output_artifact_configs: Mapping[str, Optional["ArtifactConfig"]],
-        step_run_info: "StepRunInfo",
-        cache_enabled: bool,
     ) -> None:
         """Initialize the context of the currently running step.
 
@@ -108,8 +108,6 @@ class StepContext(metaclass=SingletonMetaClass):
                 context is used in.
             output_artifact_configs: The outputs' ArtifactConfigs of the step that this
                 context is used in.
-            step_run_info: (Deprecated) info about the currently running step.
-            cache_enabled: (Deprecated) Whether caching is enabled for the step.
 
         Raises:
             StepContextError: If the keys of the output materializers and
@@ -117,13 +115,19 @@ class StepContext(metaclass=SingletonMetaClass):
         """
         from zenml.client import Client
 
+        try:
+            pipeline_run = Client().get_pipeline_run(pipeline_run.id)
+        except KeyError:
+            pass
         self.pipeline_run = pipeline_run
+        try:
+            step_run = Client().get_run_step(step_run.id)
+        except KeyError:
+            pass
         self.step_run = step_run
-        self._step_run_info = step_run_info
-        self._cache_enabled = cache_enabled
-
-        # Get the stack that we are running in
-        self._stack = Client().active_stack
+        self.model_version = (
+            step_run.model_version or pipeline_run.model_version
+        )
 
         self.step_name = self.step_run.name
 
@@ -131,7 +135,7 @@ class StepContext(metaclass=SingletonMetaClass):
         if output_materializers.keys() != output_artifact_uris.keys():
             raise StepContextError(
                 f"Mismatched keys in output materializers and output artifact "
-                f"URIs for step '{self.step_name}'. Output materializer "
+                f"URIs for step `{self.step_name}`. Output materializer "
                 f"keys: {set(output_materializers)}, output artifact URI "
                 f"keys: {set(output_artifact_uris)}"
             )
@@ -143,6 +147,7 @@ class StepContext(metaclass=SingletonMetaClass):
             )
             for key in output_materializers.keys()
         }
+        self._cleanup_registry = CallbackRegistry()
 
     @property
     def pipeline(self) -> "PipelineResponse":
@@ -157,7 +162,7 @@ class StepContext(metaclass=SingletonMetaClass):
         if self.pipeline_run.pipeline:
             return self.pipeline_run.pipeline
         raise StepContextError(
-            f"Unable to get pipeline in step '{self.step_name}' of pipeline "
+            f"Unable to get pipeline in step `{self.step_name}` of pipeline "
             f"run '{self.pipeline_run.id}': This pipeline run does not have "
             f"a pipeline associated with it."
         )
@@ -167,32 +172,27 @@ class StepContext(metaclass=SingletonMetaClass):
         """Returns configured Model.
 
         Order of resolution to search for Model is:
-            1. Model from @step
-            2. Model from @pipeline
+            1. Model from the step context
+            2. Model from the pipeline context
 
         Returns:
             The `Model` object associated with the current step.
 
         Raises:
-            StepContextError: If the `Model` object is not set in `@step` or `@pipeline`.
+            StepContextError: If no `Model` object was specified for the step
+                or pipeline.
         """
-        if self.step_run.config.model is not None:
-            model = self.step_run.config.model
-        elif self.pipeline_run.config.model is not None:
-            model = self.pipeline_run.config.model
-        else:
+        if not self.model_version:
             raise StepContextError(
-                f"Unable to get Model in step '{self.step_name}' of pipeline "
-                f"run '{self.pipeline_run.id}': it was not set in `@step` or `@pipeline`."
+                f"Unable to get Model in step `{self.step_name}` of pipeline "
+                f"run '{self.pipeline_run.id}': No model has been specified "
+                "the step or pipeline."
             )
 
-        # warm-up the model version
-        model._get_or_create_model_version()
-
-        return model
+        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:
@@ -221,13 +221,13 @@ class StepContext(metaclass=SingletonMetaClass):
         output_count = len(self._outputs)
         if output_count == 0:
             raise StepContextError(
-                f"Unable to get step output for step '{self.step_name}': "
+                f"Unable to get step output for step `{self.step_name}`: "
                 f"This step does not have any outputs."
             )
 
         if not output_name and output_count > 1:
             raise StepContextError(
-                f"Unable to get step output for step '{self.step_name}': "
+                f"Unable to get step output for step `{self.step_name}`: "
                 f"This step has multiple outputs ({set(self._outputs)}), "
                 f"please specify which output to return."
             )
@@ -236,7 +236,7 @@ class StepContext(metaclass=SingletonMetaClass):
             if output_name not in self._outputs:
                 raise StepContextError(
                     f"Unable to get step output '{output_name}' for "
-                    f"step '{self.step_name}'. This step does not have an "
+                    f"step `{self.step_name}`. This step does not have an "
                     f"output with the given name, please specify one of the "
                     f"available outputs: {set(self._outputs)}."
                 )
@@ -318,7 +318,33 @@ class StepContext(metaclass=SingletonMetaClass):
         Returns:
             Metadata for the given output.
         """
-        return self._get_output(output_name).run_metadata or {}
+        output = self._get_output(output_name)
+        custom_metadata = output.run_metadata or {}
+        if output.artifact_config:
+            custom_metadata.update(
+                **(output.artifact_config.run_metadata or {})
+            )
+        return custom_metadata
+
+    def get_output_tags(self, output_name: Optional[str] = None) -> List[str]:
+        """Returns the tags for a given step output.
+
+        Args:
+            output_name: Optional name of the output for which to get the
+                metadata. If no name is given and the step only has a single
+                output, the metadata of this output will be returned. If the
+                step has multiple outputs, an exception will be raised.
+
+        Returns:
+            Tags for the given output.
+        """
+        output = self._get_output(output_name)
+        custom_tags = set(output.tags or [])
+        if output.artifact_config:
+            return list(
+                set(output.artifact_config.tags or []).union(custom_tags)
+            )
+        return list(custom_tags)
 
     def add_output_metadata(
         self,
@@ -339,31 +365,24 @@ class StepContext(metaclass=SingletonMetaClass):
             output.run_metadata = {}
         output.run_metadata.update(**metadata)
 
-    def _set_artifact_config(
+    def add_output_tags(
         self,
-        artifact_config: "ArtifactConfig",
+        tags: List[str],
         output_name: Optional[str] = None,
     ) -> None:
-        """Adds artifact config for a given step output.
+        """Adds tags for a given step output.
 
         Args:
-            artifact_config: The artifact config of the output to set.
-            output_name: Optional name of the output for which to set the
-                output signature. If no name is given and the step only has a single
-                output, the metadata of this output will be added. If the
+            tags: The tags to add.
+            output_name: Optional name of the output for which to add the
+                tags. If no name is given and the step only has a single
+                output, the tags of this output will be added. If the
                 step has multiple outputs, an exception will be raised.
-
-        Raises:
-            EntityExistsError: If the output already has an output signature.
         """
         output = self._get_output(output_name)
-
-        if output.artifact_config is None:
-            output.artifact_config = artifact_config
-        else:
-            raise EntityExistsError(
-                f"Output with name '{output_name}' already has artifact config."
-            )
+        if not output.tags:
+            output.tags = []
+        output.tags += tags
 
 
 class StepContextOutput:
@@ -373,6 +392,7 @@ class StepContextOutput:
     artifact_uri: str
     run_metadata: Optional[Dict[str, "MetadataType"]] = None
     artifact_config: Optional["ArtifactConfig"]
+    tags: Optional[List[str]] = None
 
     def __init__(
         self,