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

zenml/orchestrators/step_runner.py

@@ -26,7 +26,6 @@ from typing import (
     Tuple,
     Type,
 )
-from uuid import UUID
 
 from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
 from zenml.artifacts.utils import save_artifact
@@ -41,20 +40,14 @@ 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.orchestrators.publish_utils import (
     publish_step_run_metadata,
     publish_successful_step_run,
 )
 from zenml.orchestrators.utils import (
-    _link_pipeline_run_to_model_from_artifacts,
-    _link_pipeline_run_to_model_from_context,
     is_setting_enabled,
 )
-from zenml.steps.step_environment import StepEnvironment
+from zenml.steps.step_context import StepContext, get_step_context
 from zenml.steps.utils import (
     OutputSignature,
     parse_return_type_annotations,
@@ -64,6 +57,7 @@ from zenml.utils import materializer_utils, source_utils
 from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
+    from zenml.artifact_stores import BaseArtifactStore
     from zenml.config.source import Source
     from zenml.config.step_configurations import Step
     from zenml.models import (
@@ -135,7 +129,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(
@@ -150,70 +145,53 @@ 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()
-                    },
-                )
 
-                # Parse the inputs for the entrypoint function.
-                function_params = self._parse_inputs(
-                    args=spec.args,
-                    annotations=spec.annotations,
-                    input_artifacts=input_artifacts,
-                )
+            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,
+                output_artifact_configs={
+                    k: v.artifact_config for k, v in output_annotations.items()
+                },
+            )
 
-                _link_pipeline_run_to_model_from_context(
-                    pipeline_run_id=pipeline_run.id
-                )
+            # 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
+            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
                     ):
-                        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:
+                        logger.info("Detected failure hook. Running...")
+                        self.load_and_run_hook(
+                            failure_hook_source,
+                            step_exception=step_exception,
+                        )
+                raise
+            finally:
+                try:
                     step_run_metadata = self._stack.get_step_run_metadata(
                         info=step_run_info,
                     )
@@ -247,7 +225,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,
@@ -255,19 +233,26 @@ 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
-                        )
-                        _link_pipeline_run_to_model_from_artifacts(
-                            pipeline_run_id=pipeline_run.id,
-                            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=output_artifacts,
+                                output_configurations=step_run.config.outputs,
+                                model_version=model_version,
+                            )
+                finally:
                     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,
@@ -338,17 +323,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/how-to/track-metrics-metadata/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
                 )
@@ -380,8 +355,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":
@@ -391,33 +364,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.model_validate(
-                    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/how-to/track-metrics-metadata/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
@@ -455,15 +408,25 @@ class StepRunner:
             )
         )
 
-        with register_artifact_store_filesystem(
-            artifact.artifact_store_id
-        ) as target_artifact_store:
+        def _load_artifact(artifact_store: "BaseArtifactStore") -> Any:
             materializer: BaseMaterializer = materializer_class(
-                uri=artifact.uri, artifact_store=target_artifact_store
+                uri=artifact.uri, artifact_store=artifact_store
             )
-            materializer.validate_type_compatibility(data_type)
+            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,
         return_values: Any,
@@ -550,7 +513,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:
@@ -568,7 +531,7 @@ class StepRunner:
             The IDs of the published output artifacts.
         """
         step_context = get_step_context()
-        output_artifacts: Dict[str, UUID] = {}
+        output_artifacts: Dict[str, "ArtifactVersionResponse"] = {}
 
         for output_name, return_value in output_data.items():
             data_type = type(return_value)
@@ -642,7 +605,7 @@ class StepRunner:
                 user_metadata=user_metadata,
                 manual_save=False,
             )
-            output_artifacts[output_name] = artifact.id
+            output_artifacts[output_name] = artifact
 
         return output_artifacts
 

zenml/orchestrators/utils.py

@@ -15,36 +15,28 @@
 
 import os
 import random
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, cast
+from typing import TYPE_CHECKING, Any, Dict, Optional, cast
 from uuid import UUID
 
 from zenml.client import Client
 from zenml.config.global_config import (
     GlobalConfiguration,
 )
-from zenml.config.source import Source
 from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
+    ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
     ENV_ZENML_SERVER,
     ENV_ZENML_STORE_PREFIX,
     PIPELINE_API_TOKEN_EXPIRES_MINUTES,
 )
-from zenml.enums import StackComponentType, StoreType
-from zenml.exceptions import StepContextError
+from zenml.enums import AuthScheme, StackComponentType, StoreType
 from zenml.logger import get_logger
-from zenml.model.utils import link_artifact_config_to_model
-from zenml.models.v2.core.step_run import StepRunRequest
-from zenml.new.steps.step_context import get_step_context
 from zenml.stack import StackComponent
 from zenml.utils.string_utils import format_name_template
 
 if TYPE_CHECKING:
     from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
-    from zenml.artifacts.external_artifact_config import (
-        ExternalArtifactConfiguration,
-    )
-    from zenml.model.model import Model
     from zenml.models import PipelineDeploymentResponse
 
 
@@ -102,32 +94,54 @@ def get_config_environment_vars(
     Returns:
         Environment variable dict.
     """
+    from zenml.login.credentials_store import get_credentials_store
     from zenml.zen_stores.rest_zen_store import RestZenStore
 
     global_config = GlobalConfiguration()
     environment_vars = global_config.get_config_environment_vars()
 
-    if deployment and global_config.store_configuration.type == StoreType.REST:
-        # When connected to a ZenML server, if a pipeline deployment is
-        # supplied, we need to fetch an API token that will be valid for the
-        # duration of the pipeline run.
-        assert isinstance(global_config.zen_store, RestZenStore)
-        pipeline_id: Optional[UUID] = None
-        if deployment.pipeline:
-            pipeline_id = deployment.pipeline.id
-        schedule_id: Optional[UUID] = None
-        expires_minutes: Optional[int] = PIPELINE_API_TOKEN_EXPIRES_MINUTES
-        if deployment.schedule:
-            schedule_id = deployment.schedule.id
-            # If a schedule is given, this is a long running pipeline that
-            # should not have an API token that expires.
-            expires_minutes = None
-        api_token = global_config.zen_store.get_api_token(
-            pipeline_id=pipeline_id,
-            schedule_id=schedule_id,
-            expires_minutes=expires_minutes,
-        )
-        environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = api_token
+    if (
+        global_config.store_configuration.type == StoreType.REST
+        and global_config.zen_store.get_store_info().auth_scheme
+        != AuthScheme.NO_AUTH
+    ):
+        credentials_store = get_credentials_store()
+        url = global_config.store_configuration.url
+        api_key = credentials_store.get_api_key(url)
+        api_token = credentials_store.get_token(url, allow_expired=False)
+        if api_key:
+            environment_vars[ENV_ZENML_STORE_PREFIX + "API_KEY"] = api_key
+        elif deployment:
+            # When connected to an authenticated ZenML server, if a pipeline
+            # deployment is supplied, we need to fetch an API token that will be
+            # valid for the duration of the pipeline run.
+            assert isinstance(global_config.zen_store, RestZenStore)
+            pipeline_id: Optional[UUID] = None
+            if deployment.pipeline:
+                pipeline_id = deployment.pipeline.id
+            schedule_id: Optional[UUID] = None
+            expires_minutes: Optional[int] = PIPELINE_API_TOKEN_EXPIRES_MINUTES
+            if deployment.schedule:
+                schedule_id = deployment.schedule.id
+                # If a schedule is given, this is a long running pipeline that
+                # should not have an API token that expires.
+                expires_minutes = None
+            new_api_token = global_config.zen_store.get_api_token(
+                pipeline_id=pipeline_id,
+                schedule_id=schedule_id,
+                expires_minutes=expires_minutes,
+            )
+            environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
+                new_api_token
+            )
+        elif api_token:
+            environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
+                api_token.access_token
+            )
+
+    # Disable credentials caching to avoid storing sensitive information
+    # in the pipeline run environment
+    environment_vars[ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING] = "true"
 
     # Make sure to use the correct active stack/workspace which might come
     # from a .zen repository and not the global config
@@ -161,154 +175,6 @@ def get_run_name(run_name_template: str) -> str:
     return run_name
 
 
-def _link_pipeline_run_to_model_from_context(
-    pipeline_run_id: "UUID", model: Optional["Model"] = None
-) -> None:
-    """Links the pipeline run to the model version using artifacts data.
-
-    Args:
-        pipeline_run_id: The ID of the current pipeline run.
-        model: Model configured in the step
-    """
-    from zenml.models import ModelVersionPipelineRunRequest
-
-    if not model:
-        model_id, model_version_id = _get_model_versions_from_config()
-    else:
-        model_id, model_version_id = model.model_id, model.id
-
-    if model_id and model_version_id:
-        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_id,
-                model_version=model_version_id,
-            )
-        )
-
-
-def _get_model_versions_from_config() -> Tuple[Optional[UUID], Optional[UUID]]:
-    """Gets the model versions from the step model version.
-
-    Returns:
-        Tuple of (model_id, model_version_id).
-    """
-    try:
-        mc = get_step_context().model
-        return mc.model_id, mc.id
-    except StepContextError:
-        return None, None
-
-
-def _link_cached_artifacts_to_model(
-    model_from_context: Optional["Model"],
-    step_run: StepRunRequest,
-    step_source: Source,
-) -> None:
-    """Links the output artifacts of the cached step to the model version in Control Plane.
-
-    Args:
-        model_from_context: The model version of the current step.
-        step_run: The step to run.
-        step_source: The source of the step.
-    """
-    from zenml.artifacts.artifact_config import ArtifactConfig
-    from zenml.steps.base_step import BaseStep
-    from zenml.steps.utils import parse_return_type_annotations
-
-    step_instance = BaseStep.load_from_source(step_source)
-    output_annotations = parse_return_type_annotations(
-        step_instance.entrypoint
-    )
-    for output_name_, output_id in step_run.outputs.items():
-        artifact_config_ = None
-        if output_name_ in output_annotations:
-            annotation = output_annotations.get(output_name_, None)
-            if annotation and annotation.artifact_config is not None:
-                artifact_config_ = annotation.artifact_config.model_copy()
-        # no artifact config found or artifact was produced by `save_artifact`
-        # inside the step body, so was never in annotations
-        if artifact_config_ is None:
-            artifact_config_ = ArtifactConfig(name=output_name_)
-
-        link_artifact_config_to_model(
-            artifact_config=artifact_config_,
-            model=model_from_context,
-            artifact_version_id=output_id,
-        )
-
-
-def _link_pipeline_run_to_model_from_artifacts(
-    pipeline_run_id: UUID,
-    artifact_names: List[str],
-    external_artifacts: List["ExternalArtifactConfiguration"],
-) -> None:
-    """Links the pipeline run to the model version using artifacts data.
-
-    Args:
-        pipeline_run_id: The ID of the 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 = _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],
-            )
-        )
-
-
-def _get_model_versions_from_artifacts(
-    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
-        )
-        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
-
-
 class register_artifact_store_filesystem:
     """Context manager for the artifact_store/filesystem_registry dependency.
 

zenml/pipelines/__init__.py

@@ -1,4 +1,4 @@
-#  Copyright (c) ZenML GmbH 2021. All Rights Reserved.
+#  Copyright (c) ZenML GmbH 2023. All Rights Reserved.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -11,21 +11,8 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""A ZenML pipeline consists of tasks that execute in order and yield artifacts.
-
-The artifacts are automatically stored within the artifact store and metadata 
-is tracked by ZenML. Each individual task within a pipeline is known as a
-step. The standard pipelines within ZenML are designed to have easy interfaces
-to add pre-decided steps, with the order also pre-decided. Other sorts of
-pipelines can be created as well from scratch, building on the `BasePipeline` class.
-
-Pipelines can be written as simple functions. They are created by using decorators appropriate to the specific use case you have. The moment it is `run`, a pipeline is compiled and passed directly to the orchestrator.
-"""
-
-
-from zenml.config import DockerSettings
 from zenml.config.schedule import Schedule
-from zenml.pipelines.base_pipeline import BasePipeline
 from zenml.pipelines.pipeline_decorator import pipeline
+from zenml.pipelines.pipeline_context import PipelineContext, get_pipeline_context
 
-__all__ = ["BasePipeline", "DockerSettings", "pipeline", "Schedule"]
+__all__ = ["pipeline", "Schedule", "PipelineContext", "get_pipeline_context"]

zenml/{new/pipelines → pipelines}/pipeline_context.py

@@ -13,12 +13,11 @@
 #  permissions and limitations under the License.
 """Pipeline context class."""
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING
 
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
-    from zenml import Model
     from zenml.config.pipeline_configurations import PipelineConfiguration
 
 logger = get_logger(__name__)
@@ -34,11 +33,11 @@ def get_pipeline_context() -> "PipelineContext":
         RuntimeError: If no active pipeline is found.
         RuntimeError: If inside a running step.
     """
-    from zenml.new.pipelines.pipeline import Pipeline
+    from zenml.pipelines.pipeline_definition import Pipeline
 
     if Pipeline.ACTIVE_PIPELINE is None:
         try:
-            from zenml.new.steps.step_context import get_step_context
+            from zenml.steps.step_context import get_step_context
 
             get_step_context()
         except RuntimeError:
@@ -113,17 +112,3 @@ class PipelineContext:
         self.settings = pipeline_configuration.settings
         self.extra = pipeline_configuration.extra
         self.model = pipeline_configuration.model
-        self._model_version = pipeline_configuration.model
-
-    # TODO: deprecate me
-    @property
-    def model_version(self) -> Optional["Model"]:
-        """DEPRECATED, use `model` instead.
-
-        Returns:
-            The `Model` object associated with the current pipeline.
-        """
-        logger.warning(
-            "Pipeline context `model_version` is deprecated. Please use `model` instead."
-        )
-        return self.model