Execution

`zenml.execution`

Step and pipeline execution.

Modules

`pipeline`

Pipeline execution.

Modules

`dynamic`

Dynamic pipeline execution.

Modules

compilation

Compilation helpers for dynamic pipelines.

Classes Functions:

compile_child_pipeline(pipeline: DynamicPipeline, args: Tuple[Any, ...], kwargs: Dict[str, Any], parent_snapshot: PipelineSnapshotResponse) -> PipelineSnapshotResponse

Compile and persist a snapshot for a dynamic child pipeline run.

The child snapshot inherits the parent build/code context so child execution runs in the same orchestration environment.

Parameters:

Name	Type	Description	Default
`pipeline`	`DynamicPipeline`	The child dynamic pipeline	required
`args`	`Tuple[Any, ...]`	Positional arguments for the child pipeline.	required
`kwargs`	`Dict[str, Any]`	Keyword arguments for the child pipeline.	required
`parent_snapshot`	`PipelineSnapshotResponse`	Snapshot of the parent dynamic run.	required

Raises:

Type	Description
`RuntimeError`	If the parent snapshot has no stack.

Returns:

Type	Description
`PipelineSnapshotResponse`	The persisted child snapshot.

Source code in src/zenml/execution/pipeline/dynamic/compilation.py

def compile_child_pipeline(
    pipeline: "DynamicPipeline",
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    parent_snapshot: "PipelineSnapshotResponse",
) -> "PipelineSnapshotResponse":
    """Compile and persist a snapshot for a dynamic child pipeline run.

    The child snapshot inherits the parent build/code context so child
    execution runs in the same orchestration environment.

    Args:
        pipeline: The child dynamic pipeline
        args: Positional arguments for the child pipeline.
        kwargs: Keyword arguments for the child pipeline.
        parent_snapshot: Snapshot of the parent dynamic run.

    Raises:
        RuntimeError: If the parent snapshot has no stack.

    Returns:
        The persisted child snapshot.
    """
    if parent_snapshot.stack is None:
        raise RuntimeError(
            "Cannot compile a child pipeline snapshot: parent snapshot has no "
            "associated stack."
        )

    inputs = convert_to_keyword_arguments(
        pipeline.entrypoint, tuple(args), kwargs
    )
    inputs = await_step_inputs(inputs)
    pipeline.prepare(**inputs)

    snapshot_base, _, _ = pipeline._compile()

    code_reference = None
    if parent_snapshot.code_reference:
        code_reference = CodeReferenceRequest(
            commit=parent_snapshot.code_reference.commit,
            subdirectory=parent_snapshot.code_reference.subdirectory,
            code_repository=parent_snapshot.code_reference.code_repository.id,
        )

    # The child pipeline will use the parent's build, so we simply copy the Docker
    # settings from the parent snapshot.
    snapshot_base.pipeline_configuration.settings[DOCKER_SETTINGS_KEY] = (
        parent_snapshot.pipeline_configuration.docker_settings.model_copy()
    )
    for _, step_configuration in snapshot_base.step_configurations.items():
        step_configuration.config.settings.pop(DOCKER_SETTINGS_KEY, None)
        step_configuration.step_config_overrides.settings.pop(
            DOCKER_SETTINGS_KEY, None
        )

    request = PipelineSnapshotRequest(
        project=parent_snapshot.project_id,
        stack=parent_snapshot.stack.id,
        pipeline=pipeline.register().id,
        build=parent_snapshot.build.id if parent_snapshot.build else None,
        code_reference=code_reference,
        code_path=parent_snapshot.code_path,
        source_code=pipeline.source_code,
        **snapshot_base.model_dump(),
    )
    return Client().zen_store.create_snapshot(snapshot=request)

compile_dynamic_step_invocation(snapshot: PipelineSnapshotResponse, pipeline: DynamicPipeline, step: BaseStep, invocation_id: str, inputs: Dict[str, Any], pipeline_docker_settings: DockerSettings, after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None] = None, config: Optional[StepConfigurationUpdate] = None) -> Step

Compile a dynamic step invocation.

Parameters:

Name	Type	Description	Default
`snapshot`	`PipelineSnapshotResponse`	The snapshot.	required
`pipeline`	`DynamicPipeline`	The dynamic pipeline.	required
`step`	`BaseStep`	The step to compile.	required
`invocation_id`	`str`	The invocation ID of the step.	required
`inputs`	`Dict[str, Any]`	The inputs for the step function.	required
`pipeline_docker_settings`	`DockerSettings`	The Docker settings of the parent pipeline.	required
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	The step run output futures to wait for.	`None`
`config`	`Optional[StepConfigurationUpdate]`	The configuration for the step.	`None`

Returns:

Type	Description
`Step`	The compiled step.

Source code in src/zenml/execution/pipeline/dynamic/compilation.py

def compile_dynamic_step_invocation(
    snapshot: "PipelineSnapshotResponse",
    pipeline: "DynamicPipeline",
    step: "BaseStep",
    invocation_id: str,
    inputs: Dict[str, Any],
    pipeline_docker_settings: "DockerSettings",
    after: Union["AnyOutputFuture", Sequence["AnyOutputFuture"], None] = None,
    config: Optional[StepConfigurationUpdate] = None,
) -> "Step":
    """Compile a dynamic step invocation.

    Args:
        snapshot: The snapshot.
        pipeline: The dynamic pipeline.
        step: The step to compile.
        invocation_id: The invocation ID of the step.
        inputs: The inputs for the step function.
        pipeline_docker_settings: The Docker settings of the parent pipeline.
        after: The step run output futures to wait for.
        config: The configuration for the step.

    Returns:
        The compiled step.
    """
    upstream_steps = set()

    for future in collect_futures(after=after, expand_map_results=True):
        _await_input_future(future, return_result=False)
        if isinstance(future, PipelineFuture):
            # A pipeline future means we're waiting for a child pipeline to
            # finish. No such step exists in our pipeline, so we can't track
            # it as an upstream step.
            continue
        upstream_steps.add(future.invocation_id)

    inputs = await_step_inputs(inputs)

    for value in inputs.values():
        if isinstance(value, OutputArtifact):
            upstream_steps.add(value.step_name)

        if (
            isinstance(value, Sequence)
            and value
            and all(isinstance(item, OutputArtifact) for item in value)
        ):
            upstream_steps.update(item.step_name for item in value)

    input_artifacts: Dict[str, Union[StepArtifact, List[StepArtifact]]] = {}
    external_artifacts = {}
    for name, value in inputs.items():
        if isinstance(value, OutputArtifact):
            input_artifacts[name] = StepArtifact(
                invocation_id=value.step_name,
                output_name=value.output_name,
                annotation=OutputSignature(resolved_annotation=Any),
                pipeline=pipeline,
                chunk_index=value.chunk_index,
                chunk_size=value.chunk_size,
            )
        elif (
            isinstance(value, list)
            and value
            and all(isinstance(item, OutputArtifact) for item in value)
        ):
            input_artifacts[name] = [
                StepArtifact(
                    invocation_id=item.step_name,
                    output_name=item.output_name,
                    annotation=OutputSignature(resolved_annotation=Any),
                    pipeline=pipeline,
                    chunk_index=item.chunk_index,
                    chunk_size=item.chunk_size,
                )
                for item in value
            ]
        elif isinstance(value, (ArtifactVersionResponse, ExternalArtifact)):
            external_artifacts[name] = value
        else:
            # TODO: should some of these be parameters?
            external_artifacts[name] = ExternalArtifact(value=value)

    if template := get_config_template(snapshot, step, pipeline):
        logger.debug(
            "Using config template `%s` for step `%s`",
            template.spec.invocation_id,
            invocation_id,
        )
        step._configuration = template.config.model_copy(
            update={"template": template.spec.invocation_id}
        )

    default_parameters = {
        key: value
        for key, value in convert_to_keyword_arguments(
            step.entrypoint, (), inputs, apply_defaults=True
        ).items()
        if key not in inputs and key not in step.configuration.parameters
    }

    step_invocation = StepInvocation(
        id=invocation_id,
        step=step,
        input_artifacts=input_artifacts,
        external_artifacts=external_artifacts,
        default_parameters=default_parameters,
        upstream_steps=upstream_steps,
        pipeline=pipeline,
        model_artifacts_or_metadata={},
        client_lazy_loaders={},
        parameters={},
    )

    compiled_step = Compiler()._compile_step_invocation(
        invocation=step_invocation,
        stack=Client().active_stack,
        step_config=config,
        pipeline=pipeline,
    )

    if not compiled_step.config.docker_settings.skip_build:
        if template:
            if (
                template.config.docker_settings
                != compiled_step.config.docker_settings
            ):
                logger.warning(
                    "Custom Docker settings specified for step %s will be "
                    "ignored. The image built for template %s will be used "
                    "instead.",
                    invocation_id,
                    template.spec.invocation_id,
                )
        elif compiled_step.config.docker_settings != pipeline_docker_settings:
            logger.warning(
                "Custom Docker settings specified for step %s will be "
                "ignored. The image built for the pipeline will be used "
                "instead.",
                invocation_id,
            )

    return compiled_step

get_config_template(snapshot: PipelineSnapshotResponse, step: BaseStep, pipeline: DynamicPipeline) -> Optional[Step]

Get the config template for a step executed in a dynamic pipeline.

Parameters:

Name	Type	Description	Default
`snapshot`	`PipelineSnapshotResponse`	The snapshot of the pipeline.	required
`step`	`BaseStep`	The step to get the config template for.	required
`pipeline`	`DynamicPipeline`	The dynamic pipeline that the step is being executed in.	required

Returns:

Type	Description
`Optional[Step]`	The config template for the step.

Source code in src/zenml/execution/pipeline/dynamic/compilation.py

def get_config_template(
    snapshot: "PipelineSnapshotResponse",
    step: "BaseStep",
    pipeline: "DynamicPipeline",
) -> Optional["Step"]:
    """Get the config template for a step executed in a dynamic pipeline.

    Args:
        snapshot: The snapshot of the pipeline.
        step: The step to get the config template for.
        pipeline: The dynamic pipeline that the step is being executed in.

    Returns:
        The config template for the step.
    """
    for index, step_ in enumerate(pipeline.depends_on):
        if step_._static_id == step._static_id:
            break
    else:
        return None

    return list(snapshot.step_configurations.values())[index]

get_step_runtime(step_config: StepConfiguration, pipeline_docker_settings: DockerSettings, orchestrator: Optional[BaseOrchestrator] = None) -> StepRuntime

Determine if a step should be run in process.

Parameters:

Name	Type	Description	Default
`step_config`	`StepConfiguration`	The step configuration.	required
`pipeline_docker_settings`	`DockerSettings`	The Docker settings of the parent pipeline.	required
`orchestrator`	`Optional[BaseOrchestrator]`	The orchestrator to use. If not provided, the orchestrator will be inferred from the active stack.	`None`

Returns:

Type	Description
`StepRuntime`	The runtime for the step.

Source code in src/zenml/execution/pipeline/dynamic/compilation.py

def get_step_runtime(
    step_config: "StepConfiguration",
    pipeline_docker_settings: "DockerSettings",
    orchestrator: Optional["BaseOrchestrator"] = None,
) -> StepRuntime:
    """Determine if a step should be run in process.

    Args:
        step_config: The step configuration.
        pipeline_docker_settings: The Docker settings of the parent pipeline.
        orchestrator: The orchestrator to use. If not provided, the
            orchestrator will be inferred from the active stack.

    Returns:
        The runtime for the step.
    """
    if step_config.step_operator:
        return StepRuntime.ISOLATED

    if not orchestrator:
        orchestrator = Client().active_stack.orchestrator

    if not orchestrator.can_run_isolated_steps:
        return StepRuntime.INLINE

    runtime = step_config.runtime

    if runtime is None:
        if not step_config.resource_settings.empty:
            runtime = StepRuntime.ISOLATED
        elif step_config.docker_settings != pipeline_docker_settings:
            runtime = StepRuntime.ISOLATED
        else:
            runtime = StepRuntime.INLINE

    return runtime

future_registry

Future registry for dynamic pipeline execution.

Classes

FutureRegistry()

Registry for user futures.

Initialize the future registry.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def __init__(self) -> None:
    """Initialize the future registry."""
    self._lock = threading.RLock()
    self._step_futures: Dict[str, StepFuture] = {}
    self._map_futures: Dict[str, MapResultsFuture] = {}
    self._pipeline_futures: Dict[str, PipelineFuture] = {}

Methods:

await_all_no_raise() -> None

Wait for all tracked futures to finish without raising.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def await_all_no_raise(self) -> None:
    """Wait for all tracked futures to finish without raising."""
    futures = self.get_all_futures()

    # Poll all futures concurrently before awaiting each of them. This
    # avoids long cumulative waits when individual futures use backoff-based
    # polling internally.
    poll_interval_seconds = 2.0
    while True:
        any_running = False
        for future in futures:
            try:
                if future.running():
                    any_running = True
                    break
            except Exception:
                # We still fall back to `wait()` below for this future.
                continue

        if not any_running:
            break

        time.sleep(poll_interval_seconds)

    for future in futures:
        try:
            future.wait()
        except Exception:
            pass

bind_map_child_futures(map_id: str, child_futures: List[StepFuture]) -> None

Bind expanded child futures to a map.

Parameters:

Name	Type	Description	Default
`map_id`	`str`	The map ID.	required
`child_futures`	`List[StepFuture]`	The child step futures created for the map.	required

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def bind_map_child_futures(
    self, map_id: str, child_futures: List[StepFuture]
) -> None:
    """Bind expanded child futures to a map.

    Args:
        map_id: The map ID.
        child_futures: The child step futures created for the map.
    """
    with self._lock:
        future = self.get_map_future(map_id=map_id)
        future._set_startup_result(child_futures)

bind_step_execution_future(invocation_id: str, future: StepExecutionFuture) -> None

Bind the step execution future to a step invocation.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	The step invocation ID.	required
`future`	`StepExecutionFuture`	The future that represents the started step execution.	required

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def bind_step_execution_future(
    self, invocation_id: str, future: StepExecutionFuture
) -> None:
    """Bind the step execution future to a step invocation.

    Args:
        invocation_id: The step invocation ID.
        future: The future that represents the started step execution.
    """
    with self._lock:
        step_future = self.get_step_future(invocation_id=invocation_id)
        step_future._set_startup_result(future)

get_all_futures() -> List[Union[StepFuture, MapResultsFuture, PipelineFuture]]

Return all tracked futures.

Returns:

Type	Description
`List[Union[StepFuture, MapResultsFuture, PipelineFuture]]`	A snapshot of all tracked futures.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def get_all_futures(
    self,
) -> List[Union[StepFuture, MapResultsFuture, PipelineFuture]]:
    """Return all tracked futures.

    Returns:
        A snapshot of all tracked futures.
    """
    with self._lock:
        return [
            *self._step_futures.values(),
            *self._map_futures.values(),
            *self._pipeline_futures.values(),
        ]

get_map_future(map_id: str) -> MapResultsFuture

Get a map future.

Parameters:

Name	Type	Description	Default
`map_id`	`str`	The map ID.	required

Raises:

Type	Description
`KeyError`	If the future does not exist.

Returns:

Type	Description
`MapResultsFuture`	The map future.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def get_map_future(self, map_id: str) -> MapResultsFuture:
    """Get a map future.

    Args:
        map_id: The map ID.

    Raises:
        KeyError: If the future does not exist.

    Returns:
        The map future.
    """
    with self._lock:
        future = self._map_futures.get(map_id)
        if future is None:
            raise KeyError(f"Unknown map future `{map_id}`.")
        return future

get_pipeline_future(node_id: str) -> PipelineFuture

Get a child pipeline future.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	Dependency-graph node ID of the child pipeline call.	required

Raises:

Type	Description
`KeyError`	If the future does not exist.

Returns:

Type	Description
`PipelineFuture`	The pipeline future.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def get_pipeline_future(self, node_id: str) -> PipelineFuture:
    """Get a child pipeline future.

    Args:
        node_id: Dependency-graph node ID of the child pipeline call.

    Raises:
        KeyError: If the future does not exist.

    Returns:
        The pipeline future.
    """
    with self._lock:
        future = self._pipeline_futures.get(node_id)
        if future is None:
            raise KeyError(f"Unknown pipeline future `{node_id}`.")
        return future

get_step_future(invocation_id: str) -> StepFuture

Get a step future.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	The step invocation ID.	required

Raises:

Type	Description
`KeyError`	If the future does not exist.

Returns:

Type	Description
`StepFuture`	The step future.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def get_step_future(self, invocation_id: str) -> StepFuture:
    """Get a step future.

    Args:
        invocation_id: The step invocation ID.

    Raises:
        KeyError: If the future does not exist.

    Returns:
        The step future.
    """
    with self._lock:
        future = self._step_futures.get(invocation_id)
        if future is None:
            raise KeyError(f"Unknown step future `{invocation_id}`.")
        return future

has_in_progress_work() -> bool

Check whether any tracked future is still running.

Returns:

Type	Description
`bool`	True if any tracked future is still running, False otherwise.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def has_in_progress_work(self) -> bool:
    """Check whether any tracked future is still running.

    Returns:
        True if any tracked future is still running, False otherwise.
    """
    return any(future.running() for future in self.get_all_futures())

rebind_step_execution_future(invocation_id: str, future: StepExecutionFuture) -> None

Rebind the step execution future to a new attempt of the step.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	The step invocation ID.	required
`future`	`StepExecutionFuture`	The future that represents the new step execution attempt.	required

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def rebind_step_execution_future(
    self, invocation_id: str, future: StepExecutionFuture
) -> None:
    """Rebind the step execution future to a new attempt of the step.

    Args:
        invocation_id: The step invocation ID.
        future: The future that represents the new step execution attempt.
    """
    with self._lock:
        step_future = self.get_step_future(invocation_id=invocation_id)
        step_future._rebind_execution_future(future)

register_map_future(map_id: str, future: MapResultsFuture) -> MapResultsFuture

Register the future for a map invocation.

Parameters:

Name	Type	Description	Default
`map_id`	`str`	The map ID.	required
`future`	`MapResultsFuture`	The map future.	required

Raises:

Type	Description
`RuntimeError`	If a future already exists for the map.

Returns:

Type	Description
`MapResultsFuture`	The registered map future.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def register_map_future(
    self, map_id: str, future: MapResultsFuture
) -> MapResultsFuture:
    """Register the future for a map invocation.

    Args:
        map_id: The map ID.
        future: The map future.

    Raises:
        RuntimeError: If a future already exists for the map.

    Returns:
        The registered map future.
    """
    with self._lock:
        if map_id in self._map_futures:
            raise RuntimeError(
                f"Map future for map `{map_id}` already exists."
            )

        self._map_futures[map_id] = future
        return future

register_pipeline_future(node_id: str, future: PipelineFuture) -> PipelineFuture

Register a child pipeline future.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	Dependency-graph node ID of the child pipeline call (e.g. `pipeline:<name>`), not a pipeline run UUID.	required
`future`	`PipelineFuture`	The pipeline future.	required

Raises:

Type	Description
`RuntimeError`	If a future already exists for the node.

Returns:

Type	Description
`PipelineFuture`	The registered pipeline future.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def register_pipeline_future(
    self, node_id: str, future: PipelineFuture
) -> PipelineFuture:
    """Register a child pipeline future.

    Args:
        node_id: Dependency-graph node ID of the child pipeline call (e.g.
            `pipeline:<name>`), not a pipeline run UUID.
        future: The pipeline future.

    Raises:
        RuntimeError: If a future already exists for the node.

    Returns:
        The registered pipeline future.
    """
    with self._lock:
        if node_id in self._pipeline_futures:
            raise RuntimeError(
                f"Pipeline future for node `{node_id}` already exists."
            )

        self._pipeline_futures[node_id] = future
        return future

register_step_future(invocation_id: str, future: StepFuture) -> StepFuture

Register a step invocation future.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	The step invocation ID.	required
`future`	`StepFuture`	The step future.	required

Raises:

Type	Description
`RuntimeError`	If a future already exists for the invocation.

Returns:

Type	Description
`StepFuture`	The registered step future.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def register_step_future(
    self,
    invocation_id: str,
    future: StepFuture,
) -> StepFuture:
    """Register a step invocation future.

    Args:
        invocation_id: The step invocation ID.
        future: The step future.

    Raises:
        RuntimeError: If a future already exists for the invocation.

    Returns:
        The registered step future.
    """
    with self._lock:
        if invocation_id in self._step_futures:
            raise RuntimeError(
                f"Step future for invocation `{invocation_id}` already exists."
            )

        self._step_futures[invocation_id] = future
        return future

set_startup_exception(invocation_id: str, exception: BaseException) -> None

Set the startup exception for any registered future.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	The invocation ID of the future.	required
`exception`	`BaseException`	The exception to record on the future.	required

Raises:

Type	Description
`KeyError`	If no future is registered for `invocation_id`.

Source code in src/zenml/execution/pipeline/dynamic/future_registry.py

def set_startup_exception(
    self, invocation_id: str, exception: BaseException
) -> None:
    """Set the startup exception for any registered future.

    Args:
        invocation_id: The invocation ID of the future.
        exception: The exception to record on the future.

    Raises:
        KeyError: If no future is registered for `invocation_id`.
    """
    with self._lock:
        for store in (
            self._step_futures,
            self._map_futures,
            self._pipeline_futures,
        ):
            future = store.get(invocation_id)
            if future is not None:
                future._set_startup_failed(exception)
                return
        raise KeyError(
            f"No future registered for invocation `{invocation_id}`."
        )

StartupCancelled

Bases: Exception

Exception raised when an invocation startup is cancelled.

Functions:

inputs

Input resolution helpers for dynamic pipeline execution.

Classes Functions:

await_step_inputs(inputs: Dict[str, Any]) -> Dict[str, Any]

Await the inputs of a step.

Parameters:

Name	Type	Description	Default
`inputs`	`Dict[str, Any]`	The inputs of the step.	required

Raises:

Type	Description
`RuntimeError`	If a step run future or a child pipeline future referring to multiple output artifacts is passed as an input.

Returns:

Type	Description
`Dict[str, Any]`	The awaited inputs.

Source code in src/zenml/execution/pipeline/dynamic/inputs.py

def await_step_inputs(inputs: Dict[str, Any]) -> Dict[str, Any]:
    """Await the inputs of a step.

    Args:
        inputs: The inputs of the step.

    Raises:
        RuntimeError: If a step run future or a child pipeline future referring
            to multiple output artifacts is passed as an input.

    Returns:
        The awaited inputs.
    """
    result = {}
    for key, value in inputs.items():
        if isinstance(value, MapResultsFuture):
            value = value.futures

        if (
            isinstance(value, (list, tuple))
            and value
            and all(isinstance(item, StepFuture) for item in value)
        ):
            if any(len(item._output_keys) != 1 for item in value):
                raise RuntimeError(
                    f"Invalid step input `{key}`: Passing a future that refers "
                    "to multiple output artifacts as an input to another step "
                    "is not allowed."
                )
            value = [_await_input_future(item) for item in value]
        elif isinstance(value, StepFuture):
            if len(value._output_keys) != 1:
                raise RuntimeError(
                    f"Invalid step input `{key}`: Passing a future that refers "
                    "to multiple output artifacts as an input to another step "
                    "is not allowed."
                )
            value = _await_input_future(value)
        elif (
            isinstance(value, (list, tuple))
            and value
            and all(isinstance(item, PipelineFuture) for item in value)
        ):
            value = [
                _resolve_single_child_pipeline_artifact(
                    _await_input_future(item), key=key
                )
                for item in value
            ]
        elif isinstance(value, PipelineFuture):
            value = _resolve_single_child_pipeline_artifact(
                _await_input_future(value), key=key
            )

        if (
            isinstance(value, (list, tuple))
            and value
            and all(isinstance(item, ArtifactFuture) for item in value)
        ):
            value = [_await_input_future(item) for item in value]

        if isinstance(value, ArtifactFuture):
            value = _await_input_future(value)

        result[key] = value

    return result

collect_upstream_node_ids(inputs: Dict[str, Any], after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]) -> List[str]

Collect upstream node IDs from step inputs and after futures.

Parameters:

Name	Type	Description	Default
`inputs`	`Dict[str, Any]`	The step inputs.	required
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	Optional upstream futures for explicit ordering.	required

Returns:

Type	Description
`List[str]`	The upstream node IDs.

Source code in src/zenml/execution/pipeline/dynamic/inputs.py

def collect_upstream_node_ids(
    inputs: Dict[str, Any],
    after: Union["AnyOutputFuture", Sequence["AnyOutputFuture"], None],
) -> List[str]:
    """Collect upstream node IDs from step inputs and `after` futures.

    Args:
        inputs: The step inputs.
        after: Optional upstream futures for explicit ordering.

    Returns:
        The upstream node IDs.
    """
    return [
        future.invocation_id
        for future in collect_futures(inputs=inputs, after=after)
    ]

convert_to_keyword_arguments(func: Callable[..., Any], args: Tuple[Any, ...], kwargs: Dict[str, Any], apply_defaults: bool = False) -> Dict[str, Any]

Convert function arguments to keyword arguments.

Parameters:

Name	Type	Description	Default
`func`	`Callable[..., Any]`	The function to convert the arguments to keyword arguments for.	required
`args`	`Tuple[Any, ...]`	The arguments to convert to keyword arguments.	required
`kwargs`	`Dict[str, Any]`	The keyword arguments to convert to keyword arguments.	required
`apply_defaults`	`bool`	Whether to apply the function default values.	`False`

Returns:

Type	Description
`Dict[str, Any]`	The keyword arguments.

Source code in src/zenml/execution/pipeline/dynamic/inputs.py

def convert_to_keyword_arguments(
    func: Callable[..., Any],
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    apply_defaults: bool = False,
) -> Dict[str, Any]:
    """Convert function arguments to keyword arguments.

    Args:
        func: The function to convert the arguments to keyword arguments for.
        args: The arguments to convert to keyword arguments.
        kwargs: The keyword arguments to convert to keyword arguments.
        apply_defaults: Whether to apply the function default values.

    Returns:
        The keyword arguments.
    """
    signature = inspect.signature(func, follow_wrapped=True)
    bound_args = signature.bind_partial(*args, **kwargs)
    if apply_defaults:
        bound_args.apply_defaults()

    return bound_args.arguments

get_running_upstream_dependencies(inputs: Dict[str, Any], after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]) -> List[str]

Get all running upstream dependencies for a step.

Parameters:

Name	Type	Description	Default
`inputs`	`Dict[str, Any]`	The inputs of the step.	required
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	The step run futures to wait for.	required

Raises:

Type	Description
`TypeError`	If an unexpected future type is passed.

Returns:

Type	Description
`List[str]`	The list of running upstream dependencies.

Source code in src/zenml/execution/pipeline/dynamic/inputs.py

def get_running_upstream_dependencies(
    inputs: Dict[str, Any],
    after: Union["AnyOutputFuture", Sequence["AnyOutputFuture"], None],
) -> List[str]:
    """Get all running upstream dependencies for a step.

    Args:
        inputs: The inputs of the step.
        after: The step run futures to wait for.

    Raises:
        TypeError: If an unexpected future type is passed.

    Returns:
        The list of running upstream dependencies.
    """
    futures = collect_futures(inputs=inputs, after=after)

    dependencies = []

    for future in futures:
        if isinstance(future, MapResultsFuture):
            if future.startup_succeeded:
                for item in future.futures:
                    if item.running():
                        dependencies.append(item.invocation_id)
            elif future.running():
                dependencies.append(future.invocation_id)
        elif isinstance(future, BaseStepFuture):
            if future.running():
                dependencies.append(future.invocation_id)
        elif isinstance(future, PipelineFuture):
            if future.running():
                dependencies.append(future.invocation_id)
        else:
            raise TypeError(f"Unexpected future type: {type(future)}")

    return dependencies

interactive_input_utils

Local terminal input helpers for dynamic pipeline wait conditions.

Classes Functions:

can_answer_wait_condition_interactively(orchestrator: BaseOrchestrator) -> bool

Check whether a wait condition can be answered interactively.

Parameters:

Name	Type	Description	Default
`orchestrator`	`BaseOrchestrator`	The orchestrator running the dynamic pipeline.	required

Returns:

Type	Description
`bool`	Whether terminal prompting should be enabled for the current process.

Source code in src/zenml/execution/pipeline/dynamic/interactive_input_utils.py

def can_answer_wait_condition_interactively(
    orchestrator: "BaseOrchestrator",
) -> bool:
    """Check whether a wait condition can be answered interactively.

    Args:
        orchestrator: The orchestrator running the dynamic pipeline.

    Returns:
        Whether terminal prompting should be enabled for the current process.
    """
    from zenml.orchestrators.local.local_orchestrator import LocalOrchestrator

    try:
        # We use select to check if stdin is ready to be read. If it isn't
        # available on the current platform, we don't support interactive input.
        select.select([sys.stdin], [], [], 0)
    except (AttributeError, OSError, ValueError):
        return False

    return bool(
        not handle_bool_env_var(
            ENV_ZENML_DISABLE_INTERACTIVE_INPUT, default=False
        )
        and isinstance(orchestrator, LocalOrchestrator)
        and sys.stdin.isatty()
        and sys.stdout.isatty()
    )

maybe_enable_interactive_wait_prompt(orchestrator: BaseOrchestrator, condition: RunWaitConditionResponse) -> Iterator[bool]

Render and clean up the terminal prompt for interactive waiting.

Parameters:

Name	Type	Description	Default
`orchestrator`	`BaseOrchestrator`	The orchestrator running the dynamic pipeline.	required
`condition`	`RunWaitConditionResponse`	The wait condition shown to the user.	required

Yields:

Type	Description
`bool`	Whether interactive input is enabled.

Source code in src/zenml/execution/pipeline/dynamic/interactive_input_utils.py

@contextmanager
def maybe_enable_interactive_wait_prompt(
    orchestrator: "BaseOrchestrator",
    condition: "RunWaitConditionResponse",
) -> Iterator[bool]:
    """Render and clean up the terminal prompt for interactive waiting.

    Args:
        orchestrator: The orchestrator running the dynamic pipeline.
        condition: The wait condition shown to the user.

    Yields:
        Whether interactive input is enabled.
    """
    enabled = can_answer_wait_condition_interactively(orchestrator)
    if not enabled:
        yield False
        return

    print_wait_condition_details(condition=condition)
    try:
        yield True
    finally:
        print()

poll_interactive_wait_condition_input(condition: RunWaitConditionResponse, poll_interval: int) -> None

Poll interactive input for a wait condition.

Parameters:

Name	Type	Description	Default
`condition`	`RunWaitConditionResponse`	The wait condition to resolve.	required
`poll_interval`	`int`	Maximum number of seconds to wait.	required

Source code in src/zenml/execution/pipeline/dynamic/interactive_input_utils.py

def poll_interactive_wait_condition_input(
    condition: "RunWaitConditionResponse",
    poll_interval: int,
) -> None:
    """Poll interactive input for a wait condition.

    Args:
        condition: The wait condition to resolve.
        poll_interval: Maximum number of seconds to wait.
    """
    has_input, raw_value = read_stdin_line_with_timeout(
        timeout=float(poll_interval),
    )
    if not has_input:
        return

    result: Optional[Any] = None

    if raw_value:
        if condition.data_schema is not None:
            try:
                result = parse_value_for_schema(
                    raw_value, condition.data_schema
                )
            except ValueError as e:
                print(f"Invalid input: {e}")
                print("> ", end="", flush=True)
                return

    try:
        Client().zen_store.resolve_run_wait_condition(
            run_wait_condition_id=condition.id,
            resolve_request=RunWaitConditionResolveRequest(
                resolution=RunWaitConditionResolution.CONTINUE,
                result=result,
            ),
        )
    except Exception as e:
        print(f"Failed to resolve: {e}")
        print("> ", end="", flush=True)
        return

print_wait_condition_details(condition: RunWaitConditionResponse) -> None

Print wait-condition details for interactive terminal input.

Parameters:

Name	Type	Description	Default
`condition`	`RunWaitConditionResponse`	The wait condition shown to the user.	required

Source code in src/zenml/execution/pipeline/dynamic/interactive_input_utils.py

def print_wait_condition_details(
    condition: "RunWaitConditionResponse",
) -> None:
    """Print wait-condition details for interactive terminal input.

    Args:
        condition: The wait condition shown to the user.
    """
    # Keep the logs neutral because they're printed as-is also in Kitaru.
    print()
    if condition.data_schema is not None:
        print("Waiting for input.")
    else:
        print("Waiting for confirmation.")
    print(f"Question: {condition.question or 'Please provide input.'}")
    if condition.data_schema is not None:
        print("Expected JSON schema:")
        print(json.dumps(condition.data_schema, indent=2, sort_keys=True))
        print("Press Enter on an empty line to submit null.")
    else:
        print("Press Enter to continue.")
    print("Press Ctrl+C to abort.")
    print("> ", end="", flush=True)

read_stdin_line_with_timeout(timeout: float) -> Tuple[bool, Optional[str]]

Read a line from stdin with a timeout.

Parameters:

Name	Type	Description	Default
`timeout`	`float`	Maximum number of seconds to wait.	required

Returns:

Type	Description
`bool`	A tuple indicating whether input was received and the entered value.
`Optional[str]`	Empty input is returned as an empty string.

Source code in src/zenml/execution/pipeline/dynamic/interactive_input_utils.py

def read_stdin_line_with_timeout(
    timeout: float,
) -> Tuple[bool, Optional[str]]:
    """Read a line from stdin with a timeout.

    Args:
        timeout: Maximum number of seconds to wait.

    Returns:
        A tuple indicating whether input was received and the entered value.
        Empty input is returned as an empty string.
    """
    if timeout <= 0:
        return False, None

    try:
        readable, _, _ = select.select([sys.stdin], [], [], timeout)
    except (AttributeError, OSError, ValueError):
        return False, None

    if not readable:
        return False, None

    raw_value = sys.stdin.readline()
    if raw_value == "":
        return False, None

    return True, raw_value.rstrip("\n")

invocation_dependency_graph

Invocation dependency graph.

Classes

BaseNode(*, node_id: str, state: NodeState = NodeState.PENDING, upstream_ids: Set[str] = set(), downstream_ids: Set[str] = set())

dataclass

Base node in the dependency graph.

Attributes

is_terminal: bool property

Whether the node is terminal.

Returns:

Type	Description
`bool`	True if the node is terminal, False otherwise.

ChildPipelineNode(*, node_id: str, state: NodeState = NodeState.PENDING, upstream_ids: Set[str] = set(), downstream_ids: Set[str] = set(), pipeline: DynamicPipeline, args: Sequence[Any] = tuple(), kwargs: Dict[str, Any] = dict())

dataclass

Bases: BaseNode

Child pipeline graph node.

InvocationDependencyGraph()

Invocation dependency graph.

Initialize the graph.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def __init__(self) -> None:
    """Initialize the graph."""
    self._lock = threading.RLock()
    self._nodes: Dict[str, AnyNode] = {}

Methods:

attach_map_children(map_node_id: str, child_node_ids: Sequence[str]) -> bool

Attach expanded child nodes to a map node.

Parameters:

Name	Type	Description	Default
`map_node_id`	`str`	The map node ID.	required
`child_node_ids`	`Sequence[str]`	The child step node IDs created by the expansion.	required

Returns:

Type	Description
`bool`	Whether the attachment caused any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def attach_map_children(
    self, map_node_id: str, child_node_ids: Sequence[str]
) -> bool:
    """Attach expanded child nodes to a map node.

    Args:
        map_node_id: The map node ID.
        child_node_ids: The child step node IDs created by the expansion.

    Returns:
        Whether the attachment caused any newly ready nodes.
    """
    with self._lock:
        map_node = self.get_map_node(node_id=map_node_id)
        should_wake_startup_loop = False

        if map_node.state in {NodeState.READY, NodeState.STARTING}:
            self._set_node_state(
                node_id=map_node_id, state=NodeState.RUNNING
            )

        for child_node_id in child_node_ids:
            child_node = self.get_step_node(node_id=child_node_id)
            map_node.child_node_ids.add(child_node_id)
            child_node.parent_id = map_node_id

        if not child_node_ids:
            should_wake_startup_loop = self._set_node_state(
                node_id=map_node_id, state=NodeState.SUCCEEDED
            )
        else:
            should_wake_startup_loop = self._maybe_finalize_map_node(
                map_node_id=map_node_id
            )

        return should_wake_startup_loop

get_child_pipeline_node(node_id: str) -> ChildPipelineNode

Get a child pipeline node by ID.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Raises:

Type	Description
`RuntimeError`	If the node does not exist or is not a child pipeline node.

Returns:

Type	Description
`ChildPipelineNode`	The child pipeline node.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def get_child_pipeline_node(self, node_id: str) -> ChildPipelineNode:
    """Get a child pipeline node by ID.

    Args:
        node_id: The node ID.

    Raises:
        RuntimeError: If the node does not exist or is not a child pipeline
            node.

    Returns:
        The child pipeline node.
    """
    with self._lock:
        node = self._get_node(node_id=node_id)
        if not isinstance(node, ChildPipelineNode):
            raise RuntimeError(
                f"Node `{node_id}` is not a child pipeline node."
            )
        return node

get_map_node(node_id: str) -> MapNode

Get a map node by ID.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Raises:

Type	Description
`RuntimeError`	If the node does not exist or is not a map node.

Returns:

Type	Description
`MapNode`	The map node.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def get_map_node(self, node_id: str) -> MapNode:
    """Get a map node by ID.

    Args:
        node_id: The node ID.

    Raises:
        RuntimeError: If the node does not exist or is not a map node.

    Returns:
        The map node.
    """
    with self._lock:
        node = self._get_node(node_id=node_id)
        if not isinstance(node, MapNode):
            raise RuntimeError(f"Node `{node_id}` is not a map node.")
        return node

get_node_state(node_id: str) -> NodeState

Get the current state of a node.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Returns:

Type	Description
`NodeState`	The current state of the node.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def get_node_state(self, node_id: str) -> NodeState:
    """Get the current state of a node.

    Args:
        node_id: The node ID.

    Returns:
        The current state of the node.
    """
    with self._lock:
        return self._get_node(node_id=node_id).state

get_ready_node() -> Optional[AnyNode]

Get one ready node in insertion order.

Step nodes are prioritized over child pipeline nodes over map nodes.

Returns:

Type	Description
`Optional[AnyNode]`	A ready node if one exists, otherwise `None`.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def get_ready_node(self) -> Optional[AnyNode]:
    """Get one ready node in insertion order.

    Step nodes are prioritized over child pipeline nodes over map nodes.

    Returns:
        A ready node if one exists, otherwise `None`.
    """
    with self._lock:
        ready_child_pipeline_node: Optional[ChildPipelineNode] = None
        ready_map_node: Optional[MapNode] = None
        for node in self._nodes.values():
            if node.state != NodeState.READY:
                continue
            if isinstance(node, StepNode):
                return node
            if (
                isinstance(node, ChildPipelineNode)
                and ready_child_pipeline_node is None
            ):
                ready_child_pipeline_node = node
            if isinstance(node, MapNode) and ready_map_node is None:
                ready_map_node = node

        return ready_child_pipeline_node or ready_map_node

get_step_node(node_id: str) -> StepNode

Get a step node by ID.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Raises:

Type	Description
`RuntimeError`	If the node does not exist or is not a step node.

Returns:

Type	Description
`StepNode`	The step node.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def get_step_node(self, node_id: str) -> StepNode:
    """Get a step node by ID.

    Args:
        node_id: The node ID.

    Raises:
        RuntimeError: If the node does not exist or is not a step node.

    Returns:
        The step node.
    """
    with self._lock:
        node = self._get_node(node_id=node_id)
        if not isinstance(node, StepNode):
            raise RuntimeError(f"Node `{node_id}` is not a step node.")
        return node

list_nodes(states: Optional[Collection[NodeState]] = None) -> List[AnyNode]

List graph nodes in insertion order.

Parameters:

Name	Type	Description	Default
`states`	`Optional[Collection[NodeState]]`	Optional node states to filter by.	`None`

Returns:

Type	Description
`List[AnyNode]`	The graph nodes in insertion order.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def list_nodes(
    self, states: Optional[Collection[NodeState]] = None
) -> List[AnyNode]:
    """List graph nodes in insertion order.

    Args:
        states: Optional node states to filter by.

    Returns:
        The graph nodes in insertion order.
    """
    with self._lock:
        return [
            node
            for node in self._nodes.values()
            if states is None or node.state in states
        ]

mark_node_failed(node_id: str) -> bool

Mark a node as failed.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Returns:

Type	Description
`bool`	Whether the transition caused any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def mark_node_failed(self, node_id: str) -> bool:
    """Mark a node as failed.

    Args:
        node_id: The node ID.

    Returns:
        Whether the transition caused any newly ready nodes.
    """
    return self._set_node_state(node_id=node_id, state=NodeState.FAILED)

mark_node_paused(node_id: str) -> List[str]

Mark a node as paused and cascade to all downstream nodes.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Returns:

Type	Description
`List[str]`	The IDs of all nodes whose state transitioned to PAUSED.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def mark_node_paused(self, node_id: str) -> List[str]:
    """Mark a node as paused and cascade to all downstream nodes.

    Args:
        node_id: The node ID.

    Returns:
        The IDs of all nodes whose state transitioned to PAUSED.
    """
    with self._lock:
        cascaded: List[str] = []
        stack: List[str] = [node_id]
        while stack:
            current_id = stack.pop()
            node = self._get_node(node_id=current_id)
            if node.state.is_terminal:
                continue

            self._set_node_state(
                node_id=current_id, state=NodeState.PAUSED
            )
            cascaded.append(current_id)
            stack.extend(node.downstream_ids)
        return cascaded

mark_node_running(node_id: str) -> bool

Mark a node as running.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Returns:

Type	Description
`bool`	Whether the transition caused any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def mark_node_running(self, node_id: str) -> bool:
    """Mark a node as running.

    Args:
        node_id: The node ID.

    Returns:
        Whether the transition caused any newly ready nodes.
    """
    return self._set_node_state(node_id=node_id, state=NodeState.RUNNING)

mark_node_starting(node_id: str) -> bool

Mark a node as starting.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Returns:

Type	Description
`bool`	Whether the transition caused any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def mark_node_starting(self, node_id: str) -> bool:
    """Mark a node as starting.

    Args:
        node_id: The node ID.

    Returns:
        Whether the transition caused any newly ready nodes.
    """
    return self._set_node_state(node_id=node_id, state=NodeState.STARTING)

mark_node_succeeded(node_id: str) -> bool

Mark a node as successful.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Returns:

Type	Description
`bool`	Whether the transition caused any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def mark_node_succeeded(self, node_id: str) -> bool:
    """Mark a node as successful.

    Args:
        node_id: The node ID.

    Returns:
        Whether the transition caused any newly ready nodes.
    """
    return self._set_node_state(node_id=node_id, state=NodeState.SUCCEEDED)

register_child_pipeline_node(node_id: str, pipeline: DynamicPipeline, args: Optional[Sequence[Any]] = None, kwargs: Optional[Dict[str, Any]] = None, upstream_ids: Optional[Sequence[str]] = None, state: Optional[NodeState] = None) -> Tuple[ChildPipelineNode, bool]

Register a child pipeline node.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required
`pipeline`	`DynamicPipeline`	The child pipeline payload for startup.	required
`args`	`Optional[Sequence[Any]]`	Optional positional payload for startup.	`None`
`kwargs`	`Optional[Dict[str, Any]]`	Optional keyword payload for startup.	`None`
`upstream_ids`	`Optional[Sequence[str]]`	Optional upstream node IDs.	`None`
`state`	`Optional[NodeState]`	Optional initial state for the node.	`None`

Returns:

Type	Description
`ChildPipelineNode`	The registered child pipeline node and whether the registration
`bool`	caused any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def register_child_pipeline_node(
    self,
    node_id: str,
    pipeline: "DynamicPipeline",
    args: Optional[Sequence[Any]] = None,
    kwargs: Optional[Dict[str, Any]] = None,
    upstream_ids: Optional[Sequence[str]] = None,
    state: Optional[NodeState] = None,
) -> Tuple[ChildPipelineNode, bool]:
    """Register a child pipeline node.

    Args:
        node_id: The node ID.
        pipeline: The child pipeline payload for startup.
        args: Optional positional payload for startup.
        kwargs: Optional keyword payload for startup.
        upstream_ids: Optional upstream node IDs.
        state: Optional initial state for the node.

    Returns:
        The registered child pipeline node and whether the registration
        caused any newly ready nodes.
    """
    node = ChildPipelineNode(
        node_id=node_id,
        state=state or NodeState.PENDING,
        pipeline=pipeline,
        args=args or (),
        kwargs=kwargs or {},
    )
    registered, newly_ready = self._register_node(
        node=node, upstream_ids=upstream_ids
    )
    assert isinstance(registered, ChildPipelineNode)
    return registered, newly_ready

register_map_node(node_id: str, step: BaseStep, inputs: Dict[str, Any], product: bool, upstream_ids: Optional[Sequence[str]] = None, state: Optional[NodeState] = None, after: Optional[Union[AnyOutputFuture, Sequence[AnyOutputFuture]]] = None) -> Tuple[MapNode, bool]

Register a map aggregate node.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The graph node ID.	required
`upstream_ids`	`Optional[Sequence[str]]`	Optional upstream node IDs.	`None`
`state`	`Optional[NodeState]`	Optional initial state for the node.	`None`
`step`	`BaseStep`	The mapped step payload for startup.	required
`inputs`	`Dict[str, Any]`	The input payload for startup.	required
`after`	`Optional[Union[AnyOutputFuture, Sequence[AnyOutputFuture]]]`	Optional `after` payload for startup.	`None`
`product`	`bool`	The map expansion mode.	required

Returns:

Type	Description
`MapNode`	The registered map node and whether the registration caused
`bool`	any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def register_map_node(
    self,
    node_id: str,
    step: BaseStep,
    inputs: Dict[str, Any],
    product: bool,
    upstream_ids: Optional[Sequence[str]] = None,
    state: Optional[NodeState] = None,
    after: Optional[
        Union[AnyOutputFuture, Sequence[AnyOutputFuture]]
    ] = None,
) -> Tuple[MapNode, bool]:
    """Register a map aggregate node.

    Args:
        node_id: The graph node ID.
        upstream_ids: Optional upstream node IDs.
        state: Optional initial state for the node.
        step: The mapped step payload for startup.
        inputs: The input payload for startup.
        after: Optional `after` payload for startup.
        product: The map expansion mode.

    Returns:
        The registered map node and whether the registration caused
        any newly ready nodes.
    """
    node = MapNode(
        node_id=node_id,
        state=state or NodeState.PENDING,
        step=step,
        inputs=inputs,
        after=after,
        product=product,
    )
    registered, newly_ready = self._register_node(
        node=node, upstream_ids=upstream_ids
    )
    assert isinstance(registered, MapNode)
    return registered, newly_ready

register_step_node(node_id: str, upstream_ids: Optional[Sequence[str]] = None, state: Optional[NodeState] = None, step: Optional[BaseStep] = None, inputs: Optional[Dict[str, Any]] = None, after: Optional[Union[AnyOutputFuture, Sequence[AnyOutputFuture]]] = None, config_overrides: Optional[StepConfigurationUpdate] = None) -> Tuple[StepNode, bool]

Register a step node.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required
`upstream_ids`	`Optional[Sequence[str]]`	Optional upstream node IDs.	`None`
`state`	`Optional[NodeState]`	Optional initial state for the node.	`None`
`step`	`Optional[BaseStep]`	Optional step payload for startup.	`None`
`inputs`	`Optional[Dict[str, Any]]`	Optional input payload for startup.	`None`
`after`	`Optional[Union[AnyOutputFuture, Sequence[AnyOutputFuture]]]`	Optional `after` payload for startup.	`None`
`config_overrides`	`Optional[StepConfigurationUpdate]`	Optional config overrides for startup.	`None`

Returns:

Type	Description
`StepNode`	The registered step node and whether the registration caused
`bool`	any newly ready nodes.

Source code in src/zenml/execution/pipeline/dynamic/invocation_dependency_graph.py

def register_step_node(
    self,
    node_id: str,
    upstream_ids: Optional[Sequence[str]] = None,
    state: Optional[NodeState] = None,
    step: Optional[BaseStep] = None,
    inputs: Optional[Dict[str, Any]] = None,
    after: Optional[
        Union[AnyOutputFuture, Sequence[AnyOutputFuture]]
    ] = None,
    config_overrides: Optional["StepConfigurationUpdate"] = None,
) -> Tuple[StepNode, bool]:
    """Register a step node.

    Args:
        node_id: The node ID.
        upstream_ids: Optional upstream node IDs.
        state: Optional initial state for the node.
        step: Optional step payload for startup.
        inputs: Optional input payload for startup.
        after: Optional `after` payload for startup.
        config_overrides: Optional config overrides for startup.

    Returns:
        The registered step node and whether the registration caused
        any newly ready nodes.
    """
    node = StepNode(
        node_id=node_id,
        state=state or NodeState.PENDING,
        step=step,
        inputs=inputs,
        after=after,
        config_overrides=config_overrides,
    )
    registered, newly_ready = self._register_node(
        node=node, upstream_ids=upstream_ids
    )
    assert isinstance(registered, StepNode)
    return registered, newly_ready

MapNode(*, node_id: str, state: NodeState = NodeState.PENDING, upstream_ids: Set[str] = set(), downstream_ids: Set[str] = set(), child_node_ids: Set[str] = set(), step: BaseStep, inputs: Dict[str, Any], after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None], product: bool)

dataclass

Bases: BaseNode

Map graph node.

NodeState

Bases: StrEnum

Invocation dependency graph node state.

Attributes

is_terminal: bool property

Whether the state is terminal.

Returns:

Type	Description
`bool`	True if the state is terminal, False otherwise.

StepNode(*, node_id: str, state: NodeState = NodeState.PENDING, upstream_ids: Set[str] = set(), downstream_ids: Set[str] = set(), parent_id: Optional[str] = None, step: Optional[BaseStep] = None, inputs: Optional[Dict[str, Any]] = None, after: Optional[Union[AnyOutputFuture, Sequence[AnyOutputFuture]]] = None, config_overrides: Optional[StepConfigurationUpdate] = None)

dataclass

Bases: BaseNode

Step graph node.

Functions:

outputs

Dynamic pipeline execution outputs.

Classes

ArtifactFuture(parent: StepFuture, index: int)

Bases: BaseStepFuture

Future for a step run output artifact.

Initialize the future.

Parameters:

Name	Type	Description	Default
`parent`	`StepFuture`	The parent step future object.	required
`index`	`int`	The index of the output artifact.	required

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def __init__(
    self,
    parent: "StepFuture",
    index: int,
) -> None:
    """Initialize the future.

    Args:
        parent: The parent step future object.
        index: The index of the output artifact.
    """
    super().__init__(invocation_id=parent.invocation_id)
    self._index = index
    self._parent = parent

Methods:

chunk(index: int) -> OutputArtifact

Get a chunk of the output artifact.

This method will wait for the future to complete and then return the artifact chunk.

Parameters:

Name	Type	Description	Default
`index`	`int`	The index of the chunk.	required

Returns:

Type	Description
`OutputArtifact`	The artifact chunk.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def chunk(self, index: int) -> "OutputArtifact":
    """Get a chunk of the output artifact.

    This method will wait for the future to complete and then return the
    artifact chunk.

    Args:
        index: The index of the chunk.

    Returns:
        The artifact chunk.
    """
    return self.result().chunk(index=index)

load(disable_cache: bool = False) -> Any

Load the step run output artifact data.

Parameters:

Name	Type	Description	Default
`disable_cache`	`bool`	Whether to disable the artifact cache.	`False`

Returns:

Type	Description
`Any`	The step run output artifact data.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def load(self, disable_cache: bool = False) -> Any:
    """Load the step run output artifact data.

    Args:
        disable_cache: Whether to disable the artifact cache.

    Returns:
        The step run output artifact data.
    """
    return self.result().load(disable_cache=disable_cache)

result() -> OutputArtifact

Get the output artifact this future represents.

Raises:

Type	Description
`RuntimeError`	If the future returned an invalid output.

Returns:

Type	Description
`OutputArtifact`	The output artifact.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def result(self) -> OutputArtifact:
    """Get the output artifact this future represents.

    Raises:
        RuntimeError: If the future returned an invalid output.

    Returns:
        The output artifact.
    """
    from zenml.execution.pipeline.dynamic.utils import (
        load_step_run_outputs,
    )

    step_run = self._parent._wait()
    result = load_step_run_outputs(step_run.id)

    if isinstance(result, OutputArtifact):
        return result
    elif isinstance(result, tuple):
        return result[self._index]
    else:
        raise RuntimeError(
            f"Step {self.invocation_id} returned an invalid output: "
            f"{result}."
        )

running() -> bool

Check if the artifact future is running.

Returns:

Type	Description
`bool`	True if the artifact future is running, False otherwise.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def running(self) -> bool:
    """Check if the artifact future is running.

    Returns:
        True if the artifact future is running, False otherwise.
    """
    return self._parent.running()

wait() -> None

Wait for the artifact future to complete.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def wait(self) -> None:
    """Wait for the artifact future to complete."""
    self._parent.wait()

BaseFuture

Bases: ABC

Base future.

Methods:

result() -> Any abstractmethod

Get the result of the future.

Returns:

Type	Description
`Any`	The result of the future.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

@abstractmethod
def result(self) -> Any:
    """Get the result of the future.

    Returns:
        The result of the future.
    """

running() -> bool abstractmethod

Check if the future is running.

Returns:

Type	Description
`bool`	True if the future is running, False otherwise.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

@abstractmethod
def running(self) -> bool:
    """Check if the future is running.

    Returns:
        True if the future is running, False otherwise.
    """

BaseStepFuture(invocation_id: str, **kwargs: Any)

Bases: BaseFuture, ABC

Base step future.

Initialize the dynamic step run future.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	The invocation ID of the future.	required
`**kwargs`	`Any`	Additional keyword arguments.	`{}`

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def __init__(
    self,
    invocation_id: str,
    **kwargs: Any,
) -> None:
    """Initialize the dynamic step run future.

    Args:
        invocation_id: The invocation ID of the future.
        **kwargs: Additional keyword arguments.
    """
    self._invocation_id = invocation_id

Attributes

invocation_id: str property

The step run invocation ID.

Returns:

Type	Description
`str`	The step run invocation ID.

Methods:

wait() -> None abstractmethod

Wait for the future to finish.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

@abstractmethod
def wait(self) -> None:
    """Wait for the future to finish."""

MapResultsFuture(invocation_id: str)

Bases: BaseFuture

Future that represents the results of a step.map/product(...) call.

Initialize an empty map results future.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	Stable invocation ID for the map expansion.	required

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def __init__(self, invocation_id: str) -> None:
    """Initialize an empty map results future.

    Args:
        invocation_id: Stable invocation ID for the map expansion.
    """
    self._invocation_id = invocation_id
    self._startup = _StartupResult[List[StepFuture]]()

Attributes

futures: List[StepFuture] property

Get the child step futures for the map.

Returns:

Type	Description
`List[StepFuture]`	The child step futures.

invocation_id: str property

Stable invocation ID for this map expansion.

Returns:

Type	Description
`str`	The invocation ID.

startup_failed: bool property

Whether the map startup failed.

Returns:

Type	Description
`bool`	Whether the map startup failed.

startup_succeeded: bool property

Whether map startup completed successfully with child futures.

Returns:

Type	Description
`bool`	Whether the startup completed successfully.

Methods:

load(disable_cache: bool = False) -> List[Any]

Load the step run output artifacts.

Parameters:

Name	Type	Description	Default
`disable_cache`	`bool`	Whether to disable the artifact cache.	`False`

Returns:

Type	Description
`List[Any]`	The step run output artifacts.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def load(self, disable_cache: bool = False) -> List[Any]:
    """Load the step run output artifacts.

    Args:
        disable_cache: Whether to disable the artifact cache.

    Returns:
        The step run output artifacts.
    """
    return [
        future.load(disable_cache=disable_cache) for future in self.futures
    ]

result() -> List[StepRunOutputs]

Get the step run outputs this future represents.

Returns:

Type	Description
`List[StepRunOutputs]`	The step run outputs.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def result(self) -> List[StepRunOutputs]:
    """Get the step run outputs this future represents.

    Returns:
        The step run outputs.
    """
    return [future.result() for future in self.futures]

running() -> bool

Check if the map results future is running.

Returns:

Type	Description
`bool`	True if the map results future is running, False otherwise.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def running(self) -> bool:
    """Check if the map results future is running.

    Returns:
        True if the map results future is running, False otherwise.
    """
    if not self._startup.done():
        return True
    if self._startup.failed():
        return False
    return any(future.running() for future in self.futures)

unpack() -> Tuple[List[ArtifactFuture], ...]

Unpack the map results future.

This method can be used to get lists of artifact futures that represent the outputs of all the step runs that are part of this map result.

Example:

from zenml import pipeline, step

@step
def create_int_list() -> list[int]:
    return [1, 2]

@step
def do_something(a: int) -> Tuple[int, int]:
    return a * 2, a * 3

@pipeline
def map_pipeline():
    int_list = create_int_list()
    results = do_something.map(a=int_list)
    double, triple = results.unpack()

    # [future.load() for future in double] will return [2, 4]
    # [future.load() for future in triple] will return [3, 6]

Returns:

Type	Description
`Tuple[List[ArtifactFuture], ...]`	The unpacked map results.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def unpack(self) -> Tuple[List[ArtifactFuture], ...]:
    """Unpack the map results future.

    This method can be used to get lists of artifact futures that represent
    the outputs of all the step runs that are part of this map result.

    Example:
    ```python
    from zenml import pipeline, step

    @step
    def create_int_list() -> list[int]:
        return [1, 2]

    @step
    def do_something(a: int) -> Tuple[int, int]:
        return a * 2, a * 3

    @pipeline
    def map_pipeline():
        int_list = create_int_list()
        results = do_something.map(a=int_list)
        double, triple = results.unpack()

        # [future.load() for future in double] will return [2, 4]
        # [future.load() for future in triple] will return [3, 6]
    ```

    Returns:
        The unpacked map results.
    """
    return tuple(map(list, zip(*self.futures)))

wait() -> None

Wait for the map results future to complete.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def wait(self) -> None:
    """Wait for the map results future to complete."""
    for future in self.futures:
        future.wait()

OutputArtifact

Bases: ArtifactVersionResponse

Dynamic step run output artifact.

Methods:

chunk(index: int) -> OutputArtifact

Get a chunk of the output artifact.

Parameters:

Name	Type	Description	Default
`index`	`int`	The index of the chunk.	required

Raises:

Type	Description
`ValueError`	If the output artifact can not be chunked or the index is out of range.

Returns:

Type	Description
`OutputArtifact`	The artifact chunk.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def chunk(self, index: int) -> "OutputArtifact":
    """Get a chunk of the output artifact.

    Args:
        index: The index of the chunk.

    Raises:
        ValueError: If the output artifact can not be chunked or the index
            is out of range.

    Returns:
        The artifact chunk.
    """
    if not self.item_count:
        raise ValueError(
            f"Output artifact `{self.output_name}` of step "
            f"`{self.step_name}` can not be chunked."
        )

    if index < 0 or index >= self.item_count:
        raise ValueError(
            f"Chunk index `{index}` out of range for output artifact "
            f"`{self.output_name}` of step `{self.step_name}`."
        )

    if self.chunk_index is not None and self.chunk_index != index:
        raise ValueError(
            f"Output artifact `{self.output_name}` of step "
            f"`{self.step_name}` is already referring to a "
            "different chunk."
        )

    return self.model_copy(update={"chunk_index": index, "chunk_size": 1})

PipelineFuture(invocation_id: str, declared_output_names: Optional[List[str]] = None)

Bases: BaseFuture

Future for a child pipeline run output.

Initialize the future.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	Invocation ID of the child pipeline node.	required
`declared_output_names`	`Optional[List[str]]`	Output names declared on the child pipeline entrypoint (from its return-type annotation). Empty if the entrypoint is unannotated; in that case effective names are derived from the actual outputs after execution.	`None`

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def __init__(
    self,
    invocation_id: str,
    declared_output_names: Optional[List[str]] = None,
) -> None:
    """Initialize the future.

    Args:
        invocation_id: Invocation ID of the child pipeline node.
        declared_output_names: Output names declared on the child pipeline
            entrypoint (from its return-type annotation). Empty if the
            entrypoint is unannotated; in that case effective names are
            derived from the actual outputs after execution.
    """
    self._invocation_id = invocation_id
    self._startup = _StartupResult[Future[PipelineRunResponse]]()
    self._declared_output_names = declared_output_names or []

Attributes

invocation_id: str property

Invocation ID of this child pipeline future.

Returns:

Type	Description
`str`	The invocation ID.

Methods:

artifacts() -> PipelineRunOutputs

Get the child pipeline output artifacts.

Returns:

Type	Description
`PipelineRunOutputs`	The child pipeline output artifacts.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def artifacts(self) -> PipelineRunOutputs:
    """Get the child pipeline output artifacts.

    Returns:
        The child pipeline output artifacts.
    """
    return self.result()

get_artifact(key: str) -> ArtifactVersionResponse

Get an output artifact by output name.

Parameters:

Name	Type	Description	Default
`key`	`str`	The output name.	required

Raises:

Type	Description
`KeyError`	If no output exists for the key.

Returns:

Type	Description
`ArtifactVersionResponse`	The output artifact.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def get_artifact(self, key: str) -> ArtifactVersionResponse:
    """Get an output artifact by output name.

    Args:
        key: The output name.

    Raises:
        KeyError: If no output exists for the key.

    Returns:
        The output artifact.
    """
    outputs = self._output_tuple()
    names = self._effective_output_names(count=len(outputs))
    if key not in names:
        raise KeyError(
            f"Child pipeline `{self.invocation_id}` does not have an output "
            f"with the name `{key}`."
        )
    return outputs[names.index(key)]

result() -> PipelineRunOutputs

Get the child pipeline outputs.

Returns:

Type	Description
`PipelineRunOutputs`	The child pipeline outputs.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def result(self) -> PipelineRunOutputs:
    """Get the child pipeline outputs.

    Returns:
        The child pipeline outputs.
    """
    from zenml.execution.pipeline.dynamic.utils import (
        load_pipeline_run_outputs,
    )

    with _maybe_release_pipeline_thread():
        run = self._startup.result().result()
    return load_pipeline_run_outputs(run=run)

running() -> bool

Check if the child pipeline future is running.

Returns:

Type	Description
`bool`	True if the future is running, False otherwise.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def running(self) -> bool:
    """Check if the child pipeline future is running.

    Returns:
        True if the future is running, False otherwise.
    """
    if not self._startup.done():
        return True
    if self._startup.failed():
        return False

    return not self._startup.result().done()

wait() -> None

Wait for the child pipeline to finish.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def wait(self) -> None:
    """Wait for the child pipeline to finish."""
    with _maybe_release_pipeline_thread():
        self._startup.result().result()

StepFuture(invocation_id: str, output_keys: List[str], execution_future: Optional[StepExecutionFuture] = None)

Bases: BaseStepFuture

Future for a step run output.

Initialize the future.

Parameters:

Name	Type	Description	Default
`invocation_id`	`str`	The invocation ID of the step run.	required
`output_keys`	`List[str]`	The output keys of the step run.	required
`execution_future`	`Optional[StepExecutionFuture]`	Optional execution future if the startup has already completed.	`None`

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def __init__(
    self,
    invocation_id: str,
    output_keys: List[str],
    execution_future: Optional[StepExecutionFuture] = None,
) -> None:
    """Initialize the future.

    Args:
        invocation_id: The invocation ID of the step run.
        output_keys: The output keys of the step run.
        execution_future: Optional execution future if the startup has
            already completed.
    """
    super().__init__(invocation_id=invocation_id)
    self._startup = _StartupResult[StepExecutionFuture]()
    if execution_future is not None:
        self._startup.set_result(execution_future)
    self._output_keys = output_keys

Methods:

artifacts() -> StepRunOutputs

Get the step run output artifacts.

Returns:

Type	Description
`StepRunOutputs`	The step run output artifacts.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def artifacts(self) -> StepRunOutputs:
    """Get the step run output artifacts.

    Returns:
        The step run output artifacts.
    """
    return self.result()

get_artifact(key: str) -> ArtifactFuture

Get an artifact future by key.

Parameters:

Name	Type	Description	Default
`key`	`str`	The key of the artifact future.	required

Raises:

Type	Description
`KeyError`	If no artifact for the given name exists.

Returns:

Type	Description
`ArtifactFuture`	The artifact future.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def get_artifact(self, key: str) -> ArtifactFuture:
    """Get an artifact future by key.

    Args:
        key: The key of the artifact future.

    Raises:
        KeyError: If no artifact for the given name exists.

    Returns:
        The artifact future.
    """
    if key not in self._output_keys:
        raise KeyError(
            f"Step run {self.invocation_id} does not have an output with "
            f"the name: {key}."
        )

    return ArtifactFuture(
        parent=self,
        index=self._output_keys.index(key),
    )

load(disable_cache: bool = False) -> Any

Get the step run output artifact data.

Parameters:

Name	Type	Description	Default
`disable_cache`	`bool`	Whether to disable the artifact cache.	`False`

Raises:

Type	Description
`ValueError`	If the step run output is invalid.

Returns:

Type	Description
`Any`	The step run output artifact data.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def load(self, disable_cache: bool = False) -> Any:
    """Get the step run output artifact data.

    Args:
        disable_cache: Whether to disable the artifact cache.

    Raises:
        ValueError: If the step run output is invalid.

    Returns:
        The step run output artifact data.
    """
    result = self.artifacts()

    if result is None:
        return None
    elif isinstance(result, ArtifactVersionResponse):
        return result.load(disable_cache=disable_cache)
    elif isinstance(result, tuple):
        return tuple(
            item.load(disable_cache=disable_cache) for item in result
        )
    else:
        raise ValueError(f"Invalid step run output: {result}")

result() -> StepRunOutputs

Get the step run outputs this future represents.

Returns:

Type	Description
`StepRunOutputs`	The step run outputs.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def result(self) -> StepRunOutputs:
    """Get the step run outputs this future represents.

    Returns:
        The step run outputs.
    """
    from zenml.execution.pipeline.dynamic.utils import (
        load_step_run_outputs,
    )

    step_run = self._wait()
    return load_step_run_outputs(step_run.id)

running() -> bool

Check if the step future is running.

Returns:

Type	Description
`bool`	True if the step future is running, False otherwise.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def running(self) -> bool:
    """Check if the step future is running.

    Returns:
        True if the step future is running, False otherwise.
    """
    if not self._startup.done():
        return True

    if self._startup.failed():
        return False

    return self._startup.result().running()

wait() -> None

Wait for the step to finish.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def wait(self) -> None:
    """Wait for the step to finish."""
    self._wait()

Functions:

wrap_step_failure(exception: BaseException, invocation_id: Optional[str] = None) -> StepExecutionException

Wrap a step failure that was not explicitly awaited.

Parameters:

Name	Type	Description	Default
`exception`	`BaseException`	The original step failure.	required
`invocation_id`	`Optional[str]`	The invocation ID of the failed step, if known.	`None`

Returns:

Type	Description
`StepExecutionException`	The exception unchanged if it already is a `StepExecutionException`,
`StepExecutionException`	otherwise a new `StepExecutionException` chaining the original.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py

def wrap_step_failure(
    exception: BaseException, invocation_id: Optional[str] = None
) -> StepExecutionException:
    """Wrap a step failure that was not explicitly awaited.

    Args:
        exception: The original step failure.
        invocation_id: The invocation ID of the failed step, if known.

    Returns:
        The exception unchanged if it already is a `StepExecutionException`,
        otherwise a new `StepExecutionException` chaining the original.
    """
    if isinstance(exception, StepExecutionException):
        return exception

    if invocation_id:
        message = f"Step `{invocation_id}` failed."
    else:
        message = "A step failed during pipeline execution."

    wrapped = StepExecutionException(message, original_exception=exception)
    wrapped.__cause__ = exception
    wrapped.__suppress_context__ = True
    return wrapped

Modules

pipeline_output_utils

Helpers for dynamic pipeline output validation and normalization.

Classes Functions:

get_pipeline_entrypoint_output_names(pipeline_entrypoint: Callable[..., Any]) -> List[str]

Output names from the entrypoint's return-type annotation.

Returns the Empty if the entrypoint has no annotated returns.

Parameters:

Name	Type	Description	Default
`pipeline_entrypoint`	`Callable[..., Any]`	The dynamic pipeline entrypoint function.	required

Returns:

Type	Description
`List[str]`	Output names from the entrypoint annotation.

Source code in src/zenml/execution/pipeline/dynamic/pipeline_output_utils.py

def get_pipeline_entrypoint_output_names(
    pipeline_entrypoint: Callable[..., Any],
) -> List[str]:
    """Output names from the entrypoint's return-type annotation.

    Returns the
    Empty if the entrypoint has no annotated returns.

    Args:
        pipeline_entrypoint: The dynamic pipeline entrypoint function.

    Returns:
        Output names from the entrypoint annotation.
    """
    return list(parse_return_type_annotations(pipeline_entrypoint).keys())

prepare_pipeline_output_artifacts(value: Any, pipeline_entrypoint: Callable[..., Any]) -> Dict[str, UUID]

Build the output ID payload for a pipeline run status update.

Parameters:

Name	Type	Description	Default
`value`	`Any`	The dynamic pipeline return value.	required
`pipeline_entrypoint`	`Callable[..., Any]`	Dynamic pipeline entrypoint for output naming.	required

Returns:

Type	Description
`Dict[str, UUID]`	Output artifact version IDs keyed by effective output name.

Source code in src/zenml/execution/pipeline/dynamic/pipeline_output_utils.py

def prepare_pipeline_output_artifacts(
    value: Any,
    pipeline_entrypoint: Callable[..., Any],
) -> Dict[str, UUID]:
    """Build the output ID payload for a pipeline run status update.

    Args:
        value: The dynamic pipeline return value.
        pipeline_entrypoint: Dynamic pipeline entrypoint for output naming.

    Returns:
        Output artifact version IDs keyed by effective output name.
    """
    artifacts = validate_pipeline_outputs(value=value)
    output_names = resolve_pipeline_output_names(
        declared_names=get_pipeline_entrypoint_output_names(
            pipeline_entrypoint
        ),
        count=len(artifacts),
    )

    return {
        output_name: artifact.id
        for output_name, artifact in zip(output_names, artifacts)
    }

resolve_pipeline_output_names(declared_names: List[str], count: int) -> List[str]

Resolve effective output names for persisted dynamic pipeline outputs.

Used by both server-side persistence and client-side PipelineFuture so that the names always agree. Falls back to SINGLE_RETURN_OUT_NAME for a single unnamed output and output_0, output_1, ... for unnamed multi-output pipelines.

Parameters:

Name	Type	Description	Default
`declared_names`	`List[str]`	Output names declared on the pipeline entrypoint, if any.	required
`count`	`int`	Number of output artifacts actually produced.	required

Returns:

Type	Description
`List[str]`	Output names in deterministic order.

Source code in src/zenml/execution/pipeline/dynamic/pipeline_output_utils.py

def resolve_pipeline_output_names(
    declared_names: List[str], count: int
) -> List[str]:
    """Resolve effective output names for persisted dynamic pipeline outputs.

    Used by both server-side persistence and client-side `PipelineFuture` so
    that the names always agree. Falls back to `SINGLE_RETURN_OUT_NAME` for a
    single unnamed output and `output_0`, `output_1`, ... for unnamed
    multi-output pipelines.

    Args:
        declared_names: Output names declared on the pipeline entrypoint, if
            any.
        count: Number of output artifacts actually produced.

    Returns:
        Output names in deterministic order.
    """
    if count == 0:
        return []

    if count == 1:
        if declared_names:
            return [declared_names[0]]
        return [SINGLE_RETURN_OUT_NAME]

    if len(declared_names) == count:
        return declared_names

    return [f"output_{i}" for i in range(count)]

validate_pipeline_outputs(value: Any) -> List[ArtifactVersionResponse]

Validate a child pipeline return value.

Parameters:

Name	Type	Description	Default
`value`	`Any`	The pipeline return value.	required

Raises:

Type	Description
`ValueError`	If outputs are not valid.

Returns:

Type	Description
`List[ArtifactVersionResponse]`	A list of artifact responses.

Source code in src/zenml/execution/pipeline/dynamic/pipeline_output_utils.py

def validate_pipeline_outputs(
    value: Any,
) -> List[ArtifactVersionResponse]:
    """Validate a child pipeline return value.

    Args:
        value: The pipeline return value.

    Raises:
        ValueError: If outputs are not valid.

    Returns:
        A list of artifact responses.
    """
    error_message = (
        "Only artifacts (or tuples of artifacts) are currently supported as "
        "pipeline outputs. Please return artifacts generated by steps or "
        "manually store your data as an artifact. Got "
        f"{type(value).__name__}."
    )

    if value is None:
        artifacts = []
    elif isinstance(value, ArtifactVersionResponse):
        artifacts = [value]
    elif isinstance(value, dict):
        raise ValueError(error_message)
    elif isinstance(value, tuple):
        if not all(
            isinstance(item, ArtifactVersionResponse) for item in value
        ):
            raise ValueError(error_message)
        artifacts = list(value)
    else:
        raise ValueError(error_message)

    return [_to_plain_artifact_version(artifact) for artifact in artifacts]

run_context

Dynamic pipeline run context.

Classes

DynamicPipelineRunContext(pipeline: DynamicPipeline, snapshot: PipelineSnapshotResponse, run: PipelineRunResponse, runner: DynamicPipelineRunner)

Bases: BaseContext

Dynamic pipeline run context.

Initialize the dynamic pipeline run context.

Parameters:

Name	Type	Description	Default
`pipeline`	`DynamicPipeline`	The dynamic pipeline that is being executed.	required
`snapshot`	`PipelineSnapshotResponse`	The snapshot of the pipeline.	required
`run`	`PipelineRunResponse`	The pipeline run.	required
`runner`	`DynamicPipelineRunner`	The dynamic pipeline runner.	required

Source code in src/zenml/execution/pipeline/dynamic/run_context.py

def __init__(
    self,
    pipeline: "DynamicPipeline",
    snapshot: "PipelineSnapshotResponse",
    run: "PipelineRunResponse",
    runner: "DynamicPipelineRunner",
) -> None:
    """Initialize the dynamic pipeline run context.

    Args:
        pipeline: The dynamic pipeline that is being executed.
        snapshot: The snapshot of the pipeline.
        run: The pipeline run.
        runner: The dynamic pipeline runner.
    """
    super().__init__()
    self._pipeline = pipeline
    self._snapshot = snapshot
    self._run = run
    self._runner = runner
    self._wait_condition_counter = 0
    self._wait_condition_counter_lock = threading.Lock()

Attributes

pipeline: DynamicPipeline property

The pipeline that is being executed.

Returns:

Type	Description
`DynamicPipeline`	The pipeline that is being executed.

run: PipelineRunResponse property

The pipeline run.

Returns:

Type	Description
`PipelineRunResponse`	The pipeline run.

runner: DynamicPipelineRunner property

The runner executing the pipeline.

Returns:

Type	Description
`DynamicPipelineRunner`	The runner executing the pipeline.

snapshot: PipelineSnapshotResponse property

The snapshot of the pipeline.

Returns:

Type	Description
`PipelineSnapshotResponse`	The snapshot of the pipeline.

Methods:

next_wait_condition_name() -> str

Get the next deterministic wait condition name for this run.

Returns:

Type	Description
`str`	A deterministic wait condition name.

Source code in src/zenml/execution/pipeline/dynamic/run_context.py

def next_wait_condition_name(self) -> str:
    """Get the next deterministic wait condition name for this run.

    Returns:
        A deterministic wait condition name.
    """
    with self._wait_condition_counter_lock:
        counter = self._wait_condition_counter
        self._wait_condition_counter += 1
    return f"wait_condition:{counter}"

Modules

runner

Dynamic pipeline runner.

Classes

DynamicPipelineRunner(snapshot: PipelineSnapshotResponse, run: Optional[PipelineRunResponse], orchestrator: Optional[BaseOrchestrator] = None, pause_coordinator: Optional[_PauseCoordinator] = None, parent_runner: Optional[DynamicPipelineRunner] = None)

Dynamic pipeline runner.

Initialize the dynamic pipeline runner.

Parameters:

Name	Type	Description	Default
`snapshot`	`PipelineSnapshotResponse`	The snapshot of the pipeline.	required
`run`	`Optional[PipelineRunResponse]`	The pipeline run.	required
`orchestrator`	`Optional[BaseOrchestrator]`	The orchestrator to use. If not provided, the orchestrator will be inferred from the snapshot stack.	`None`
`pause_coordinator`	`Optional[_PauseCoordinator]`	The pause coordinator for the run tree. The root runner passes `None` (creates a new coordinator); child runners inherit the parent's coordinator so tree-wide queries see the entire run tree.	`None`
`parent_runner`	`Optional[DynamicPipelineRunner]`	The parent runner, if this is a child pipeline. Used to cascade `_is_paused` up the ancestor chain when a descendant run pauses.	`None`

Raises:

Type	Description
`RuntimeError`	If the snapshot has no associated stack.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def __init__(
    self,
    snapshot: "PipelineSnapshotResponse",
    run: Optional["PipelineRunResponse"],
    orchestrator: Optional["BaseOrchestrator"] = None,
    pause_coordinator: Optional[_PauseCoordinator] = None,
    parent_runner: Optional["DynamicPipelineRunner"] = None,
) -> None:
    """Initialize the dynamic pipeline runner.

    Args:
        snapshot: The snapshot of the pipeline.
        run: The pipeline run.
        orchestrator: The orchestrator to use. If not provided, the
            orchestrator will be inferred from the snapshot stack.
        pause_coordinator: The pause coordinator for the run tree. The
            root runner passes `None` (creates a new coordinator);
            child runners inherit the parent's coordinator so tree-wide
            queries see the entire run tree.
        parent_runner: The parent runner, if this is a child pipeline.
            Used to cascade `_is_paused` up the ancestor chain when a
            descendant run pauses.

    Raises:
        RuntimeError: If the snapshot has no associated stack.
    """
    if not snapshot.stack:
        raise RuntimeError("Missing stack for snapshot.")

    if (
        snapshot.pipeline_configuration.execution_mode
        == ExecutionMode.CONTINUE_ON_FAILURE
    ):
        logger.warning(
            "The `%s` execution mode is not supported for "
            "dynamic pipelines right now. "
            "The `%s` execution mode will be used instead.",
            snapshot.pipeline_configuration.execution_mode,
            ExecutionMode.STOP_ON_FAILURE,
        )

    self._parent_runner = parent_runner
    self._snapshot = snapshot
    self._pipeline: Optional["DynamicPipeline"] = None
    self._fail_fast = (
        snapshot.pipeline_configuration.execution_mode
        == ExecutionMode.FAIL_FAST
    )

    worker_count = handle_int_env_var(
        ENV_ZENML_DYNAMIC_PIPELINE_WORKER_COUNT, default=10
    )
    self._executor = ThreadPoolExecutor(max_workers=worker_count)

    stack = Stack.from_model(snapshot.stack)
    if orchestrator:
        self._orchestrator = orchestrator
    else:
        self._orchestrator = stack.orchestrator

    self._step_operator = stack.step_operator
    self._invocation_id_lock = threading.Lock()
    self._invocation_ids: Set[str] = set()

    self._run, self._orchestrator_run_id = self._prepare_run(run)
    self._existing_step_runs = {}
    self._existing_child_runs = {}

    if self._run.status not in {
        ExecutionStatus.INITIALIZING,
        ExecutionStatus.PROVISIONING,
    }:
        self._existing_step_runs = self._run.steps
        self._existing_child_runs = self._load_existing_child_runs()

    self._steps_to_monitor: Dict[str, "StepRunResponse"] = {}

    self._shutdown_requested = False
    self._failure_detected = False
    self._exception: Optional[BaseException] = None
    self._lifecycle_lock = threading.RLock()
    self._monitoring_event = threading.Event()
    self._startup_event = threading.Event()
    self._dependency_graph = InvocationDependencyGraph()
    self._future_registry = FutureRegistry()
    self._child_runners: Dict[str, "DynamicPipelineRunner"] = {}
    self._state = _PipelineThreadState()

    self._is_paused = False
    self._pause_coordinator = pause_coordinator or _PauseCoordinator()
    self._pause_coordinator.register(self)

Attributes

orchestrator: BaseOrchestrator property

The orchestrator used by this runner.

Returns:

Type	Description
`BaseOrchestrator`	The orchestrator.

orchestrator_run_id: str property

The orchestrator run ID associated with this runner.

Returns:

Type	Description
`str`	The orchestrator run ID.

pipeline: DynamicPipeline property

The pipeline that the runner is executing.

Raises:

Type	Description
`RuntimeError`	If the pipeline can't be loaded.

Returns:

Type	Description
`DynamicPipeline`	The pipeline that the runner is executing.

run: PipelineRunResponse property

The run executed by this runner.

Returns:

Type	Description
`PipelineRunResponse`	The pipeline run.

snapshot: PipelineSnapshotResponse property

The snapshot executed by this runner.

Returns:

Type	Description
`PipelineSnapshotResponse`	The pipeline snapshot.

Methods:

allocate_invocation_id(base_name: str, allow_suffix: bool = True) -> str

Allocate a new invocation ID with the given base name.

Parameters:

Name	Type	Description	Default
`base_name`	`str`	Base name for the invocation ID.	required
`allow_suffix`	`bool`	Whether to allow suffixing the invocation ID.	`True`

Returns:

Type	Description
`str`	The allocated invocation ID.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def allocate_invocation_id(
    self,
    base_name: str,
    allow_suffix: bool = True,
) -> str:
    """Allocate a new invocation ID with the given base name.

    Args:
        base_name: Base name for the invocation ID.
        allow_suffix: Whether to allow suffixing the invocation ID.

    Returns:
        The allocated invocation ID.
    """
    # TODO: maybe prevent `map:` prefixes for invocation IDs
    with self._invocation_id_lock:
        invocation_id = compute_invocation_id(
            existing_invocations=self._invocation_ids,
            base_name=base_name,
            allow_suffix=allow_suffix,
        )
        self._invocation_ids.add(invocation_id)

    return invocation_id

has_active_work() -> bool

Whether this runner is currently doing active work.

Returns:

Type	Description
`bool`	True if this runner has any active work, False otherwise.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def has_active_work(self) -> bool:
    """Whether this runner is currently doing active work.

    Returns:
        True if this runner has any active work, False otherwise.
    """
    if self._state.is_working():
        # Pipeline function is doing active work.
        return True
    for node in self._dependency_graph.list_nodes():
        if node.state in {NodeState.STARTING, NodeState.READY}:
            # A node in the graph is waiting to be executed.
            return True
        if node.state == NodeState.RUNNING:
            if isinstance(node, (ChildPipelineNode, MapNode)):
                continue

            # A step node is in progress -> we're doing active work.
            return True
    return False

launch_step(step: BaseStep, id: Optional[str], args: Tuple[Any, ...], kwargs: Dict[str, Any], after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None] = None, group: Optional[GroupInfo] = None, concurrent: bool = False) -> Union[StepRunOutputs, StepFuture]

launch_step(
    step: BaseStep,
    id: Optional[str],
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = None,
    group: Optional[GroupInfo] = None,
    concurrent: Literal[False] = False,
) -> StepRunOutputs

launch_step(
    step: BaseStep,
    id: Optional[str],
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = None,
    group: Optional[GroupInfo] = None,
    concurrent: Literal[True] = True,
) -> StepFuture

Launch a step.

Parameters:

Name	Type	Description	Default
`step`	`BaseStep`	The step to launch.	required
`id`	`Optional[str]`	The invocation ID of the step.	required
`args`	`Tuple[Any, ...]`	The arguments for the step function.	required
`kwargs`	`Dict[str, Any]`	The keyword arguments for the step function.	required
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	The step run output futures to wait for.	`None`
`group`	`Optional[GroupInfo]`	The group information for this step.	`None`
`concurrent`	`bool`	Whether to launch the step concurrently.	`False`

Raises:

Type	Description
`BaseException`	If the step failed.

Returns:

Type	Description
`Union[StepRunOutputs, StepFuture]`	The step run outputs or a future for the step run outputs.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def launch_step(
    self,
    step: "BaseStep",
    id: Optional[str],
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    after: Union[
        "AnyOutputFuture", Sequence["AnyOutputFuture"], None
    ] = None,
    group: Optional["GroupInfo"] = None,
    concurrent: bool = False,
) -> Union[StepRunOutputs, "StepFuture"]:
    """Launch a step.

    Args:
        step: The step to launch.
        id: The invocation ID of the step.
        args: The arguments for the step function.
        kwargs: The keyword arguments for the step function.
        after: The step run output futures to wait for.
        group: The group information for this step.
        concurrent: Whether to launch the step concurrently.

    Raises:
        BaseException: If the step failed.

    Returns:
        The step run outputs or a future for the step run outputs.
    """  # noqa: DOC502, DOC503
    step = step.copy()

    invocation_id = self.allocate_invocation_id(
        base_name=id or step.name, allow_suffix=not id
    )

    remaining_retries = None
    if step_run := self._existing_step_runs.get(invocation_id):
        runtime = get_step_runtime(
            step_config=step_run.config,
            pipeline_docker_settings=self._snapshot.pipeline_configuration.docker_settings,
            orchestrator=self._orchestrator,
        )

        if step_run.status.is_successful:
            # The step finished successfully, but we still need to return
            # a future in case the step was launched concurrently so the
            # caller gets the correct object back.
            if concurrent:
                execution_future = _IsolatedStepFuture(
                    pipeline_run_id=self._run.id,
                    invocation_id=invocation_id,
                )
                future = StepFuture(
                    invocation_id=invocation_id,
                    execution_future=execution_future,
                    output_keys=list(step_run.config.outputs),
                )
                self._register_concurrent_step_invocation(
                    future=future,
                    initial_state=NodeState.SUCCEEDED,
                )
                self.mark_node_succeeded(node_id=invocation_id)
                return future
            else:
                return load_step_run_outputs(step_run.id)

        if (
            runtime == StepRuntime.INLINE
            and step_run.status == ExecutionStatus.RUNNING
        ):
            # Inline steps that are in running state didn't have the
            # chance to report their failure back to ZenML before the
            # orchestration environment was shut down. But there is no
            # way that they're actually still running if we're in a new
            # orchestration environment, so we mark them as failed and
            # potentially restart them depending on the retry config.
            step_run = publish_failed_step_run(step_run.id)
            # Store the updated step run so we can compute the remaining
            # retries in the async submission flow.
            self._existing_step_runs[invocation_id] = step_run

        if step_run.status.is_failed:
            exception = self._get_step_exception(step_run=step_run)

            # If the step is running concurrently, we only raise the
            # exception once the future is awaited.
            if concurrent:
                execution_future = _IsolatedStepFuture(
                    pipeline_run_id=self._run.id,
                    invocation_id=invocation_id,
                )
                future = StepFuture(
                    invocation_id=invocation_id,
                    execution_future=execution_future,
                    output_keys=list(step_run.config.outputs),
                )
                self._register_concurrent_step_invocation(
                    future=future,
                    initial_state=NodeState.FAILED,
                )
                self.mark_node_failed(node_id=invocation_id)
                self.record_failure(
                    exception=wrap_step_failure(
                        exception, invocation_id=invocation_id
                    )
                )
                return future
            else:
                raise exception

        remaining_retries = get_remaining_retries(step_run=step_run)

        if step_run.status in {
            ExecutionStatus.PROVISIONING,
            ExecutionStatus.RUNNING,
        }:
            logger.info(
                "Restarting the monitoring of existing step `%s` "
                "(ID: %s). Remaining retries: %d",
                step_run.name,
                step_run.id,
                remaining_retries,
            )
            execution_future = _IsolatedStepFuture(
                pipeline_run_id=self._run.id,
                invocation_id=invocation_id,
            )
            monitoring_future = StepFuture(
                invocation_id=invocation_id,
                execution_future=execution_future,
                output_keys=list(step_run.config.outputs),
            )
            self._register_concurrent_step_invocation(
                future=monitoring_future,
                initial_state=NodeState.RUNNING,
            )
            self._register_isolated_step_for_monitoring(
                invocation_id=invocation_id,
                step_run=step_run,
            )
            self.mark_node_running(node_id=invocation_id)
            if concurrent:
                return monitoring_future
            else:
                return monitoring_future.result()

    inputs = convert_to_keyword_arguments(step.entrypoint, args, kwargs)

    config_overrides = None
    if self._run and self._run.triggered_by_deployment:
        # Deployment-specific step overrides
        config_overrides = StepConfigurationUpdate(
            enable_cache=False,
            step_operator=None,
            parameters={},
            runtime=StepRuntime.INLINE,
            group=group,
        )
    elif group:
        config_overrides = StepConfigurationUpdate(
            group=group,
        )

    if concurrent:
        future = StepFuture(
            invocation_id=invocation_id,
            output_keys=list(step.entrypoint_definition.outputs),
        )
        self._register_concurrent_step_invocation(
            future=future,
            step=step,
            inputs=inputs,
            after=after,
            config_overrides=config_overrides,
        )
        return future
    else:
        if (
            running_upstream_dependencies
            := get_running_upstream_dependencies(inputs, after)
        ):
            logger.info(
                "Waiting for upstream dependencies `%s` to finish before "
                "executing step `%s`.",
                ", ".join(running_upstream_dependencies),
                invocation_id,
            )

        compiled_step = self._compile_or_reuse(
            step=step,
            invocation_id=invocation_id,
            inputs=inputs,
            after=after,
            config=config_overrides,
        )

        step_run = self._run_sync_step(
            step=compiled_step, remaining_retries=remaining_retries
        )
        return load_step_run_outputs(step_run.id)

map(step: BaseStep, args: Tuple[Any, ...], kwargs: Dict[str, Any], after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None] = None, product: bool = False) -> MapResultsFuture

Map over step inputs.

Parameters:

Name	Type	Description	Default
`step`	`BaseStep`	The step to run.	required
`args`	`Tuple[Any, ...]`	The arguments for the step function.	required
`kwargs`	`Dict[str, Any]`	The keyword arguments for the step function.	required
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	The step run output futures to wait for before executing the steps.	`None`
`product`	`bool`	Whether to produce a cartesian product of the mapped inputs.	`False`

Returns:

Type	Description
`MapResultsFuture`	A future that represents the map results.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def map(
    self,
    step: "BaseStep",
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    after: Union[
        "AnyOutputFuture", Sequence["AnyOutputFuture"], None
    ] = None,
    product: bool = False,
) -> "MapResultsFuture":
    """Map over step inputs.

    Args:
        step: The step to run.
        args: The arguments for the step function.
        kwargs: The keyword arguments for the step function.
        after: The step run output futures to wait for before executing the
            steps.
        product: Whether to produce a cartesian product of the mapped
            inputs.

    Returns:
        A future that represents the map results.
    """
    step = step.copy()
    inputs = convert_to_keyword_arguments(step.entrypoint, args, kwargs)
    invocation_id = self.allocate_invocation_id(
        base_name=step.name, allow_suffix=True
    )
    map_future = MapResultsFuture(invocation_id=invocation_id)
    self._register_map_invocation(
        future=map_future,
        step=step,
        inputs=inputs,
        after=after,
        product=product,
    )
    return map_future

mark_node_failed(node_id: str) -> None

Mark a graph node as failed and propagate readiness changes.

This only updates the graph state. Use record_failure(exception) to also trigger the pipeline-level failure cascade.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def mark_node_failed(self, node_id: str) -> None:
    """Mark a graph node as failed and propagate readiness changes.

    This only updates the graph state. Use `record_failure(exception)` to
    also trigger the pipeline-level failure cascade.

    Args:
        node_id: The node ID.
    """
    self.notify_graph_changed(
        self._dependency_graph.mark_node_failed(node_id=node_id)
    )

mark_node_running(node_id: str) -> None

Mark a graph node as running and propagate readiness changes.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def mark_node_running(self, node_id: str) -> None:
    """Mark a graph node as running and propagate readiness changes.

    Args:
        node_id: The node ID.
    """
    self.notify_graph_changed(
        self._dependency_graph.mark_node_running(node_id=node_id)
    )

mark_node_starting(node_id: str) -> None

Mark a graph node as starting and propagate readiness changes.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def mark_node_starting(self, node_id: str) -> None:
    """Mark a graph node as starting and propagate readiness changes.

    Args:
        node_id: The node ID.
    """
    self.notify_graph_changed(
        self._dependency_graph.mark_node_starting(node_id=node_id)
    )

mark_node_succeeded(node_id: str) -> None

Mark a graph node as succeeded and propagate readiness changes.

Parameters:

Name	Type	Description	Default
`node_id`	`str`	The node ID.	required

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def mark_node_succeeded(self, node_id: str) -> None:
    """Mark a graph node as succeeded and propagate readiness changes.

    Args:
        node_id: The node ID.
    """
    self.notify_graph_changed(
        self._dependency_graph.mark_node_succeeded(node_id=node_id)
    )

notify_graph_changed(nodes_ready: bool) -> None

Wake up the startup loop if new nodes became ready.

Parameters:

Name	Type	Description	Default
`nodes_ready`	`bool`	Whether any new nodes are ready.	required

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def notify_graph_changed(self, nodes_ready: bool) -> None:
    """Wake up the startup loop if new nodes became ready.

    Args:
        nodes_ready: Whether any new nodes are ready.
    """
    if nodes_ready:
        self._startup_event.set()

raise_if_startup_cancelled() -> None

Abort startup work once a failure has been detected.

Raises:

Type	Description
`StartupCancelled`	If a failure has been detected.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def raise_if_startup_cancelled(self) -> None:
    """Abort startup work once a failure has been detected.

    Raises:
        StartupCancelled: If a failure has been detected.
    """
    if self._failure_detected:
        raise StartupCancelled(
            "Dynamic pipeline runner is not starting new work."
        )

record_failure(exception: BaseException) -> None

Trigger the pipeline-level failure cascade.

Marks the runner as failed (idempotent), cancels in-flight startup work, stops isolated steps in fail-fast mode and forwards shutdown to running child pipelines.

Parameters:

Name	Type	Description	Default
`exception`	`BaseException`	The failure exception.	required

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def record_failure(self, exception: BaseException) -> None:
    """Trigger the pipeline-level failure cascade.

    Marks the runner as failed (idempotent), cancels in-flight startup
    work, stops isolated steps in fail-fast mode and forwards shutdown to
    running child pipelines.

    Args:
        exception: The failure exception.
    """
    # RunPaused must be handled separately and not recorded as a failure.
    assert not isinstance(exception, RunPaused)
    self._on_failure_detected(exception=exception)

request_shutdown(reason: BaseException) -> None

Ask the runner to shut down.

Parameters:

Name	Type	Description	Default
`reason`	`BaseException`	Exception describing why shutdown is being requested.	required

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def request_shutdown(self, reason: BaseException) -> None:
    """Ask the runner to shut down.

    Args:
        reason: Exception describing why shutdown is being requested.
    """
    self._on_failure_detected(exception=reason)

run_pipeline() -> None

Run the pipeline.

Raises:

Type	Description
`BaseException`	If the init hook fails.
`Exception`	If loading the pipeline fails.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def run_pipeline(self) -> None:
    """Run the pipeline.

    Raises:
        BaseException: If the init hook fails.
        Exception: If loading the pipeline fails.
    """
    self._state.id = threading.get_ident()
    logs_context: ContextManager[Any] = nullcontext()
    if is_pipeline_logging_enabled(self._snapshot.pipeline_configuration):
        logs_context = setup_logging_context(
            source="orchestrator", pipeline_run=self._run
        )

    with logs_context:
        run_start_hook = False
        run_resume_hook = self._run.status in {
            ExecutionStatus.RESUMING,
            ExecutionStatus.PAUSED,
        }

        if self._run.status.is_finished:
            logger.info("Run `%s` is already finished.", str(self._run.id))
            return
        elif self._run.status in {
            ExecutionStatus.RESUMING,
            # TODO: We should probably not include paused here, as we want
            # to require the user to move a run to resuming first.
            # Otherwise, the same run might be resumed multiple times.
            ExecutionStatus.PAUSED,
            ExecutionStatus.RETRYING,
        }:
            self._run = Client().zen_store.update_run(
                run_id=self._run.id,
                run_update=PipelineRunUpdate(
                    status=ExecutionStatus.RUNNING,
                ),
            )
            logger.info("Resuming run `%s`.", str(self._run.id))
        elif self._run.status == ExecutionStatus.RUNNING:
            logger.info("Continuing existing run `%s`.", str(self._run.id))
        elif self._run.status in {
            ExecutionStatus.INITIALIZING,
            ExecutionStatus.PROVISIONING,
        }:
            # Set the run status to running already, in case no steps start
            # immediately which would otherwise cause the run to be stuck in
            # some init state.
            self._run = Client().zen_store.update_run(
                run_id=self._run.id,
                run_update=PipelineRunUpdate(
                    status=ExecutionStatus.RUNNING,
                ),
            )
            run_start_hook = True

        assert self._snapshot.stack

        try:
            pipeline = self.pipeline
        except Exception as e:
            self._publish_run_status(exception=e)
            raise

        with (
            InMemoryArtifactCache(),
            env_utils.temporary_runtime_environment(
                self._snapshot.pipeline_configuration, self._snapshot.stack
            ),
            DynamicPipelineRunContext(
                pipeline=pipeline,
                run=self._run,
                snapshot=self._snapshot,
                runner=self,
            ),
        ):
            monitoring_thread = self._start_monitoring_loop()
            startup_thread = self._start_startup_loop()

            if run_start_hook:
                run_lifecycle_hook(
                    self._snapshot.pipeline_configuration.start_hook_source,
                    HookType.RUN_START,
                )
            elif run_resume_hook:
                run_lifecycle_hook(
                    self._snapshot.pipeline_configuration.resume_hook_source,
                    HookType.RUN_RESUME,
                )

            if not self._run.triggered_by_deployment:
                # Only run the init hook if the run is not triggered by
                # a deployment, as the deployment service will have
                # already run the init hook.
                try:
                    self._orchestrator.run_init_hook(
                        snapshot=self._snapshot
                    )
                except BaseException as e:
                    # The init hook itself is not tracked, but its failure
                    # ends the run, so the run end and failure hooks fire
                    # before the failure propagates.
                    self._fire_terminal_run_hooks(
                        ExecutionStatus.FAILED, e
                    )
                    raise

            try:
                self._run_entrypoint_and_finalize()
            finally:
                if not self._run.triggered_by_deployment:
                    # Only run the cleanup hook if the run is not
                    # triggered by a deployment, as the deployment
                    # service will have already run the cleanup hook.
                    self._orchestrator.run_cleanup_hook(
                        snapshot=self._snapshot
                    )

                self._executor.shutdown(wait=True, cancel_futures=True)

                self._shutdown_requested = True
                self._startup_event.set()
                self._monitoring_event.set()
                monitoring_thread.join()
                startup_thread.join()
                self._pause_coordinator.unregister(self)

submit_child_pipeline(pipeline: DynamicPipeline, args: Sequence[Any], kwargs: Dict[str, Any], after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None] = None, *, concurrent: bool = False) -> Union[PipelineFuture, PipelineRunOutputs]

submit_child_pipeline(
    pipeline: DynamicPipeline,
    args: Sequence[Any],
    kwargs: Dict[str, Any],
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = None,
    *,
    concurrent: Literal[True],
) -> PipelineFuture

submit_child_pipeline(
    pipeline: DynamicPipeline,
    args: Sequence[Any],
    kwargs: Dict[str, Any],
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = None,
    *,
    concurrent: Literal[False] = False,
) -> PipelineRunOutputs

Submit a child pipeline.

Parameters:

Name	Type	Description	Default
`pipeline`	`DynamicPipeline`	Child pipeline to execute.	required
`args`	`Sequence[Any]`	Positional pipeline arguments.	required
`kwargs`	`Dict[str, Any]`	Keyword pipeline arguments.	required
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	Optional dependency futures.	`None`
`concurrent`	`bool`	Whether to run the child pipeline concurrently.	`False`

Returns:

Type	Description
`Union[PipelineFuture, PipelineRunOutputs]`	A `PipelineFuture` for concurrent calls or the resolved
`Union[PipelineFuture, PipelineRunOutputs]`	`PipelineRunOutputs` for synchronous calls.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def submit_child_pipeline(
    self,
    pipeline: "DynamicPipeline",
    args: Sequence[Any],
    kwargs: Dict[str, Any],
    after: Union[
        "AnyOutputFuture", Sequence["AnyOutputFuture"], None
    ] = None,
    *,
    concurrent: bool = False,
) -> Union[PipelineFuture, PipelineRunOutputs]:
    """Submit a child pipeline.

    Args:
        pipeline: Child pipeline to execute.
        args: Positional pipeline arguments.
        kwargs: Keyword pipeline arguments.
        after: Optional dependency futures.
        concurrent: Whether to run the child pipeline concurrently.

    Returns:
        A `PipelineFuture` for concurrent calls or the resolved
        `PipelineRunOutputs` for synchronous calls.
    """
    pipeline = pipeline.copy()
    # Fail early for invalid args/kwargs.
    convert_to_keyword_arguments(pipeline.entrypoint, tuple(args), kwargs)
    node_id = self.allocate_invocation_id(
        base_name=f"{CHILD_PIPELINE_INVOCATION_ID_PREFIX}{pipeline.name}"
    )

    if concurrent:
        pipeline_future = PipelineFuture(
            invocation_id=node_id,
            declared_output_names=get_pipeline_entrypoint_output_names(
                pipeline.entrypoint
            ),
        )
        self._register_concurrent_child_pipeline_invocation(
            node_id=node_id,
            pipeline=pipeline,
            args=tuple(args),
            kwargs=kwargs,
            after=after,
            future=pipeline_future,
        )
        return pipeline_future
    else:
        child_run = self._prepare_child_run(
            pipeline=pipeline,
            args=args,
            kwargs=kwargs,
            after=after,
            child_invocation_id=node_id,
        )
        child_runner = self._build_child_runner(child_run=child_run)

        with self._lifecycle_lock:
            if self._failure_detected:
                self._mark_child_run_placeholder_failed(
                    child_run=child_run,
                    reason=(
                        "Parent run shut down before child run started."
                    ),
                )
                self.raise_if_startup_cancelled()
            self._child_runners[node_id] = child_runner
        logger.info(
            "Launching child pipeline `%s`", child_runner.pipeline.name
        )
        try:
            # Yield the pipeline thread while the child runner is running
            # sync. This is necessary to avoid the following deadlock:
            # - The child run waits on a wait condition
            # - Us (= the parent runner) reports active work, which causes
            #   the wait condition to never timeout.
            with self._state.release():
                child_runner.run_pipeline()
            terminal_run = self._validate_successful_child_run(
                child_runner.run
            )
        finally:
            self._unregister_child_runner(node_id=node_id)
        logger.info(
            "Child pipeline `%s` completed", child_runner.pipeline.name
        )
        return load_pipeline_run_outputs(terminal_run)

wait(schema: Optional[Any] = None, type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT, timeout: int = 600, poll_interval: int = 5, question: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, name: Optional[str] = None, after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None] = None) -> Any

wait(
    schema: Type[T],
    type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT,
    timeout: int = 600,
    poll_interval: int = 5,
    question: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    name: Optional[str] = None,
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = None,
) -> T

wait(
    schema: object = None,
    type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT,
    timeout: int = 600,
    poll_interval: int = 5,
    question: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    name: Optional[str] = None,
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = None,
) -> Any

Create and poll a run wait condition.

Parameters:

Name	Type	Description	Default
`schema`	`Optional[Any]`	Optional expected output type for the resolved result.	`None`
`type`	`RunWaitConditionType`	Wait condition type.	`EXTERNAL_INPUT`
`timeout`	`int`	Earliest time in seconds at which polling may give up and pause the run. The actual pause is deferred until all active work in this pipeline (or any parent/child pipelines) has finished.	`600`
`poll_interval`	`int`	Poll interval in seconds.	`5`
`question`	`Optional[str]`	Optional question shown to external actors.	`None`
`metadata`	`Optional[Dict[str, Any]]`	Optional metadata attached to the condition.	`None`
`name`	`Optional[str]`	Optional deterministic wait condition name.	`None`
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	Optional upstream futures that must finish before waiting.	`None`

Raises:

Type	Description
`RuntimeError`	If called outside the dynamic pipeline function.
`_WaitConditionAborted`	If the wait condition was aborted.
`RunPaused`	If the wait condition polling timed out and the run was transitioned to PAUSED.
`KeyboardInterrupt`	If interrupted while waiting.
`BaseException`	If polling fails after the lease is abandoned.

Returns:

Type	Description
`Any`	The resolved wait condition value.

Source code in src/zenml/execution/pipeline/dynamic/runner.py

def wait(
    self,
    schema: Optional[Any] = None,
    type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT,
    timeout: int = 600,
    poll_interval: int = 5,
    question: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    name: Optional[str] = None,
    after: Union[
        "AnyOutputFuture", Sequence["AnyOutputFuture"], None
    ] = None,
) -> Any:
    """Create and poll a run wait condition.

    Args:
        schema: Optional expected output type for the resolved result.
        type: Wait condition type.
        timeout: Earliest time in seconds at which polling may give up
            and pause the run. The actual pause is deferred until all
            active work in this pipeline (or any parent/child pipelines) has
            finished.
        poll_interval: Poll interval in seconds.
        question: Optional question shown to external actors.
        metadata: Optional metadata attached to the condition.
        name: Optional deterministic wait condition name.
        after: Optional upstream futures that must finish before waiting.

    Raises:
        RuntimeError: If called outside the dynamic pipeline function.
        _WaitConditionAborted: If the wait condition was aborted.
        RunPaused: If the wait condition polling timed out and the run
            was transitioned to PAUSED.
        KeyboardInterrupt: If interrupted while waiting.
        BaseException: If polling fails after the lease is abandoned.

    Returns:
        The resolved wait condition value.
    """  # noqa: DOC503
    from zenml.execution.pipeline.dynamic.run_context import (
        DynamicPipelineRunContext,
    )
    from zenml.steps.step_context import StepContext
    from zenml.utils.time_utils import utc_now

    context = DynamicPipelineRunContext.get()
    if not context:
        raise RuntimeError(
            "`zenml.wait(...)` can only be used inside dynamic pipelines."
        )

    if StepContext.is_active():
        raise RuntimeError(
            "`zenml.wait(...)` cannot be called inside a step function. "
            "Use it only in the pipeline function."
        )

    if threading.get_ident() != self._state.id:
        # `wait(...)` is currently only allowed as a sync call in the
        # pipeline function.
        raise RuntimeError(
            "`zenml.wait(...)` must be called from the pipeline "
            "thread, not from a worker thread spawned inside the "
            "pipeline function."
        )

    wait_condition_name = name or context.next_wait_condition_name()

    for future in collect_futures(after=after, expand_map_results=True):
        future.wait()

    condition = Client().zen_store.create_run_wait_condition(
        RunWaitConditionRequest(
            project=self._run.project_id,
            run=self._run.id,
            name=wait_condition_name,
            type=type,
            question=question,
            metadata=metadata or {},
            data_schema=pydantic_utils.get_json_schema_for_type(schema)
            if schema
            else None,
        )
    )
    state = self._handle_wait_condition_state(
        condition=condition, schema=schema
    )
    if state.is_terminal:
        return state.value

    if poll_interval <= 0:
        logger.debug(
            "Non-positive poll interval provided, falling back to 5 seconds."
        )
        poll_interval = 5

    deadline = time.time() + timeout
    logger.info(
        "Waiting on wait condition `%s` (type=%s, timeout=%ss, poll=%ss).",
        wait_condition_name,
        type.value,
        timeout,
        poll_interval,
    )

    with self._state.claim(deadline):
        try:
            with maybe_enable_interactive_wait_prompt(
                orchestrator=self._orchestrator,
                condition=condition,
            ) as interactive_prompting_enabled:
                # Poll while own deadline hasn't expired or any runner
                # in the tree still has active work. Once both flip
                # false the loop exits naturally; the lease finalize
                # below transitions the run to PAUSED on the server.
                while (
                    time.time() < deadline
                    or self._pause_coordinator.has_active_work()
                ):
                    lease_now = utc_now()
                    status = Client().zen_store.update_run_wait_condition_lease(
                        run_wait_condition_id=condition.id,
                        lease_update=RunWaitConditionLeaseUpdate(
                            poller_instance_id=self._orchestrator_run_id,
                            poller_lease_expires_at=lease_now
                            + timedelta(
                                seconds=max(15, poll_interval * 2)
                            ),
                        ),
                    )
                    if status != RunWaitConditionStatus.PENDING:
                        condition = (
                            Client().zen_store.get_run_wait_condition(
                                condition.id, hydrate=True
                            )
                        )
                    state = self._handle_wait_condition_state(
                        condition=condition, schema=schema
                    )
                    if state.is_terminal:
                        return state.value

                    if interactive_prompting_enabled:
                        # If we're running interactively, sleep until
                        # the polling interval is reached or the user
                        # submitted input in their terminal.
                        poll_interactive_wait_condition_input(
                            condition=condition,
                            poll_interval=poll_interval,
                        )
                    else:
                        time.sleep(poll_interval)
        except _WaitConditionAborted:
            raise
        except KeyboardInterrupt as keyboard_interrupt:
            try:
                Client().zen_store.resolve_run_wait_condition(
                    run_wait_condition_id=condition.id,
                    resolve_request=RunWaitConditionResolveRequest(
                        resolution=RunWaitConditionResolution.ABORT,
                    ),
                )
            except Exception as e:
                logger.warning(
                    "Failed to abort wait condition `%s` after "
                    "keyboard interrupt: %s",
                    condition.id,
                    e,
                )
                # When resolving a wait condition with `ABORT` resolution,
                # the server will transition the run to a `STOPPED` status.
                # This failed, which means we re-raise the
                # KeyboardInterrupt, which causes downstream code to publish
                # a `FAILED` run status.
                raise keyboard_interrupt
            else:
                # Raise _WaitConditionAborted to signal that the wait
                # condition was aborted. Downstream code will not publish
                # a `FAILED` run status in this case.
                raise _WaitConditionAborted(
                    f"Wait condition `{condition.name}` was aborted."
                ) from keyboard_interrupt
        except BaseException:
            try:
                Client().zen_store.update_run_wait_condition_lease(
                    run_wait_condition_id=condition.id,
                    lease_update=RunWaitConditionLeaseUpdate(
                        poller_instance_id=self._orchestrator_run_id,
                        poller_lease_expires_at=utc_now(),
                        mode=RunWaitConditionLeaseMode.ABANDON,
                    ),
                )
            except Exception as e:
                logger.warning(
                    "Failed to abandon wait condition `%s` after wait "
                    "loop failure: %s",
                    condition.id,
                    e,
                )
            raise

    status = Client().zen_store.update_run_wait_condition_lease(
        run_wait_condition_id=condition.id,
        lease_update=RunWaitConditionLeaseUpdate(
            poller_instance_id=self._orchestrator_run_id,
            poller_lease_expires_at=utc_now(),
            mode=RunWaitConditionLeaseMode.FINALIZE,
        ),
    )
    if status != RunWaitConditionStatus.PENDING:
        condition = Client().zen_store.get_run_wait_condition(
            condition.id, hydrate=True
        )
    state = self._handle_wait_condition_state(
        condition=condition, schema=schema
    )
    if state.is_terminal:
        return state.value

    raise RunPaused()

RunPaused

Bases: Exception

Raised when a run (or one of its descendants) entered a paused state.

Functions: Modules

utils

Dynamic pipeline execution utilities.

Classes Functions:

collect_futures(inputs: Optional[Dict[str, Any]] = None, after: Union[AnyOutputFuture, Sequence[AnyOutputFuture], None] = None, expand_map_results: bool = False) -> List[AnyOutputFuture]

collect_futures(
    inputs: Optional[Dict[str, Any]] = ...,
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = ...,
    expand_map_results: Literal[True] = ...,
) -> List[AnyOutputFuture]

collect_futures(
    inputs: Optional[Dict[str, Any]] = ...,
    after: Union[
        AnyOutputFuture, Sequence[AnyOutputFuture], None
    ] = ...,
    expand_map_results: Literal[False] = ...,
) -> List[AnyOutputFuture]

Collect futures referenced in step inputs and after.

Parameters:

Name	Type	Description	Default
`inputs`	`Optional[Dict[str, Any]]`	Optional step inputs to inspect for futures.	`None`
`after`	`Union[AnyOutputFuture, Sequence[AnyOutputFuture], None]`	Optional explicit upstream dependencies. Must be a future or a sequence containing only futures.	`None`
`expand_map_results`	`bool`	Whether map futures should be expanded into their child step futures.	`False`

Raises:

Type	Description
`TypeError`	If `after` is not a future or a sequence containing only futures.

Returns:

Type	Description
`List[AnyOutputFuture]`	The collected futures.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def collect_futures(
    inputs: Optional[Dict[str, Any]] = None,
    after: Union["AnyOutputFuture", Sequence["AnyOutputFuture"], None] = None,
    expand_map_results: bool = False,
) -> List["AnyOutputFuture"]:
    """Collect futures referenced in step inputs and `after`.

    Args:
        inputs: Optional step inputs to inspect for futures.
        after: Optional explicit upstream dependencies. Must be a future or a
            sequence containing only futures.
        expand_map_results: Whether map futures should be expanded into their
            child step futures.

    Raises:
        TypeError: If `after` is not a future or a sequence containing only
            futures.

    Returns:
        The collected futures.
    """
    from zenml.execution.pipeline.dynamic.outputs import (
        ArtifactFuture,
        MapResultsFuture,
        PipelineFuture,
        StepFuture,
    )

    VALID_OUTPUT_FUTURE_CLASSES = (
        ArtifactFuture,
        StepFuture,
        MapResultsFuture,
        PipelineFuture,
    )

    futures: List["AnyOutputFuture"] = []

    def _append_future(future: "AnyOutputFuture") -> None:
        if expand_map_results and isinstance(future, MapResultsFuture):
            futures.extend(future.futures)
        else:
            futures.append(future)

    def _collect_input_value(value: Any) -> None:
        if isinstance(value, VALID_OUTPUT_FUTURE_CLASSES):
            _append_future(value)
            return

        if isinstance(value, Sequence) and all(
            isinstance(item, VALID_OUTPUT_FUTURE_CLASSES) for item in value
        ):
            for item in value:
                _append_future(item)
            return

    if inputs:
        for value in inputs.values():
            _collect_input_value(value=value)

    if after is not None:
        if isinstance(after, VALID_OUTPUT_FUTURE_CLASSES):
            _append_future(after)
        elif isinstance(after, Sequence) and all(
            isinstance(item, VALID_OUTPUT_FUTURE_CLASSES) for item in after
        ):
            for item in after:
                _append_future(item)
        else:
            raise TypeError(
                "`after` must be a future or a sequence of futures."
            )

    return futures

expand_mapped_inputs(inputs: Dict[str, Any], product: bool = False) -> List[Dict[str, Any]]

Find the mapped and unmapped inputs of a step.

Parameters:

Name	Type	Description	Default
`inputs`	`Dict[str, Any]`	The step function inputs.	required
`product`	`bool`	Whether to produce a cartesian product of the mapped inputs.	`False`

Raises:

Type	Description
`RuntimeError`	If no mapped inputs are found or the input combinations are not valid.

Returns:

Type	Description
`List[Dict[str, Any]]`	The step inputs.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def expand_mapped_inputs(
    inputs: Dict[str, Any],
    product: bool = False,
) -> List[Dict[str, Any]]:
    """Find the mapped and unmapped inputs of a step.

    Args:
        inputs: The step function inputs.
        product: Whether to produce a cartesian product of the mapped inputs.

    Raises:
        RuntimeError: If no mapped inputs are found or the input combinations
            are not valid.

    Returns:
        The step inputs.
    """
    from zenml.execution.pipeline.dynamic.outputs import OutputArtifact

    static_inputs: Dict[str, Any] = {}
    mapped_input_names: List[str] = []
    mapped_inputs: List[Tuple["OutputArtifact", ...]] = []

    for key, value in inputs.items():
        if isinstance(value, _Unmapped):
            static_inputs[key] = value.value
        elif isinstance(value, OutputArtifact):
            if value.item_count is None:
                static_inputs[key] = value
            elif value.item_count == 0:
                raise RuntimeError(
                    f"Artifact `{value.id}` has 0 items and cannot be mapped "
                    "over. Wrap it with the `unmapped(...)` function to pass "
                    "the artifact without mapping over it."
                )
            else:
                mapped_input_names.append(key)
                mapped_inputs.append(
                    tuple(
                        value.chunk(index=i) for i in range(value.item_count)
                    )
                )
        elif (
            isinstance(value, ArtifactVersionResponse)
            and value.item_count is not None
        ):
            static_inputs[key] = value
            logger.warning(
                "Received sequence-like artifact for step input `%s`. Mapping "
                "over artifacts that are not step output artifacts is "
                "currently not supported, and the complete artifact will be "
                "passed to all steps. If you want to silence this warning, "
                "wrap your input with the `unmapped(...)` function.",
                key,
            )
        elif (
            isinstance(value, Sequence)
            and value
            and all(isinstance(item, OutputArtifact) for item in value)
        ):
            # List of step output artifacts, in this case the mapping is over
            # the items of the list
            mapped_input_names.append(key)
            mapped_inputs.append(tuple(value))
        elif isinstance(value, Sequence):
            logger.warning(
                "Received sequence-like data for step input `%s`. Mapping over "
                "data that is not a step output artifact is currently not "
                "supported, and the complete data will be passed to all steps. "
                "If you want to silence this warning, wrap your input with the "
                "`unmapped(...)` function.",
                key,
            )
            static_inputs[key] = value
        else:
            static_inputs[key] = value

    if len(mapped_inputs) == 0:
        raise RuntimeError(
            "No inputs to map over found. When calling `.map(...)` or "
            "`.product(...)` on a step, you need to pass at least one "
            "sequence-like step output of a previous step as input."
        )

    step_inputs = []

    if product:
        for input_combination in itertools.product(*mapped_inputs):
            all_inputs = copy.deepcopy(static_inputs)
            for name, value in zip(mapped_input_names, input_combination):
                all_inputs[name] = value
            step_inputs.append(all_inputs)
    else:
        item_counts = [len(inputs) for inputs in mapped_inputs]
        if not all(count == item_counts[0] for count in item_counts):
            raise RuntimeError(
                f"All mapped input artifacts must have the same "
                "item counts, but you passed artifacts with item counts "
                f"{item_counts}. If you want "
                "to pass sequence-like artifacts without mapping over "
                "them, wrap them with the `unmapped(...)` function."
            )

        for i in range(item_counts[0]):
            all_inputs = copy.deepcopy(static_inputs)
            for name, artifact in zip(
                mapped_input_names,
                [artifact_list[i] for artifact_list in mapped_inputs],
            ):
                all_inputs[name] = artifact
            step_inputs.append(all_inputs)

    return step_inputs

get_latest_step_run(pipeline_run_id: UUID, invocation_id: str, hydrate: bool = False) -> StepRunResponse

Get the latest step run for a step.

Parameters:

Name	Type	Description	Default
`pipeline_run_id`	`UUID`	The ID of the pipeline run.	required
`invocation_id`	`str`	The invocation ID of the step.	required
`hydrate`	`bool`	Whether to hydrate the step run.	`False`

Raises:

Type	Description
`RuntimeError`	If no step run exists for the given invocation ID.

Returns:

Type	Description
`StepRunResponse`	The latest step run.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def get_latest_step_run(
    pipeline_run_id: UUID, invocation_id: str, hydrate: bool = False
) -> "StepRunResponse":
    """Get the latest step run for a step.

    Args:
        pipeline_run_id: The ID of the pipeline run.
        invocation_id: The invocation ID of the step.
        hydrate: Whether to hydrate the step run.

    Raises:
        RuntimeError: If no step run exists for the given invocation ID.

    Returns:
        The latest step run.
    """
    step_runs = Client().list_run_steps(
        pipeline_run_id=pipeline_run_id,
        name=invocation_id,
        exclude_retried=True,
        size=1,
        hydrate=hydrate,
    )

    if not step_runs:
        raise RuntimeError(
            f"Step `{invocation_id}` not found in pipeline run "
            f"`{pipeline_run_id}`."
        )

    return step_runs.items[0]

get_remaining_retries(step_run: StepRunResponse) -> int

Get the remaining retries for a step run.

Parameters:

Name	Type	Description	Default
`step_run`	`StepRunResponse`	The step run to get the remaining retries for.	required

Returns:

Type	Description
`int`	The remaining retries for the step run.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def get_remaining_retries(step_run: "StepRunResponse") -> int:
    """Get the remaining retries for a step run.

    Args:
        step_run: The step run to get the remaining retries for.

    Returns:
        The remaining retries for the step run.
    """
    max_retries = (
        step_run.config.retry.max_retries if step_run.config.retry else 0
    )
    return max(0, 1 + max_retries - step_run.version)

load_pipeline_run_outputs(run: PipelineRunResponse) -> PipelineRunOutputs

Load output artifacts from a finished pipeline run.

Parameters:

Name	Type	Description	Default
`run`	`PipelineRunResponse`	The finished pipeline run.	required

Returns:

Type	Description
`PipelineRunOutputs`	The pipeline outputs.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def load_pipeline_run_outputs(
    run: PipelineRunResponse,
) -> "PipelineRunOutputs":
    """Load output artifacts from a finished pipeline run.

    Args:
        run: The finished pipeline run.

    Returns:
        The pipeline outputs.
    """
    outputs = tuple(run.outputs.values())
    if not outputs:
        return None
    if len(outputs) == 1:
        return outputs[0]
    return outputs

load_step_run_outputs(step_run_id: UUID) -> StepRunOutputs

Load the outputs of a step run.

Parameters:

Name	Type	Description	Default
`step_run_id`	`UUID`	The ID of the step run.	required

Returns:

Type	Description
`StepRunOutputs`	The outputs of the step run.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def load_step_run_outputs(step_run_id: UUID) -> "StepRunOutputs":
    """Load the outputs of a step run.

    Args:
        step_run_id: The ID of the step run.

    Returns:
        The outputs of the step run.
    """
    from zenml.execution.pipeline.dynamic.outputs import OutputArtifact

    step_run = Client().zen_store.get_run_step(step_run_id)

    def _convert_output_artifact(
        output_name: str, artifact: "ArtifactVersionResponse"
    ) -> "OutputArtifact":
        return OutputArtifact(
            output_name=output_name,
            step_name=step_run.name,
            **artifact.model_dump(),
        )

    output_artifacts = step_run.regular_outputs
    if len(output_artifacts) == 0:
        return None
    elif len(output_artifacts) == 1:
        name, artifact = next(iter(output_artifacts.items()))
        return _convert_output_artifact(output_name=name, artifact=artifact)
    else:
        # Make sure we return them in the same order as they're defined in the
        # step configuration, as we don't enforce any ordering in the DB.
        outputs = []

        for template_name in step_run.config.outputs.keys():
            name = string_utils.format_name_template(
                template_name,
                substitutions=step_run.config.substitutions,
            )
            outputs.append(
                _convert_output_artifact(
                    output_name=template_name, artifact=output_artifacts[name]
                )
            )

        return tuple(outputs)

unmapped(value: T) -> _Unmapped[T]

Helper function to pass an input without mapping over it.

Wrap any step input with this function and then pass it to step.map(...) to pass the full value to all steps.

Parameters:

Name	Type	Description	Default
`value`	`T`	The value to wrap.	required

Returns:

Type	Description
`_Unmapped[T]`	The wrapped value.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def unmapped(value: T) -> _Unmapped[T]:
    """Helper function to pass an input without mapping over it.

    Wrap any step input with this function and then pass it to `step.map(...)`
    to pass the full value to all steps.

    Args:
        value: The value to wrap.

    Returns:
        The wrapped value.
    """
    return _Unmapped(value)

wait(schema: Optional[Any] = None, type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT, timeout: int = 600, poll_interval: int = 5, question: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, after: Optional[Sequence[AnyOutputFuture]] = None, name: Optional[str] = None) -> Any

wait(
    schema: Type[T],
    type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT,
    timeout: int = 600,
    poll_interval: int = 5,
    question: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    after: Optional[Sequence[AnyOutputFuture]] = None,
    name: Optional[str] = None,
) -> T

wait(
    schema: object = None,
    type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT,
    timeout: int = 600,
    poll_interval: int = 5,
    question: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    after: Optional[Sequence[AnyOutputFuture]] = None,
    name: Optional[str] = None,
) -> Any

Pause dynamic execution on an external wait condition.

Parameters:

Name	Type	Description	Default
`schema`	`Optional[Any]`	Optional expected output type for the resolved result.	`None`
`type`	`RunWaitConditionType`	Wait condition type.	`EXTERNAL_INPUT`
`timeout`	`int`	Maximum time in seconds to poll before pausing.	`600`
`poll_interval`	`int`	Poll interval in seconds.	`5`
`question`	`Optional[str]`	Optional question shown to external actors.	`None`
`metadata`	`Optional[Dict[str, Any]]`	Optional metadata attached to the condition.	`None`
`after`	`Optional[Sequence[AnyOutputFuture]]`	Optional upstream futures that must finish before waiting.	`None`
`name`	`Optional[str]`	Optional deterministic wait condition name.	`None`

Raises:

Type	Description
`RuntimeError`	If called outside of dynamic pipeline execution.

Returns:

Type	Description
`Any`	The resolved wait condition value.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def wait(
    schema: Optional[Any] = None,
    type: RunWaitConditionType = RunWaitConditionType.EXTERNAL_INPUT,
    timeout: int = 600,
    poll_interval: int = 5,
    question: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    after: Optional[Sequence["AnyOutputFuture"]] = None,
    name: Optional[str] = None,
) -> Any:
    """Pause dynamic execution on an external wait condition.

    Args:
        schema: Optional expected output type for the resolved result.
        type: Wait condition type.
        timeout: Maximum time in seconds to poll before pausing.
        poll_interval: Poll interval in seconds.
        question: Optional question shown to external actors.
        metadata: Optional metadata attached to the condition.
        after: Optional upstream futures that must finish before waiting.
        name: Optional deterministic wait condition name.

    Raises:
        RuntimeError: If called outside of dynamic pipeline execution.

    Returns:
        The resolved wait condition value.
    """
    from zenml.execution.pipeline.dynamic.run_context import (
        DynamicPipelineRunContext,
    )

    context = DynamicPipelineRunContext.get()
    if not context:
        raise RuntimeError(
            "`zenml.wait(...)` can only be used inside dynamic pipelines."
        )

    return context.runner.wait(
        schema=schema,
        type=type,
        timeout=timeout,
        poll_interval=poll_interval,
        question=question,
        after=after,
        metadata=metadata,
        name=name,
    )

wait_for_step_run_to_finish(step_run_id: UUID) -> StepRunResponse

Wait until a step run is finished.

Parameters:

Name	Type	Description	Default
`step_run_id`	`UUID`	The ID of the step run.	required

Returns:

Type	Description
`StepRunResponse`	The finished step run.

Source code in src/zenml/execution/pipeline/dynamic/utils.py

def wait_for_step_run_to_finish(step_run_id: UUID) -> "StepRunResponse":
    """Wait until a step run is finished.

    Args:
        step_run_id: The ID of the step run.

    Returns:
        The finished step run.
    """
    sleep_interval = 1
    max_sleep_interval = 64

    while True:
        step_run = Client().zen_store.get_run_step(step_run_id)

        if step_run.status != ExecutionStatus.RUNNING:
            return step_run

        logger.debug(
            "Waiting for step run with ID %s to finish (current status: %s)",
            step_run_id,
            step_run.status,
        )
        time.sleep(sleep_interval)
        if sleep_interval < max_sleep_interval:
            sleep_interval *= 2

Modules

`utils`

Pipeline execution utilities.

Classes Functions:

compute_invocation_id(existing_invocations: Set[str], base_name: str, allow_suffix: bool = True) -> str

Compute the invocation ID.

Parameters:

Name	Type	Description	Default
`existing_invocations`	`Set[str]`	The existing invocation IDs.	required
`base_name`	`str`	Base name for the invocation. Used as the ID directly when unique, otherwise suffixed with `_2`, `_3`, ... until a free slot is found (when `allow_suffix=True`).	required
`allow_suffix`	`bool`	Whether a suffix can be appended to the invocation ID.	`True`

Raises:

Type	Description
`RuntimeError`	If no ID suffix is allowed and an invocation for the same ID already exists, or if no unique invocation ID can be found.

Returns:

Type	Description
`str`	The invocation ID.

Source code in src/zenml/execution/pipeline/utils.py

def compute_invocation_id(
    existing_invocations: Set[str],
    base_name: str,
    allow_suffix: bool = True,
) -> str:
    """Compute the invocation ID.

    Args:
        existing_invocations: The existing invocation IDs.
        base_name: Base name for the invocation. Used as the ID directly when
            unique, otherwise suffixed with `_2`, `_3`, ... until a free slot
            is found (when `allow_suffix=True`).
        allow_suffix: Whether a suffix can be appended to the invocation
            ID.

    Raises:
        RuntimeError: If no ID suffix is allowed and an invocation for the
            same ID already exists, or if no unique invocation ID can be
            found.

    Returns:
        The invocation ID.
    """
    base_id = id_ = base_name

    if id_ not in existing_invocations:
        return id_

    if not allow_suffix:
        raise RuntimeError(f"Duplicate invocation ID `{id_}`")

    for index in range(2, 10000):
        id_ = f"{base_id}_{index}"
        if id_ not in existing_invocations:
            return id_

    raise RuntimeError("Unable to find invocation ID")

prevent_pipeline_execution() -> Generator[None, None, None]

Context manager to prevent pipeline execution.

Yields:

Type	Description
`None`	None.

Source code in src/zenml/execution/pipeline/utils.py

@contextmanager
def prevent_pipeline_execution() -> Generator[None, None, None]:
    """Context manager to prevent pipeline execution.

    Yields:
        None.
    """
    token = _prevent_pipeline_execution.set(True)
    try:
        yield
    finally:
        _prevent_pipeline_execution.reset(token)

should_prevent_pipeline_execution() -> bool

Whether to prevent pipeline execution.

Returns:

Type	Description
`bool`	Whether to prevent pipeline execution.

Source code in src/zenml/execution/pipeline/utils.py

def should_prevent_pipeline_execution() -> bool:
    """Whether to prevent pipeline execution.

    Returns:
        Whether to prevent pipeline execution.
    """
    return _prevent_pipeline_execution.get()

skip_steps_and_prune_snapshot(snapshot: PipelineSnapshotResponse, pipeline_run: PipelineRunResponse) -> bool

Skip steps and prune the snapshot.

Parameters:

Name	Type	Description	Default
`snapshot`	`PipelineSnapshotResponse`	The snapshot to prune.	required
`pipeline_run`	`PipelineRunResponse`	The pipeline run to skip steps for.	required

Raises:

Type	Description
`RuntimeError`	If the pipeline run is not a replayed run, if a step has an upstream step that is not skipped, or if a step run request cannot be populated.

Returns:

Type	Description
`bool`	Whether a pipeline run is still required.

Source code in src/zenml/execution/pipeline/utils.py

def skip_steps_and_prune_snapshot(
    snapshot: "PipelineSnapshotResponse",
    pipeline_run: "PipelineRunResponse",
) -> bool:
    """Skip steps and prune the snapshot.

    Args:
        snapshot: The snapshot to prune.
        pipeline_run: The pipeline run to skip steps for.

    Raises:
        RuntimeError: If the pipeline run is not a replayed run, if a step
            has an upstream step that is not skipped, or if a step run
            request cannot be populated.

    Returns:
        Whether a pipeline run is still required.
    """
    from zenml.orchestrators.step_run_utils import StepRunRequestFactory

    if snapshot.is_dynamic:
        # In dynamic pipelines, the steps will be skipped at runtime.
        return True

    if not pipeline_run.original_run:
        raise RuntimeError(
            "Unable to skip steps because the pipeline run is not a "
            "replayed run."
        )

    logger.debug("Skipping steps and pruning snapshot.")

    client = Client()
    request_factory = StepRunRequestFactory(
        snapshot=snapshot,
        pipeline_run=pipeline_run,
        stack=client.active_stack,
    )

    explicitly_skipped_steps = set(pipeline_run.config.steps_to_skip)
    skipped_invocations: Set[str] = set()

    for invocation_id, step in snapshot.step_configurations.items():
        explicitly_skipped = invocation_id in explicitly_skipped_steps
        should_skip = request_factory.should_skip_step(invocation_id)

        if not should_skip:
            continue

        unskipped_upstream_steps = (
            set(step.spec.upstream_steps) - skipped_invocations
        )

        if unskipped_upstream_steps:
            if not explicitly_skipped:
                logger.debug(
                    "Not skipping successful step `%s` because upstream "
                    "steps `%s` are not skipped.",
                    invocation_id,
                    ", ".join(unskipped_upstream_steps),
                )
                continue

            raise RuntimeError(
                f"Unable to skip step `{invocation_id}` because it has "
                f"upstream steps `{', '.join(unskipped_upstream_steps)}` that "
                "are not skipped."
            )

        request = request_factory.create_request(invocation_id)
        try:
            request_factory.populate_request(request)
        except Exception as e:
            # We failed to populate the step run request. This might be due
            # to some input resolution error, or an error importing the step
            # source (there might be some missing dependencies). We do not want
            # the orchestrator to spin up an environment for this step, so we
            # fail early here.
            raise RuntimeError(
                "Failed to populate step run request for step "
                f"`{invocation_id}`: {str(e)}"
            ) from e

        if request.status != ExecutionStatus.SKIPPED:
            # This shouldn't happen, but just in case.
            raise RuntimeError(
                f"Expected step request `{invocation_id}` to have status "
                f"`{ExecutionStatus.SKIPPED}`, but got `{request.status}`."
            )

        client.zen_store.create_run_step(request)
        skipped_invocations.add(invocation_id)
        logger.info("Skipping step `%s`.", invocation_id)

    for invocation_id in skipped_invocations:
        # Remove the skipped step invocations from the snapshot so
        # the orchestrator does not try to run them
        snapshot.step_configurations.pop(invocation_id)

    for step in snapshot.step_configurations.values():
        for invocation_id in skipped_invocations:
            if invocation_id in step.spec.upstream_steps:
                step.spec.upstream_steps.remove(invocation_id)

    if len(snapshot.step_configurations) == 0:
        logger.info("All steps were skipped.")
        return False

    return True

submit_pipeline(snapshot: PipelineSnapshotResponse, stack: Stack, placeholder_run: Optional[PipelineRunResponse] = None) -> None

Submit a snapshot for execution.

Parameters:

Name	Type	Description	Default
`snapshot`	`PipelineSnapshotResponse`	The snapshot to submit.	required
`stack`	`Stack`	The stack on which to submit the snapshot.	required
`placeholder_run`	`Optional[PipelineRunResponse]`	An optional placeholder run for the snapshot.	`None`

Raises:

Type	Description
`BaseException`	Any exception that happened while submitting or running (in case it happens synchronously) the pipeline.

Source code in src/zenml/execution/pipeline/utils.py

def submit_pipeline(
    snapshot: "PipelineSnapshotResponse",
    stack: "Stack",
    placeholder_run: Optional["PipelineRunResponse"] = None,
) -> None:
    """Submit a snapshot for execution.

    Args:
        snapshot: The snapshot to submit.
        stack: The stack on which to submit the snapshot.
        placeholder_run: An optional placeholder run for the snapshot.

    Raises:
        BaseException: Any exception that happened while submitting or running
            (in case it happens synchronously) the pipeline.
    """  # noqa: DOC502, DOC503
    # Prevent execution of nested pipelines which might lead to
    # unexpected behavior
    with prevent_pipeline_execution():
        try:
            stack.prepare_pipeline_submission(snapshot=snapshot)
            stack.submit_pipeline(
                snapshot=snapshot,
                placeholder_run=placeholder_run,
            )
        except RunMonitoringError as e:
            # Don't mark the run as failed if the error happened during
            # monitoring of the run.
            raise e.original_exception from None
        except BaseException as e:
            if (
                placeholder_run
                and not Client()
                .get_pipeline_run(placeholder_run.id, hydrate=False)
                .status.is_finished
            ):
                # We failed during/before the submission of the run, so we mark
                # the run as failed if it's still in an unfinished state.
                publish_failed_pipeline_run(placeholder_run.id)

            raise e

`step`

Step execution.

Modules

`utils`

Step execution utilities.

Classes Functions:

launch_step(snapshot: PipelineSnapshotResponse, step: Step, orchestrator_run_id: str, retry: bool = False, remaining_retries: Optional[int] = None, wait: bool = True) -> StepRunResponse

Launch a step.

Parameters:

Name	Type	Description	Default
`snapshot`	`PipelineSnapshotResponse`	The snapshot.	required
`step`	`Step`	The step to run.	required
`orchestrator_run_id`	`str`	The orchestrator run ID.	required
`retry`	`bool`	Whether to retry the step if it fails.	`False`
`remaining_retries`	`Optional[int]`	The number of remaining retries. If not passed, this will be read from the step configuration.	`None`
`wait`	`bool`	Whether to wait for the step to complete.	`True`

Raises:

Type	Description
`RunStoppedException`	If the run was stopped.
`BaseException`	If the step failed all retries.

Returns:

Type	Description
`StepRunResponse`	The step run response.

Source code in src/zenml/execution/step/utils.py

def launch_step(
    snapshot: "PipelineSnapshotResponse",
    step: "Step",
    orchestrator_run_id: str,
    retry: bool = False,
    remaining_retries: Optional[int] = None,
    wait: bool = True,
) -> StepRunResponse:
    """Launch a step.

    Args:
        snapshot: The snapshot.
        step: The step to run.
        orchestrator_run_id: The orchestrator run ID.
        retry: Whether to retry the step if it fails.
        remaining_retries: The number of remaining retries. If not passed, this
            will be read from the step configuration.
        wait: Whether to wait for the step to complete.

    Raises:
        RunStoppedException: If the run was stopped.
        BaseException: If the step failed all retries.

    Returns:
        The step run response.
    """

    def _launch_without_retry() -> StepRunResponse:
        launcher = StepLauncher(
            snapshot=snapshot,
            step=step,
            orchestrator_run_id=orchestrator_run_id,
            wait=wait,
        )
        return launcher.launch()

    if not retry:
        step_run = _launch_without_retry()
    else:
        retries = 0
        retry_config = step.config.retry
        if remaining_retries is None:
            max_retries = retry_config.max_retries if retry_config else 0
        else:
            max_retries = remaining_retries
        delay = retry_config.delay if retry_config else 0
        backoff = retry_config.backoff if retry_config else 1

        while retries <= max_retries:
            try:
                step_run = _launch_without_retry()
            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

    return step_run

`utils`

Execution utilities.

Classes

`DebugModeContext()`

Bases: BaseContext

Context manager for enabling debug mode.

When debug mode is enabled, a local orchestrator is used instead of the actual orchestrator.

Source code in src/zenml/utils/context_utils.py

def __init__(self) -> None:
    """Initialize the context."""
    self._token: Optional[contextvars.Token[Any]] = None