zenml-nightly 0.83.1.dev20250708__py3-none-any.whl → 0.83.1.dev20250710__py3-none-any.whl

zenml/login/web_login.py

@@ -21,6 +21,7 @@ from typing import Optional, Union
 import requests
 
 from zenml import __version__
+from zenml.cli import utils as cli_utils
 from zenml.config.global_config import GlobalConfiguration
 from zenml.constants import (
     API,
@@ -175,11 +176,17 @@ def web_login(
         # URL to it
         verification_uri = base_url + verification_uri
     webbrowser.open(verification_uri)
-    logger.info(
-        f"If your browser did not open automatically, please open the "
-        f"following URL into your browser to proceed with the authentication:"
-        f"\n\n{verification_uri}\n"
+
+    # Display the verification URL without panel styling
+    from zenml.console import console
+
+    console.print()
+    console.print(
+        "If your browser did not open automatically, please open the following URL into your browser to proceed with the authentication:",
+        style="white",
     )
+    console.print(verification_uri, style="bright_blue underline")
+    console.print()
 
     # Poll the OAuth2 server until the user has authorized the device
     token_request = OAuthDeviceTokenRequest(
@@ -201,9 +208,9 @@ def web_login(
             # The user has authorized the device, so we can extract the access token
             token_response = OAuthTokenResponse(**response.json())
             if zenml_pro:
-                logger.info("Successfully logged in to ZenML Pro.")
+                cli_utils.success("✔ Successfully logged in to ZenML Pro.")
             else:
-                logger.info(f"Successfully logged in to {url}.")
+                cli_utils.success(f"✔ Successfully logged in to {url}.")
             break
         elif response.status_code == 400:
             try:

zenml/models/v2/core/model_version.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing model versions."""
 
+import json
 from typing import (
     TYPE_CHECKING,
     ClassVar,
@@ -411,10 +412,17 @@ class ModelVersionResponse(
 
         from zenml.client import Client
 
+        data_artifact_types = [
+            value
+            for value in ArtifactType.values()
+            if value
+            not in [ArtifactType.MODEL.value, ArtifactType.SERVICE.value]
+        ]
+
         artifact_versions = pagination_utils.depaginate(
             Client().list_artifact_versions,
             model_version_id=self.id,
-            type=ArtifactType.DATA,
+            type="oneof:" + json.dumps(data_artifact_types),
             project=self.project_id,
         )
 

zenml/models/v2/core/pipeline_run.py

@@ -497,6 +497,7 @@ class PipelineRunResponse(
                 Client().list_run_steps,
                 pipeline_run_id=self.id,
                 project=self.project_id,
+                exclude_retried=True,
             )
         }
 

zenml/models/v2/core/step_run.py

@@ -177,6 +177,12 @@ class StepRunResponseBody(ProjectScopedResponseBody):
     """Response body for step runs."""
 
     status: ExecutionStatus = Field(title="The status of the step.")
+    version: int = Field(
+        title="The version of the step run.",
+    )
+    is_retriable: bool = Field(
+        title="Whether the step run is retriable.",
+    )
     start_time: Optional[datetime] = Field(
         title="The start time of the step run.",
         default=None,
@@ -420,6 +426,24 @@ class StepRunResponse(
         """
         return self.get_body().status
 
+    @property
+    def version(self) -> int:
+        """The `version` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().version
+
+    @property
+    def is_retriable(self) -> bool:
+        """The `is_retriable` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().is_retriable
+
     @property
     def inputs(self) -> Dict[str, List[StepRunInputResponse]]:
         """The `inputs` property.
@@ -602,6 +626,7 @@ class StepRunFilter(ProjectScopedFilter, RunMetadataFilterMixin):
         *ProjectScopedFilter.FILTER_EXCLUDE_FIELDS,
         *RunMetadataFilterMixin.FILTER_EXCLUDE_FIELDS,
         "model",
+        "exclude_retried",
     ]
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *ProjectScopedFilter.CLI_EXCLUDE_FIELDS,
@@ -666,6 +691,10 @@ class StepRunFilter(ProjectScopedFilter, RunMetadataFilterMixin):
         default=None,
         description="Name/ID of the model associated with the step run.",
     )
+    exclude_retried: Optional[bool] = Field(
+        default=None,
+        description="Whether to exclude retried step runs.",
+    )
     model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
@@ -681,7 +710,7 @@ class StepRunFilter(ProjectScopedFilter, RunMetadataFilterMixin):
         """
         custom_filters = super().get_custom_filters(table)
 
-        from sqlmodel import and_
+        from sqlmodel import and_, col
 
         from zenml.zen_stores.schemas import (
             ModelSchema,
@@ -699,4 +728,9 @@ class StepRunFilter(ProjectScopedFilter, RunMetadataFilterMixin):
             )
             custom_filters.append(model_filter)
 
+        if self.exclude_retried:
+            custom_filters.append(
+                col(StepRunSchema.status) != ExecutionStatus.RETRIED.value
+            )
+
         return custom_filters

zenml/orchestrators/base_orchestrator.py

@@ -13,6 +13,7 @@
 #  permissions and limitations under the License.
 """Base orchestrator class."""
 
+import time
 from abc import ABC, abstractmethod
 from typing import (
     TYPE_CHECKING,
@@ -33,7 +34,7 @@ from zenml.constants import (
     handle_bool_env_var,
 )
 from zenml.enums import ExecutionStatus, StackComponentType
-from zenml.exceptions import RunMonitoringError
+from zenml.exceptions import RunMonitoringError, RunStoppedException
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
 from zenml.orchestrators.publish_utils import (
@@ -127,6 +128,15 @@ class BaseOrchestratorConfig(StackComponentConfig):
         """
         return True
 
+    @property
+    def handles_step_retries(self) -> bool:
+        """Whether the orchestrator handles step retries.
+
+        Returns:
+            Whether the orchestrator handles step retries.
+        """
+        return False
+
 
 class BaseOrchestrator(StackComponent, ABC):
     """Base class for all orchestrators."""
@@ -346,14 +356,59 @@ class BaseOrchestrator(StackComponent, ABC):
 
         Args:
             step: The step to run.
+
+        Raises:
+            RunStoppedException: If the run was stopped.
+            BaseException: If the step failed all retries.
         """
-        assert self._active_deployment
-        launcher = StepLauncher(
-            deployment=self._active_deployment,
-            step=step,
-            orchestrator_run_id=self.get_orchestrator_run_id(),
-        )
-        launcher.launch()
+
+        def _launch_step() -> None:
+            assert self._active_deployment
+
+            launcher = StepLauncher(
+                deployment=self._active_deployment,
+                step=step,
+                orchestrator_run_id=self.get_orchestrator_run_id(),
+            )
+            launcher.launch()
+
+        if self.config.handles_step_retries:
+            _launch_step()
+        else:
+            # The orchestrator subclass doesn't handle step retries, so we
+            # handle it in-process instead
+            retries = 0
+            retry_config = step.config.retry
+            max_retries = retry_config.max_retries if retry_config else 0
+            delay = retry_config.delay if retry_config else 0
+            backoff = retry_config.backoff if retry_config else 1
+
+            while retries <= max_retries:
+                try:
+                    _launch_step()
+                except RunStoppedException:
+                    # Don't retry if the run was stopped
+                    raise
+                except BaseException:
+                    retries += 1
+                    if retries <= max_retries:
+                        logger.info(
+                            "Sleeping for %d seconds before retrying step `%s`.",
+                            delay,
+                            step.config.name,
+                        )
+                        time.sleep(delay)
+                        delay *= backoff
+                    else:
+                        if max_retries > 0:
+                            logger.error(
+                                "Failed to run step `%s` after %d retries.",
+                                step.config.name,
+                                max_retries,
+                            )
+                        raise
+                else:
+                    break
 
     @staticmethod
     def requires_resources_in_orchestration_environment(

zenml/orchestrators/dag_runner.py

@@ -217,7 +217,9 @@ class ThreadedDagRunner:
             self.node_states[node] = NodeStatus.PENDING
 
         # Run node in new thread.
-        thread = threading.Thread(target=self._run_node, args=(node,))
+        thread = threading.Thread(
+            name=node, target=self._run_node, args=(node,)
+        )
         thread.start()
         return thread
 

zenml/orchestrators/publish_utils.py

@@ -190,7 +190,10 @@ def get_pipeline_run_status(
         return ExecutionStatus.FAILED
 
     # If there is a running step, the run is running
-    elif ExecutionStatus.RUNNING in step_statuses:
+    elif (
+        ExecutionStatus.RUNNING in step_statuses
+        or ExecutionStatus.RETRYING in step_statuses
+    ):
         return ExecutionStatus.RUNNING
 
     # If there are less steps than the total number of steps, it is running

zenml/orchestrators/step_launcher.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Class to launch (run directly or using a step operator) steps."""
 
-import os
 import signal
 import time
 from contextlib import nullcontext
@@ -25,7 +24,6 @@ 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,
     handle_bool_env_var,
 )
 from zenml.enums import ExecutionStatus
@@ -245,141 +243,93 @@ class StepLauncher:
                 artifact_store_id=self._stack.artifact_store.id,
             )
 
-        try:
-            with logs_context:
-                if run_was_created:
-                    pipeline_run_metadata = (
-                        self._stack.get_pipeline_run_metadata(
-                            run_id=pipeline_run.id
-                        )
+        with logs_context:
+            if run_was_created:
+                pipeline_run_metadata = self._stack.get_pipeline_run_metadata(
+                    run_id=pipeline_run.id
+                )
+                publish_utils.publish_pipeline_run_metadata(
+                    pipeline_run_id=pipeline_run.id,
+                    pipeline_run_metadata=pipeline_run_metadata,
+                )
+                if model_version := pipeline_run.model_version:
+                    step_run_utils.log_model_version_dashboard_url(
+                        model_version=model_version
                     )
-                    publish_utils.publish_pipeline_run_metadata(
-                        pipeline_run_id=pipeline_run.id,
-                        pipeline_run_metadata=pipeline_run_metadata,
+
+            request_factory = step_run_utils.StepRunRequestFactory(
+                deployment=self._deployment,
+                pipeline_run=pipeline_run,
+                stack=self._stack,
+            )
+            step_run_request = request_factory.create_request(
+                invocation_id=self._step_name
+            )
+            step_run_request.logs = logs_model
+
+            try:
+                request_factory.populate_request(request=step_run_request)
+            except:
+                logger.exception(f"Failed preparing step `{self._step_name}`.")
+                step_run_request.status = ExecutionStatus.FAILED
+                step_run_request.end_time = utc_now()
+                raise
+            finally:
+                step_run = Client().zen_store.create_run_step(step_run_request)
+                self._step_run = step_run
+                if model_version := step_run.model_version:
+                    step_run_utils.log_model_version_dashboard_url(
+                        model_version=model_version
                     )
-                    if model_version := pipeline_run.model_version:
-                        step_run_utils.log_model_version_dashboard_url(
-                            model_version=model_version
-                        )
 
-                request_factory = step_run_utils.StepRunRequestFactory(
-                    deployment=self._deployment,
-                    pipeline_run=pipeline_run,
-                    stack=self._stack,
-                )
-                step_run_request = request_factory.create_request(
-                    invocation_id=self._step_name
-                )
-                step_run_request.logs = logs_model
+            if not step_run.status.is_finished:
+                logger.info(f"Step `{self._step_name}` has started.")
 
                 try:
-                    request_factory.populate_request(request=step_run_request)
-                except:
-                    logger.exception(
-                        f"Failed preparing step `{self._step_name}`."
-                    )
-                    step_run_request.status = ExecutionStatus.FAILED
-                    step_run_request.end_time = utc_now()
-                    raise
-                finally:
-                    step_run = Client().zen_store.create_run_step(
-                        step_run_request
-                    )
-                    # Store step run ID for signal handler
-                    self._step_run = step_run
-                    if model_version := step_run.model_version:
-                        step_run_utils.log_model_version_dashboard_url(
-                            model_version=model_version
+                    # here pass a forced save_to_file callable to be
+                    # used as a dump function to use before starting
+                    # the external jobs in step operators
+                    if isinstance(
+                        logs_context,
+                        step_logging.PipelineLogsStorageContext,
+                    ):
+                        force_write_logs = partial(
+                            logs_context.storage.save_to_file,
+                            force=True,
                         )
+                    else:
 
-                if not step_run.status.is_finished:
-                    logger.info(f"Step `{self._step_name}` has started.")
-                    retries = 0
-                    last_retry = True
-                    max_retries = (
-                        step_run.config.retry.max_retries
-                        if step_run.config.retry
-                        else 1
-                    )
-                    delay = (
-                        step_run.config.retry.delay
-                        if step_run.config.retry
-                        else 0
+                        def _bypass() -> None:
+                            return None
+
+                        force_write_logs = _bypass
+                    self._run_step(
+                        pipeline_run=pipeline_run,
+                        step_run=step_run,
+                        force_write_logs=force_write_logs,
                     )
-                    backoff = (
-                        step_run.config.retry.backoff
-                        if step_run.config.retry
-                        else 1
+                except RunStoppedException as e:
+                    raise e
+                except BaseException as e:  # noqa: E722
+                    logger.error(
+                        "Failed to run step `%s`: %s",
+                        self._step_name,
+                        e,
                     )
-
-                    while retries < max_retries:
-                        last_retry = retries == max_retries - 1
-                        try:
-                            # here pass a forced save_to_file callable to be
-                            # used as a dump function to use before starting
-                            # the external jobs in step operators
-                            if isinstance(
-                                logs_context,
-                                step_logging.PipelineLogsStorageContext,
-                            ):
-                                force_write_logs = partial(
-                                    logs_context.storage.save_to_file,
-                                    force=True,
-                                )
-                            else:
-
-                                def _bypass() -> None:
-                                    return None
-
-                                force_write_logs = _bypass
-                            self._run_step(
-                                pipeline_run=pipeline_run,
-                                step_run=step_run,
-                                last_retry=last_retry,
-                                force_write_logs=force_write_logs,
-                            )
-                            break
-                        except RunStoppedException as e:
-                            raise e
-                        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."
-                                )
-                                logger.exception(e)
-                                publish_utils.publish_failed_step_run(
-                                    step_run.id
-                                )
-                                raise
-                else:
-                    logger.info(
-                        f"Using cached version of step `{self._step_name}`."
+                    publish_utils.publish_failed_step_run(step_run.id)
+                    raise
+            else:
+                logger.info(
+                    f"Using cached version of step `{self._step_name}`."
+                )
+                if (
+                    model_version := step_run.model_version
+                    or pipeline_run.model_version
+                ):
+                    step_run_utils.link_output_artifacts_to_model_version(
+                        artifacts=step_run.outputs,
+                        model_version=model_version,
                     )
-                    if (
-                        model_version := step_run.model_version
-                        or pipeline_run.model_version
-                    ):
-                        step_run_utils.link_output_artifacts_to_model_version(
-                            artifacts=step_run.outputs,
-                            model_version=model_version,
-                        )
-        except RunStoppedException:
-            logger.info(f"Pipeline run `{pipeline_run.name}` stopped.")
-            raise
-        except:  # noqa: E722
-            logger.error(f"Pipeline run `{pipeline_run.name}` failed.")
-            raise
 
     def _create_or_reuse_run(self) -> Tuple[PipelineRunResponse, bool]:
         """Creates a pipeline run or reuses an existing one.
@@ -421,7 +371,6 @@ class StepLauncher:
         pipeline_run: PipelineRunResponse,
         step_run: StepRunResponse,
         force_write_logs: Callable[..., Any],
-        last_retry: bool = True,
     ) -> None:
         """Runs the current step.
 
@@ -429,7 +378,6 @@ class StepLauncher:
             pipeline_run: The model of the current pipeline run.
             step_run: The model of the current step run.
             force_write_logs: The context for the step logs.
-            last_retry: Whether this is the last retry of the step.
         """
         # Prepare step run information.
         step_run_info = StepRunInfo(
@@ -457,7 +405,6 @@ class StepLauncher:
                 self._run_step_with_step_operator(
                     step_operator_name=step_operator_name,
                     step_run_info=step_run_info,
-                    last_retry=last_retry,
                 )
             else:
                 self._run_step_without_step_operator(
@@ -466,7 +413,6 @@ class StepLauncher:
                     step_run_info=step_run_info,
                     input_artifacts=step_run.regular_inputs,
                     output_artifact_uris=output_artifact_uris,
-                    last_retry=last_retry,
                 )
         except:  # noqa: E722
             output_utils.remove_artifact_dirs(
@@ -484,14 +430,12 @@ class StepLauncher:
         self,
         step_operator_name: Optional[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,
@@ -509,8 +453,6 @@ class StepLauncher:
         environment = orchestrator_utils.get_config_environment_vars(
             pipeline_run_id=step_run_info.run_id,
         )
-        if last_retry:
-            environment[ENV_ZENML_IGNORE_FAILURE_HOOK] = str(False)
         logger.info(
             "Using step operator `%s` to run step `%s`.",
             step_operator.name,
@@ -529,7 +471,6 @@ class StepLauncher:
         step_run_info: StepRunInfo,
         input_artifacts: Dict[str, StepRunInputResponse],
         output_artifact_uris: Dict[str, str],
-        last_retry: bool,
     ) -> None:
         """Runs the current step without a step operator.
 
@@ -539,10 +480,7 @@ 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_run_utils.py

@@ -59,6 +59,22 @@ class StepRunRequestFactory:
         self.pipeline_run = pipeline_run
         self.stack = stack
 
+    def has_caching_enabled(self, invocation_id: str) -> bool:
+        """Check if the step has caching enabled.
+
+        Args:
+            invocation_id: The invocation ID for which to check if caching is
+                enabled.
+
+        Returns:
+            Whether the step has caching enabled.
+        """
+        step = self.deployment.step_configurations[invocation_id]
+        return utils.is_setting_enabled(
+            is_enabled_on_step=step.config.enable_cache,
+            is_enabled_on_pipeline=self.deployment.pipeline_configuration.enable_cache,
+        )
+
     def create_request(self, invocation_id: str) -> StepRunRequest:
         """Create a step run request.
 

zenml/orchestrators/step_runner.py

@@ -34,7 +34,6 @@ 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.enums import ArtifactSaveType
@@ -194,9 +193,7 @@ class StepRunner:
                 )
             except BaseException as step_exception:  # noqa: E722
                 step_failed = True
-                if not handle_bool_env_var(
-                    ENV_ZENML_IGNORE_FAILURE_HOOK, False
-                ):
+                if not step_run.is_retriable:
                     if (
                         failure_hook_source
                         := self.configuration.failure_hook_source

zenml/pipelines/pipeline_decorator.py

@@ -29,6 +29,7 @@ from zenml.logger import get_logger
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
+    from zenml.config.retry_config import StepRetryConfig
     from zenml.model.model import Model
     from zenml.pipelines.pipeline_definition import Pipeline
     from zenml.types import HookSpecification
@@ -57,6 +58,7 @@ def pipeline(
     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"], "Pipeline"]: ...
 
@@ -75,6 +77,7 @@ def pipeline(
     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["Pipeline", Callable[["F"], "Pipeline"]]:
     """Decorator to create a pipeline.
@@ -97,6 +100,7 @@ def pipeline(
             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: Retry configuration for the pipeline steps.
         substitutions: Extra placeholders to use in the name templates.
 
     Returns:
@@ -108,6 +112,7 @@ def pipeline(
 
         p = Pipeline(
             name=name or func.__name__,
+            entrypoint=func,
             enable_cache=enable_cache,
             enable_artifact_metadata=enable_artifact_metadata,
             enable_step_logs=enable_step_logs,
@@ -118,7 +123,7 @@ def pipeline(
             on_failure=on_failure,
             on_success=on_success,
             model=model,
-            entrypoint=func,
+            retry=retry,
             substitutions=substitutions,
         )