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

zenml/orchestrators/step_launcher.py

@@ -13,47 +13,40 @@
 #  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
-from typing import TYPE_CHECKING, Dict, Optional, Tuple
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple
 
 from zenml.client import Client
 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,
-    STEP_SOURCE_PARAMETER_NAME,
-    TEXT_FIELD_MAX_LENGTH,
+    ENV_ZENML_IGNORE_FAILURE_HOOK,
     handle_bool_env_var,
 )
 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,
     PipelineDeploymentResponse,
     PipelineRunRequest,
     PipelineRunResponse,
-    StepRunRequest,
     StepRunResponse,
 )
-from zenml.orchestrators import (
-    cache_utils,
-    input_utils,
-    output_utils,
-    publish_utils,
-)
+from zenml.models.v2.core.step_run import StepRunInputResponse
+from zenml.orchestrators import output_utils, publish_utils, step_run_utils
 from zenml.orchestrators import utils as orchestrator_utils
 from zenml.orchestrators.step_runner import StepRunner
 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__)
@@ -142,7 +135,7 @@ class StepLauncher:
         """Launches the step.
 
         Raises:
-            Exception: If the step failed to launch, run, or publish.
+            BaseException: If the step failed to launch, run, or publish.
         """
         pipeline_run, run_was_created = self._create_or_reuse_run()
 
@@ -166,7 +159,7 @@ class StepLauncher:
             )
 
             logs_context = step_logging.StepLogsStorageContext(
-                logs_uri=logs_uri
+                logs_uri=logs_uri, artifact_store=self._stack.artifact_store
             )  # type: ignore[assignment]
 
             logs_model = LogsRequest(
@@ -186,88 +179,125 @@ class StepLauncher:
                         pipeline_run_id=pipeline_run.id,
                         pipeline_run_metadata=pipeline_run_metadata,
                     )
-                client = Client()
-                (
-                    docstring,
-                    source_code,
-                ) = self._get_step_docstring_and_source_code()
-
-                code_hash = self._deployment.step_configurations[
-                    self._step_name
-                ].config.caching_parameters.get(STEP_SOURCE_PARAMETER_NAME)
-                step_run = StepRunRequest(
-                    name=self._step_name,
-                    pipeline_run_id=pipeline_run.id,
-                    deployment=self._deployment.id,
-                    code_hash=code_hash,
-                    status=ExecutionStatus.RUNNING,
-                    docstring=docstring,
-                    source_code=source_code,
-                    start_time=datetime.utcnow(),
-                    user=client.active_user.id,
-                    workspace=client.active_workspace.id,
-                    logs=logs_model,
+                    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
+
                 try:
-                    execution_needed, step_run = self._prepare(
-                        step_run=step_run
-                    )
+                    request_factory.populate_request(request=step_run_request)
                 except:
                     logger.exception(
-                        f"Failed preparing run step `{self._step_name}`."
+                        f"Failed preparing step `{self._step_name}`."
                     )
-                    step_run.status = ExecutionStatus.FAILED
-                    step_run.end_time = datetime.utcnow()
+                    step_run_request.status = ExecutionStatus.FAILED
+                    step_run_request.end_time = datetime.utcnow()
                     raise
                 finally:
-                    step_run_response = Client().zen_store.create_run_step(
-                        step_run
+                    step_run = Client().zen_store.create_run_step(
+                        step_run_request
                     )
-
-                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 Exception as e:  # noqa: E722
-                        logger.error(
-                            f"Failed to run step `{self._step_name}`."
+                    if model_version := step_run.model_version:
+                        step_run_utils.log_model_version_dashboard_url(
+                            model_version=model_version
                         )
-                        logger.exception(e)
-                        publish_utils.publish_failed_step_run(
-                            step_run_response.id
+
+                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
+                    )
+                    backoff = (
+                        step_run.config.retry.backoff
+                        if step_run.config.retry
+                        else 1
+                    )
+
+                    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.StepLogsStorageContext,
+                            ):
+                                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 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}`."
+                    )
+                    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,
                         )
-                        raise
 
         except:  # noqa: E722
             logger.error(f"Pipeline run `{pipeline_run.name}` failed.")
             publish_utils.publish_failed_pipeline_run(pipeline_run.id)
             raise
 
-    def _get_step_docstring_and_source_code(self) -> Tuple[Optional[str], str]:
-        """Gets the docstring and source code of the step.
-
-        If any of the two is longer than 1000 characters, it will be truncated.
-
-        Returns:
-            The docstring and source code of the step.
-        """
-        from zenml.steps.base_step import BaseStep
-
-        step_instance = BaseStep.load_from_source(self._step.spec.source)
-
-        docstring = step_instance.docstring
-        if docstring and len(docstring) > TEXT_FIELD_MAX_LENGTH:
-            docstring = docstring[: (TEXT_FIELD_MAX_LENGTH - 3)] + "..."
-
-        source_code = step_instance.source_code
-        if source_code and len(source_code) > TEXT_FIELD_MAX_LENGTH:
-            source_code = source_code[: (TEXT_FIELD_MAX_LENGTH - 3)] + "..."
-
-        return docstring, source_code
-
     def _create_or_reuse_run(self) -> Tuple[PipelineRunResponse, bool]:
         """Creates a pipeline run or reuses an existing one.
 
@@ -275,8 +305,12 @@ class StepLauncher:
             The created or existing pipeline run,
             and a boolean indicating whether the run was created or reused.
         """
-        run_name = orchestrator_utils.get_run_name(
-            run_name_template=self._deployment.run_name_template
+        start_time = datetime.utcnow()
+        run_name = string_utils.format_name_template(
+            name_template=self._deployment.run_name_template,
+            substitutions=self._deployment.pipeline_configuration._get_full_substitutions(
+                start_time
+            ),
         )
 
         logger.debug("Creating pipeline run %s", run_name)
@@ -288,130 +322,32 @@ class StepLauncher:
             user=client.active_user.id,
             workspace=client.active_workspace.id,
             deployment=self._deployment.id,
-            pipeline=self._deployment.pipeline.id
-            if self._deployment.pipeline
-            else None,
+            pipeline=(
+                self._deployment.pipeline.id
+                if self._deployment.pipeline
+                else None
+            ),
             status=ExecutionStatus.RUNNING,
             orchestrator_environment=get_run_environment_dict(),
-            start_time=datetime.utcnow(),
+            start_time=start_time,
+            tags=self._deployment.pipeline_configuration.tags,
         )
         return client.zen_store.get_or_create_run(pipeline_run)
 
-    def _prepare(
-        self,
-        step_run: StepRunRequest,
-    ) -> Tuple[bool, StepRunRequest]:
-        """Prepares running the step.
-
-        Args:
-            step_run: The step to run.
-
-        Returns:
-            Tuple that specifies whether the step needs to be executed as
-            well as the response model of the registered step run.
-        """
-        model = (
-            self._deployment.step_configurations[step_run.name].config.model
-            or self._deployment.pipeline_configuration.model
-        )
-        input_artifacts, parent_step_ids = input_utils.resolve_step_inputs(
-            step=self._step,
-            run_id=step_run.pipeline_run_id,
-        )
-        input_artifact_ids = {
-            input_name: artifact.id
-            for input_name, artifact in input_artifacts.items()
-        }
-
-        cache_key = cache_utils.generate_cache_key(
-            step=self._step,
-            input_artifact_ids=input_artifact_ids,
-            artifact_store=self._stack.artifact_store,
-            workspace_id=Client().active_workspace.id,
-        )
-
-        step_run.inputs = input_artifact_ids
-        step_run.parent_step_ids = parent_step_ids
-        step_run.cache_key = cache_key
-
-        cache_enabled = orchestrator_utils.is_setting_enabled(
-            is_enabled_on_step=self._step.config.enable_cache,
-            is_enabled_on_pipeline=self._deployment.pipeline_configuration.enable_cache,
-        )
-
-        step_cache = self._step.config.enable_cache
-        if step_cache is not None:
-            logger.info(
-                f"Caching {'`enabled`' if step_cache else '`disabled`'} "
-                f"explicitly for `{self._step_name}`."
-            )
-
-        execution_needed = True
-        if cache_enabled:
-            cached_step_run = cache_utils.get_cached_step_run(
-                cache_key=cache_key
-            )
-            if cached_step_run:
-                logger.info(f"Using cached version of `{self._step_name}`.")
-                execution_needed = False
-                cached_outputs = cached_step_run.outputs
-                step_run.original_step_run_id = cached_step_run.id
-                step_run.outputs = {
-                    output_name: artifact.id
-                    for output_name, artifact in cached_outputs.items()
-                }
-                self._link_cached_artifacts_to_model(
-                    model_from_context=model,
-                    step_run=step_run,
-                )
-                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():
-            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()
-                else:
-                    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,
+        force_write_logs: Callable[..., Any],
+        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.
+            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(
@@ -421,6 +357,7 @@ class StepLauncher:
             pipeline_step_name=self._step_name,
             run_id=pipeline_run.id,
             step_run_id=step_run.id,
+            force_write_logs=force_write_logs,
         )
 
         output_artifact_uris = output_utils.prepare_output_artifact_uris(
@@ -434,6 +371,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(
@@ -442,6 +380,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(
@@ -459,12 +398,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,
@@ -480,8 +421,11 @@ class StepLauncher:
             )
         )
         environment = orchestrator_utils.get_config_environment_vars(
-            deployment=self._deployment
+            pipeline_run_id=step_run_info.run_id,
+            step_run_id=step_run_info.step_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,
@@ -498,8 +442,9 @@ class StepLauncher:
         pipeline_run: PipelineRunResponse,
         step_run: StepRunResponse,
         step_run_info: StepRunInfo,
-        input_artifacts: Dict[str, ArtifactVersionResponse],
+        input_artifacts: Dict[str, StepRunInputResponse],
         output_artifact_uris: Dict[str, str],
+        last_retry: bool,
     ) -> None:
         """Runs the current step without a step operator.
 
@@ -509,7 +454,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,