Support ArrayNode mapping over Launch Plans

[pvditt]

Tracking issue

fixes: https://linear.app/unionai/issue/COR-1385/add-flytekit-support-for-mapping-over-launch-plans

Why are the changes needed?

ArrayNode currently only supports mapping over a subset of flyte entities (PythonFunctionTasks). To map over other flyte entities such as LaunchPlans we need a new Flytekit node type as the ArrayNodeMapTask implementation can only support Python tasks.

What changes were proposed in this pull request?

Create new ArrayNode class to map over launch plans. Will have a follow up PR to have this new ArrayNode class support Python tasks (as per discussion with @eapolinario and @wild-endeavor).

How was this patch tested?

Ran in sandbox with private fork branch: https://github.com/unionai/flyte/pull/255

@hamersaw scale tested as well

Setup process

Need to create custom image to run. I updated Dockerfile.dev to pulldown the OSS idl changes:

# Use a future version of SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FLYTEIDL such that uv resolution works.
RUN SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FLYTEKIT=$PSEUDO_VERSION \
    SETUPTOOLS_SCM_PRETEND_VERSION_FOR_FLYTEIDL=3.0.0dev0 \
        uv pip install --system --no-cache-dir -U \
            "git+https://github.com/flyteorg/flyte.git@exec-version-array-node#subdirectory=flyteidl" \
            -e /flytekit \
            -e /flytekit/plugins/flytekit-k8s-pod \
            -e /flytekit/plugins/flytekit-deck-standard \
            -e /flytekit/plugins/flytekit-flyteinteractive \
            scikit-learn \
    && apt-get clean autoclean \
    && apt-get autoremove --yes \
    && rm -rf /var/lib/{apt,dpkg,cache,log}/ \
    && useradd -u 1000 flytekit \
    && chown flytekit: /root \
    && chown flytekit: /home \
    && :

@task
def multiply(val: int, val1: int) -> int:
    # if val == 1:
    #     raise Exception("This is a test exception")
    return val * val1


@workflow
def parent_wf(num: int, num1: int) -> int:
    return multiply(val=num, val1=num1)


ex_lp1 = LaunchPlan.get_default_launch_plan(current_context(), parent_wf)


@workflow
def mapped_lp1() -> list[int]:
    a = map_task(ex_lp1, min_success_ratio=0.6)(num=[1, 3, 5], num1=[2, 4, 6])
    return a

Screenshots

Check all the applicable boxes

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

Related PRs

https://github.com/unionai/flyte/pull/255

Docs link

[eapolinario]

There's a fair amount of code duplication between array_node.py and array_node_map_task.py. Have we thought about how to unify some of that code?

Also, I see mentions to execution_version but no correspondent field in the ArrayNode proto, maybe I'm missing a flyteidl change?

I'm slightly confused about the use of mapped_entity. I'll do another pass on the PR later today.

[pvditt]

@eapolinario how/where do we bump flyteidl version for these PR builds?

[eapolinario]

@eapolinario how/where do we bump flyteidl version for these PR builds?

We run tests using master of flyteidl.

[eapolinario]

Reformat this docstring to use the reST format (e.g. [1] or [2]).

[eapolinario]

why do we need to special case min_successes?

[pvditt]

Since we default min_success_ratio to 1.0, it's always set. When passing in both into the proto the one_of will take the last assigned value - which would be min_success_ratio if both are set.

Could also handle this by nulling out min_success_ratio if min_successes is set

[eapolinario]

what do you think of the idea of changing the type of min_success_ratio in map_task? From float to Optional[float] and set its default value to None but at the same time check if both min_successes and min_success_ratio are set to None we "favor" min_success_ratio and set that to 1.0?

This way we centralize this weird interaction between two different parameters in a single function and we don't really validate anything in ArrayNodeModel (so that instantiating objects of that type will always take the arguments as-is) and no more special-casing.

[eapolinario]

@eapolinario how/where do we bump flyteidl version for these PR builds?

We run tests using master of flyteidl.

Actually, plugin tests require a small change to make this a reality. PR incoming.

[eapolinario]

@pvditt , this is getting very close, I have only one small comment.

Also, you mentioned that there was an issue running the sandbox with the launchplan watch api enabled. Is this being tracked somewhere else?

[eapolinario]

what do you think of the idea of changing the type of min_success_ratio in map_task? From float to Optional[float] and set its default value to None but at the same time check if both min_successes and min_success_ratio are set to None we "favor" min_success_ratio and set that to 1.0?

This way we centralize this weird interaction between two different parameters in a single function and we don't really validate anything in ArrayNodeModel (so that instantiating objects of that type will always take the arguments as-is) and no more special-casing.

[pvditt]

@eapolinario just made a ticket for that issue. This isn't an issue when running in any of our tenants. Work around locally is to just disable the config.

[pvditt]

@eapolinario yup - I can implement that real quick

[eapolinario]

Amazing. Thank you for your persistence, this looks great!