Eager cleanup #3148

wild-endeavor · 2025-02-20T18:43:00Z

Why are the changes needed?

Eager tasks operate outside of the fundamental assumptions made when Flyte was designed. This leads to potentially orphaned executions - executions that were launched by an eager task that aren't terminated when a parent eager task fails.

What changes were proposed in this pull request?

This PR adds an imperative workflow for each eager task, similar to what Admin does for single task executions. However it also attaches a failure node handler in cases where an eager task fails.

Added a new task type EagerFailureHandlerTask that's similar to the EchoTask. It is used to mimic the input interface of an eager task since the inputs to the failure node have to match that of the parent workflow/eager task. This failure task leverages the execution tags functionality as well as the tags that eager tasks add to child executions to make sure they're all cleaned up.
Added a new task resolver for that task type that always returns an EagerFailureHandlerTask instance with no input interface at runtime. Since we don't need to transform the inputs at failure handling runtime, we don't need to store a path to the original eager task to retrieve the Python types.
Added failure node for imperative workflows. (mostly takes care of add on failure parameter to imperative workflow flyte#5913)
Noticed a corresponding ValueNotIn filter object was missing but ended up not using it.

Todo: Before merging, two changes need to be reverted: requesting of a default secret, and setting of the failure task container image.

How was this patch tested?

Setup process

This was tested on Union internal clusters. (Testing on Flyte clusters is not yet possible due to missing single-task-execution labels.) Wrote an eager task that created some downstream executions that take varying amounts of time. One of the downstream tasks is set to fail, triggering a cleanup of the remaining ones.

@task(container_image=image)
def add_one(x: int) -> int:
    print(os.environ)
    return x + 1

@task(container_image=image)
async def double(x: int) -> int:
    print_env()
    if x == 2:
        return 42
    if x == 5:
        raise AssertionError("x is 5")
    await asyncio.sleep(20 + x)
    return x * 2

@eager(container_image=image)
async def level_1() -> typing.Tuple[int, int]:
    a1 = asyncio.create_task(double(x=1))
    a2 = asyncio.create_task(double(x=2))
    a3 = asyncio.create_task(double(x=3))
    a4 = asyncio.create_task(double(x=4))
    a5 = asyncio.create_task(double(x=5))

    (i1, i2, i3, i4, i5) = await asyncio.gather(a1, a2, a3, a4, a5)
    x, y = i1 + i2, i3 + i4 + i5
    print(f"Final result: {x}, {y}")
    return x, y

Screenshots

Failure node tasks do not run if everything succeeds.

Check all the applicable boxes

I updated the documentation accordingly.
All new and existing tests passed.
All commits are signed-off.

Related PRs

Docs link

Summary by Bito

Implementation of eager task cleanup mechanism with EagerFailureHandlerTask, introducing configurable delay settings and failure handler system. The changes include updating failure node IDs from 'nfail' to 'efn', adding task resolver and failure node support in imperative workflows. These modifications ensure proper cleanup of downstream executions and orphaned tasks when parent eager tasks fail.

Unit tests added: True

Estimated effort to review (1-5, lower is better): 3

Signed-off-by: Yee Hing Tong <[email protected]>

flyte-bot · 2025-02-20T18:43:15Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

codecov · 2025-02-20T18:45:24Z

Codecov Report

Attention: Patch coverage is 25.33333% with 56 lines in your changes missing coverage. Please review.

Project coverage is 41.23%. Comparing base (66d4aed) to head (be37cd5).
Report is 11 commits behind head on master.

Files with missing lines	Patch %	Lines
flytekit/core/python_function_task.py	42.50%	23 Missing ⚠️
flytekit/core/workflow.py	5.00%	19 Missing ⚠️
flytekit/tools/serialize_helpers.py	0.00%	7 Missing ⚠️
flytekit/core/node_creation.py	0.00%	5 Missing ⚠️
flytekit/tools/translator.py	0.00%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3148       +/-   ##
===========================================
- Coverage   76.85%   41.23%   -35.63%     
===========================================
  Files         206      211        +5     
  Lines       21851    22060      +209     
  Branches     2837     2868       +31     
===========================================
- Hits        16794     9096     -7698     
- Misses       4269    12815     +8546     
+ Partials      788      149      -639

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

flyte-bot · 2025-02-21T20:50:53Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: Yee Hing Tong <[email protected]>

flyte-bot · 2025-02-21T23:10:53Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: Yee Hing Tong <[email protected]>

flyte-bot · 2025-02-22T02:30:53Z

Code Review Agent Run Status

Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

Signed-off-by: Yee Hing Tong <[email protected]>

flyte-bot · 2025-02-22T03:40:56Z

Code Review Agent Run #ce446d

Actionable Suggestions - 3

flytekit/tools/serialize_helpers.py - 1
- Consider extracting task workflow creation logic · Line 64-67
tests/flytekit/unit/core/test_imperative.py - 1
- Consider assigning default DC object creation · Line 159-160
flytekit/core/node_creation.py - 1
- Consider adding error handling for async execution · Line 82-85

Review Details

Files reviewed - 10 · Commit Range: f31e823..cdb5189
- flytekit/core/constants.py
- flytekit/core/node_creation.py
- flytekit/core/python_function_task.py
- flytekit/core/workflow.py
- flytekit/models/filters.py
- flytekit/tools/serialize_helpers.py
- flytekit/tools/translator.py
- tests/flytekit/unit/core/test_eager_cleanup.py
- tests/flytekit/unit/core/test_imperative.py
- tests/flytekit/unit/core/test_serialization.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

flyte-bot · 2025-02-22T03:49:26Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
New Feature - Eager Task Cleanup System Implementation	- `constants.py` - Added default failure node ID constant - `node_creation.py` - Added async support for entity execution - `python_function_task.py` - Implemented EagerFailureHandlerTask and task resolver - `workflow.py` - Added failure handler support in ImperativeWorkflow - `filters.py` - Added ValueNotIn filter class - `serialize_helpers.py` - Added serialization support for eager async tasks - `translator.py` - Added validation for failure node ID
Testing - Test Coverage for Eager Cleanup System	- `test_eager_cleanup.py` - Added tests for EagerFailureHandlerTask and resolver - `test_imperative.py` - Added tests for failure handling in imperative workflows - `test_serialization.py` - Added tests for failure node naming validation

flytekit/tools/serialize_helpers.py

tests/flytekit/unit/core/test_imperative.py

flyte-bot · 2025-02-22T03:49:30Z

flytekit/core/node_creation.py

+        if inspect.iscoroutinefunction(entity.__call__):
+            outputs = run_sync(entity, **kwargs)
+        else:
+            outputs = entity(**kwargs)


Consider adding error handling for async execution

Consider handling potential exceptions from run_sync when executing coroutine functions. The current implementation may silently fail if the async execution encounters issues.

Code suggestion

Check the AI-generated fix before applying

Suggested change

if inspect.iscoroutinefunction(entity.__call__):

outputs = run_sync(entity, **kwargs)

else:

outputs = entity(**kwargs)

if inspect.iscoroutinefunction(entity.__call__):

try:

outputs = run_sync(entity, **kwargs)

except Exception as e:

raise RuntimeError(f"Async execution failed for {entity.name}: {str(e)}") from e

else:

outputs = entity(**kwargs)

Code Review Run #ce446d

Should Bito avoid suggestions like this for future reviews? (Manage Rules)

Yes, avoid them

eapolinario

just a few nits.

flytekit/core/constants.py

eapolinario · 2025-02-25T17:22:59Z

flytekit/core/python_function_task.py

+        # todo: remove this before merging
+        #   this is actually bad, but useful for developing
+        cleanup._container_image = self._container_image


flytekit/core/python_function_task.py

tests/flytekit/unit/core/test_imperative.py

flytekit/core/workflow.py

Signed-off-by: Yee Hing Tong <[email protected]>

flyte-bot · 2025-02-27T06:57:33Z

Code Review Agent Run #e76f06

Actionable Suggestions - 0

Review Details

Files reviewed - 4 · Commit Range: cdb5189..3db4249
- flytekit/core/constants.py
- flytekit/core/python_function_task.py
- tests/flytekit/unit/core/test_imperative.py
- tests/flytekit/unit/core/test_serialization.py
Files skipped - 0
Tools
- Whispers (Secret Scanner) - ✔︎ Successful
- Detect-secrets (Secret Scanner) - ✔︎ Successful
- MyPy (Static Code Analysis) - ✔︎ Successful
- Astral Ruff (Static Code Analysis) - ✔︎ Successful

AI Code Review powered by

wild-endeavor added 3 commits February 19, 2025 11:29

testing

f31e823

Signed-off-by: Yee Hing Tong <[email protected]>

add failure handler

5770a5e

Signed-off-by: Yee Hing Tong <[email protected]>

nit

943ca5c

Signed-off-by: Yee Hing Tong <[email protected]>

comments and unit test, add failure node name

be37cd5

wild-endeavor added 3 commits February 21, 2025 14:03

debug

b6dbc73

Signed-off-by: Yee Hing Tong <[email protected]>

debug worker

3937e53

Signed-off-by: Yee Hing Tong <[email protected]>

more print

54a56e5

Signed-off-by: Yee Hing Tong <[email protected]>

wild-endeavor added 3 commits February 21, 2025 17:24

add valuenotin filter, add snippet to cancel tasks

367a393

Signed-off-by: Yee Hing Tong <[email protected]>

add name

a7e1bdf

Signed-off-by: Yee Hing Tong <[email protected]>

debug

8f4cd24

Signed-off-by: Yee Hing Tong <[email protected]>

wild-endeavor added 4 commits February 21, 2025 18:33

remove debugging, add metaclass

da401fc

Signed-off-by: Yee Hing Tong <[email protected]>

change to warning

d5be81d

Signed-off-by: Yee Hing Tong <[email protected]>

nit

9d221fc

Signed-off-by: Yee Hing Tong <[email protected]>

fix test

f12b1f4

Signed-off-by: Yee Hing Tong <[email protected]>

wild-endeavor marked this pull request as ready for review February 22, 2025 03:00

wild-endeavor requested review from kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla, thomasjpfan and Future-Outlier as code owners February 22, 2025 03:00

nit

cdb5189

Signed-off-by: Yee Hing Tong <[email protected]>

flyte-bot reviewed Feb 22, 2025

View reviewed changes

flytekit/tools/serialize_helpers.py Show resolved Hide resolved

flyte-bot reviewed Feb 22, 2025

View reviewed changes

tests/flytekit/unit/core/test_imperative.py Outdated Show resolved Hide resolved

flyte-bot reviewed Feb 22, 2025

View reviewed changes

eapolinario reviewed Feb 25, 2025

View reviewed changes

wild-endeavor added 2 commits February 26, 2025 21:47

change test

6ac8010

Signed-off-by: Yee Hing Tong <[email protected]>

use constant and fix test

3db4249

Signed-off-by: Yee Hing Tong <[email protected]>

eapolinario approved these changes Feb 27, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Eager cleanup #3148

Eager cleanup #3148

wild-endeavor commented Feb 20, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Feb 20, 2025

codecov bot commented Feb 20, 2025 •

edited

Loading

flyte-bot commented Feb 21, 2025

flyte-bot commented Feb 21, 2025

flyte-bot commented Feb 22, 2025

flyte-bot commented Feb 22, 2025 •

edited

Loading

Code Review Agent Run #ce446d

flyte-bot commented Feb 22, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Feb 22, 2025

eapolinario left a comment

eapolinario Feb 25, 2025

flyte-bot commented Feb 27, 2025 •

edited

Loading

Code Review Agent Run #e76f06

Eager cleanup #3148

Are you sure you want to change the base?

Eager cleanup #3148

Conversation

wild-endeavor commented Feb 20, 2025 • edited by flyte-bot Loading

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

Summary by Bito

flyte-bot commented Feb 20, 2025

Code Review Agent Run Status

codecov bot commented Feb 20, 2025 • edited Loading

Codecov Report

flyte-bot commented Feb 21, 2025

Code Review Agent Run Status

flyte-bot commented Feb 21, 2025

Code Review Agent Run Status

flyte-bot commented Feb 22, 2025

Code Review Agent Run Status

flyte-bot commented Feb 22, 2025 • edited Loading

Code Review Agent Run #ce446d

flyte-bot commented Feb 22, 2025 • edited Loading

Changelist by Bito

flyte-bot Feb 22, 2025

Choose a reason for hiding this comment

eapolinario left a comment

Choose a reason for hiding this comment

eapolinario Feb 25, 2025

Choose a reason for hiding this comment

flyte-bot commented Feb 27, 2025 • edited Loading

Code Review Agent Run #e76f06

wild-endeavor commented Feb 20, 2025 •

edited by flyte-bot

Loading

codecov bot commented Feb 20, 2025 •

edited

Loading

flyte-bot commented Feb 22, 2025 •

edited

Loading

flyte-bot commented Feb 22, 2025 •

edited

Loading

flyte-bot commented Feb 27, 2025 •

edited

Loading