From efedefd7182a78edaaa21f34b07b235d9bd09491 Mon Sep 17 00:00:00 2001 From: Alex Wu <66259759+popojk@users.noreply.github.com> Date: Wed, 5 Feb 2025 13:07:39 +0800 Subject: [PATCH] fix failure node example code (#1778) Signed-off-by: Alex Wu --- .../data_types_and_io/data_types_and_io/pickle_type.py | 2 +- .../development_lifecycle/failure_node.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/examples/data_types_and_io/data_types_and_io/pickle_type.py b/examples/data_types_and_io/data_types_and_io/pickle_type.py index f0122aedf..5fc5a25d0 100644 --- a/examples/data_types_and_io/data_types_and_io/pickle_type.py +++ b/examples/data_types_and_io/data_types_and_io/pickle_type.py @@ -34,7 +34,7 @@ def superhero_wf(name: str = "Thor", power: str = "Flight") -> str: # The following example demonstrates how to set the batch size. from typing import Iterator -from flytekit.types.pickle.pickle import BatchSize +from flytekit.core.type_engine import BatchSize from typing_extensions import Annotated diff --git a/examples/development_lifecycle/development_lifecycle/failure_node.py b/examples/development_lifecycle/development_lifecycle/failure_node.py index 6823b1fc0..7c9332e23 100644 --- a/examples/development_lifecycle/development_lifecycle/failure_node.py +++ b/examples/development_lifecycle/development_lifecycle/failure_node.py @@ -1,4 +1,7 @@ +import typing + from flytekit import WorkflowFailurePolicy, task, workflow +from flytekit.types.error.error import FlyteError @task @@ -20,8 +23,8 @@ def delete_cluster(name: str): # Create a task that will be executed if any of the tasks in the workflow fail @task -def clean_up(name: str): - print(f"Cleaning up cluster {name}") +def clean_up(name: str, err: typing.Optional[FlyteError] = None): + print(f"Deleting cluster {name} due to {err}") # Specify the `on_failure` to a cleanup task.