[Sovietaced]

Tracking issue

Closes #8282

Why are the changes needed?

The ListChildActions operations enumerates over every single TaskAction in a run to determine the child actions for a parent which can explode memory usage for high scale + fanout runs.

What changes were proposed in this pull request?

Add a feature flag to filter TaskActions based on a new parent task action label such that filtering is done on the server side instead of the client side in application code. This reduces CPU usage, memory usage, and network I/O.

I chose to use a feature flag for this because I don't think there is any truly safe way to understand when to use the optimized query pattern vs the old one because of progressive rollouts, rollbacks, etc. As such, I've left it up to the user to enable the feature but I think at some point we should try and make it the default behavior (ie. the next minor release).

How was this patch tested?

We tested this optimization in our internal for albeit with a Redis based TaskAction store implementation we have. I think performance improvements should be similar.

Example run with total tasks set to 5000.

"""Flyte 2 scale workload with two levels of fanout over square tasks."""

import asyncio
import logging
from functools import partial
from typing import TYPE_CHECKING, Final

import flyte

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable, Coroutine, Iterable

map_env: Final = flyte.TaskEnvironment(
    name="scale_map_env",
    resources=flyte.Resources(cpu="300m", memory="300Mi"),
)
env: Final = flyte.TaskEnvironment(name="scale_fanout_env", depends_on=[map_env])
_FIRST_LEVEL_FANOUT: Final = 100
_CHILD_RETRY_DELAYS: Final = (5, 10, 20)
_CHILD_LAUNCH_DELAY: Final = 0.05
_logger: Final = logging.getLogger(__name__)


async def _gather_staggered(calls: "Iterable[Coroutine[None, None, int]]") -> list[int]:
    """Launch children 50 ms apart and collect their results in input order."""
    tasks: list[asyncio.Task[int]] = []
    try:
        for call in calls:
            tasks.append(asyncio.create_task(call))
            await asyncio.sleep(_CHILD_LAUNCH_DELAY)
        return await asyncio.gather(*tasks)
    finally:
        for task in tasks:
            if not task.done():
                task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)


@map_env.task
def simple_map_task(a: int) -> int:
    """Square one input in a separate task."""
    return a * a


async def _with_child_retries(call: "Callable[[], Awaitable[int]]") -> int:
    """Retry a failed child up to three times with exponential backoff."""
    for delay in _CHILD_RETRY_DELAYS:
        try:
            return await call()
        except Exception:
            _logger.warning("Child task failed; retrying in %s seconds", delay, exc_info=True)
            await asyncio.sleep(delay)
    return await call()


@env.task
def simple_reduce(b: list[int]) -> int:
    """Sum the completed map task outputs."""
    return sum(b)


@env.task
async def fanout_task(branch_index: int, size: int) -> int:
    """Fan out square tasks over one branch's share of the inputs, then sum them."""
    output = await _gather_staggered(
        _with_child_retries(partial(simple_map_task.aio, a=value))
        for value in range(branch_index, size, _FIRST_LEVEL_FANOUT)
    )
    return await _with_child_retries(partial(simple_reduce.aio, b=output))


@env.task
async def wf(size: int = 1000) -> int:
    """Split size integers across 100 child tasks, then sum their squared results."""
    output = await _gather_staggered(
        _with_child_retries(partial(fanout_task, branch_index=branch_index, size=size))
        for branch_index in range(_FIRST_LEVEL_FANOUT)
    )
    return await _with_child_retries(partial(simple_reduce.aio, b=output))

Before (Notice replicas autoscale to 5. 2-3 total CPU usage, memory usage shoots up > 2G, and network RX is ~ 20-30 MB/s)

After (CPU usage peaks at .3 CPU, memory peaks at 400MB, network RX peaks at ~3 MB/s)

Labels

Please add one or more of the following labels to categorize your PR:

• added: For new features.
• changed: For changes in existing functionality.
• deprecated: For soon-to-be-removed features.
• removed: For features being removed.
• fixed: For any bug fixed.
• security: In case of vulnerabilities

This is important to improve the readability of release notes.

Setup process

Screenshots

Check all the applicable boxes

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

Related PRs

Stack

If you do use git town to manage PR Stacks, the stack relevant to this PR
will show below. Otherwise, you can ignore this section.

Docs link

[cloud-verifier-agent]

cloud-verifier · ⚠️ 1 finding, not blocking

🟠 parentLabelQueries hides pre-existing unlabeled children from watch snapshot · actions/k8s/client.go:431 · medium · medium likelihood · see inline comment

_{Reviewed 31920c6973 · turn 1 · re-run with /cloud-verifier review}

[cloud-verifier-agent]

🟠 parentLabelQueries hides pre-existing unlabeled children from watch snapshot

With the flag on, the spec-based parent check is skipped, so actions without the new label are dropped from the result. Consider a backfill or fallback for legacy actions.

[Sovietaced]

That's why the flag defaults to false, it requires operators to be sure that task actions have the new labels. It is documented as such.

[kumare3]

I think you shoudl not need feature flag i feel. as you always get the parent name in every request and it seems it is just dropped today right

[Sovietaced]

I think you shoudl not need feature flag i feel. as you always get the parent name in every request and it seems it is just dropped today right

There is no existing label which can be used to query the k8s with a filter to get only the TaskActions with the target parent. So there reason for the feature flag is to allow folks to run the system long enough such that every TaskAction in the system is guaranteed to have the parent label available, making the new query filter safe to use.