zenml-nightly 0.66.0.dev20240927__py3-none-any.whl → 0.68.1.dev20241102__py3-none-any.whl

zenml/steps/base_step.py

@@ -26,12 +26,10 @@ from typing import (
     Mapping,
     Optional,
     Sequence,
-    Set,
     Tuple,
     Type,
     TypeVar,
     Union,
-    cast,
 )
 
 from pydantic import BaseModel, ConfigDict, ValidationError
@@ -40,18 +38,16 @@ from zenml.client_lazy_loader import ClientLazyLoader
 from zenml.config.retry_config import StepRetryConfig
 from zenml.config.source import Source
 from zenml.constants import (
+    CODE_HASH_PARAMETER_NAME,
     ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK,
-    STEP_SOURCE_PARAMETER_NAME,
     handle_bool_env_var,
 )
-from zenml.exceptions import MissingStepParameterError, StepInterfaceError
+from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.materializers.base_materializer import BaseMaterializer
 from zenml.materializers.materializer_registry import materializer_registry
-from zenml.steps.base_parameters import BaseParameters
 from zenml.steps.entrypoint_function_utils import (
     StepArtifact,
-    get_step_entrypoint_signature,
     validate_entrypoint_function,
 )
 from zenml.steps.utils import (
@@ -82,9 +78,9 @@ if TYPE_CHECKING:
     )
     from zenml.model.lazy_load import ModelVersionDataLazyLoader
     from zenml.model.model import Model
+    from zenml.models import ArtifactVersionResponse
     from zenml.types import HookSpecification
 
-    ParametersOrDict = Union["BaseParameters", Dict[str, Any]]
     MaterializerClassOrSource = Union[str, Source, Type["BaseMaterializer"]]
     OutputMaterializersSpecification = Union[
         "MaterializerClassOrSource",
@@ -95,43 +91,14 @@ if TYPE_CHECKING:
 
 logger = get_logger(__name__)
 
-
-class BaseStepMeta(type):
-    """Metaclass for `BaseStep`.
-
-    Makes sure that the entrypoint function has valid parameters and type
-    annotations.
-    """
-
-    def __new__(
-        mcs, name: str, bases: Tuple[Type[Any], ...], dct: Dict[str, Any]
-    ) -> "BaseStepMeta":
-        """Set up a new class with a qualified spec.
-
-        Args:
-            name: The name of the class.
-            bases: The base classes of the class.
-            dct: The attributes of the class.
-
-        Returns:
-            The new class.
-        """
-        cls = cast(Type["BaseStep"], super().__new__(mcs, name, bases, dct))
-        if name not in {"BaseStep", "_DecoratedStep"}:
-            validate_entrypoint_function(cls.entrypoint)
-
-        return cls
-
-
 T = TypeVar("T", bound="BaseStep")
 
 
-class BaseStep(metaclass=BaseStepMeta):
+class BaseStep:
     """Abstract base class for all ZenML steps."""
 
     def __init__(
         self,
-        *args: Any,
         name: Optional[str] = None,
         enable_cache: Optional[bool] = None,
         enable_artifact_metadata: Optional[bool] = None,
@@ -139,7 +106,7 @@ class BaseStep(metaclass=BaseStepMeta):
         enable_step_logs: Optional[bool] = None,
         experiment_tracker: Optional[str] = None,
         step_operator: Optional[str] = None,
-        parameters: Optional["ParametersOrDict"] = None,
+        parameters: Optional[Dict[str, Any]] = None,
         output_materializers: Optional[
             "OutputMaterializersSpecification"
         ] = None,
@@ -149,12 +116,10 @@ class BaseStep(metaclass=BaseStepMeta):
         on_success: Optional["HookSpecification"] = None,
         model: Optional["Model"] = None,
         retry: Optional[StepRetryConfig] = None,
-        **kwargs: Any,
     ) -> None:
         """Initializes a step.
 
         Args:
-            *args: Positional arguments passed to the step.
             name: The name of the step.
             enable_cache: If caching should be enabled for this step.
             enable_artifact_metadata: If artifact metadata should be enabled
@@ -179,30 +144,15 @@ class BaseStep(metaclass=BaseStepMeta):
                 function (e.g. `module.my_function`).
             model: configuration of the model version in the Model Control Plane.
             retry: Configuration for retrying the step in case of failure.
-            **kwargs: Keyword arguments passed to the step.
         """
         from zenml.config.step_configurations import PartialStepConfiguration
 
-        self._upstream_steps: Set["BaseStep"] = set()
         self.entrypoint_definition = validate_entrypoint_function(
             self.entrypoint, reserved_arguments=["after", "id"]
         )
 
         name = name or self.__class__.__name__
 
-        requires_context = self.entrypoint_definition.context is not None
-        if enable_cache is None:
-            if requires_context:
-                # Using the StepContext inside a step provides access to
-                # external resources which might influence the step execution.
-                # We therefore disable caching unless it is explicitly enabled
-                enable_cache = False
-                logger.debug(
-                    "Step `%s`: Step context required and caching not "
-                    "explicitly enabled.",
-                    name,
-                )
-
         logger.debug(
             "Step `%s`: Caching %s.",
             name,
@@ -254,7 +204,6 @@ class BaseStep(metaclass=BaseStepMeta):
             model=model,
             retry=retry,
         )
-        self._verify_and_apply_init_params(*args, **kwargs)
 
         notebook_utils.try_to_save_notebook_cell_code(self.source_object)
 
@@ -300,45 +249,6 @@ class BaseStep(metaclass=BaseStepMeta):
         """
         return source_utils.resolve(self.__class__)
 
-    @property
-    def upstream_steps(self) -> Set["BaseStep"]:
-        """Names of the upstream steps of this step.
-
-        This property will only contain the full set of upstream steps once
-        it's parent pipeline `connect(...)` method was called.
-
-        Returns:
-            Set of upstream step names.
-        """
-        return self._upstream_steps
-
-    def after(self, step: "BaseStep") -> None:
-        """Adds an upstream step to this step.
-
-        Calling this method makes sure this step only starts running once the
-        given step has successfully finished executing.
-
-        **Note**: This can only be called inside the pipeline connect function
-        which is decorated with the `@pipeline` decorator. Any calls outside
-        this function will be ignored.
-
-        Example:
-        The following pipeline will run its steps sequentially in the following
-        order: step_2 -> step_1 -> step_3
-
-        ```python
-        @pipeline
-        def example_pipeline(step_1, step_2, step_3):
-            step_1.after(step_2)
-            step_3(step_1(), step_2())
-        ```
-
-        Args:
-            step: A step which should finish executing before this step is
-                started.
-        """
-        self._upstream_steps.add(step)
-
     @property
     def source_object(self) -> Any:
         """The source object of this step.
@@ -374,7 +284,7 @@ class BaseStep(metaclass=BaseStepMeta):
             A dictionary containing the caching parameters
         """
         parameters = {
-            STEP_SOURCE_PARAMETER_NAME: source_code_utils.get_hashed_source_code(
+            CODE_HASH_PARAMETER_NAME: source_code_utils.get_hashed_source_code(
                 self.source_object
             )
         }
@@ -394,70 +304,11 @@ class BaseStep(metaclass=BaseStepMeta):
 
         return parameters
 
-    def _verify_and_apply_init_params(self, *args: Any, **kwargs: Any) -> None:
-        """Verifies the initialization args and kwargs of this step.
-
-        This method makes sure that there is only one parameters object passed
-        at initialization and that it was passed using the correct name and
-        type specified in the step declaration.
-
-        Args:
-            *args: The args passed to the init method of this step.
-            **kwargs: The kwargs passed to the init method of this step.
-
-        Raises:
-            StepInterfaceError: If there are too many arguments or arguments
-                with a wrong name/type.
-        """
-        maximum_arg_count = (
-            1 if self.entrypoint_definition.legacy_params else 0
-        )
-        arg_count = len(args) + len(kwargs)
-        if arg_count > maximum_arg_count:
-            raise StepInterfaceError(
-                f"Too many arguments ({arg_count}, expected: "
-                f"{maximum_arg_count}) passed when creating a "
-                f"'{self.name}' step."
-            )
-
-        if self.entrypoint_definition.legacy_params:
-            if args:
-                config = args[0]
-            elif kwargs:
-                key, config = kwargs.popitem()
-
-                if key != self.entrypoint_definition.legacy_params.name:
-                    raise StepInterfaceError(
-                        f"Unknown keyword argument '{key}' when creating a "
-                        f"'{self.name}' step, only expected a single "
-                        "argument with key "
-                        f"'{self.entrypoint_definition.legacy_params.name}'."
-                    )
-            else:
-                # This step requires configuration parameters but no parameters
-                # object was passed as an argument. The parameters might be
-                # set via default values in the parameters class or in a
-                # configuration file, so we continue for now and verify
-                # that all parameters are set before running the step
-                return
-
-            if not isinstance(
-                config, self.entrypoint_definition.legacy_params.annotation
-            ):
-                raise StepInterfaceError(
-                    f"`{config}` object passed when creating a "
-                    f"'{self.name}' step is not a "
-                    f"`{self.entrypoint_definition.legacy_params.annotation.__name__} "
-                    "` instance."
-                )
-
-            self.configure(parameters=config)
-
     def _parse_call_args(
         self, *args: Any, **kwargs: Any
     ) -> Tuple[
         Dict[str, "StepArtifact"],
-        Dict[str, "ExternalArtifact"],
+        Dict[str, Union["ExternalArtifact", "ArtifactVersionResponse"]],
         Dict[str, "ModelVersionDataLazyLoader"],
         Dict[str, "ClientLazyLoader"],
         Dict[str, Any],
@@ -478,11 +329,12 @@ class BaseStep(metaclass=BaseStepMeta):
         from zenml.artifacts.external_artifact import ExternalArtifact
         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 = get_step_entrypoint_signature(step=self)
+        signature = inspect.signature(self.entrypoint, follow_wrapped=True)
 
         try:
             bound_args = signature.bind_partial(*args, **kwargs)
@@ -492,7 +344,9 @@ class BaseStep(metaclass=BaseStepMeta):
             ) from e
 
         artifacts = {}
-        external_artifacts = {}
+        external_artifacts: Dict[
+            str, Union["ExternalArtifact", "ArtifactVersionResponse"]
+        ] = {}
         model_artifacts_or_metadata = {}
         client_lazy_loaders = {}
         parameters = {}
@@ -522,14 +376,18 @@ class BaseStep(metaclass=BaseStepMeta):
                     )
             elif isinstance(value, LazyArtifactVersionResponse):
                 model_artifacts_or_metadata[key] = ModelVersionDataLazyLoader(
-                    model=value.lazy_load_model,
+                    model_name=value.lazy_load_model_name,
+                    model_version=value.lazy_load_model_version,
                     artifact_name=value.lazy_load_name,
                     artifact_version=value.lazy_load_version,
                     metadata_name=None,
                 )
+            elif isinstance(value, ArtifactVersionResponse):
+                external_artifacts[key] = value
             elif isinstance(value, LazyRunMetadataResponse):
                 model_artifacts_or_metadata[key] = ModelVersionDataLazyLoader(
-                    model=value.lazy_load_model,
+                    model_name=value.lazy_load_model_name,
+                    model_version=value.lazy_load_model_version,
                     artifact_name=value.lazy_load_artifact_name,
                     artifact_version=value.lazy_load_artifact_version,
                     metadata_name=value.lazy_load_metadata_name,
@@ -589,7 +447,7 @@ class BaseStep(metaclass=BaseStepMeta):
         Returns:
             The outputs of the entrypoint function call.
         """
-        from zenml.new.pipelines.pipeline import Pipeline
+        from zenml.pipelines.pipeline_definition import Pipeline
 
         if not Pipeline.ACTIVE_PIPELINE:
             from zenml import constants, get_step_context
@@ -720,14 +578,13 @@ class BaseStep(metaclass=BaseStepMeta):
 
     def configure(
         self: T,
-        name: Optional[str] = None,
         enable_cache: Optional[bool] = None,
         enable_artifact_metadata: Optional[bool] = None,
         enable_artifact_visualization: Optional[bool] = None,
         enable_step_logs: Optional[bool] = None,
         experiment_tracker: Optional[str] = None,
         step_operator: Optional[str] = None,
-        parameters: Optional["ParametersOrDict"] = None,
+        parameters: Optional[Dict[str, Any]] = None,
         output_materializers: Optional[
             "OutputMaterializersSpecification"
         ] = None,
@@ -752,7 +609,6 @@ class BaseStep(metaclass=BaseStepMeta):
             step.configuration.extra # {"key2": 2}
 
         Args:
-            name: DEPRECATED: The name of the step.
             enable_cache: If caching should be enabled for this step.
             enable_artifact_metadata: If artifact metadata should be enabled
                 for this step.
@@ -788,9 +644,6 @@ class BaseStep(metaclass=BaseStepMeta):
         from zenml.config.step_configurations import StepConfigurationUpdate
         from zenml.hooks.hook_validators import resolve_and_validate_hook
 
-        if name:
-            logger.warning("Configuring the name of a step is deprecated.")
-
         def _resolve_if_necessary(
             value: Union[str, Source, Type[Any]],
         ) -> Source:
@@ -832,9 +685,6 @@ class BaseStep(metaclass=BaseStepMeta):
             # string of on_success hook function to be used for this step
             success_hook_source = resolve_and_validate_hook(on_success)
 
-        if isinstance(parameters, BaseParameters):
-            parameters = parameters.model_dump()
-
         values = dict_utils.remove_none_values(
             {
                 "enable_cache": enable_cache,
@@ -865,7 +715,7 @@ class BaseStep(metaclass=BaseStepMeta):
         enable_step_logs: Optional[bool] = None,
         experiment_tracker: Optional[str] = None,
         step_operator: Optional[str] = None,
-        parameters: Optional["ParametersOrDict"] = None,
+        parameters: Optional[Dict[str, Any]] = None,
         output_materializers: Optional[
             "OutputMaterializersSpecification"
         ] = None,
@@ -1006,8 +856,7 @@ class BaseStep(metaclass=BaseStepMeta):
                     conflicting_parameters[key] = (value, runtime_value)
             if key in self.entrypoint_definition.inputs:
                 self.entrypoint_definition.validate_input(key=key, value=value)
-
-            elif not self.entrypoint_definition.legacy_params:
+            else:
                 raise StepInterfaceError(
                     f"Unable to find parameter '{key}' in step function "
                     "signature."
@@ -1143,9 +992,7 @@ To avoid this consider setting step parameters only in one place (config or code
             StepConfigurationUpdate,
         )
 
-        outputs: Dict[str, Dict[str, Union[Source, Tuple[Source, ...]]]] = (
-            defaultdict(dict)
-        )
+        outputs: Dict[str, Dict[str, Any]] = defaultdict(dict)
 
         for (
             output_name,
@@ -1154,6 +1001,8 @@ To avoid this consider setting step parameters only in one place (config or code
             output = self._configuration.outputs.get(
                 output_name, PartialArtifactConfiguration()
             )
+            if artifact_config := output_annotation.artifact_config:
+                outputs[output_name]["artifact_config"] = artifact_config
 
             if output.materializer_source:
                 # The materializer source was configured by the user. We
@@ -1264,95 +1113,4 @@ To avoid this consider setting step parameters only in one place (config or code
             else:
                 params[key] = value
 
-        if self.entrypoint_definition.legacy_params:
-            legacy_params = self._finalize_legacy_parameters()
-            params[self.entrypoint_definition.legacy_params.name] = (
-                legacy_params
-            )
-
         return params
-
-    def _finalize_legacy_parameters(self) -> Dict[str, Any]:
-        """Verifies and prepares the config parameters for running this step.
-
-        When the step requires config parameters, this method:
-            - checks if config parameters were set via a config object or file
-            - tries to set missing config parameters from default values of the
-              config class
-
-        Returns:
-            Values for the previously unconfigured function parameters.
-
-        Raises:
-            MissingStepParameterError: If no value could be found for one or
-                more config parameters.
-            StepInterfaceError: If the parameter class validation failed.
-        """
-        if not self.entrypoint_definition.legacy_params:
-            return {}
-
-        logger.warning(
-            "The `BaseParameters` class to define step parameters is "
-            "deprecated. Check out our docs "
-            "https://docs.zenml.io/how-to/use-configuration-files/how-to-use-config "
-            "for information on how to parameterize your steps. As a quick "
-            "fix to get rid of this warning, make sure your parameter class "
-            "inherits from `pydantic.BaseModel` instead of the "
-            "`BaseParameters` class."
-        )
-
-        # parameters for the `BaseParameters` class specified in the "new" way
-        # by specifying a dict of parameters for the corresponding key
-        params_defined_in_new_way = (
-            self.configuration.parameters.get(
-                self.entrypoint_definition.legacy_params.name
-            )
-            or {}
-        )
-
-        values = {}
-        missing_keys = []
-        for (
-            name,
-            field,
-        ) in self.entrypoint_definition.legacy_params.annotation.model_fields.items():
-            if name in self.configuration.parameters:
-                # a value for this parameter has been set already
-                values[name] = self.configuration.parameters[name]
-            elif name in params_defined_in_new_way:
-                # a value for this parameter has been set in the "new" way
-                # already
-                values[name] = params_defined_in_new_way[name]
-            elif field.is_required():
-                # this field has no default value set and therefore needs
-                # to be passed via an initialized config object
-                missing_keys.append(name)
-            else:
-                # use default value from the pydantic config class
-                values[name] = field.default
-
-        if missing_keys:
-            raise MissingStepParameterError(
-                self.name,
-                missing_keys,
-                self.entrypoint_definition.legacy_params.annotation,
-            )
-
-        if (
-            getattr(
-                self.entrypoint_definition.legacy_params.annotation.model_config,
-                "extra",
-                None,
-            )
-            == "allow"
-        ):
-            # Add all parameters for the config class for backwards
-            # compatibility if the config class allows extra attributes
-            values.update(self.configuration.parameters)
-
-        try:
-            self.entrypoint_definition.legacy_params.annotation(**values)
-        except ValidationError:
-            raise StepInterfaceError("Failed to validate function parameters.")
-
-        return values

zenml/steps/entrypoint_function_utils.py

@@ -21,7 +21,6 @@ from typing import (
     Dict,
     NamedTuple,
     NoReturn,
-    Optional,
     Sequence,
     Type,
     Union,
@@ -42,45 +41,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."""
 
@@ -143,15 +110,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.
@@ -260,36 +222,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(
@@ -309,29 +255,7 @@ 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
@@ -340,6 +264,4 @@ def validate_entrypoint_function(
     return EntrypointFunctionDefinition(
         inputs=inputs,
         outputs=outputs,
-        context=context,
-        legacy_params=legacy_params,
     )