Function find_all_downstream_steps

src/zenml/utils/run_utils.py:188–206

Find all downstream steps of a given step. Args: step_name: The name of the step to find the downstream steps of. dag: The DAG. Returns: The set of downstream steps.

(
    step_name: str, dag: Dict[str, Set[str]]
)

def find_all_downstream_steps(
    step_name: str, dag: Dict[str, Set[str]]
) -> Set[str]:
    """Find all downstream steps of a given step.

    Args:
        step_name: The name of the step to find the downstream steps of.
        dag: The DAG.

    Returns:
        The set of downstream steps.
    """
    downstream_steps = set()
    for downstream_step in dag[step_name]:
        downstream_steps.add(downstream_step)
        downstream_steps.update(
            find_all_downstream_steps(downstream_step, dag)
        )
    return downstream_steps