Skip to content

Execution

zenml.execution

Step and pipeline execution.

Modules

pipeline

Pipeline execution.

Modules
dynamic

Dynamic pipeline execution.

Modules
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 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
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 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
@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 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
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 = json.loads(raw_value)
            except json.JSONDecodeError as e:
                print(f"Invalid JSON 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 
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
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 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
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")
outputs

Dynamic pipeline execution outputs.

Classes
ArtifactFuture(wrapped: Union[_InlineStepFuture, _IsolatedStepFuture], index: int)

Bases: BaseStepFuture

Future for a step run output artifact.

Initialize the future.

Parameters:

Name Type Description Default
wrapped Union[_InlineStepFuture, _IsolatedStepFuture]

The wrapped future object.

 required
index int

The index of the output artifact.

 required
Source code in src/zenml/execution/pipeline/dynamic/outputs.py 
236
237
238
239
240
241
242
243
244
245
246
247
248
def __init__(
    self,
    wrapped: Union[_InlineStepFuture, _IsolatedStepFuture],
    index: int,
) -> None:
    """Initialize the future.

    Args:
        wrapped: The wrapped future object.
        index: The index of the output artifact.
    """
    super().__init__(wrapped=wrapped)
    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 
287
288
289
290
291
292
293
294
295
296
297
298
299
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 
276
277
278
279
280
281
282
283
284
285
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 
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def result(self) -> OutputArtifact:
    """Get the output artifact this future represents.

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

    Returns:
        The output artifact.
    """
    step_run = self._wrapped.result()
    from zenml.execution.pipeline.dynamic.utils import (
        load_step_run_outputs,
    )

    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}."
        )
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 
85
86
87
88
89
90
91
@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 
77
78
79
80
81
82
83
@abstractmethod
def running(self) -> bool:
    """Check if the future is running.

    Returns:
        True if the future is running, False otherwise.
    """
BaseStepFuture(wrapped: Union[_InlineStepFuture, _IsolatedStepFuture], **kwargs: Any)

Bases: BaseFuture

Base step future.

Initialize the dynamic step run future.

Parameters:

Name Type Description Default
wrapped Union[_InlineStepFuture, _IsolatedStepFuture]

The wrapped future object.

 required
**kwargs Any

Additional keyword arguments.

 {}
Source code in src/zenml/execution/pipeline/dynamic/outputs.py 
202
203
204
205
206
207
208
209
210
211
212
213
def __init__(
    self,
    wrapped: Union[_InlineStepFuture, _IsolatedStepFuture],
    **kwargs: Any,
) -> None:
    """Initialize the dynamic step run future.

    Args:
        wrapped: The wrapped future object.
        **kwargs: Additional keyword arguments.
    """
    self._wrapped = wrapped
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 
224
225
226
227
228
229
230
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[StepFuture])

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[StepFuture]

The step run futures.

 required
Source code in src/zenml/execution/pipeline/dynamic/outputs.py 
463
464
465
466
467
468
469
def __init__(self, futures: List[StepFuture]) -> 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 
487
488
489
490
491
492
493
494
495
496
497
498
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 
479
480
481
482
483
484
485
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 
471
472
473
474
475
476
477
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 
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
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 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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})
StepFuture(wrapped: Union[_InlineStepFuture, _IsolatedStepFuture], output_keys: List[str])

Bases: BaseStepFuture

Future for a step run output.

Initialize the future.

Parameters:

Name Type Description Default
wrapped Union[_InlineStepFuture, _IsolatedStepFuture]

The wrapped future object.

 required
output_keys List[str]

The output keys of the step run.

 required
Source code in src/zenml/execution/pipeline/dynamic/outputs.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
def __init__(
    self,
    wrapped: Union[_InlineStepFuture, _IsolatedStepFuture],
    output_keys: List[str],
) -> None:
    """Initialize the future.

    Args:
        wrapped: The wrapped future object.
        output_keys: The output keys of the step run.
    """
    super().__init__(wrapped=wrapped)
    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 
346
347
348
349
350
351
352
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 
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
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,
        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 
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
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 
354
355
356
357
358
359
360
361
362
363
364
365
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._wrapped.result()
    return load_step_run_outputs(step_run.id)
wait() -> None

Wait for the step to finish.

Source code in src/zenml/execution/pipeline/dynamic/outputs.py 
342
343
344
def wait(self) -> None:
    """Wait for the step to finish."""
    self._wrapped.result()
Functions Modules
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 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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.
Functions
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 
 94
 95
 96
 97
 98
 99
100
101
102
103
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)

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 
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
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.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._snapshot = snapshot
    self._pipeline: Optional["DynamicPipeline"] = None

    worker_count = handle_int_env_var(
        ENV_ZENML_DYNAMIC_PIPELINE_WORKER_COUNT, default=10
    )
    self._executor = ThreadPoolExecutor(max_workers=worker_count)
    if orchestrator:
        self._orchestrator = orchestrator
    else:
        self._orchestrator = Stack.from_model(snapshot.stack).orchestrator

    self._step_operator = Stack.from_model(snapshot.stack).step_operator
    self._futures: Dict[str, "StepFuture"] = {}
    self._invocation_ids: Set[str] = set()

    self._run, self._orchestrator_run_id = self._prepare_run(run)
    self._existing_step_runs = self._run.steps

    self._steps_to_monitor: Dict[str, "StepRunResponse"] = {}
    self._shutdown_event = threading.Event()
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_futures() -> None

Await all step futures.

Source code in src/zenml/execution/pipeline/dynamic/runner.py 
1245
1246
1247
1248
1249
def await_all_step_futures(self) -> None:
    """Await all step futures."""
    for future in list(self._futures.values()):
        future.wait()
    self._futures = {}
has_in_progress_steps() -> bool

Check if there are any in-progress steps.

Returns:

Type Description
bool

True if there are any in-progress steps, False otherwise.
Source code in src/zenml/execution/pipeline/dynamic/runner.py 
1251
1252
1253
1254
1255
1256
1257
def has_in_progress_steps(self) -> bool:
    """Check if there are any in-progress steps.

    Returns:
        True if there are any in-progress steps, False otherwise.
    """
    return any(future.running() for future in list(self._futures.values()))
launch_step(step: BaseStep, id: Optional[str], args: Tuple[Any, ...], kwargs: Dict[str, Any], after: Union[AnyStepFuture, Sequence[AnyStepFuture], 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[
        AnyStepFuture, Sequence[AnyStepFuture], 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[
        AnyStepFuture, Sequence[AnyStepFuture], 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[AnyStepFuture, Sequence[AnyStepFuture], 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
noqa: DAR401

Raises: 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 
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
def launch_step(
    self,
    step: "BaseStep",
    id: Optional[str],
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    after: Union["AnyStepFuture", Sequence["AnyStepFuture"], 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.

    # noqa: DAR401
    Raises:
        BaseException: If the step failed.

    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,
            group=group,
        )
    elif group:
        step_config = StepConfigurationUpdate(
            group=group,
        )

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

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

    if step_run:
        old_config = step_run.config.model_copy(deep=True)
        # The server includes the date/time substitutions based on the
        # actual DB start time of the pipeline run. We want to compare this
        # to the newly compiled step config, so we remove them.
        old_config.substitutions.pop("date")
        old_config.substitutions.pop("time")

        if (
            old_config.model_dump_json()
            != compiled_step.config.model_dump_json()
        ):
            # We use the old config here to keep the behavior aligned across
            # all step retries: When the orchestration environment is
            # restarted while an isolated step is running, the code below
            # will start monitoring that existing execution. If that step
            # then fails, the monitoring loop will restart it with the
            # configuration of the step run that just failed, which is the
            # old (=compiled in previous orchestration environment) config.
            compiled_step = Step(
                spec=step_run.spec,
                config=old_config,
                step_config_overrides=old_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 old configuration.",
                step_run.name,
            )

        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:
                future = StepFuture(
                    wrapped=_IsolatedStepFuture(
                        pipeline_run_id=self._run.id,
                        invocation_id=invocation_id,
                    ),
                    output_keys=list(compiled_step.config.outputs),
                )
                self._futures[invocation_id] = future
                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)

        if step_run.status.is_failed:
            # If the step is running concurrently, we only raise the
            # exception once the future is awaited.
            if concurrent:
                future = StepFuture(
                    wrapped=_IsolatedStepFuture(
                        pipeline_run_id=self._run.id,
                        invocation_id=invocation_id,
                    ),
                    output_keys=list(compiled_step.config.outputs),
                )
                self._futures[invocation_id] = future
                return future
            else:
                raise exception_utils.reconstruct_exception(
                    exception_info=step_run.exception_info,
                    fallback_message=(
                        f"Step `{invocation_id}` failed with status "
                        f"`{step_run.status}`."
                    ),
                )

        remaining_retries = get_remaining_retries(step_run=step_run)

        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,
            )
            self._steps_to_monitor[invocation_id] = step_run
            monitoring_future = StepFuture(
                wrapped=_IsolatedStepFuture(
                    pipeline_run_id=self._run.id,
                    invocation_id=invocation_id,
                ),
                output_keys=list(compiled_step.config.outputs),
            )
            self._futures[invocation_id] = monitoring_future
            return monitoring_future

    if not concurrent:
        step_run = self._run_sync_step(
            step=compiled_step, remaining_retries=remaining_retries
        )
        return load_step_run_outputs(step_run.id)
    elif runtime == StepRuntime.INLINE:
        return self._queue_concurrent_inline_step(
            step=compiled_step, remaining_retries=remaining_retries
        )
    else:
        return self._queue_concurrent_isolated_step(step=compiled_step)
map(step: BaseStep, args: Tuple[Any, ...], kwargs: Dict[str, Any], after: Union[AnyStepFuture, Sequence[AnyStepFuture], 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[AnyStepFuture, Sequence[AnyStepFuture], 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 
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
def map(
    self,
    step: "BaseStep",
    args: Tuple[Any, ...],
    kwargs: Dict[str, Any],
    after: Union["AnyStepFuture", Sequence["AnyStepFuture"], 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)

    # This will overwrite any user-configured groups for the step, but
    # capturing the mapping information is more important until we introduce
    # a more flexible group system.
    group_info = GroupInfo(
        id=str(uuid4()),
        name=step.name,
        type=GroupType.MAP,
    )

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

    return MapResultsFuture(futures=step_futures)
run_pipeline() -> None

Run the pipeline.

Raises:

Type Description
Exception

If the pipeline run failed.
Source code in src/zenml/execution/pipeline/dynamic/runner.py 
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
def run_pipeline(self) -> None:
    """Run the pipeline.

    Raises:
        Exception: If the pipeline run failed.
    """
    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:
        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,
            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,
                ),
            )

        assert self._snapshot.stack

        with (
            InMemoryArtifactCache(),
            env_utils.temporary_runtime_environment(
                self._snapshot.pipeline_configuration, self._snapshot.stack
            ),
            DynamicPipelineRunContext(
                pipeline=self.pipeline,
                run=self._run,
                snapshot=self._snapshot,
                runner=self,
            ),
        ):
            monitoring_thread = threading.Thread(
                name="DynamicPipelineRunner-Monitoring-Loop",
                target=lambda: context_utils.run_in_current_context(
                    self._monitoring_loop
                ),
                daemon=True,
            )
            monitoring_thread.start()
            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.
                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_futures()
            except _WaitConditionPollTimeout:
                logger.info("Pausing pipeline run `%s`.", self._run.id)
            except _WaitConditionAborted:
                logger.info(
                    "Stopping pipeline run `%s` because a wait condition "
                    "was aborted.",
                    self._run.id,
                )
            except Exception as e:
                exception_info = (
                    exception_utils.collect_exception_information(
                        exception=e,
                        user_func=self.pipeline.entrypoint,
                    )
                )
                # 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(
                    self._run.id, exception_info=exception_info
                )
                raise
            finally:
                self._shutdown_event.set()

                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._maybe_stop_isolated_steps()
                self._executor.shutdown(wait=True, cancel_futures=True)
                monitoring_thread.join()

            self._run = Client().zen_store.get_run(
                self._run.id, hydrate=False
            )
            if self._run.status == ExecutionStatus.RUNNING:
                publish_successful_pipeline_run(self._run.id)
                logger.info("Pipeline completed successfully.")
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[AnyStepFuture, Sequence[AnyStepFuture], 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[
        AnyStepFuture, Sequence[AnyStepFuture], 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[
        AnyStepFuture, Sequence[AnyStepFuture], 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

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
name Optional[str]

Optional deterministic wait condition name.

 None
after Union[AnyStepFuture, Sequence[AnyStepFuture], 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.
_WaitConditionPollTimeout

If the wait condition polling timed out.
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 
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
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["AnyStepFuture", Sequence["AnyStepFuture"], 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: 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.
        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.
        _WaitConditionPollTimeout: If the wait condition polling timed out.
        KeyboardInterrupt: If interrupted while waiting.
        BaseException: If polling fails after the lease is abandoned.

    Returns:
        The resolved wait condition value.
    """
    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."
        )

    wait_condition_name = name or context.next_wait_condition_name()

    if isinstance(after, BaseStepFuture):
        after.result()
    elif isinstance(after, MapResultsFuture):
        for future in after:
            future.result()
    elif isinstance(after, Sequence):
        for item in after:
            if isinstance(item, BaseStepFuture):
                item.result()
            elif isinstance(item, MapResultsFuture):
                for future in item:
                    future.result()

    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

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

    try:
        with maybe_enable_interactive_wait_prompt(
            orchestrator=self._orchestrator,
            condition=condition,
        ) as interactive_prompting_enabled:
            # We keep polling until the deadline is reached and all ongoing
            # steps have finished.
            while time.time() < deadline or self.has_in_progress_steps():
                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, we 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:
        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,
            )
        raise
    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 _WaitConditionPollTimeout(
        f"Wait condition `{condition.name}` polling timed out."
    )
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 
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
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, 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 = [item.artifacts() 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 = 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[AnyStepFuture, Sequence[AnyStepFuture], 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[AnyStepFuture, Sequence[AnyStepFuture], 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 
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
def compile_dynamic_step_invocation(
    snapshot: "PipelineSnapshotResponse",
    pipeline: "DynamicPipeline",
    step: "BaseStep",
    invocation_id: str,
    inputs: Dict[str, Any],
    pipeline_docker_settings: "DockerSettings",
    after: Union["AnyStepFuture", Sequence["AnyStepFuture"], 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, BaseStepFuture):
        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, BaseStepFuture):
                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)

    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):
        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
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 
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
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 
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
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 
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
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_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/runner.py 
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
1727
1728
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)
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 
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
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
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 
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
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]
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 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
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.
        return tuple(
            _convert_output_artifact(
                output_name=name, artifact=output_artifacts[name]
            )
            for name in step_run.config.outputs.keys()
        )
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 
63
64
65
66
67
68
69
70
71
72
73
74
75
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[AnyStepFuture]] = 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[AnyStepFuture]] = 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[AnyStepFuture]] = 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[AnyStepFuture]]

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 
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
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["AnyStepFuture"]] = 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 
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
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
wait_for_step_to_finish(pipeline_run_id: UUID, step_name: str) -> StepRunResponse

Wait until a step is finished.

Parameters:

Name Type Description Default
pipeline_run_id UUID

The ID of the pipeline run.

 required
step_name str

The name of the step.

 required

Returns:

Type Description
StepRunResponse

The finished step run.
Source code in src/zenml/execution/pipeline/dynamic/utils.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
def wait_for_step_to_finish(
    pipeline_run_id: UUID, step_name: str
) -> "StepRunResponse":
    """Wait until a step is finished.

    Args:
        pipeline_run_id: The ID of the pipeline run.
        step_name: The name of the step.

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

    while True:
        step_run = get_latest_step_run(
            pipeline_run_id, step_name, hydrate=False
        )
        # If a step is in `retrying` status, another step run will be
        # created and we will try to pick it up in the next iteration.
        if step_run.status not in {
            ExecutionStatus.RUNNING,
            ExecutionStatus.RETRYING,
        }:
            return step_run

        logger.debug(
            "Waiting for step `%s` to finish (current status: %s)",
            step_name,
            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 
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
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 
55
56
57
58
59
60
61
62
63
64
65
66
@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 
46
47
48
49
50
51
52
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.
RuntimeError

If a step has an upstream step that is not skipped.
RuntimeError

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 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
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.
        RuntimeError: If a step has an upstream step that is not skipped.
        RuntimeError: 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
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 
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
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, 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 
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
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 
42
43
44
def __init__(self) -> None:
    """Initialize the context."""
    self._token: Optional[contextvars.Token[Any]] = None
Modules