zenml-nightly 0.57.1.dev20240522__py3-none-any.whl → 0.57.1.dev20240527__py3-none-any.whl

zenml/new/steps/step_decorator.py

@@ -33,6 +33,7 @@ if TYPE_CHECKING:
     from types import FunctionType
 
     from zenml.config.base_settings import SettingsOrDict
+    from zenml.config.retry_config import StepRetryConfig
     from zenml.config.source import Source
     from zenml.materializers.base_materializer import BaseMaterializer
     from zenml.model.model import Model
@@ -72,6 +73,7 @@ def step(
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
+    retry: Optional["StepRetryConfig"] = None,
     model_version: Optional["Model"] = None,  # TODO: deprecate me
 ) -> Callable[["F"], "BaseStep"]: ...
 
@@ -92,6 +94,7 @@ def step(
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
     model: Optional["Model"] = None,
+    retry: Optional["StepRetryConfig"] = None,
     model_version: Optional["Model"] = None,  # TODO: deprecate me
 ) -> Union["BaseStep", Callable[["F"], "BaseStep"]]:
     """Decorator to create a ZenML step.
@@ -123,6 +126,7 @@ def step(
             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.
         model_version: DEPRECATED, please use `model` instead.
 
     Returns:
@@ -162,6 +166,7 @@ def step(
             on_failure=on_failure,
             on_success=on_success,
             model=model or model_version,
+            retry=retry,
         )
 
         return step_instance

zenml/orchestrators/step_launcher.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Class to launch (run directly or using a step operator) steps."""
 
+import os
 import time
 from contextlib import nullcontext
 from datetime import datetime
@@ -23,6 +24,7 @@ from zenml.config.step_configurations import Step
 from zenml.config.step_run_info import StepRunInfo
 from zenml.constants import (
     ENV_ZENML_DISABLE_STEP_LOGS_STORAGE,
+    ENV_ZENML_IGNORE_FAILURE_HOOK,
     STEP_SOURCE_PARAMETER_NAME,
     TEXT_FIELD_MAX_LENGTH,
     handle_bool_env_var,
@@ -31,7 +33,6 @@ from zenml.enums import ExecutionStatus
 from zenml.environment import get_run_environment_dict
 from zenml.logger import get_logger
 from zenml.logging import step_logging
-from zenml.model.utils import link_artifact_config_to_model
 from zenml.models import (
     ArtifactVersionResponse,
     LogsRequest,
@@ -53,7 +54,6 @@ from zenml.stack import Stack
 from zenml.utils import string_utils
 
 if TYPE_CHECKING:
-    from zenml.model.model import Model
     from zenml.step_operators import BaseStepOperator
 
 logger = get_logger(__name__)
@@ -226,20 +226,56 @@ class StepLauncher:
 
                 logger.info(f"Step `{self._step_name}` has started.")
                 if execution_needed:
-                    try:
-                        self._run_step(
-                            pipeline_run=pipeline_run,
-                            step_run=step_run_response,
-                        )
-                    except BaseException as e:  # noqa: E722
-                        logger.error(
-                            f"Failed to run step `{self._step_name}`."
-                        )
-                        logger.exception(e)
-                        publish_utils.publish_failed_step_run(
-                            step_run_response.id
-                        )
-                        raise
+                    retries = 0
+                    last_retry = True
+                    max_retries = (
+                        step_run_response.config.retry.max_retries
+                        if step_run_response.config.retry
+                        else 1
+                    )
+                    delay = (
+                        step_run_response.config.retry.delay
+                        if step_run_response.config.retry
+                        else 0
+                    )
+                    backoff = (
+                        step_run_response.config.retry.backoff
+                        if step_run_response.config.retry
+                        else 1
+                    )
+
+                    while retries < max_retries:
+                        last_retry = retries == max_retries - 1
+                        try:
+                            self._run_step(
+                                pipeline_run=pipeline_run,
+                                step_run=step_run_response,
+                                last_retry=last_retry,
+                            )
+                            logger.info(
+                                f"Step '{self._step_name}' completed successfully."
+                            )
+                            break
+                        except BaseException as e:  # noqa: E722
+                            retries += 1
+                            if retries < max_retries:
+                                logger.error(
+                                    f"Failed to run step '{self._step_name}'. Retrying..."
+                                )
+                                logger.exception(e)
+                                logger.info(
+                                    f"Sleeping for {delay} seconds before retrying."
+                                )
+                                time.sleep(delay)
+                                delay *= backoff
+                            else:
+                                logger.error(
+                                    f"Failed to run step '{self._step_name}' after {max_retries} retries. Exiting."
+                                )
+                                publish_utils.publish_failed_step_run(
+                                    step_run_response.id
+                                )
+                                raise
 
         except:  # noqa: E722
             logger.error(f"Pipeline run `{pipeline_run.name}` failed.")
@@ -360,61 +396,33 @@ class StepLauncher:
                     output_name: artifact.id
                     for output_name, artifact in cached_outputs.items()
                 }
-                self._link_cached_artifacts_to_model(
+                orchestrator_utils._link_cached_artifacts_to_model(
                     model_from_context=model,
                     step_run=step_run,
+                    step_source=self._step.spec.source,
                 )
+                if self._step.config.model:
+                    orchestrator_utils._link_pipeline_run_to_model_from_context(
+                        pipeline_run_id=step_run.pipeline_run_id,
+                        model=self._step.config.model,
+                    )
                 step_run.status = ExecutionStatus.CACHED
                 step_run.end_time = step_run.start_time
 
         return execution_needed, step_run
 
-    def _link_cached_artifacts_to_model(
-        self,
-        model_from_context: Optional["Model"],
-        step_run: StepRunRequest,
-    ) -> 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.
-        """
-        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(self._step.spec.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.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 _run_step(
         self,
         pipeline_run: PipelineRunResponse,
         step_run: StepRunResponse,
+        last_retry: bool = True,
     ) -> None:
         """Runs the current step.
 
         Args:
             pipeline_run: The model of the current pipeline run.
             step_run: The model of the current step run.
+            last_retry: Whether this is the last retry of the step.
         """
         # Prepare step run information.
         step_run_info = StepRunInfo(
@@ -437,6 +445,7 @@ class StepLauncher:
                 self._run_step_with_step_operator(
                     step_operator_name=self._step.config.step_operator,
                     step_run_info=step_run_info,
+                    last_retry=last_retry,
                 )
             else:
                 self._run_step_without_step_operator(
@@ -445,6 +454,7 @@ class StepLauncher:
                     step_run_info=step_run_info,
                     input_artifacts=step_run.inputs,
                     output_artifact_uris=output_artifact_uris,
+                    last_retry=last_retry,
                 )
         except:  # noqa: E722
             output_utils.remove_artifact_dirs(
@@ -462,12 +472,14 @@ class StepLauncher:
         self,
         step_operator_name: str,
         step_run_info: StepRunInfo,
+        last_retry: bool,
     ) -> None:
         """Runs the current step with a step operator.
 
         Args:
             step_operator_name: The name of the step operator to use.
             step_run_info: Additional information needed to run the step.
+            last_retry: Whether this is the last retry of the step.
         """
         step_operator = _get_step_operator(
             stack=self._stack,
@@ -485,6 +497,8 @@ class StepLauncher:
         environment = orchestrator_utils.get_config_environment_vars(
             deployment=self._deployment
         )
+        if last_retry:
+            environment[ENV_ZENML_IGNORE_FAILURE_HOOK] = str(False)
         logger.info(
             "Using step operator `%s` to run step `%s`.",
             step_operator.name,
@@ -503,6 +517,7 @@ class StepLauncher:
         step_run_info: StepRunInfo,
         input_artifacts: Dict[str, ArtifactVersionResponse],
         output_artifact_uris: Dict[str, str],
+        last_retry: bool,
     ) -> None:
         """Runs the current step without a step operator.
 
@@ -512,7 +527,10 @@ class StepLauncher:
             step_run_info: Additional information needed to run the step.
             input_artifacts: The input artifact versions of the current step.
             output_artifact_uris: The output artifact URIs of the current step.
+            last_retry: Whether this is the last retry of the step.
         """
+        if last_retry:
+            os.environ[ENV_ZENML_IGNORE_FAILURE_HOOK] = "false"
         runner = StepRunner(step=self._step, stack=self._stack)
         runner.run(
             pipeline_run=pipeline_run,

zenml/orchestrators/step_runner.py

@@ -23,7 +23,6 @@ from typing import (
     Dict,
     List,
     Optional,
-    Set,
     Tuple,
     Type,
 )
@@ -33,11 +32,11 @@ from pydantic.typing import get_origin, is_union
 
 from zenml.artifacts.unmaterialized_artifact import UnmaterializedArtifact
 from zenml.artifacts.utils import save_artifact
-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
@@ -52,7 +51,11 @@ from zenml.orchestrators.publish_utils import (
     publish_step_run_metadata,
     publish_successful_step_run,
 )
-from zenml.orchestrators.utils import is_setting_enabled
+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.utils import (
     OutputSignature,
@@ -62,9 +65,6 @@ from zenml.steps.utils import (
 from zenml.utils import materializer_utils, source_utils
 
 if TYPE_CHECKING:
-    from zenml.artifacts.external_artifact_config import (
-        ExternalArtifactConfiguration,
-    )
     from zenml.config.source import Source
     from zenml.config.step_configurations import Step
     from zenml.models import (
@@ -192,8 +192,8 @@ class StepRunner:
                     input_artifacts=input_artifacts,
                 )
 
-                self._link_pipeline_run_to_model_from_context(
-                    pipeline_run=pipeline_run
+                _link_pipeline_run_to_model_from_context(
+                    pipeline_run_id=pipeline_run.id
                 )
 
                 step_failed = False
@@ -203,15 +203,18 @@ class StepRunner:
                     )
                 except BaseException as step_exception:  # noqa: E722
                     step_failed = True
-                    failure_hook_source = (
-                        self.configuration.failure_hook_source
-                    )
-                    if failure_hook_source:
-                        logger.info("Detected failure hook. Running...")
-                        self.load_and_run_hook(
-                            failure_hook_source,
-                            step_exception=step_exception,
-                        )
+                    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:
                     step_run_metadata = self._stack.get_step_run_metadata(
@@ -225,10 +228,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,
@@ -258,8 +261,8 @@ class StepRunner:
                         link_step_artifacts_to_model(
                             artifact_version_ids=output_artifact_ids
                         )
-                        self._link_pipeline_run_to_model_from_artifacts(
-                            pipeline_run=pipeline_run,
+                        _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()
@@ -645,115 +648,6 @@ class StepRunner:
         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
-            )
-            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
-
-    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],
-                )
-            )
-
     def load_and_run_hook(
         self,
         hook_source: "Source",

zenml/orchestrators/utils.py

@@ -15,13 +15,14 @@
 
 import random
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple
 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,
@@ -29,8 +30,16 @@ from zenml.constants import (
     PIPELINE_API_TOKEN_EXPIRES_MINUTES,
 )
 from zenml.enums import StoreType
+from zenml.exceptions import StepContextError
+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
 
 if TYPE_CHECKING:
+    from zenml.artifacts.external_artifact_config import (
+        ExternalArtifactConfiguration,
+    )
+    from zenml.model.model import Model
     from zenml.models import PipelineDeploymentResponse
 
 
@@ -148,3 +157,151 @@ def get_run_name(run_name_template: str) -> str:
         raise ValueError("Empty run names are not allowed.")
 
     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.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