Execution

zenml.execution

Step and pipeline execution.

Modules

pipeline

Pipeline execution.

Modules

dynamic

Dynamic pipeline execution.

Modules

outputs

Dynamic pipeline execution outputs.

Classes

ArtifactFuture(wrapped: Future[StepRunOutputs], invocation_id: str, index: int)

Bases: BaseStepRunFuture

Future for a step run output artifact.

Initialize the future.

Parameters:

Name	Type	Description	Default
wrapped	Future[StepRunOutputs]	The wrapped future object.	required
invocation_id	str	The invocation ID of the step run.	required
index	int	The index of the output artifact.	required

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

def __init__(
    self, wrapped: Future[StepRunOutputs], invocation_id: str, index: int
) -> None:
    """Initialize the future.

    Args:
        wrapped: The wrapped future object.
        invocation_id: The invocation ID of the step run.
        index: The index of the output artifact.
    """
    super().__init__(wrapped=wrapped, invocation_id=invocation_id)
    self._index = index

Functions

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.
    """
    result = self._wrapped.result()
    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}."
        )

BaseFuture

Bases: ABC

Base future.

Functions

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.
    """

BaseStepRunFuture(wrapped: Future[StepRunOutputs], invocation_id: str, **kwargs: Any)

Bases: BaseFuture

Base step run future.

Initialize the dynamic step run future.

Parameters:

Name	Type	Description	Default
wrapped	Future[StepRunOutputs]	The wrapped future object.	required
invocation_id	str	The invocation ID of the step run.	required
**kwargs	Any	Additional keyword arguments.	{}

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

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

    Args:
        wrapped: The wrapped future object.
        invocation_id: The invocation ID of the step run.
        **kwargs: Additional keyword arguments.
    """
    self._wrapped = wrapped
    self._invocation_id = invocation_id

Attributes

invocation_id: str property

The step run invocation ID.

Returns:

Type	Description
str	The step run invocation ID.

Functions

running() -> bool

Check if the step run future is running.

Returns:

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

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

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

    Returns:
        True if the step run future is running, False otherwise.
    """
    return self._wrapped.running()

MapResultsFuture(futures: List[StepRunOutputsFuture])

Bases: BaseFuture

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

Initialize the map results future.

Parameters:

Name	Type	Description	Default
futures	List[StepRunOutputsFuture]	The step run futures.	required

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

def __init__(self, futures: List[StepRunOutputsFuture]) -> None:
    """Initialize the map results future.

    Args:
        futures: The step run futures.
    """
    self.futures = futures

Functions

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.
    """
    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)))

OutputArtifact

Bases: ArtifactVersionResponse

Dynamic step run output artifact.

Functions

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})

StepRunOutputsFuture(wrapped: Future[StepRunOutputs], invocation_id: str, output_keys: List[str])

Bases: BaseStepRunFuture

Future for a step run output.

Initialize the future.

Parameters:

Name	Type	Description	Default
wrapped	Future[StepRunOutputs]	The wrapped future object.	required
invocation_id	str	The invocation ID of the step run.	required
output_keys	List[str]	The output keys of the step run.	required

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

def __init__(
    self,
    wrapped: Future[StepRunOutputs],
    invocation_id: str,
    output_keys: List[str],
) -> None:
    """Initialize the future.

    Args:
        wrapped: The wrapped future object.
        invocation_id: The invocation ID of the step run.
        output_keys: The output keys of the step run.
    """
    super().__init__(wrapped=wrapped, invocation_id=invocation_id)
    self._output_keys = output_keys

Functions

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._wrapped.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(
        wrapped=self._wrapped,
        invocation_id=self._invocation_id,
        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.
    """
    return self._wrapped.result()

Functions

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

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.

Functions

Modules

runner

Dynamic pipeline runner.

Classes

DynamicPipelineRunner(snapshot: PipelineSnapshotResponse, run: Optional[PipelineRunResponse], orchestrator: Optional[BaseOrchestrator] = 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

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,
) -> 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.

    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.STOP_ON_FAILURE
    ):
        logger.warning(
            "Only the `%s` execution mode is supported for "
            "dynamic pipelines right now. "
            "The execution mode `%s` will be ignored.",
            ExecutionMode.STOP_ON_FAILURE,
            snapshot.pipeline_configuration.execution_mode,
        )

    self._snapshot = snapshot
    self._run = run
    # TODO: make this configurable
    self._executor = ThreadPoolExecutor(max_workers=10)
    self._pipeline: Optional["DynamicPipeline"] = None
    if orchestrator:
        self._orchestrator = orchestrator
    else:
        self._orchestrator = Stack.from_model(snapshot.stack).orchestrator
    self._futures: List[StepRunOutputsFuture] = []
    self._invocation_ids: Set[str] = set()

    self._existing_step_runs: Dict[str, "StepRunResponse"] = {}
    if run and run.orchestrator_run_id:
        self._orchestrator_run_id = run.orchestrator_run_id

        if run.status == ExecutionStatus.RUNNING:
            self._existing_step_runs = run.steps.copy()
    else:
        self._orchestrator_run_id = (
            self._orchestrator.get_orchestrator_run_id()
        )

Attributes

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.

Functions

await_all_step_run_futures() -> None

Await all step run output futures.

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

def await_all_step_run_futures(self) -> None:
    """Await all step run output futures."""
    for future in self._futures:
        future.result()
    self._futures = []

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

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

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

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[AnyStepRunFuture, Sequence[AnyStepRunFuture], None]	The step run output futures to wait for.	None
concurrent	bool	Whether to launch the step concurrently.	False

Returns:

Type	Description
Union[StepRunOutputs, StepRunOutputsFuture]	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[
        "AnyStepRunFuture", Sequence["AnyStepRunFuture"], None
    ] = None,
    concurrent: bool = False,
) -> Union[StepRunOutputs, "StepRunOutputsFuture"]:
    """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.
        concurrent: Whether to launch the step concurrently.

    Returns:
        The step run outputs or a future for the step run outputs.
    """
    step = step.copy()
    step_config = None
    if self._run and self._run.triggered_by_deployment:
        # Deployment-specific step overrides
        step_config = StepConfigurationUpdate(
            enable_cache=False,
            step_operator=None,
            parameters={},
            runtime=StepRuntime.INLINE,
        )

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

    invocation_id = compute_invocation_id(
        existing_invocations=self._invocation_ids,
        step=step,
        custom_id=id,
        allow_suffix=not id,
    )
    self._invocation_ids.add(invocation_id)

    if waiting_for_steps := _get_running_upstream_steps(inputs, after):
        logger.info(
            "Waiting for step(s) `%s` to finish before executing step `%s`.",
            ", ".join(waiting_for_steps),
            invocation_id,
        )

    compiled_step = compile_dynamic_step_invocation(
        snapshot=self._snapshot,
        pipeline=self.pipeline,
        step=step,
        invocation_id=invocation_id,
        inputs=inputs,
        pipeline_docker_settings=self._snapshot.pipeline_configuration.docker_settings,
        after=after,
        config=step_config,
    )

    should_retry = _should_retry_locally(
        step=compiled_step,
        pipeline_docker_settings=self._snapshot.pipeline_configuration.docker_settings,
        orchestrator=self._orchestrator,
    )

    def _launch_step_and_load_outputs(
        remaining_retries: Optional[int] = None,
    ) -> StepRunOutputs:
        # TODO: maybe pass run here to avoid extra server requests?
        step_run = launch_step(
            snapshot=self._snapshot,
            step=compiled_step,
            orchestrator_run_id=self._orchestrator_run_id,
            retry=should_retry,
            remaining_retries=remaining_retries,
        )
        return _load_step_run_outputs(step_run.id)

    step_run = self._existing_step_runs.get(invocation_id)
    if step_run and step_run.config != compiled_step.config:
        logger.warning(
            "Configuration for step `%s` changed since the the "
            "orchestration environment was restarted. If the step "
            "needs to be retried, it will use the new configuration.",
            step_run.name,
        )

    def _run() -> StepRunOutputs:
        nonlocal step_run

        if not step_run:
            return _launch_step_and_load_outputs()

        if step_run.status.is_successful:
            return _load_step_run_outputs(step_run.id)

        runtime = get_step_runtime(
            step_config=compiled_step.config,
            pipeline_docker_settings=self._snapshot.pipeline_configuration.docker_settings,
            orchestrator=self._orchestrator,
        )
        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)

        remaining_retries = 0

        if should_retry:
            max_retries = (
                compiled_step.config.retry.max_retries
                if compiled_step.config.retry
                else 0
            )
            remaining_retries = max(0, 1 + max_retries - step_run.version)

        if step_run.status == ExecutionStatus.RUNNING:
            logger.info(
                "Restarting the monitoring of existing step `%s` "
                "(ID: %s). Remaining retries: %d",
                step_run.name,
                step_run.id,
                remaining_retries,
            )

        step_run = wait_for_step_run_to_finish(step_run.id)

        if remaining_retries > 0:
            if not step_run.status.is_successful:
                logger.error("Failed to run step `%s`.", step_run.name)
                return _launch_step_and_load_outputs(
                    remaining_retries=remaining_retries
                )
            else:
                return _load_step_run_outputs(step_run.id)
        else:
            if not step_run.status.is_successful:
                # This is the last retry, in which case we have to raise
                # an error that the step failed.
                # TODO: Make this better by raising the actual exception
                # that caused the step to fail instead of just a generic
                # runtime error.
                raise RuntimeError(
                    f"Failed to run step `{step_run.name}`."
                )
            return _load_step_run_outputs(step_run.id)

    if concurrent:
        ctx = contextvars.copy_context()
        future = self._executor.submit(ctx.run, _run)
        step_run_future = StepRunOutputsFuture(
            wrapped=future,
            invocation_id=invocation_id,
            output_keys=list(compiled_step.config.outputs),
        )
        self._futures.append(step_run_future)
        return step_run_future
    else:
        return _run()

map(step: BaseStep, args: Tuple[Any, ...], kwargs: Dict[str, Any], after: Union[AnyStepRunFuture, Sequence[AnyStepRunFuture], 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[AnyStepRunFuture, Sequence[AnyStepRunFuture], 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[
        "AnyStepRunFuture", Sequence["AnyStepRunFuture"], 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.
    """
    kwargs = convert_to_keyword_arguments(step.entrypoint, args, kwargs)
    kwargs = await_step_inputs(kwargs)
    step_inputs = expand_mapped_inputs(kwargs, product=product)

    step_run_futures = [
        self.launch_step(
            step,
            id=None,
            args=(),
            kwargs=inputs,
            after=after,
            concurrent=True,
        )
        for inputs in step_inputs
    ]

    return MapResultsFuture(futures=step_run_futures)

run_pipeline() -> None

Run the pipeline.

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

def run_pipeline(self) -> None:
    """Run the pipeline."""
    if self._run:
        if self._run.status.is_finished:
            logger.info("Run `%s` is already finished.", str(self._run.id))
            return
        if self._run.orchestrator_run_id:
            logger.info("Continuing existing run `%s`.", str(self._run.id))
            run = self._run
        else:
            run = Client().zen_store.update_run(
                run_id=self._run.id,
                run_update=PipelineRunUpdate(
                    orchestrator_run_id=self._orchestrator_run_id,
                ),
            )
    else:
        existing_runs = Client().list_pipeline_runs(
            snapshot_id=self._snapshot.id,
            orchestrator_run_id=self._orchestrator_run_id,
        )
        if existing_runs.total == 1:
            run = existing_runs.items[0]
            if run.status.is_finished:
                logger.info("Run `%s` is already finished.", str(run.id))
                return
            else:
                logger.info("Continuing existing run `%s`.", str(run.id))
        else:
            run = create_placeholder_run(
                snapshot=self._snapshot,
                orchestrator_run_id=self._orchestrator_run_id,
            )

    logging_context = nullcontext()
    if not run.triggered_by_deployment and is_pipeline_logging_enabled(
        self._snapshot.pipeline_configuration
    ):
        logging_context = setup_run_logging(
            pipeline_run=run,
            source="orchestrator",
        )

    with logging_context:
        with InMemoryArtifactCache():
            with DynamicPipelineRunContext(
                pipeline=self.pipeline,
                run=run,
                snapshot=self._snapshot,
                runner=self,
            ):
                if not 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.
                    self._orchestrator.run_init_hook(
                        snapshot=self._snapshot
                    )

                try:
                    # TODO: what should be allowed as pipeline returns?
                    #  (artifacts, json serializable, anything?)
                    #  how do we show it in the UI?
                    params = self.pipeline.configuration.parameters or {}
                    self.pipeline._call_entrypoint(**params)
                    # The pipeline function finished successfully, but some
                    # steps might still be running. We now wait for all of
                    # them and raise any exceptions that occurred.
                    self.await_all_step_run_futures()
                except:
                    # TODO: this call already invalidates the token, so
                    # the steps will keep running but won't be able to
                    # report their status back to ZenML.
                    publish_failed_pipeline_run(run.id)
                    logger.error(
                        "Pipeline run failed. All in-progress step runs "
                        "will still finish executing."
                    )
                    raise
                finally:
                    if not 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)

                publish_successful_pipeline_run(run.id)
                logger.info("Pipeline completed successfully.")

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 with 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/runner.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 with 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, Sequence)
            and value
            and all(isinstance(item, StepRunOutputsFuture) 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 = [item.artifacts() for item in value]
        elif isinstance(value, StepRunOutputsFuture):
            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 = value.artifacts()

        if (
            isinstance(value, Sequence)
            and value
            and all(isinstance(item, ArtifactFuture) for item in value)
        ):
            value = [item.result() for item in value]

        if isinstance(value, ArtifactFuture):
            value = value.result()

        result[key] = value

    return result

compile_dynamic_step_invocation(snapshot: PipelineSnapshotResponse, pipeline: DynamicPipeline, step: BaseStep, invocation_id: str, inputs: Dict[str, Any], pipeline_docker_settings: DockerSettings, after: Union[AnyStepRunFuture, Sequence[AnyStepRunFuture], 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[AnyStepRunFuture, Sequence[AnyStepRunFuture], 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/runner.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[
        "AnyStepRunFuture", Sequence["AnyStepRunFuture"], 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()

    if isinstance(after, BaseStepRunFuture):
        after.result()
        upstream_steps.add(after.invocation_id)
    elif isinstance(after, MapResultsFuture):
        for future in after:
            future.result()
            upstream_steps.add(future.invocation_id)
    elif isinstance(after, Sequence):
        for item in after:
            if isinstance(item, BaseStepRunFuture):
                item.result()
                upstream_steps.add(item.invocation_id)
            elif isinstance(item, MapResultsFuture):
                for future in item:
                    future.result()
                    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)

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

    input_artifacts = {}
    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):
        step._configuration = template.config.model_copy(
            update={"template": template.spec.invocation_id}
        )

    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

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/runner.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

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/runner.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.
    """
    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_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/runner.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/runner.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

Modules

utils

Dynamic pipeline execution utilities.

Classes Functions

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_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

utils

Pipeline execution utilities.

Classes Functions

compute_invocation_id(existing_invocations: Set[str], step: BaseStep, custom_id: Optional[str] = None, allow_suffix: bool = True) -> str

Compute the invocation ID.

Parameters:

Name	Type	Description	Default
existing_invocations	Set[str]	The existing invocation IDs.	required
step	BaseStep	The step for which to compute the ID.	required
custom_id	Optional[str]	Custom ID to use for the invocation.	None
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.
RuntimeError	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],
    step: "BaseStep",
    custom_id: Optional[str] = None,
    allow_suffix: bool = True,
) -> str:
    """Compute the invocation ID.

    Args:
        existing_invocations: The existing invocation IDs.
        step: The step for which to compute the ID.
        custom_id: Custom ID to use for the invocation.
        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.
        RuntimeError: If no unique invocation ID can be found.

    Returns:
        The invocation ID.
    """
    base_id = id_ = custom_id or step.name

    if id_ not in existing_invocations:
        return id_

    if not allow_suffix:
        raise RuntimeError(f"Duplicate step 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 step 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()

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

noqa: DAR401

Raises: 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.

    # noqa: DAR401
    Raises:
        BaseException: Any exception that happened while submitting or running
            (in case it happens synchronously) the pipeline.
    """
    # 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) -> 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

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,
) -> 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.

    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,
        )
        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