Steps

Initializer for ZenML steps.

A step is a single piece or stage of a ZenML pipeline. Think of each step as being one of the nodes of a Directed Acyclic Graph (or DAG). Steps are responsible for one aspect of processing or interacting with the data / artifacts in the pipeline.

Conceptually, a Step is a discrete and independent part of a pipeline that is responsible for one particular aspect of data manipulation inside a ZenML pipeline.

Steps can be subclassed from the BaseStep class, or used via our @step decorator.

BaseStep

Abstract base class for all ZenML steps.

Source code in src/zenml/steps/base_step.py

class BaseStep:
    """Abstract base class for all ZenML steps."""

    def __init__(
        self,
        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[Dict[str, Any]] = None,
        output_materializers: Optional[
            "OutputMaterializersSpecification"
        ] = None,
        settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
        extra: Optional[Dict[str, Any]] = None,
        on_failure: Optional["HookSpecification"] = None,
        on_success: Optional["HookSpecification"] = None,
        model: Optional["Model"] = None,
        retry: Optional[StepRetryConfig] = None,
        substitutions: Optional[Dict[str, str]] = None,
    ) -> None:
        """Initializes a step.

        Args:
            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
                for this step.
            enable_artifact_visualization: If artifact visualization should be
                enabled for this step.
            enable_step_logs: Enable step logs for this step.
            experiment_tracker: The experiment tracker to use for this step.
            step_operator: The step operator to use for this step.
            parameters: Function parameters for this step
            output_materializers: Output materializers for this step. If
                given as a dict, the keys must be a subset of the output names
                of this step. If a single value (type or string) is given, the
                materializer will be used for all outputs.
            settings: settings for this step.
            extra: Extra configurations for this step.
            on_failure: Callback function in event of failure of the step. Can
                be a function with a single argument of type `BaseException`, or
                a source path to such a function (e.g. `module.my_function`).
            on_success: Callback function in event of success of the step. Can
                be a function with no arguments, or a source path to such a
                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.
            substitutions: Extra placeholders to use in the name template.
        """
        from zenml.config.step_configurations import PartialStepConfiguration

        self.entrypoint_definition = validate_entrypoint_function(
            self.entrypoint,
            reserved_arguments=["after", "id"],
        )

        name = name or self.__class__.__name__

        logger.debug(
            "Step `%s`: Caching %s.",
            name,
            "enabled" if enable_cache is not False else "disabled",
        )
        logger.debug(
            "Step `%s`: Artifact metadata %s.",
            name,
            "enabled" if enable_artifact_metadata is not False else "disabled",
        )
        logger.debug(
            "Step `%s`: Artifact visualization %s.",
            name,
            "enabled"
            if enable_artifact_visualization is not False
            else "disabled",
        )
        logger.debug(
            "Step `%s`: logs %s.",
            name,
            "enabled" if enable_step_logs is not False else "disabled",
        )
        if model is not None:
            logger.debug(
                "Step `%s`: Is in Model context %s.",
                name,
                {
                    "model": model.name,
                    "version": model.version,
                },
            )

        self._configuration = PartialStepConfiguration(
            name=name,
            enable_cache=enable_cache,
            enable_artifact_metadata=enable_artifact_metadata,
            enable_artifact_visualization=enable_artifact_visualization,
            enable_step_logs=enable_step_logs,
        )
        self.configure(
            experiment_tracker=experiment_tracker,
            step_operator=step_operator,
            output_materializers=output_materializers,
            parameters=parameters,
            settings=settings,
            extra=extra,
            on_failure=on_failure,
            on_success=on_success,
            model=model,
            retry=retry,
            substitutions=substitutions,
        )

        notebook_utils.try_to_save_notebook_cell_code(self.source_object)

    @abstractmethod
    def entrypoint(self, *args: Any, **kwargs: Any) -> Any:
        """Abstract method for core step logic.

        Args:
            *args: Positional arguments passed to the step.
            **kwargs: Keyword arguments passed to the step.

        Returns:
            The output of the step.
        """

    @classmethod
    def load_from_source(cls, source: Union[Source, str]) -> "BaseStep":
        """Loads a step from source.

        Args:
            source: The path to the step source.

        Returns:
            The loaded step.

        Raises:
            ValueError: If the source is not a valid step source.
        """
        obj = source_utils.load(source)

        if isinstance(obj, BaseStep):
            return obj
        elif isinstance(obj, type) and issubclass(obj, BaseStep):
            return obj()
        else:
            raise ValueError("Invalid step source.")

    def resolve(self) -> Source:
        """Resolves the step.

        Returns:
            The step source.
        """
        return source_utils.resolve(self.__class__)

    @property
    def source_object(self) -> Any:
        """The source object of this step.

        Returns:
            The source object of this step.
        """
        return self.__class__

    @property
    def source_code(self) -> str:
        """The source code of this step.

        Returns:
            The source code of this step.
        """
        return inspect.getsource(self.source_object)

    @property
    def docstring(self) -> Optional[str]:
        """The docstring of this step.

        Returns:
            The docstring of this step.
        """
        return self.__doc__

    @property
    def caching_parameters(self) -> Dict[str, Any]:
        """Caching parameters for this step.

        Returns:
            A dictionary containing the caching parameters
        """
        parameters = {
            CODE_HASH_PARAMETER_NAME: source_code_utils.get_hashed_source_code(
                self.source_object
            )
        }
        for name, output in self.configuration.outputs.items():
            if output.materializer_source:
                key = f"{name}_materializer_source"
                hash_ = hashlib.md5()  # nosec

                for source in output.materializer_source:
                    materializer_class = source_utils.load(source)
                    code_hash = source_code_utils.get_hashed_source_code(
                        materializer_class
                    )
                    hash_.update(code_hash.encode())

                parameters[key] = hash_.hexdigest()

        return parameters

    def _parse_call_args(
        self, *args: Any, **kwargs: Any
    ) -> Tuple[
        Dict[str, "StepArtifact"],
        Dict[str, Union["ExternalArtifact", "ArtifactVersionResponse"]],
        Dict[str, "ModelVersionDataLazyLoader"],
        Dict[str, "ClientLazyLoader"],
        Dict[str, Any],
        Dict[str, Any],
    ]:
        """Parses the call args for the step entrypoint.

        Args:
            *args: Entrypoint function arguments.
            **kwargs: Entrypoint function keyword arguments.

        Raises:
            StepInterfaceError: If invalid function arguments were passed.

        Returns:
            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,
        )

        signature = inspect.signature(self.entrypoint, follow_wrapped=True)

        try:
            bound_args = signature.bind_partial(*args, **kwargs)
        except TypeError as e:
            raise StepInterfaceError(
                f"Wrong arguments when calling step '{self.name}': {e}"
            ) from e

        artifacts = {}
        external_artifacts: Dict[
            str, Union["ExternalArtifact", "ArtifactVersionResponse"]
        ] = {}
        model_artifacts_or_metadata = {}
        client_lazy_loaders = {}
        parameters = {}
        default_parameters = {}

        for key, value in bound_args.arguments.items():
            self.entrypoint_definition.validate_input(key=key, value=value)

            if isinstance(value, StepArtifact):
                artifacts[key] = value
                if key in self.configuration.parameters:
                    logger.warning(
                        "Got duplicate value for step input %s, using value "
                        "provided as artifact.",
                        key,
                    )
            elif isinstance(value, ExternalArtifact):
                external_artifacts[key] = value
                if not value.id:
                    # If the external artifact references a fixed artifact by
                    # ID, caching behaves as expected.
                    logger.warning(
                        "Using an external artifact as step input currently "
                        "invalidates caching for the step and all downstream "
                        "steps. Future releases will introduce hashing of "
                        "artifacts which will improve this behavior."
                    )
            elif isinstance(value, LazyArtifactVersionResponse):
                model_artifacts_or_metadata[key] = ModelVersionDataLazyLoader(
                    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_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,
                )
            elif isinstance(value, ClientLazyLoader):
                client_lazy_loaders[key] = value
            else:
                parameters[key] = value

        # Above we iterated over the provided arguments which should overwrite
        # any parameters previously defined on the step instance. Now we apply
        # the default values on the entrypoint function and add those as
        # parameters for any argument that has no value yet. If we were to do
        # that in the above loop, we would overwrite previously configured
        # parameters with the default values.
        bound_args.apply_defaults()
        for key, value in bound_args.arguments.items():
            self.entrypoint_definition.validate_input(key=key, value=value)
            if (
                key not in artifacts
                and key not in external_artifacts
                and key not in model_artifacts_or_metadata
                and key not in self.configuration.parameters
                and key not in client_lazy_loaders
            ):
                default_parameters[key] = value

        return (
            artifacts,
            external_artifacts,
            model_artifacts_or_metadata,
            client_lazy_loaders,
            parameters,
            default_parameters,
        )

    def __call__(
        self,
        *args: Any,
        id: Optional[str] = None,
        after: Union[str, Sequence[str], None] = None,
        **kwargs: Any,
    ) -> Any:
        """Handle a call of the step.

        This method does one of two things:
        * If there is an active pipeline context, it adds an invocation of the
          step instance to the pipeline.
        * If no pipeline is active, it calls the step entrypoint function.

        Args:
            *args: Entrypoint function arguments.
            id: Invocation ID to use.
            after: Upstream steps for the invocation.
            **kwargs: Entrypoint function keyword arguments.

        Returns:
            The outputs of the entrypoint function call.
        """
        from zenml.pipelines.pipeline_definition import Pipeline

        if not Pipeline.ACTIVE_PIPELINE:
            from zenml import constants, get_step_context

            # If the environment variable was set to explicitly not run on the
            # stack, we do that.
            run_without_stack = handle_bool_env_var(
                ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK, default=False
            )
            if run_without_stack:
                return self.call_entrypoint(*args, **kwargs)

            try:
                get_step_context()
            except RuntimeError:
                pass
            else:
                # We're currently inside the execution of a different step
                # -> We don't want to launch another single step pipeline here,
                # but instead just call the step function
                return self.call_entrypoint(*args, **kwargs)

            if constants.SHOULD_PREVENT_PIPELINE_EXECUTION:
                logger.info(
                    "Preventing execution of step '%s'.",
                    self.name,
                )
                return

            return run_as_single_step_pipeline(self, *args, **kwargs)

        (
            input_artifacts,
            external_artifacts,
            model_artifacts_or_metadata,
            client_lazy_loaders,
            parameters,
            default_parameters,
        ) = self._parse_call_args(*args, **kwargs)

        upstream_steps = {
            artifact.invocation_id for artifact in input_artifacts.values()
        }
        if isinstance(after, str):
            upstream_steps.add(after)
        elif isinstance(after, Sequence):
            upstream_steps = upstream_steps.union(after)

        invocation_id = Pipeline.ACTIVE_PIPELINE.add_step_invocation(
            step=self,
            input_artifacts=input_artifacts,
            external_artifacts=external_artifacts,
            model_artifacts_or_metadata=model_artifacts_or_metadata,
            client_lazy_loaders=client_lazy_loaders,
            parameters=parameters,
            default_parameters=default_parameters,
            upstream_steps=upstream_steps,
            custom_id=id,
            allow_id_suffix=not id,
        )

        outputs = []
        for key, annotation in self.entrypoint_definition.outputs.items():
            output = StepArtifact(
                invocation_id=invocation_id,
                output_name=key,
                annotation=annotation,
                pipeline=Pipeline.ACTIVE_PIPELINE,
            )
            outputs.append(output)
        return outputs[0] if len(outputs) == 1 else outputs

    def call_entrypoint(self, *args: Any, **kwargs: Any) -> Any:
        """Calls the entrypoint function of the step.

        Args:
            *args: Entrypoint function arguments.
            **kwargs: Entrypoint function keyword arguments.

        Returns:
            The return value of the entrypoint function.

        Raises:
            StepInterfaceError: If the arguments to the entrypoint function are
                invalid.
        """
        try:
            validated_args = pydantic_utils.validate_function_args(
                self.entrypoint,
                ConfigDict(arbitrary_types_allowed=True),
                *args,
                **kwargs,
            )
        except ValidationError as e:
            raise StepInterfaceError(
                "Invalid step function entrypoint arguments. Check out the "
                "pydantic error above for more details."
            ) from e

        return self.entrypoint(**validated_args)

    @property
    def name(self) -> str:
        """The name of the step.

        Returns:
            The name of the step.
        """
        return self.configuration.name

    @property
    def enable_cache(self) -> Optional[bool]:
        """If caching is enabled for the step.

        Returns:
            If caching is enabled for the step.
        """
        return self.configuration.enable_cache

    @property
    def configuration(self) -> "PartialStepConfiguration":
        """The configuration of the step.

        Returns:
            The configuration of the step.
        """
        return self._configuration

    def configure(
        self: T,
        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[Dict[str, Any]] = None,
        output_materializers: Optional[
            "OutputMaterializersSpecification"
        ] = None,
        settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
        extra: Optional[Dict[str, Any]] = None,
        on_failure: Optional["HookSpecification"] = None,
        on_success: Optional["HookSpecification"] = None,
        model: Optional["Model"] = None,
        merge: bool = True,
        retry: Optional[StepRetryConfig] = None,
        substitutions: Optional[Dict[str, str]] = None,
    ) -> T:
        """Configures the step.

        Configuration merging example:
        * `merge==True`:
            step.configure(extra={"key1": 1})
            step.configure(extra={"key2": 2}, merge=True)
            step.configuration.extra # {"key1": 1, "key2": 2}
        * `merge==False`:
            step.configure(extra={"key1": 1})
            step.configure(extra={"key2": 2}, merge=False)
            step.configuration.extra # {"key2": 2}

        Args:
            enable_cache: If caching should be enabled for this step.
            enable_artifact_metadata: If artifact metadata should be enabled
                for this step.
            enable_artifact_visualization: If artifact visualization should be
                enabled for this step.
            enable_step_logs: If step logs should be enabled for this step.
            experiment_tracker: The experiment tracker to use for this step.
            step_operator: The step operator to use for this step.
            parameters: Function parameters for this step
            output_materializers: Output materializers for this step. If
                given as a dict, the keys must be a subset of the output names
                of this step. If a single value (type or string) is given, the
                materializer will be used for all outputs.
            settings: settings for this step.
            extra: Extra configurations for this step.
            on_failure: Callback function in event of failure of the step. Can
                be a function with a single argument of type `BaseException`, or
                a source path to such a function (e.g. `module.my_function`).
            on_success: Callback function in event of success of the step. Can
                be a function with no arguments, or a source path to such a
                function (e.g. `module.my_function`).
            model: configuration of the model version in the Model Control Plane.
            merge: If `True`, will merge the given dictionary configurations
                like `parameters` and `settings` with existing
                configurations. If `False` the given configurations will
                overwrite all existing ones. See the general description of this
                method for an example.
            retry: Configuration for retrying the step in case of failure.
            substitutions: Extra placeholders to use in the name template.

        Returns:
            The step instance that this method was called on.
        """
        from zenml.config.step_configurations import StepConfigurationUpdate
        from zenml.hooks.hook_validators import resolve_and_validate_hook

        def _resolve_if_necessary(
            value: Union[str, Source, Type[Any]],
        ) -> Source:
            if isinstance(value, str):
                return Source.from_import_path(value)
            elif isinstance(value, Source):
                return value
            else:
                return source_utils.resolve(value)

        def _convert_to_tuple(value: Any) -> Tuple[Source, ...]:
            if isinstance(value, str) or not isinstance(value, Sequence):
                return (_resolve_if_necessary(value),)
            else:
                return tuple(_resolve_if_necessary(v) for v in value)

        outputs: Dict[str, Dict[str, Tuple[Source, ...]]] = defaultdict(dict)
        allowed_output_names = set(self.entrypoint_definition.outputs)

        if output_materializers:
            if not isinstance(output_materializers, Mapping):
                sources = _convert_to_tuple(output_materializers)
                output_materializers = {
                    output_name: sources
                    for output_name in allowed_output_names
                }

            for output_name, materializer in output_materializers.items():
                sources = _convert_to_tuple(materializer)
                outputs[output_name]["materializer_source"] = sources

        failure_hook_source = None
        if on_failure:
            # string of on_failure hook function to be used for this step
            failure_hook_source = resolve_and_validate_hook(on_failure)

        success_hook_source = None
        if on_success:
            # string of on_success hook function to be used for this step
            success_hook_source = resolve_and_validate_hook(on_success)

        values = dict_utils.remove_none_values(
            {
                "enable_cache": enable_cache,
                "enable_artifact_metadata": enable_artifact_metadata,
                "enable_artifact_visualization": enable_artifact_visualization,
                "enable_step_logs": enable_step_logs,
                "experiment_tracker": experiment_tracker,
                "step_operator": step_operator,
                "parameters": parameters,
                "settings": settings,
                "outputs": outputs or None,
                "extra": extra,
                "failure_hook_source": failure_hook_source,
                "success_hook_source": success_hook_source,
                "model": model,
                "retry": retry,
                "substitutions": substitutions,
            }
        )
        config = StepConfigurationUpdate(**values)
        self._apply_configuration(config, merge=merge)
        return self

    def with_options(
        self,
        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[Dict[str, Any]] = None,
        output_materializers: Optional[
            "OutputMaterializersSpecification"
        ] = None,
        settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
        extra: Optional[Dict[str, Any]] = None,
        on_failure: Optional["HookSpecification"] = None,
        on_success: Optional["HookSpecification"] = None,
        model: Optional["Model"] = None,
        merge: bool = True,
        substitutions: Optional[Dict[str, str]] = None,
    ) -> "BaseStep":
        """Copies the step and applies the given configurations.

        Args:
            enable_cache: If caching should be enabled for this step.
            enable_artifact_metadata: If artifact metadata should be enabled
                for this step.
            enable_artifact_visualization: If artifact visualization should be
                enabled for this step.
            enable_step_logs: If step logs should be enabled for this step.
            experiment_tracker: The experiment tracker to use for this step.
            step_operator: The step operator to use for this step.
            parameters: Function parameters for this step
            output_materializers: Output materializers for this step. If
                given as a dict, the keys must be a subset of the output names
                of this step. If a single value (type or string) is given, the
                materializer will be used for all outputs.
            settings: settings for this step.
            extra: Extra configurations for this step.
            on_failure: Callback function in event of failure of the step. Can
                be a function with a single argument of type `BaseException`, or
                a source path to such a function (e.g. `module.my_function`).
            on_success: Callback function in event of success of the step. Can
                be a function with no arguments, or a source path to such a
                function (e.g. `module.my_function`).
            model: configuration of the model version in the Model Control Plane.
            merge: If `True`, will merge the given dictionary configurations
                like `parameters` and `settings` with existing
                configurations. If `False` the given configurations will
                overwrite all existing ones. See the general description of this
                method for an example.
            substitutions: Extra placeholders for the step name.

        Returns:
            The copied step instance.
        """
        step_copy = self.copy()
        step_copy.configure(
            enable_cache=enable_cache,
            enable_artifact_metadata=enable_artifact_metadata,
            enable_artifact_visualization=enable_artifact_visualization,
            enable_step_logs=enable_step_logs,
            experiment_tracker=experiment_tracker,
            step_operator=step_operator,
            parameters=parameters,
            output_materializers=output_materializers,
            settings=settings,
            extra=extra,
            on_failure=on_failure,
            on_success=on_success,
            model=model,
            merge=merge,
            substitutions=substitutions,
        )
        return step_copy

    def copy(self) -> "BaseStep":
        """Copies the step.

        Returns:
            The step copy.
        """
        return copy.deepcopy(self)

    def _apply_configuration(
        self,
        config: "StepConfigurationUpdate",
        merge: bool = True,
        runtime_parameters: Dict[str, Any] = {},
    ) -> None:
        """Applies an update to the step configuration.

        Args:
            config: The configuration update.
            runtime_parameters: Dictionary of parameters passed to a step from runtime
            merge: Whether to merge the updates with the existing configuration
                or not. See the `BaseStep.configure(...)` method for a detailed
                explanation.
        """
        self._validate_configuration(config, runtime_parameters)

        self._configuration = pydantic_utils.update_model(
            self._configuration, update=config, recursive=merge
        )

        logger.debug("Updated step configuration:")
        logger.debug(self._configuration)

    def _validate_configuration(
        self,
        config: "StepConfigurationUpdate",
        runtime_parameters: Dict[str, Any],
    ) -> None:
        """Validates a configuration update.

        Args:
            config: The configuration update to validate.
            runtime_parameters: Dictionary of parameters passed to a step from runtime
        """
        settings_utils.validate_setting_keys(list(config.settings))
        self._validate_function_parameters(
            parameters=config.parameters, runtime_parameters=runtime_parameters
        )
        self._validate_outputs(outputs=config.outputs)

    def _validate_function_parameters(
        self,
        parameters: Dict[str, Any],
        runtime_parameters: Dict[str, Any],
    ) -> None:
        """Validates step function parameters.

        Args:
            parameters: The parameters to validate.
            runtime_parameters: Dictionary of parameters passed to a step from runtime

        Raises:
            StepInterfaceError: If the step requires no function parameters but
                parameters were configured.
            RuntimeError: If the step has parameters configured differently in
                configuration file and code.
        """
        if not parameters:
            return

        conflicting_parameters = {}
        for key, value in parameters.items():
            if key in runtime_parameters:
                runtime_value = runtime_parameters[key]
                if runtime_value != value:
                    conflicting_parameters[key] = (value, runtime_value)
            if key in self.entrypoint_definition.inputs:
                self.entrypoint_definition.validate_input(key=key, value=value)
            else:
                raise StepInterfaceError(
                    f"Unable to find parameter '{key}' in step function "
                    "signature."
                )
        if conflicting_parameters:
            is_plural = "s" if len(conflicting_parameters) > 1 else ""
            msg = f"Configured parameter{is_plural} for the step '{self.name}' conflict{'' if not is_plural else 's'} with parameter{is_plural} passed in runtime:\n"
            for key, values in conflicting_parameters.items():
                msg += (
                    f"`{key}`: config=`{values[0]}` | runtime=`{values[1]}`\n"
                )
            msg += """This happens, if you define values for step parameters in configuration file and pass same parameters from the code. Example:
```
# config.yaml

steps:
    step_name:
        parameters:
            param_name: value1


# pipeline.py

@pipeline
def pipeline_():
    step_name(param_name="other_value")
```
To avoid this consider setting step parameters only in one place (config or code).
"""
            raise RuntimeError(msg)

    def _validate_outputs(
        self, outputs: Mapping[str, "PartialArtifactConfiguration"]
    ) -> None:
        """Validates the step output configuration.

        Args:
            outputs: The configured step outputs.

        Raises:
            StepInterfaceError: If an output for a non-existent name is
                configured of an output artifact/materializer source does not
                resolve to the correct class.
        """
        allowed_output_names = set(self.entrypoint_definition.outputs)
        for output_name, output in outputs.items():
            if output_name not in allowed_output_names:
                raise StepInterfaceError(
                    f"Got unexpected materializers for non-existent "
                    f"output '{output_name}' in step '{self.name}'. "
                    f"Only materializers for the outputs "
                    f"{allowed_output_names} of this step can"
                    f" be registered."
                )

            if output.materializer_source:
                for source in output.materializer_source:
                    if not source_utils.validate_source_class(
                        source, expected_class=BaseMaterializer
                    ):
                        raise StepInterfaceError(
                            f"Materializer source `{source}` "
                            f"for output '{output_name}' of step '{self.name}' "
                            "does not resolve to a `BaseMaterializer` subclass."
                        )

    def _validate_inputs(
        self,
        input_artifacts: Dict[str, "StepArtifact"],
        external_artifacts: Dict[str, "ExternalArtifactConfiguration"],
        model_artifacts_or_metadata: Dict[str, "ModelVersionDataLazyLoader"],
        client_lazy_loaders: Dict[str, "ClientLazyLoader"],
    ) -> None:
        """Validates the step inputs.

        This method makes sure that all inputs are provided either as an
        artifact or parameter.

        Args:
            input_artifacts: The input artifacts.
            external_artifacts: The external input artifacts.
            model_artifacts_or_metadata: The model artifacts or metadata.
            client_lazy_loaders: The client lazy loaders.

        Raises:
            StepInterfaceError: If an entrypoint input is missing.
        """
        for key in self.entrypoint_definition.inputs.keys():
            if (
                key in input_artifacts
                or key in self.configuration.parameters
                or key in external_artifacts
                or key in model_artifacts_or_metadata
                or key in client_lazy_loaders
            ):
                continue
            raise StepInterfaceError(
                f"Missing entrypoint input '{key}' in step '{self.name}'."
            )

    def _finalize_configuration(
        self,
        input_artifacts: Dict[str, "StepArtifact"],
        external_artifacts: Dict[str, "ExternalArtifactConfiguration"],
        model_artifacts_or_metadata: Dict[str, "ModelVersionDataLazyLoader"],
        client_lazy_loaders: Dict[str, "ClientLazyLoader"],
    ) -> "StepConfiguration":
        """Finalizes the configuration after the step was called.

        Once the step was called, we know the outputs of previous steps
        and that no additional user configurations will be made. That means
        we can now collect the remaining artifact and materializer types
        as well as check for the completeness of the step function parameters.

        Args:
            input_artifacts: The input artifacts of this step.
            external_artifacts: The external artifacts of this step.
            model_artifacts_or_metadata: The model artifacts or metadata of
                this step.
            client_lazy_loaders: The client lazy loaders of this step.

        Raises:
            StepInterfaceError: If explicit materializers were specified for an
                output but they do not work for the data type(s) defined by
                the type annotation.

        Returns:
            The finalized step configuration.
        """
        from zenml.config.step_configurations import (
            PartialArtifactConfiguration,
            StepConfiguration,
            StepConfigurationUpdate,
        )

        outputs: Dict[str, Dict[str, Any]] = defaultdict(dict)

        for (
            output_name,
            output_annotation,
        ) in self.entrypoint_definition.outputs.items():
            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
                # validate that their configured materializer supports the
                # output type. If the output annotation is a Union, we check
                # that at least one of the specified materializers works with at
                # least one of the types in the Union. If that's not the case,
                # it would be a guaranteed failure at runtime and we fail early
                # here.
                if output_annotation.resolved_annotation is Any:
                    continue

                materializer_classes: List[Type["BaseMaterializer"]] = [
                    source_utils.load(materializer_source)
                    for materializer_source in output.materializer_source
                ]

                for data_type in output_annotation.get_output_types():
                    try:
                        materializer_utils.select_materializer(
                            data_type=data_type,
                            materializer_classes=materializer_classes,
                        )
                        break
                    except RuntimeError:
                        pass
                else:
                    materializer_strings = [
                        materializer_source.import_path
                        for materializer_source in output.materializer_source
                    ]
                    raise StepInterfaceError(
                        "Invalid materializers specified for output "
                        f"{output_name} of step {self.name}. None of the "
                        f"materializers ({materializer_strings}) are "
                        "able to save or load data of the type that is defined "
                        "for the output "
                        f"({output_annotation.resolved_annotation})."
                    )
            else:
                if output_annotation.resolved_annotation is Any:
                    outputs[output_name]["materializer_source"] = ()
                    outputs[output_name]["default_materializer_source"] = (
                        source_utils.resolve(
                            materializer_registry.get_default_materializer()
                        )
                    )
                    continue

                materializer_sources = []

                for output_type in output_annotation.get_output_types():
                    materializer_class = materializer_registry[output_type]
                    materializer_sources.append(
                        source_utils.resolve(materializer_class)
                    )

                outputs[output_name]["materializer_source"] = tuple(
                    materializer_sources
                )

        parameters = self._finalize_parameters()
        self.configure(parameters=parameters, merge=False)
        self._validate_inputs(
            input_artifacts=input_artifacts,
            external_artifacts=external_artifacts,
            model_artifacts_or_metadata=model_artifacts_or_metadata,
            client_lazy_loaders=client_lazy_loaders,
        )

        values = dict_utils.remove_none_values({"outputs": outputs or None})
        config = StepConfigurationUpdate(**values)
        self._apply_configuration(config)

        self._configuration = self._configuration.model_copy(
            update={
                "caching_parameters": self.caching_parameters,
                "external_input_artifacts": external_artifacts,
                "model_artifacts_or_metadata": model_artifacts_or_metadata,
                "client_lazy_loaders": client_lazy_loaders,
            }
        )

        return StepConfiguration.model_validate(
            self._configuration.model_dump()
        )

    def _finalize_parameters(self) -> Dict[str, Any]:
        """Finalizes the config parameters for running this step.

        Returns:
            All parameter values for running this step.
        """
        params = {}
        for key, value in self.configuration.parameters.items():
            if key not in self.entrypoint_definition.inputs:
                continue

            annotation = self.entrypoint_definition.inputs[key].annotation
            annotation = resolve_type_annotation(annotation)
            if inspect.isclass(annotation) and issubclass(
                annotation, BaseModel
            ):
                # Make sure we have all necessary values to instantiate the
                # pydantic model later
                model = annotation(**value)
                params[key] = model.model_dump()
            else:
                params[key] = value

        return params

caching_parameters property

Caching parameters for this step.

Returns:

Type	Description
Dict[str, Any]	A dictionary containing the caching parameters

configuration property

The configuration of the step.

Returns:

Type	Description
PartialStepConfiguration	The configuration of the step.

docstring property

The docstring of this step.

Returns:

Type	Description
Optional[str]	The docstring of this step.

enable_cache property

If caching is enabled for the step.

Returns:

Type	Description
Optional[bool]	If caching is enabled for the step.

name property

The name of the step.

Returns:

Type	Description
str	The name of the step.

source_code property

The source code of this step.

Returns:

Type	Description
str	The source code of this step.

source_object property

The source object of this step.

Returns:

Type	Description
Any	The source object of this step.

__call__(*args, id=None, after=None, **kwargs)

Handle a call of the step.

This method does one of two things: * If there is an active pipeline context, it adds an invocation of the step instance to the pipeline. * If no pipeline is active, it calls the step entrypoint function.

Parameters:

Name	Type	Description	Default
*args	Any	Entrypoint function arguments.	()
id	Optional[str]	Invocation ID to use.	None
after	Union[str, Sequence[str], None]	Upstream steps for the invocation.	None
**kwargs	Any	Entrypoint function keyword arguments.	{}

Returns:

Type	Description
Any	The outputs of the entrypoint function call.

Source code in src/zenml/steps/base_step.py

def __call__(
    self,
    *args: Any,
    id: Optional[str] = None,
    after: Union[str, Sequence[str], None] = None,
    **kwargs: Any,
) -> Any:
    """Handle a call of the step.

    This method does one of two things:
    * If there is an active pipeline context, it adds an invocation of the
      step instance to the pipeline.
    * If no pipeline is active, it calls the step entrypoint function.

    Args:
        *args: Entrypoint function arguments.
        id: Invocation ID to use.
        after: Upstream steps for the invocation.
        **kwargs: Entrypoint function keyword arguments.

    Returns:
        The outputs of the entrypoint function call.
    """
    from zenml.pipelines.pipeline_definition import Pipeline

    if not Pipeline.ACTIVE_PIPELINE:
        from zenml import constants, get_step_context

        # If the environment variable was set to explicitly not run on the
        # stack, we do that.
        run_without_stack = handle_bool_env_var(
            ENV_ZENML_RUN_SINGLE_STEPS_WITHOUT_STACK, default=False
        )
        if run_without_stack:
            return self.call_entrypoint(*args, **kwargs)

        try:
            get_step_context()
        except RuntimeError:
            pass
        else:
            # We're currently inside the execution of a different step
            # -> We don't want to launch another single step pipeline here,
            # but instead just call the step function
            return self.call_entrypoint(*args, **kwargs)

        if constants.SHOULD_PREVENT_PIPELINE_EXECUTION:
            logger.info(
                "Preventing execution of step '%s'.",
                self.name,
            )
            return

        return run_as_single_step_pipeline(self, *args, **kwargs)

    (
        input_artifacts,
        external_artifacts,
        model_artifacts_or_metadata,
        client_lazy_loaders,
        parameters,
        default_parameters,
    ) = self._parse_call_args(*args, **kwargs)

    upstream_steps = {
        artifact.invocation_id for artifact in input_artifacts.values()
    }
    if isinstance(after, str):
        upstream_steps.add(after)
    elif isinstance(after, Sequence):
        upstream_steps = upstream_steps.union(after)

    invocation_id = Pipeline.ACTIVE_PIPELINE.add_step_invocation(
        step=self,
        input_artifacts=input_artifacts,
        external_artifacts=external_artifacts,
        model_artifacts_or_metadata=model_artifacts_or_metadata,
        client_lazy_loaders=client_lazy_loaders,
        parameters=parameters,
        default_parameters=default_parameters,
        upstream_steps=upstream_steps,
        custom_id=id,
        allow_id_suffix=not id,
    )

    outputs = []
    for key, annotation in self.entrypoint_definition.outputs.items():
        output = StepArtifact(
            invocation_id=invocation_id,
            output_name=key,
            annotation=annotation,
            pipeline=Pipeline.ACTIVE_PIPELINE,
        )
        outputs.append(output)
    return outputs[0] if len(outputs) == 1 else outputs

__init__(name=None, enable_cache=None, enable_artifact_metadata=None, enable_artifact_visualization=None, enable_step_logs=None, experiment_tracker=None, step_operator=None, parameters=None, output_materializers=None, settings=None, extra=None, on_failure=None, on_success=None, model=None, retry=None, substitutions=None)

Initializes a step.

Parameters:

Name	Type	Description	Default
name	Optional[str]	The name of the step.	None
enable_cache	Optional[bool]	If caching should be enabled for this step.	None
enable_artifact_metadata	Optional[bool]	If artifact metadata should be enabled for this step.	None
enable_artifact_visualization	Optional[bool]	If artifact visualization should be enabled for this step.	None
enable_step_logs	Optional[bool]	Enable step logs for this step.	None
experiment_tracker	Optional[str]	The experiment tracker to use for this step.	None
step_operator	Optional[str]	The step operator to use for this step.	None
parameters	Optional[Dict[str, Any]]	Function parameters for this step	None
output_materializers	Optional[OutputMaterializersSpecification]	Output materializers for this step. If given as a dict, the keys must be a subset of the output names of this step. If a single value (type or string) is given, the materializer will be used for all outputs.	None
settings	Optional[Mapping[str, SettingsOrDict]]	settings for this step.	None
extra	Optional[Dict[str, Any]]	Extra configurations for this step.	None
on_failure	Optional[HookSpecification]	Callback function in event of failure of the step. Can be a function with a single argument of type BaseException, or a source path to such a function (e.g. module.my_function).	None
on_success	Optional[HookSpecification]	Callback function in event of success of the step. Can be a function with no arguments, or a source path to such a function (e.g. module.my_function).	None
model	Optional[Model]	configuration of the model version in the Model Control Plane.	None
retry	Optional[StepRetryConfig]	Configuration for retrying the step in case of failure.	None
substitutions	Optional[Dict[str, str]]	Extra placeholders to use in the name template.	None

Source code in src/zenml/steps/base_step.py

def __init__(
    self,
    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[Dict[str, Any]] = None,
    output_materializers: Optional[
        "OutputMaterializersSpecification"
    ] = None,
    settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
    extra: Optional[Dict[str, Any]] = None,
    on_failure: Optional["HookSpecification"] = None,
    on_success: Optional["HookSpecification"] = None,
    model: Optional["Model"] = None,
    retry: Optional[StepRetryConfig] = None,
    substitutions: Optional[Dict[str, str]] = None,
) -> None:
    """Initializes a step.

    Args:
        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
            for this step.
        enable_artifact_visualization: If artifact visualization should be
            enabled for this step.
        enable_step_logs: Enable step logs for this step.
        experiment_tracker: The experiment tracker to use for this step.
        step_operator: The step operator to use for this step.
        parameters: Function parameters for this step
        output_materializers: Output materializers for this step. If
            given as a dict, the keys must be a subset of the output names
            of this step. If a single value (type or string) is given, the
            materializer will be used for all outputs.
        settings: settings for this step.
        extra: Extra configurations for this step.
        on_failure: Callback function in event of failure of the step. Can
            be a function with a single argument of type `BaseException`, or
            a source path to such a function (e.g. `module.my_function`).
        on_success: Callback function in event of success of the step. Can
            be a function with no arguments, or a source path to such a
            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.
        substitutions: Extra placeholders to use in the name template.
    """
    from zenml.config.step_configurations import PartialStepConfiguration

    self.entrypoint_definition = validate_entrypoint_function(
        self.entrypoint,
        reserved_arguments=["after", "id"],
    )

    name = name or self.__class__.__name__

    logger.debug(
        "Step `%s`: Caching %s.",
        name,
        "enabled" if enable_cache is not False else "disabled",
    )
    logger.debug(
        "Step `%s`: Artifact metadata %s.",
        name,
        "enabled" if enable_artifact_metadata is not False else "disabled",
    )
    logger.debug(
        "Step `%s`: Artifact visualization %s.",
        name,
        "enabled"
        if enable_artifact_visualization is not False
        else "disabled",
    )
    logger.debug(
        "Step `%s`: logs %s.",
        name,
        "enabled" if enable_step_logs is not False else "disabled",
    )
    if model is not None:
        logger.debug(
            "Step `%s`: Is in Model context %s.",
            name,
            {
                "model": model.name,
                "version": model.version,
            },
        )

    self._configuration = PartialStepConfiguration(
        name=name,
        enable_cache=enable_cache,
        enable_artifact_metadata=enable_artifact_metadata,
        enable_artifact_visualization=enable_artifact_visualization,
        enable_step_logs=enable_step_logs,
    )
    self.configure(
        experiment_tracker=experiment_tracker,
        step_operator=step_operator,
        output_materializers=output_materializers,
        parameters=parameters,
        settings=settings,
        extra=extra,
        on_failure=on_failure,
        on_success=on_success,
        model=model,
        retry=retry,
        substitutions=substitutions,
    )

    notebook_utils.try_to_save_notebook_cell_code(self.source_object)

call_entrypoint(*args, **kwargs)

Calls the entrypoint function of the step.

Parameters:

Name	Type	Description	Default
*args	Any	Entrypoint function arguments.	()
**kwargs	Any	Entrypoint function keyword arguments.	{}

Returns:

Type	Description
Any	The return value of the entrypoint function.

Raises:

Type	Description
StepInterfaceError	If the arguments to the entrypoint function are invalid.

Source code in src/zenml/steps/base_step.py

def call_entrypoint(self, *args: Any, **kwargs: Any) -> Any:
    """Calls the entrypoint function of the step.

    Args:
        *args: Entrypoint function arguments.
        **kwargs: Entrypoint function keyword arguments.

    Returns:
        The return value of the entrypoint function.

    Raises:
        StepInterfaceError: If the arguments to the entrypoint function are
            invalid.
    """
    try:
        validated_args = pydantic_utils.validate_function_args(
            self.entrypoint,
            ConfigDict(arbitrary_types_allowed=True),
            *args,
            **kwargs,
        )
    except ValidationError as e:
        raise StepInterfaceError(
            "Invalid step function entrypoint arguments. Check out the "
            "pydantic error above for more details."
        ) from e

    return self.entrypoint(**validated_args)

configure(enable_cache=None, enable_artifact_metadata=None, enable_artifact_visualization=None, enable_step_logs=None, experiment_tracker=None, step_operator=None, parameters=None, output_materializers=None, settings=None, extra=None, on_failure=None, on_success=None, model=None, merge=True, retry=None, substitutions=None)

Configures the step.

Configuration merging example: * merge==True: step.configure(extra={"key1": 1}) step.configure(extra={"key2": 2}, merge=True) step.configuration.extra # {"key1": 1, "key2": 2} * merge==False: step.configure(extra={"key1": 1}) step.configure(extra={"key2": 2}, merge=False) step.configuration.extra # {"key2": 2}

Parameters:

Name	Type	Description	Default
enable_cache	Optional[bool]	If caching should be enabled for this step.	None
enable_artifact_metadata	Optional[bool]	If artifact metadata should be enabled for this step.	None
enable_artifact_visualization	Optional[bool]	If artifact visualization should be enabled for this step.	None
enable_step_logs	Optional[bool]	If step logs should be enabled for this step.	None
experiment_tracker	Optional[str]	The experiment tracker to use for this step.	None
step_operator	Optional[str]	The step operator to use for this step.	None
parameters	Optional[Dict[str, Any]]	Function parameters for this step	None
output_materializers	Optional[OutputMaterializersSpecification]	Output materializers for this step. If given as a dict, the keys must be a subset of the output names of this step. If a single value (type or string) is given, the materializer will be used for all outputs.	None
settings	Optional[Mapping[str, SettingsOrDict]]	settings for this step.	None
extra	Optional[Dict[str, Any]]	Extra configurations for this step.	None
on_failure	Optional[HookSpecification]	Callback function in event of failure of the step. Can be a function with a single argument of type BaseException, or a source path to such a function (e.g. module.my_function).	None
on_success	Optional[HookSpecification]	Callback function in event of success of the step. Can be a function with no arguments, or a source path to such a function (e.g. module.my_function).	None
model	Optional[Model]	configuration of the model version in the Model Control Plane.	None
merge	bool	If True, will merge the given dictionary configurations like parameters and settings with existing configurations. If False the given configurations will overwrite all existing ones. See the general description of this method for an example.	True
retry	Optional[StepRetryConfig]	Configuration for retrying the step in case of failure.	None
substitutions	Optional[Dict[str, str]]	Extra placeholders to use in the name template.	None

Returns:

Type	Description
T	The step instance that this method was called on.

Source code in src/zenml/steps/base_step.py

def configure(
    self: T,
    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[Dict[str, Any]] = None,
    output_materializers: Optional[
        "OutputMaterializersSpecification"
    ] = None,
    settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
    extra: Optional[Dict[str, Any]] = None,
    on_failure: Optional["HookSpecification"] = None,
    on_success: Optional["HookSpecification"] = None,
    model: Optional["Model"] = None,
    merge: bool = True,
    retry: Optional[StepRetryConfig] = None,
    substitutions: Optional[Dict[str, str]] = None,
) -> T:
    """Configures the step.

    Configuration merging example:
    * `merge==True`:
        step.configure(extra={"key1": 1})
        step.configure(extra={"key2": 2}, merge=True)
        step.configuration.extra # {"key1": 1, "key2": 2}
    * `merge==False`:
        step.configure(extra={"key1": 1})
        step.configure(extra={"key2": 2}, merge=False)
        step.configuration.extra # {"key2": 2}

    Args:
        enable_cache: If caching should be enabled for this step.
        enable_artifact_metadata: If artifact metadata should be enabled
            for this step.
        enable_artifact_visualization: If artifact visualization should be
            enabled for this step.
        enable_step_logs: If step logs should be enabled for this step.
        experiment_tracker: The experiment tracker to use for this step.
        step_operator: The step operator to use for this step.
        parameters: Function parameters for this step
        output_materializers: Output materializers for this step. If
            given as a dict, the keys must be a subset of the output names
            of this step. If a single value (type or string) is given, the
            materializer will be used for all outputs.
        settings: settings for this step.
        extra: Extra configurations for this step.
        on_failure: Callback function in event of failure of the step. Can
            be a function with a single argument of type `BaseException`, or
            a source path to such a function (e.g. `module.my_function`).
        on_success: Callback function in event of success of the step. Can
            be a function with no arguments, or a source path to such a
            function (e.g. `module.my_function`).
        model: configuration of the model version in the Model Control Plane.
        merge: If `True`, will merge the given dictionary configurations
            like `parameters` and `settings` with existing
            configurations. If `False` the given configurations will
            overwrite all existing ones. See the general description of this
            method for an example.
        retry: Configuration for retrying the step in case of failure.
        substitutions: Extra placeholders to use in the name template.

    Returns:
        The step instance that this method was called on.
    """
    from zenml.config.step_configurations import StepConfigurationUpdate
    from zenml.hooks.hook_validators import resolve_and_validate_hook

    def _resolve_if_necessary(
        value: Union[str, Source, Type[Any]],
    ) -> Source:
        if isinstance(value, str):
            return Source.from_import_path(value)
        elif isinstance(value, Source):
            return value
        else:
            return source_utils.resolve(value)

    def _convert_to_tuple(value: Any) -> Tuple[Source, ...]:
        if isinstance(value, str) or not isinstance(value, Sequence):
            return (_resolve_if_necessary(value),)
        else:
            return tuple(_resolve_if_necessary(v) for v in value)

    outputs: Dict[str, Dict[str, Tuple[Source, ...]]] = defaultdict(dict)
    allowed_output_names = set(self.entrypoint_definition.outputs)

    if output_materializers:
        if not isinstance(output_materializers, Mapping):
            sources = _convert_to_tuple(output_materializers)
            output_materializers = {
                output_name: sources
                for output_name in allowed_output_names
            }

        for output_name, materializer in output_materializers.items():
            sources = _convert_to_tuple(materializer)
            outputs[output_name]["materializer_source"] = sources

    failure_hook_source = None
    if on_failure:
        # string of on_failure hook function to be used for this step
        failure_hook_source = resolve_and_validate_hook(on_failure)

    success_hook_source = None
    if on_success:
        # string of on_success hook function to be used for this step
        success_hook_source = resolve_and_validate_hook(on_success)

    values = dict_utils.remove_none_values(
        {
            "enable_cache": enable_cache,
            "enable_artifact_metadata": enable_artifact_metadata,
            "enable_artifact_visualization": enable_artifact_visualization,
            "enable_step_logs": enable_step_logs,
            "experiment_tracker": experiment_tracker,
            "step_operator": step_operator,
            "parameters": parameters,
            "settings": settings,
            "outputs": outputs or None,
            "extra": extra,
            "failure_hook_source": failure_hook_source,
            "success_hook_source": success_hook_source,
            "model": model,
            "retry": retry,
            "substitutions": substitutions,
        }
    )
    config = StepConfigurationUpdate(**values)
    self._apply_configuration(config, merge=merge)
    return self

copy()

Copies the step.

Returns:

Type	Description
BaseStep	The step copy.

Source code in src/zenml/steps/base_step.py

def copy(self) -> "BaseStep":
    """Copies the step.

    Returns:
        The step copy.
    """
    return copy.deepcopy(self)

entrypoint(*args, **kwargs) abstractmethod

Abstract method for core step logic.

Parameters:

Name	Type	Description	Default
*args	Any	Positional arguments passed to the step.	()
**kwargs	Any	Keyword arguments passed to the step.	{}

Returns:

Type	Description
Any	The output of the step.

Source code in src/zenml/steps/base_step.py

@abstractmethod
def entrypoint(self, *args: Any, **kwargs: Any) -> Any:
    """Abstract method for core step logic.

    Args:
        *args: Positional arguments passed to the step.
        **kwargs: Keyword arguments passed to the step.

    Returns:
        The output of the step.
    """

load_from_source(source) classmethod

Loads a step from source.

Parameters:

Name	Type	Description	Default
source	Union[Source, str]	The path to the step source.	required

Returns:

Type	Description
BaseStep	The loaded step.

Raises:

Type	Description
ValueError	If the source is not a valid step source.

Source code in src/zenml/steps/base_step.py

@classmethod
def load_from_source(cls, source: Union[Source, str]) -> "BaseStep":
    """Loads a step from source.

    Args:
        source: The path to the step source.

    Returns:
        The loaded step.

    Raises:
        ValueError: If the source is not a valid step source.
    """
    obj = source_utils.load(source)

    if isinstance(obj, BaseStep):
        return obj
    elif isinstance(obj, type) and issubclass(obj, BaseStep):
        return obj()
    else:
        raise ValueError("Invalid step source.")

resolve()

Resolves the step.

Returns:

Type	Description
Source	The step source.

Source code in src/zenml/steps/base_step.py

def resolve(self) -> Source:
    """Resolves the step.

    Returns:
        The step source.
    """
    return source_utils.resolve(self.__class__)

with_options(enable_cache=None, enable_artifact_metadata=None, enable_artifact_visualization=None, enable_step_logs=None, experiment_tracker=None, step_operator=None, parameters=None, output_materializers=None, settings=None, extra=None, on_failure=None, on_success=None, model=None, merge=True, substitutions=None)

Copies the step and applies the given configurations.

Parameters:

Name	Type	Description	Default
enable_cache	Optional[bool]	If caching should be enabled for this step.	None
enable_artifact_metadata	Optional[bool]	If artifact metadata should be enabled for this step.	None
enable_artifact_visualization	Optional[bool]	If artifact visualization should be enabled for this step.	None
enable_step_logs	Optional[bool]	If step logs should be enabled for this step.	None
experiment_tracker	Optional[str]	The experiment tracker to use for this step.	None
step_operator	Optional[str]	The step operator to use for this step.	None
parameters	Optional[Dict[str, Any]]	Function parameters for this step	None
output_materializers	Optional[OutputMaterializersSpecification]	Output materializers for this step. If given as a dict, the keys must be a subset of the output names of this step. If a single value (type or string) is given, the materializer will be used for all outputs.	None
settings	Optional[Mapping[str, SettingsOrDict]]	settings for this step.	None
extra	Optional[Dict[str, Any]]	Extra configurations for this step.	None
on_failure	Optional[HookSpecification]	Callback function in event of failure of the step. Can be a function with a single argument of type BaseException, or a source path to such a function (e.g. module.my_function).	None
on_success	Optional[HookSpecification]	Callback function in event of success of the step. Can be a function with no arguments, or a source path to such a function (e.g. module.my_function).	None
model	Optional[Model]	configuration of the model version in the Model Control Plane.	None
merge	bool	If True, will merge the given dictionary configurations like parameters and settings with existing configurations. If False the given configurations will overwrite all existing ones. See the general description of this method for an example.	True
substitutions	Optional[Dict[str, str]]	Extra placeholders for the step name.	None

Returns:

Type	Description
BaseStep	The copied step instance.

Source code in src/zenml/steps/base_step.py

def with_options(
    self,
    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[Dict[str, Any]] = None,
    output_materializers: Optional[
        "OutputMaterializersSpecification"
    ] = None,
    settings: Optional[Mapping[str, "SettingsOrDict"]] = None,
    extra: Optional[Dict[str, Any]] = None,
    on_failure: Optional["HookSpecification"] = None,
    on_success: Optional["HookSpecification"] = None,
    model: Optional["Model"] = None,
    merge: bool = True,
    substitutions: Optional[Dict[str, str]] = None,
) -> "BaseStep":
    """Copies the step and applies the given configurations.

    Args:
        enable_cache: If caching should be enabled for this step.
        enable_artifact_metadata: If artifact metadata should be enabled
            for this step.
        enable_artifact_visualization: If artifact visualization should be
            enabled for this step.
        enable_step_logs: If step logs should be enabled for this step.
        experiment_tracker: The experiment tracker to use for this step.
        step_operator: The step operator to use for this step.
        parameters: Function parameters for this step
        output_materializers: Output materializers for this step. If
            given as a dict, the keys must be a subset of the output names
            of this step. If a single value (type or string) is given, the
            materializer will be used for all outputs.
        settings: settings for this step.
        extra: Extra configurations for this step.
        on_failure: Callback function in event of failure of the step. Can
            be a function with a single argument of type `BaseException`, or
            a source path to such a function (e.g. `module.my_function`).
        on_success: Callback function in event of success of the step. Can
            be a function with no arguments, or a source path to such a
            function (e.g. `module.my_function`).
        model: configuration of the model version in the Model Control Plane.
        merge: If `True`, will merge the given dictionary configurations
            like `parameters` and `settings` with existing
            configurations. If `False` the given configurations will
            overwrite all existing ones. See the general description of this
            method for an example.
        substitutions: Extra placeholders for the step name.

    Returns:
        The copied step instance.
    """
    step_copy = self.copy()
    step_copy.configure(
        enable_cache=enable_cache,
        enable_artifact_metadata=enable_artifact_metadata,
        enable_artifact_visualization=enable_artifact_visualization,
        enable_step_logs=enable_step_logs,
        experiment_tracker=experiment_tracker,
        step_operator=step_operator,
        parameters=parameters,
        output_materializers=output_materializers,
        settings=settings,
        extra=extra,
        on_failure=on_failure,
        on_success=on_success,
        model=model,
        merge=merge,
        substitutions=substitutions,
    )
    return step_copy

ResourceSettings

Bases: BaseSettings

Hardware resource settings.

Attributes:

Name	Type	Description
cpu_count	Optional[PositiveFloat]	The amount of CPU cores that should be configured.
gpu_count	Optional[NonNegativeInt]	The amount of GPUs that should be configured.
memory	Optional[str]	The amount of memory that should be configured.

Source code in src/zenml/config/resource_settings.py

class ResourceSettings(BaseSettings):
    """Hardware resource settings.

    Attributes:
        cpu_count: The amount of CPU cores that should be configured.
        gpu_count: The amount of GPUs that should be configured.
        memory: The amount of memory that should be configured.
    """

    cpu_count: Optional[PositiveFloat] = None
    gpu_count: Optional[NonNegativeInt] = None
    memory: Optional[str] = Field(pattern=MEMORY_REGEX, default=None)

    @property
    def empty(self) -> bool:
        """Returns if this object is "empty" (=no values configured) or not.

        Returns:
            `True` if no values were configured, `False` otherwise.
        """
        # To detect whether this config is empty (= no values specified), we
        # check if there are any attributes which are explicitly set to any
        # value other than `None`.
        return len(self.model_dump(exclude_unset=True, exclude_none=True)) == 0

    def get_memory(
        self, unit: Union[str, ByteUnit] = ByteUnit.GB
    ) -> Optional[float]:
        """Gets the memory configuration in a specific unit.

        Args:
            unit: The unit to which the memory should be converted.

        Raises:
            ValueError: If the memory string is invalid.

        Returns:
            The memory configuration converted to the requested unit, or None
            if no memory was configured.
        """
        if not self.memory:
            return None

        if isinstance(unit, str):
            unit = ByteUnit(unit)

        memory = self.memory
        for memory_unit in ByteUnit:
            if memory.endswith(memory_unit.value):
                memory_value = int(memory[: -len(memory_unit.value)])
                return memory_value * memory_unit.byte_value / unit.byte_value
        else:
            # Should never happen due to the regex validation
            raise ValueError(f"Unable to parse memory unit from '{memory}'.")

    model_config = SettingsConfigDict(
        # public attributes are immutable
        frozen=True,
        # prevent extra attributes during model initialization
        extra="forbid",
    )

empty property

Returns if this object is "empty" (=no values configured) or not.

Returns:

Type	Description
bool	True if no values were configured, False otherwise.

get_memory(unit=ByteUnit.GB)

Gets the memory configuration in a specific unit.

Parameters:

Name	Type	Description	Default
unit	Union[str, ByteUnit]	The unit to which the memory should be converted.	GB

Raises:

Type	Description
ValueError	If the memory string is invalid.

Returns:

Type	Description
Optional[float]	The memory configuration converted to the requested unit, or None
Optional[float]	if no memory was configured.

Source code in src/zenml/config/resource_settings.py

def get_memory(
    self, unit: Union[str, ByteUnit] = ByteUnit.GB
) -> Optional[float]:
    """Gets the memory configuration in a specific unit.

    Args:
        unit: The unit to which the memory should be converted.

    Raises:
        ValueError: If the memory string is invalid.

    Returns:
        The memory configuration converted to the requested unit, or None
        if no memory was configured.
    """
    if not self.memory:
        return None

    if isinstance(unit, str):
        unit = ByteUnit(unit)

    memory = self.memory
    for memory_unit in ByteUnit:
        if memory.endswith(memory_unit.value):
            memory_value = int(memory[: -len(memory_unit.value)])
            return memory_value * memory_unit.byte_value / unit.byte_value
    else:
        # Should never happen due to the regex validation
        raise ValueError(f"Unable to parse memory unit from '{memory}'.")

StepContext

Provides additional context inside a step function.

This singleton class is used to access information about the current run, step run, or its outputs inside a step function.

Usage example:

from zenml.steps import get_step_context

@step
def my_trainer_step() -> Any:
    context = get_step_context()

    # get info about the current pipeline run
    current_pipeline_run = context.pipeline_run

    # get info about the current step run
    current_step_run = context.step_run

    # get info about the future output artifacts of this step
    output_artifact_uri = context.get_output_artifact_uri()

    ...

Source code in src/zenml/steps/step_context.py

class StepContext(metaclass=SingletonMetaClass):
    """Provides additional context inside a step function.

    This singleton class is used to access information about the current run,
    step run, or its outputs inside a step function.

    Usage example:

    ```python
    from zenml.steps import get_step_context

    @step
    def my_trainer_step() -> Any:
        context = get_step_context()

        # get info about the current pipeline run
        current_pipeline_run = context.pipeline_run

        # get info about the current step run
        current_step_run = context.step_run

        # get info about the future output artifacts of this step
        output_artifact_uri = context.get_output_artifact_uri()

        ...
    ```
    """

    def __init__(
        self,
        pipeline_run: "PipelineRunResponse",
        step_run: "StepRunResponse",
        output_materializers: Mapping[str, Sequence[Type["BaseMaterializer"]]],
        output_artifact_uris: Mapping[str, str],
        output_artifact_configs: Mapping[str, Optional["ArtifactConfig"]],
    ) -> None:
        """Initialize the context of the currently running step.

        Args:
            pipeline_run: The model of the current pipeline run.
            step_run: The model of the current step run.
            output_materializers: The output materializers of the step that
                this context is used in.
            output_artifact_uris: The output artifacts of the step that this
                context is used in.
            output_artifact_configs: The outputs' ArtifactConfigs of the step that this
                context is used in.

        Raises:
            StepContextError: If the keys of the output materializers and
                output artifacts do not match.
        """
        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.model_version = (
            step_run.model_version or pipeline_run.model_version
        )

        self.step_name = self.step_run.name

        # set outputs
        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"keys: {set(output_materializers)}, output artifact URI "
                f"keys: {set(output_artifact_uris)}"
            )
        self._outputs = {
            key: StepContextOutput(
                materializer_classes=output_materializers[key],
                artifact_uri=output_artifact_uris[key],
                artifact_config=output_artifact_configs[key],
            )
            for key in output_materializers.keys()
        }
        self._cleanup_registry = CallbackRegistry()

    @property
    def pipeline(self) -> "PipelineResponse":
        """Returns the current pipeline.

        Returns:
            The current pipeline or None.

        Raises:
            StepContextError: If the pipeline run does not have a pipeline.
        """
        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"run '{self.pipeline_run.id}': This pipeline run does not have "
            f"a pipeline associated with it."
        )

    @property
    def model(self) -> "Model":
        """Returns configured Model.

        Order of resolution to search for Model is:
            1. Model from the step context
            2. Model from the pipeline context

        Returns:
            The `Model` object associated with the current step.

        Raises:
            StepContextError: If no `Model` object was specified for the step
                or pipeline.
        """
        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}': No model has been specified "
                "the step or pipeline."
            )

        return self.model_version.to_model_class()

    @property
    def inputs(self) -> Dict[str, "StepRunInputResponse"]:
        """Returns the input artifacts of the current step.

        Returns:
            The input artifacts of the current step.
        """
        return self.step_run.inputs

    def _get_output(
        self, output_name: Optional[str] = None
    ) -> "StepContextOutput":
        """Returns the materializer and artifact URI for a given step output.

        Args:
            output_name: Optional name of the output for which to get the
                materializer and URI.

        Returns:
            Tuple containing the materializer and artifact URI for the
                given output.

        Raises:
            StepContextError: If the step has no outputs, no output for
                the given `output_name` or if no `output_name` was given but
                the step has multiple outputs.
        """
        output_count = len(self._outputs)
        if output_count == 0:
            raise StepContextError(
                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"This step has multiple outputs ({set(self._outputs)}), "
                f"please specify which output to return."
            )

        if output_name:
            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"output with the given name, please specify one of the "
                    f"available outputs: {set(self._outputs)}."
                )
            return self._outputs[output_name]
        else:
            return next(iter(self._outputs.values()))

    def get_output_materializer(
        self,
        output_name: Optional[str] = None,
        custom_materializer_class: Optional[Type["BaseMaterializer"]] = None,
        data_type: Optional[Type[Any]] = None,
    ) -> "BaseMaterializer":
        """Returns a materializer for a given step output.

        Args:
            output_name: Optional name of the output for which to get the
                materializer. If no name is given and the step only has a
                single output, the materializer of this output will be
                returned. If the step has multiple outputs, an exception
                will be raised.
            custom_materializer_class: If given, this `BaseMaterializer`
                subclass will be initialized with the output artifact instead
                of the materializer that was registered for this step output.
            data_type: If the output annotation is of type `Union` and the step
                therefore has multiple materializers configured, you can provide
                a data type for the output which will be used to select the
                correct materializer. If not provided, the first materializer
                will be used.

        Returns:
            A materializer initialized with the output artifact for
            the given output.
        """
        from zenml.utils import materializer_utils

        output = self._get_output(output_name)
        materializer_classes = output.materializer_classes
        artifact_uri = output.artifact_uri

        if custom_materializer_class:
            materializer_class = custom_materializer_class
        elif len(materializer_classes) == 1 or not data_type:
            materializer_class = materializer_classes[0]
        else:
            materializer_class = materializer_utils.select_materializer(
                data_type=data_type, materializer_classes=materializer_classes
            )

        return materializer_class(artifact_uri)

    def get_output_artifact_uri(
        self, output_name: Optional[str] = None
    ) -> str:
        """Returns the artifact URI for a given step output.

        Args:
            output_name: Optional name of the output for which to get the URI.
                If no name is given and the step only has a single output,
                the URI of this output will be returned. If the step has
                multiple outputs, an exception will be raised.

        Returns:
            Artifact URI for the given output.
        """
        return self._get_output(output_name).artifact_uri

    def get_output_metadata(
        self, output_name: Optional[str] = None
    ) -> Dict[str, "MetadataType"]:
        """Returns the metadata 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:
            Metadata for the given output.
        """
        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,
        metadata: Dict[str, "MetadataType"],
        output_name: Optional[str] = None,
    ) -> None:
        """Adds metadata for a given step output.

        Args:
            metadata: The metadata to add.
            output_name: Optional name of the output for which to add the
                metadata. If no name is given and the step only has a single
                output, the metadata of this output will be added. If the
                step has multiple outputs, an exception will be raised.
        """
        output = self._get_output(output_name)
        if not output.run_metadata:
            output.run_metadata = {}
        output.run_metadata.update(**metadata)

    def add_output_tags(
        self,
        tags: List[str],
        output_name: Optional[str] = None,
    ) -> None:
        """Adds tags for a given step output.

        Args:
            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.
        """
        output = self._get_output(output_name)
        if not output.tags:
            output.tags = []
        output.tags += tags

inputs property

Returns the input artifacts of the current step.

Returns:

Type	Description
Dict[str, StepRunInputResponse]	The input artifacts of the current step.

model property

Returns configured Model.

Order of resolution to search for Model is

1. Model from the step context
2. Model from the pipeline context

Returns:

Type	Description
Model	The Model object associated with the current step.

Raises:

Type	Description
StepContextError	If no Model object was specified for the step or pipeline.

pipeline property

Returns the current pipeline.

Returns:

Type	Description
PipelineResponse	The current pipeline or None.

Raises:

Type	Description
StepContextError	If the pipeline run does not have a pipeline.

__init__(pipeline_run, step_run, output_materializers, output_artifact_uris, output_artifact_configs)

Initialize the context of the currently running step.

Parameters:

Name	Type	Description	Default
pipeline_run	PipelineRunResponse	The model of the current pipeline run.	required
step_run	StepRunResponse	The model of the current step run.	required
output_materializers	Mapping[str, Sequence[Type[BaseMaterializer]]]	The output materializers of the step that this context is used in.	required
output_artifact_uris	Mapping[str, str]	The output artifacts of the step that this context is used in.	required
output_artifact_configs	Mapping[str, Optional[ArtifactConfig]]	The outputs' ArtifactConfigs of the step that this context is used in.	required

Raises:

Type	Description
StepContextError	If the keys of the output materializers and output artifacts do not match.

Source code in src/zenml/steps/step_context.py

def __init__(
    self,
    pipeline_run: "PipelineRunResponse",
    step_run: "StepRunResponse",
    output_materializers: Mapping[str, Sequence[Type["BaseMaterializer"]]],
    output_artifact_uris: Mapping[str, str],
    output_artifact_configs: Mapping[str, Optional["ArtifactConfig"]],
) -> None:
    """Initialize the context of the currently running step.

    Args:
        pipeline_run: The model of the current pipeline run.
        step_run: The model of the current step run.
        output_materializers: The output materializers of the step that
            this context is used in.
        output_artifact_uris: The output artifacts of the step that this
            context is used in.
        output_artifact_configs: The outputs' ArtifactConfigs of the step that this
            context is used in.

    Raises:
        StepContextError: If the keys of the output materializers and
            output artifacts do not match.
    """
    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.model_version = (
        step_run.model_version or pipeline_run.model_version
    )

    self.step_name = self.step_run.name

    # set outputs
    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"keys: {set(output_materializers)}, output artifact URI "
            f"keys: {set(output_artifact_uris)}"
        )
    self._outputs = {
        key: StepContextOutput(
            materializer_classes=output_materializers[key],
            artifact_uri=output_artifact_uris[key],
            artifact_config=output_artifact_configs[key],
        )
        for key in output_materializers.keys()
    }
    self._cleanup_registry = CallbackRegistry()

add_output_metadata(metadata, output_name=None)

Adds metadata for a given step output.

Parameters:

Name	Type	Description	Default
metadata	Dict[str, MetadataType]	The metadata to add.	required
output_name	Optional[str]	Optional name of the output for which to add the metadata. If no name is given and the step only has a single output, the metadata of this output will be added. If the step has multiple outputs, an exception will be raised.	None

Source code in src/zenml/steps/step_context.py

def add_output_metadata(
    self,
    metadata: Dict[str, "MetadataType"],
    output_name: Optional[str] = None,
) -> None:
    """Adds metadata for a given step output.

    Args:
        metadata: The metadata to add.
        output_name: Optional name of the output for which to add the
            metadata. If no name is given and the step only has a single
            output, the metadata of this output will be added. If the
            step has multiple outputs, an exception will be raised.
    """
    output = self._get_output(output_name)
    if not output.run_metadata:
        output.run_metadata = {}
    output.run_metadata.update(**metadata)

add_output_tags(tags, output_name=None)

Adds tags for a given step output.

Parameters:

Name	Type	Description	Default
tags	List[str]	The tags to add.	required
output_name	Optional[str]	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.	None

Source code in src/zenml/steps/step_context.py

def add_output_tags(
    self,
    tags: List[str],
    output_name: Optional[str] = None,
) -> None:
    """Adds tags for a given step output.

    Args:
        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.
    """
    output = self._get_output(output_name)
    if not output.tags:
        output.tags = []
    output.tags += tags

get_output_artifact_uri(output_name=None)

Returns the artifact URI for a given step output.

Parameters:

Name	Type	Description	Default
output_name	Optional[str]	Optional name of the output for which to get the URI. If no name is given and the step only has a single output, the URI of this output will be returned. If the step has multiple outputs, an exception will be raised.	None

Returns:

Type	Description
str	Artifact URI for the given output.

Source code in src/zenml/steps/step_context.py

def get_output_artifact_uri(
    self, output_name: Optional[str] = None
) -> str:
    """Returns the artifact URI for a given step output.

    Args:
        output_name: Optional name of the output for which to get the URI.
            If no name is given and the step only has a single output,
            the URI of this output will be returned. If the step has
            multiple outputs, an exception will be raised.

    Returns:
        Artifact URI for the given output.
    """
    return self._get_output(output_name).artifact_uri

get_output_materializer(output_name=None, custom_materializer_class=None, data_type=None)

Returns a materializer for a given step output.

Parameters:

Name	Type	Description	Default
output_name	Optional[str]	Optional name of the output for which to get the materializer. If no name is given and the step only has a single output, the materializer of this output will be returned. If the step has multiple outputs, an exception will be raised.	None
custom_materializer_class	Optional[Type[BaseMaterializer]]	If given, this BaseMaterializer subclass will be initialized with the output artifact instead of the materializer that was registered for this step output.	None
data_type	Optional[Type[Any]]	If the output annotation is of type Union and the step therefore has multiple materializers configured, you can provide a data type for the output which will be used to select the correct materializer. If not provided, the first materializer will be used.	None

Returns:

Type	Description
BaseMaterializer	A materializer initialized with the output artifact for
BaseMaterializer	the given output.

Source code in src/zenml/steps/step_context.py

def get_output_materializer(
    self,
    output_name: Optional[str] = None,
    custom_materializer_class: Optional[Type["BaseMaterializer"]] = None,
    data_type: Optional[Type[Any]] = None,
) -> "BaseMaterializer":
    """Returns a materializer for a given step output.

    Args:
        output_name: Optional name of the output for which to get the
            materializer. If no name is given and the step only has a
            single output, the materializer of this output will be
            returned. If the step has multiple outputs, an exception
            will be raised.
        custom_materializer_class: If given, this `BaseMaterializer`
            subclass will be initialized with the output artifact instead
            of the materializer that was registered for this step output.
        data_type: If the output annotation is of type `Union` and the step
            therefore has multiple materializers configured, you can provide
            a data type for the output which will be used to select the
            correct materializer. If not provided, the first materializer
            will be used.

    Returns:
        A materializer initialized with the output artifact for
        the given output.
    """
    from zenml.utils import materializer_utils

    output = self._get_output(output_name)
    materializer_classes = output.materializer_classes
    artifact_uri = output.artifact_uri

    if custom_materializer_class:
        materializer_class = custom_materializer_class
    elif len(materializer_classes) == 1 or not data_type:
        materializer_class = materializer_classes[0]
    else:
        materializer_class = materializer_utils.select_materializer(
            data_type=data_type, materializer_classes=materializer_classes
        )

    return materializer_class(artifact_uri)

get_output_metadata(output_name=None)

Returns the metadata for a given step output.

Parameters:

Name	Type	Description	Default
output_name	Optional[str]	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.	None

Returns:

Type	Description
Dict[str, MetadataType]	Metadata for the given output.

Source code in src/zenml/steps/step_context.py

def get_output_metadata(
    self, output_name: Optional[str] = None
) -> Dict[str, "MetadataType"]:
    """Returns the metadata 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:
        Metadata for the given output.
    """
    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

get_output_tags(output_name=None)

Returns the tags for a given step output.

Parameters:

Name	Type	Description	Default
output_name	Optional[str]	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.	None

Returns:

Type	Description
List[str]	Tags for the given output.

Source code in src/zenml/steps/step_context.py

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)

get_step_context()

Get the context of the currently running step.

Returns:

Type	Description
StepContext	The context of the currently running step.

Raises:

Type	Description
RuntimeError	If no step is currently running.

Source code in src/zenml/steps/step_context.py

def get_step_context() -> "StepContext":
    """Get the context of the currently running step.

    Returns:
        The context of the currently running step.

    Raises:
        RuntimeError: If no step is currently running.
    """
    if StepContext._exists():
        return StepContext()  # type: ignore
    raise RuntimeError(
        "The step context is only available inside a step function."
    )

step(_func=None, *, name=None, enable_cache=None, enable_artifact_metadata=None, enable_artifact_visualization=None, enable_step_logs=None, experiment_tracker=None, step_operator=None, output_materializers=None, settings=None, extra=None, on_failure=None, on_success=None, model=None, retry=None, substitutions=None)

step(_func: F) -> BaseStep

step(
    *,
    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,
    output_materializers: Optional[
        OutputMaterializersSpecification
    ] = None,
    settings: Optional[Dict[str, SettingsOrDict]] = None,
    extra: Optional[Dict[str, Any]] = None,
    on_failure: Optional[HookSpecification] = None,
    on_success: Optional[HookSpecification] = None,
    model: Optional[Model] = None,
    retry: Optional[StepRetryConfig] = None,
    substitutions: Optional[Dict[str, str]] = None,
) -> Callable[[F], BaseStep]

Decorator to create a ZenML step.

Parameters:

Name	Type	Description	Default
_func	Optional[F]	The decorated function.	None
name	Optional[str]	The name of the step. If left empty, the name of the decorated function will be used as a fallback.	None
enable_cache	Optional[bool]	Specify whether caching is enabled for this step. If no value is passed, caching is enabled by default.	None
enable_artifact_metadata	Optional[bool]	Specify whether metadata is enabled for this step. If no value is passed, metadata is enabled by default.	None
enable_artifact_visualization	Optional[bool]	Specify whether visualization is enabled for this step. If no value is passed, visualization is enabled by default.	None
enable_step_logs	Optional[bool]	Specify whether step logs are enabled for this step.	None
experiment_tracker	Optional[str]	The experiment tracker to use for this step.	None
step_operator	Optional[str]	The step operator to use for this step.	None
output_materializers	Optional[OutputMaterializersSpecification]	Output materializers for this step. If given as a dict, the keys must be a subset of the output names of this step. If a single value (type or string) is given, the materializer will be used for all outputs.	None
settings	Optional[Dict[str, SettingsOrDict]]	Settings for this step.	None
extra	Optional[Dict[str, Any]]	Extra configurations for this step.	None
on_failure	Optional[HookSpecification]	Callback function in event of failure of the step. Can be a function with a single argument of type BaseException, or a source path to such a function (e.g. module.my_function).	None
on_success	Optional[HookSpecification]	Callback function in event of success of the step. Can be a function with no arguments, or a source path to such a function (e.g. module.my_function).	None
model	Optional[Model]	configuration of the model in the Model Control Plane.	None
retry	Optional[StepRetryConfig]	configuration of step retry in case of step failure.	None
substitutions	Optional[Dict[str, str]]	Extra placeholders for the step name.	None

Returns:

Type	Description
Union[BaseStep, Callable[[F], BaseStep]]	The step instance.

Source code in src/zenml/steps/step_decorator.py

def step(
    _func: Optional["F"] = None,
    *,
    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,
    output_materializers: Optional["OutputMaterializersSpecification"] = None,
    settings: Optional[Dict[str, "SettingsOrDict"]] = None,
    extra: Optional[Dict[str, Any]] = None,
    on_failure: Optional["HookSpecification"] = None,
    on_success: Optional["HookSpecification"] = None,
    model: Optional["Model"] = None,
    retry: Optional["StepRetryConfig"] = None,
    substitutions: Optional[Dict[str, str]] = None,
) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
    """Decorator to create a ZenML step.

    Args:
        _func: The decorated function.
        name: The name of the step. If left empty, the name of the decorated
            function will be used as a fallback.
        enable_cache: Specify whether caching is enabled for this step. If no
            value is passed, caching is enabled by default.
        enable_artifact_metadata: Specify whether metadata is enabled for this
            step. If no value is passed, metadata is enabled by default.
        enable_artifact_visualization: Specify whether visualization is enabled
            for this step. If no value is passed, visualization is enabled by
            default.
        enable_step_logs: Specify whether step logs are enabled for this step.
        experiment_tracker: The experiment tracker to use for this step.
        step_operator: The step operator to use for this step.
        output_materializers: Output materializers for this step. If
            given as a dict, the keys must be a subset of the output names
            of this step. If a single value (type or string) is given, the
            materializer will be used for all outputs.
        settings: Settings for this step.
        extra: Extra configurations for this step.
        on_failure: Callback function in event of failure of the step. Can be a
            function with a single argument of type `BaseException`, or a source
            path to such a function (e.g. `module.my_function`).
        on_success: Callback function in event of success of the step. Can be a
            function with no arguments, or a source path to such a function
            (e.g. `module.my_function`).
        model: configuration of the model in the Model Control Plane.
        retry: configuration of step retry in case of step failure.
        substitutions: Extra placeholders for the step name.

    Returns:
        The step instance.
    """

    def inner_decorator(func: "F") -> "BaseStep":
        from zenml.steps.decorated_step import _DecoratedStep

        class_: Type["BaseStep"] = type(
            func.__name__,
            (_DecoratedStep,),
            {
                "entrypoint": staticmethod(func),
                "__module__": func.__module__,
                "__doc__": func.__doc__,
            },
        )

        step_instance = class_(
            name=name or func.__name__,
            enable_cache=enable_cache,
            enable_artifact_metadata=enable_artifact_metadata,
            enable_artifact_visualization=enable_artifact_visualization,
            enable_step_logs=enable_step_logs,
            experiment_tracker=experiment_tracker,
            step_operator=step_operator,
            output_materializers=output_materializers,
            settings=settings,
            extra=extra,
            on_failure=on_failure,
            on_success=on_success,
            model=model,
            retry=retry,
            substitutions=substitutions,
        )

        return step_instance

    if _func is None:
        return inner_decorator
    else:
        return inner_decorator(_func)