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

zenml/orchestrators/step_runner.py

@@ -23,46 +23,44 @@ from typing import (
     Dict,
     List,
     Optional,
-    Set,
     Tuple,
     Type,
 )
-from uuid import UUID
-
-from pydantic.typing import get_origin, is_union
 
 from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
-from zenml.artifacts.utils import save_artifact
+from zenml.artifacts.utils import _store_artifact_data_and_prepare_request
 from zenml.client import Client
 from zenml.config.step_configurations import StepConfiguration
 from zenml.config.step_run_info import StepRunInfo
 from zenml.constants import (
     ENV_ZENML_DISABLE_STEP_LOGS_STORAGE,
+    ENV_ZENML_IGNORE_FAILURE_HOOK,
     handle_bool_env_var,
 )
-from zenml.exceptions import StepContextError, StepInterfaceError
+from zenml.enums import ArtifactSaveType
+from zenml.exceptions import StepInterfaceError
 from zenml.logger import get_logger
 from zenml.logging.step_logging import StepLogsStorageContext, redirected
 from zenml.materializers.base_materializer import BaseMaterializer
-from zenml.model.utils import link_step_artifacts_to_model
-from zenml.new.steps.step_context import StepContext, get_step_context
+from zenml.models.v2.core.step_run import StepRunInputResponse
 from zenml.orchestrators.publish_utils import (
     publish_step_run_metadata,
     publish_successful_step_run,
 )
-from zenml.orchestrators.utils import is_setting_enabled
-from zenml.steps.step_environment import StepEnvironment
+from zenml.orchestrators.utils import (
+    is_setting_enabled,
+)
+from zenml.steps.step_context import StepContext, get_step_context
 from zenml.steps.utils import (
     OutputSignature,
     parse_return_type_annotations,
     resolve_type_annotation,
 )
-from zenml.utils import materializer_utils, source_utils
+from zenml.utils import materializer_utils, source_utils, string_utils
+from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
-    from zenml.artifacts.external_artifact_config import (
-        ExternalArtifactConfiguration,
-    )
+    from zenml.artifact_stores import BaseArtifactStore
     from zenml.config.source import Source
     from zenml.config.step_configurations import Step
     from zenml.models import (
@@ -103,7 +101,7 @@ class StepRunner:
         self,
         pipeline_run: "PipelineRunResponse",
         step_run: "StepRunResponse",
-        input_artifacts: Dict[str, "ArtifactVersionResponse"],
+        input_artifacts: Dict[str, StepRunInputResponse],
         output_artifact_uris: Dict[str, str],
         step_run_info: StepRunInfo,
     ) -> None:
@@ -134,7 +132,8 @@ class StepRunner:
         if step_logging_enabled and not redirected.get():
             if step_run.logs:
                 logs_context = StepLogsStorageContext(  # type: ignore[assignment]
-                    logs_uri=step_run.logs.uri
+                    logs_uri=step_run.logs.uri,
+                    artifact_store=self._stack.artifact_store,
                 )
             else:
                 logger.debug(
@@ -149,69 +148,62 @@ class StepRunner:
                 inspect.unwrap(step_instance.entrypoint)
             )
 
-            # (Deprecated) Wrap the execution of the step function in a step
-            # environment that the step function code can access to retrieve
-            # information about the pipeline runtime, such as the current step
-            # name and the current pipeline run ID
-            cache_enabled = is_setting_enabled(
-                is_enabled_on_step=step_run_info.config.enable_cache,
-                is_enabled_on_pipeline=step_run_info.pipeline.enable_cache,
-            )
             output_annotations = parse_return_type_annotations(
                 func=step_instance.entrypoint
             )
-            with StepEnvironment(
-                step_run_info=step_run_info,
-                cache_enabled=cache_enabled,
-            ):
-                self._stack.prepare_step_run(info=step_run_info)
-
-                # Initialize the step context singleton
-                StepContext._clear()
-                StepContext(
-                    pipeline_run=pipeline_run,
-                    step_run=step_run,
-                    output_materializers=output_materializers,
-                    output_artifact_uris=output_artifact_uris,
-                    step_run_info=step_run_info,
-                    cache_enabled=cache_enabled,
-                    output_artifact_configs={
-                        k: v.artifact_config
-                        for k, v in output_annotations.items()
-                    },
-                )
-                # Prepare Model Context
-                self._prepare_model_context_for_step()
-
-                # Parse the inputs for the entrypoint function.
-                function_params = self._parse_inputs(
-                    args=spec.args,
-                    annotations=spec.annotations,
-                    input_artifacts=input_artifacts,
-                )
 
-                self._link_pipeline_run_to_model_from_context(
-                    pipeline_run=pipeline_run
-                )
+            self._evaluate_artifact_names_in_collections(
+                step_run,
+                output_annotations,
+                [
+                    output_artifact_uris,
+                    output_materializers,
+                ],
+            )
 
-                step_failed = False
-                try:
-                    return_values = step_instance.call_entrypoint(
-                        **function_params
-                    )
-                except BaseException as step_exception:  # noqa: E722
-                    step_failed = True
-                    failure_hook_source = (
-                        self.configuration.failure_hook_source
-                    )
-                    if failure_hook_source:
+            self._stack.prepare_step_run(info=step_run_info)
+
+            # Initialize the step context singleton
+            StepContext._clear()
+            step_context = StepContext(
+                pipeline_run=pipeline_run,
+                step_run=step_run,
+                output_materializers=output_materializers,
+                output_artifact_uris=output_artifact_uris,
+                output_artifact_configs={
+                    k: v.artifact_config for k, v in output_annotations.items()
+                },
+            )
+
+            # Parse the inputs for the entrypoint function.
+            function_params = self._parse_inputs(
+                args=spec.args,
+                annotations=spec.annotations,
+                input_artifacts=input_artifacts,
+            )
+
+            step_failed = False
+            try:
+                return_values = step_instance.call_entrypoint(
+                    **function_params
+                )
+            except BaseException as step_exception:  # noqa: E722
+                step_failed = True
+                if not handle_bool_env_var(
+                    ENV_ZENML_IGNORE_FAILURE_HOOK, False
+                ):
+                    if (
+                        failure_hook_source
+                        := self.configuration.failure_hook_source
+                    ):
                         logger.info("Detected failure hook. Running...")
                         self.load_and_run_hook(
                             failure_hook_source,
                             step_exception=step_exception,
                         )
-                    raise
-                finally:
+                raise
+            finally:
+                try:
                     step_run_metadata = self._stack.get_step_run_metadata(
                         info=step_run_info,
                     )
@@ -223,10 +215,10 @@ class StepRunner:
                         info=step_run_info, step_failed=step_failed
                     )
                     if not step_failed:
-                        success_hook_source = (
-                            self.configuration.success_hook_source
-                        )
-                        if success_hook_source:
+                        if (
+                            success_hook_source
+                            := self.configuration.success_hook_source
+                        ):
                             logger.info("Detected success hook. Running...")
                             self.load_and_run_hook(
                                 success_hook_source,
@@ -245,7 +237,7 @@ class StepRunner:
                             is_enabled_on_step=step_run_info.config.enable_artifact_visualization,
                             is_enabled_on_pipeline=step_run_info.pipeline.enable_artifact_visualization,
                         )
-                        output_artifact_ids = self._store_output_artifacts(
+                        output_artifacts = self._store_output_artifacts(
                             output_data=output_data,
                             output_artifact_uris=output_artifact_uris,
                             output_materializers=output_materializers,
@@ -253,24 +245,63 @@ class StepRunner:
                             artifact_metadata_enabled=artifact_metadata_enabled,
                             artifact_visualization_enabled=artifact_visualization_enabled,
                         )
-                        link_step_artifacts_to_model(
-                            artifact_version_ids=output_artifact_ids
-                        )
-                        self._link_pipeline_run_to_model_from_artifacts(
-                            pipeline_run=pipeline_run,
-                            artifact_names=list(output_artifact_ids.keys()),
-                            external_artifacts=list(
-                                step_run.config.external_input_artifacts.values()
-                            ),
-                        )
+
+                        if (
+                            model_version := step_run.model_version
+                            or pipeline_run.model_version
+                        ):
+                            from zenml.orchestrators import step_run_utils
+
+                            step_run_utils.link_output_artifacts_to_model_version(
+                                artifacts={
+                                    k: [v] for k, v in output_artifacts.items()
+                                },
+                                model_version=model_version,
+                            )
+                finally:
+                    step_context._cleanup_registry.execute_callbacks(
+                        raise_on_exception=False
+                    )
                     StepContext._clear()  # Remove the step context singleton
 
             # Update the status and output artifacts of the step run.
+            output_artifact_ids = {
+                output_name: [
+                    artifact.id,
+                ]
+                for output_name, artifact in output_artifacts.items()
+            }
             publish_successful_step_run(
                 step_run_id=step_run_info.step_run_id,
                 output_artifact_ids=output_artifact_ids,
             )
 
+    def _evaluate_artifact_names_in_collections(
+        self,
+        step_run: "StepRunResponse",
+        output_annotations: Dict[str, OutputSignature],
+        collections: List[Dict[str, Any]],
+    ) -> None:
+        """Evaluates the artifact names in the collections.
+
+        Args:
+            step_run: The step run.
+            output_annotations: The output annotations of the step function
+                (also evaluated).
+            collections: The collections to evaluate.
+        """
+        collections.append(output_annotations)
+        for k, v in list(output_annotations.items()):
+            name = k
+            if v.artifact_config and v.artifact_config.name:
+                name = string_utils.format_name_template(
+                    v.artifact_config.name,
+                    substitutions=step_run.config.substitutions,
+                )
+
+            for d in collections:
+                d[name] = d.pop(k)
+
     def _load_step(self) -> "BaseStep":
         """Load the step instance.
 
@@ -297,10 +328,10 @@ class StepRunner:
             output_materializers = []
 
             for source in output.materializer_source:
-                materializer_class: Type[
-                    BaseMaterializer
-                ] = source_utils.load_and_validate_class(
-                    source, expected_class=BaseMaterializer
+                materializer_class: Type[BaseMaterializer] = (
+                    source_utils.load_and_validate_class(
+                        source, expected_class=BaseMaterializer
+                    )
                 )
                 output_materializers.append(materializer_class)
 
@@ -312,7 +343,7 @@ class StepRunner:
         self,
         args: List[str],
         annotations: Dict[str, Any],
-        input_artifacts: Dict[str, "ArtifactVersionResponse"],
+        input_artifacts: Dict[str, StepRunInputResponse],
     ) -> Dict[str, Any]:
         """Parses the inputs for a step entrypoint function.
 
@@ -336,17 +367,7 @@ class StepRunner:
             arg_type = annotations.get(arg, None)
             arg_type = resolve_type_annotation(arg_type)
 
-            if inspect.isclass(arg_type) and issubclass(arg_type, StepContext):
-                step_name = self.configuration.name
-                logger.warning(
-                    "Passing a `StepContext` as an argument to a step function "
-                    "is deprecated and will be removed in a future release. "
-                    f"Please adjust your '{step_name}' step to instead import "
-                    "the `StepContext` inside your step, as shown here: "
-                    "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/fetch-metadata-within-steps"
-                )
-                function_params[arg] = get_step_context()
-            elif arg in input_artifacts:
+            if arg in input_artifacts:
                 function_params[arg] = self._load_input_artifact(
                     input_artifacts[arg], arg_type
                 )
@@ -378,8 +399,6 @@ class StepRunner:
         Raises:
             TypeError: If hook function is passed a wrong parameter type.
         """
-        from zenml.steps import BaseParameters
-
         function_params: Dict[str, Any] = {}
 
         if args and args[0] == "self":
@@ -389,33 +408,13 @@ class StepRunner:
             arg_type = annotations.get(arg, None)
             arg_type = resolve_type_annotation(arg_type)
 
-            # Parse the parameters
-            if issubclass(arg_type, BaseParameters):
-                step_params = arg_type.parse_obj(
-                    self.configuration.parameters[arg]
-                )
-                function_params[arg] = step_params
-
-            # Parse the step context
-            elif issubclass(arg_type, StepContext):
-                step_name = self.configuration.name
-                logger.warning(
-                    "Passing a `StepContext` as an argument to a hook function "
-                    "is deprecated and will be removed in a future release. "
-                    f"Please adjust your '{step_name}' hook to instead import "
-                    "the `StepContext` inside your hook, as shown here: "
-                    "https://docs.zenml.io/user-guide/advanced-guide/pipelining-features/fetch-metadata-within-steps"
-                )
-                function_params[arg] = get_step_context()
-
-            elif issubclass(arg_type, BaseException):
+            if issubclass(arg_type, BaseException):
                 function_params[arg] = step_exception
-
             else:
                 # It should not be of any other type
                 raise TypeError(
                     "Hook functions can only take arguments of type "
-                    f"`BaseParameters`, or `BaseException`, not {arg_type}"
+                    f"`BaseException`, not {arg_type}"
                 )
 
         return function_params
@@ -434,21 +433,43 @@ class StepRunner:
         """
         # Skip materialization for `UnmaterializedArtifact`.
         if data_type == UnmaterializedArtifact:
-            return UnmaterializedArtifact.parse_obj(artifact)
+            return UnmaterializedArtifact(
+                **artifact.get_hydrated_version().model_dump()
+            )
 
-        if data_type is Any or is_union(get_origin(data_type)):
+        if data_type in (None, Any) or is_union(get_origin(data_type)):
             # Entrypoint function does not define a specific type for the input,
             # we use the datatype of the stored artifact
             data_type = source_utils.load(artifact.data_type)
 
-        materializer_class: Type[
-            BaseMaterializer
-        ] = source_utils.load_and_validate_class(
-            artifact.materializer, expected_class=BaseMaterializer
+        from zenml.orchestrators.utils import (
+            register_artifact_store_filesystem,
         )
-        materializer: BaseMaterializer = materializer_class(artifact.uri)
-        materializer.validate_type_compatibility(data_type)
-        return materializer.load(data_type=data_type)
+
+        materializer_class: Type[BaseMaterializer] = (
+            source_utils.load_and_validate_class(
+                artifact.materializer, expected_class=BaseMaterializer
+            )
+        )
+
+        def _load_artifact(artifact_store: "BaseArtifactStore") -> Any:
+            materializer: BaseMaterializer = materializer_class(
+                uri=artifact.uri, artifact_store=artifact_store
+            )
+            materializer.validate_load_type_compatibility(data_type)
+            return materializer.load(data_type=data_type)
+
+        if artifact.artifact_store_id == self._stack.artifact_store.id:
+            # Register the artifact store of the active stack here to avoid
+            # unnecessary component/flavor calls when using
+            # `register_artifact_store_filesystem(...)`
+            self._stack.artifact_store._register()
+            return _load_artifact(artifact_store=self._stack.artifact_store)
+        else:
+            with register_artifact_store_filesystem(
+                artifact.artifact_store_id
+            ) as target_artifact_store:
+                return _load_artifact(artifact_store=target_artifact_store)
 
     def _validate_outputs(
         self,
@@ -474,7 +495,7 @@ class StepRunner:
         if len(output_annotations) == 0:
             if return_values is not None:
                 raise StepInterfaceError(
-                    f"Wrong step function output type for step '{step_name}': "
+                    f"Wrong step function output type for step `{step_name}`: "
                     f"Expected no outputs but the function returned something: "
                     f"{return_values}."
                 )
@@ -490,7 +511,7 @@ class StepRunner:
         # or tuple.
         if not isinstance(return_values, (list, tuple)):
             raise StepInterfaceError(
-                f"Wrong step function output type for step '{step_name}': "
+                f"Wrong step function output type for step `{step_name}`: "
                 f"Expected multiple outputs ({output_annotations}) but "
                 f"the function did not return a list or tuple "
                 f"(actual return value: {return_values})."
@@ -506,8 +527,6 @@ class StepRunner:
                 f"(return values: {return_values})."
             )
 
-        from pydantic.typing import get_origin, is_union
-
         from zenml.steps.utils import get_args
 
         validated_outputs: Dict[str, Any] = {}
@@ -538,7 +557,7 @@ class StepRunner:
         output_annotations: Dict[str, OutputSignature],
         artifact_metadata_enabled: bool,
         artifact_visualization_enabled: bool,
-    ) -> Dict[str, UUID]:
+    ) -> Dict[str, "ArtifactVersionResponse"]:
         """Stores the output artifacts of the step.
 
         Args:
@@ -556,7 +575,7 @@ class StepRunner:
             The IDs of the published output artifacts.
         """
         step_context = get_step_context()
-        output_artifacts: Dict[str, UUID] = {}
+        artifact_requests = []
 
         for output_name, return_value in output_data.items():
             data_type = type(return_value)
@@ -579,11 +598,11 @@ class StepRunner:
                 ].default_materializer_source
 
                 if default_materializer_source:
-                    default_materializer_class: Type[
-                        BaseMaterializer
-                    ] = source_utils.load_and_validate_class(
-                        default_materializer_source,
-                        expected_class=BaseMaterializer,
+                    default_materializer_class: Type[BaseMaterializer] = (
+                        source_utils.load_and_validate_class(
+                            default_materializer_source,
+                            expected_class=BaseMaterializer,
+                        )
                     )
                     materializer_registry.default_materializer = (
                         default_materializer_class
@@ -594,12 +613,13 @@ class StepRunner:
             uri = output_artifact_uris[output_name]
             artifact_config = output_annotations[output_name].artifact_config
 
+            artifact_type = None
             if artifact_config is not None:
                 has_custom_name = bool(artifact_config.name)
                 version = artifact_config.version
-                tags = artifact_config.tags
+                artifact_type = artifact_config.artifact_type
             else:
-                has_custom_name, version, tags = False, None, None
+                has_custom_name, version = False, None
 
             # Override the artifact name if it is not a custom name.
             if has_custom_name:
@@ -615,138 +635,29 @@ class StepRunner:
             # Get metadata that the user logged manually
             user_metadata = step_context.get_output_metadata(output_name)
 
-            artifact = save_artifact(
+            # Get full set of tags
+            tags = step_context.get_output_tags(output_name)
+
+            artifact_request = _store_artifact_data_and_prepare_request(
                 name=artifact_name,
                 data=return_value,
-                materializer=materializer_class,
+                materializer_class=materializer_class,
                 uri=uri,
-                extract_metadata=artifact_metadata_enabled,
-                include_visualizations=artifact_visualization_enabled,
+                artifact_type=artifact_type,
+                store_metadata=artifact_metadata_enabled,
+                store_visualizations=artifact_visualization_enabled,
                 has_custom_name=has_custom_name,
                 version=version,
                 tags=tags,
-                user_metadata=user_metadata,
-                manual_save=False,
-            )
-            output_artifacts[output_name] = artifact.id
-
-        return output_artifacts
-
-    def _prepare_model_context_for_step(self) -> None:
-        try:
-            model = get_step_context().model
-            model._get_or_create_model_version()
-        except StepContextError:
-            return
-
-    def _get_model_versions_from_artifacts(
-        self,
-        artifact_names: List[str],
-    ) -> Set[Tuple[UUID, UUID]]:
-        """Gets the model versions from the artifacts.
-
-        Args:
-            artifact_names: The names of the published output artifacts.
-
-        Returns:
-            Set of tuples of (model_id, model_version_id).
-        """
-        models = set()
-        for artifact_name in artifact_names:
-            artifact_config = (
-                get_step_context()._get_output(artifact_name).artifact_config
+                save_type=ArtifactSaveType.STEP_OUTPUT,
+                metadata=user_metadata,
             )
-            if artifact_config is not None:
-                if (model := artifact_config._model) is not None:
-                    model_version_response = (
-                        model._get_or_create_model_version()
-                    )
-                    models.add(
-                        (
-                            model_version_response.model.id,
-                            model_version_response.id,
-                        )
-                    )
-                else:
-                    break
-        return models
+            artifact_requests.append(artifact_request)
 
-    def _get_model_versions_from_config(self) -> Set[Tuple[UUID, UUID]]:
-        """Gets the model versions from the step model version.
-
-        Returns:
-            Set of tuples of (model_id, model_version_id).
-        """
-        try:
-            mc = get_step_context().model
-            model_version = mc._get_or_create_model_version()
-            return {(model_version.model.id, model_version.id)}
-        except StepContextError:
-            return set()
-
-    def _link_pipeline_run_to_model_from_context(
-        self,
-        pipeline_run: "PipelineRunResponse",
-    ) -> None:
-        """Links the pipeline run to the model version using artifacts data.
-
-        Args:
-            pipeline_run: The response model of current pipeline run.
-        """
-        from zenml.models import ModelVersionPipelineRunRequest
-
-        models = self._get_model_versions_from_config()
-
-        client = Client()
-        for model in models:
-            client.zen_store.create_model_version_pipeline_run_link(
-                ModelVersionPipelineRunRequest(
-                    user=Client().active_user.id,
-                    workspace=Client().active_workspace.id,
-                    pipeline_run=pipeline_run.id,
-                    model=model[0],
-                    model_version=model[1],
-                )
-            )
-
-    def _link_pipeline_run_to_model_from_artifacts(
-        self,
-        pipeline_run: "PipelineRunResponse",
-        artifact_names: List[str],
-        external_artifacts: List["ExternalArtifactConfiguration"],
-    ) -> None:
-        """Links the pipeline run to the model version using artifacts data.
-
-        Args:
-            pipeline_run: The response model of current pipeline run.
-            artifact_names: The name of the published output artifacts.
-            external_artifacts: The external artifacts of the step.
-        """
-        from zenml.models import ModelVersionPipelineRunRequest
-
-        models = self._get_model_versions_from_artifacts(artifact_names)
-        client = Client()
-
-        # Add models from external artifacts
-        for external_artifact in external_artifacts:
-            if external_artifact.model:
-                models.add(
-                    (
-                        external_artifact.model.model_id,
-                        external_artifact.model.id,
-                    )
-                )
-
-        for model in models:
-            client.zen_store.create_model_version_pipeline_run_link(
-                ModelVersionPipelineRunRequest(
-                    user=client.active_user.id,
-                    workspace=client.active_workspace.id,
-                    pipeline_run=pipeline_run.id,
-                    model=model[0],
-                    model_version=model[1],
-                )
-            )
+        responses = Client().zen_store.batch_create_artifact_versions(
+            artifact_requests
+        )
+        return dict(zip(output_data.keys(), responses))
 
     def load_and_run_hook(
         self,