Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support overriding node metadata for array node #2865

Open
wants to merge 4 commits into
base: master
Choose a base branch
from

Conversation

pvditt
Copy link
Contributor

@pvditt pvditt commented Oct 25, 2024

Tracking issue

fixes: https://linear.app/unionai/issue/COR-2090/enable-caching-when-mapping-over-launchplans

Why are the changes needed?

with_overrides doesn't work for array node.

solution:
map_task(my_wf_lp)(a=[1, 2, 3], b=[4, 5, 6]).with_overrides(cache=True, cache_version="3.0")

overrides node metadata for the node that encompasses the array node. This is then copied down into the array node's subnode causing for the node metadata to then be utilized when the nodes are executed. Note, the parent node's node metadata isn't utilized in propeller.

class NodeMetadata(_common.FlyteIdlEntity):
    def __init__(
        self,
        name,
        timeout=None,
        retries=None,
        interruptible: typing.Optional[bool] = None,
        cacheable: typing.Optional[bool] = None,
        cache_version: typing.Optional[str] = None,
        cache_serializable: typing.Optional[bool] = None,

NodeMetadata also doesn't really apply to the parent ArrayNode

Alternative could be to set something like:
map_task(my_wf_lp.with_overrides(cache=True, cache_version="3.0"))(a=[1, 2, 3], b=[4, 5, 6])
but that would require some substantial changes with LaunchPlan and differs from the current understanding of with_overrides

What changes were proposed in this pull request?

  • pass in parent node metadata to be set for subnode

cleaned up a little code:

  • clean up setting passed in metadata for array node (not used atm)

How was this patch tested?

@task
def add(a: int, b: int) -> int:
    return a + b


@workflow
def my_wf(a: int, b: int) -> int:
    return add(a=a, b=b)


my_wf_lp = LaunchPlan.get_default_launch_plan(current_context(), my_wf)


@workflow
def parent_wf_3():
    map_task(my_wf_lp)(a=[1, 2, 3], b=[4, 5, 6]).with_overrides(cache=True, cache_version="3.0")
remote = FlyteRemote(config=Config.auto(), default_project="flytesnacks", default_domain="development")
lp1 = remote.fetch_launch_plan(name="advanced_composition.subworkflow.regression_line_wf",
                               version="k9cgEFiW1sCkKqQhLYbn0Q")


@workflow
def test_wf1_remote():

    val = [1, 2, 3]
    x = [[-3, 0, 3], [-8, 2, 4], [7, 3, 1]]
    y = [[7, 4, -2], [-2, 4, 7], [3, 6, 4]]
    map_task(lp1)(val=val, x=x, y=y).with_overrides(cache=True, cache_version="2.0")

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

Docs link

else:
raise ValueError(f"Only LaunchPlans are supported for now, but got {type(target)}")

def construct_node_metadata(self) -> _workflow_model.NodeMetadata:
# Part of SupportsNodeCreation interface
# TODO - include passed in metadata
return _workflow_model.NodeMetadata(name=self.target.name)
return self.metadata or _workflow_model.NodeMetadata(name=self.target.name)
Copy link
Contributor Author

@pvditt pvditt Oct 25, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

these changes aren't really necessary - just cleaning up code. Don't even support directly passing in node metadata atm.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant