feat(sdk): support collecting outputs from conditional branches using `dsl.OneOf` #10067

connor-mccarthy · 2023-10-06T22:53:55Z

Description of your changes:
Supports collecting outputs from conditional branches in a pipeline using dsl.OneOf. For example:

from kfp import dsl

@dsl.component
def flip_coin() -> str:
    import random
    return 'heads' if random.randint(0, 1) == 0 else 'tails'

@dsl.component
def print_and_return(text: str) -> str:
    print(text)
    return text

@dsl.pipeline
def flip_coin_pipeline() -> str:
    flip_coin_task = flip_coin()
    with dsl.If(flip_coin_task.output == 'heads'):
        print_task_1 = print_and_return(text='Got heads!')
    with dsl.Else():
        print_task_2 = print_and_return(text='Got tails!')

    # use the output from the branch that gets executed
    oneof = dsl.OneOf(print_task_1.output, print_task_2.output)

    # consume it
    print_and_return(text=oneof)

    # return it
    return oneof

There are three future work items to consider following this PR:

Migrate the dsl.Collected implementation to the dsl.OneOf-style implementation. The abstraction for the dsl.OneOf is much cleaner and permits more implementation alignment and code reuse between dsl.OneOf and dsl.Collected. I left TODOs for this, opting not to implement in this CL to simplify the already large diff.
After 1, we can consider (but may choose not to) support composing dsl.Collected into dsl.OneOf. This composability increases the complexity of the compiler code + authorable pipelines considerably (high implementation + maintenance cost). We should weigh this against the suspected benefit. Leaving this note here to document the choice and the likely dependency on 1.
We can consider supporting dsl.OneOf and dsl.Collected in f-strings. dsl.Collected support in f-strings follows naturally from 1 above. dsl.OneOf support in f-strings requires a new way of injecting a channel into a user string in the compiler code such that all the member channels of a dsl.OneOf are not lost. The approach which I suspect will be most successful (included in a code comment) is to maintain a pipeline-level map of unique key to pipeline channel, then injecting/extracting that key from the user string, then looking up in the map. Similar level of effort consideration to 2 above.

Checklist:

The title for your pull request (PR) should follow our title convention. Learn more about the pull request title convention used in this
repository.

connor-mccarthy · 2023-10-12T22:40:11Z

/assign @chensun

sdk/python/kfp/compiler/compiler_test.py

chensun

/lgtm

Only a few nitpicks and questions, nothing blocking.
Great work!

chensun · 2023-10-17T21:31:24Z

sdk/python/kfp/dsl/pipeline_channel.py

+            if not isinstance(channel, expected_type):
+                raise TypeError(
+                    f'Task outputs passed to dsl.{OneOf.__name__} must be the same type. Got two channels with different types: {readable_name_map[expected_type]} at index 0 and {readable_name_map[type(channel)]} at index {i}.'
+                )


nit: this function reads a bit complex, while the essential check is all channels are of the same type. Is the complexity to allow specific error message? Not sure if user would understand what index # refers to.

I simplified the logic a bit and the error message. Please feel free to re-open if there are more changes you suggest.

chensun · 2023-10-17T21:42:58Z

sdk/python/kfp/dsl/pipeline_channel.py

+    ) -> None:
+        # avoid circular imports
+        from kfp.dsl import tasks_group
+        if not isinstance(parent_group.groups[-1], tasks_group.Else):


Does this mean users have to have an ending dsl.Else?
So the following is illegal?

dsl.If(a == True): f = foo() dsl.Elif(a == False): b = bar() oneof = dsl.OneOf(f.output, b.output)

Yes, exactly.

For posterity, I'll leave a few notes to document my thinking:

This is aligned with the current treatment of value_from_oneof field in PipelineSpec on Vertex Pipelines, which is a backend that supports such collection.

If we want different behavior... We can, of course, prescribe new behavior for any backend and drop this validation in the KFP SDK. Since the current implementation is forward compatible with looser validation, my preference is to proceed with the strong validation and listen for user demands to loosen. This helps us get this feature in the hands of users faster with effectively no added cost for users or maintainers.

If we like this behavior, but want to permit more flexibility... It's technically possible that a user might want to use only dsl.If and dsl.Elif and ensure the conditions are mutually exclusive themselves (no dsl.Else). Since this comes with the cost of shifting errors right (to runtime, instead of compilation time as in the current implementation) and because this seems like it wouldn't be the typical case, my preference is again to start conservatively with the current approach.

chensun · 2023-10-17T23:05:17Z

sdk/python/kfp/dsl/pipeline_channel.py

+        # the combination of OneOf and an f-string is not common, so prefer
+        # deferring implementation
+        raise NotImplementedError(
+            f'dsl.{OneOf.__name__} is not yet supported in f-strings.')


nit: Technically f-string is just one special syntax for string concatenation, while this limitation would apply to all formats, for example, "{} world".format(str1), "Hello" + str1, etc.

Nice catch. Updated.

chensun · 2023-10-17T23:05:48Z

sdk/python/kfp/dsl/pipeline_channel.py

+# isinstance(<some channel>, PipelineParameterChannel/PipelineArtifactChannel)
+# checks continue to behave in desirable ways
+class OneOfParameter(PipelineParameterChannel, OneOfMixin):
+    """OneOf that results in an parameter channel for all downstream tasks."""


nit: a parameter channel

chensun · 2023-10-17T23:18:34Z

sdk/python/kfp/dsl/tasks_group.py

@@ -150,6 +153,16 @@ def __init__(self) -> None:
            is_root=False,
        )

+    def _get_oneof_id(self) -> int:


nit: this seems to be an public method, so maybe no underscore prefix?

Yes, makes sense.

connor-mccarthy

Thanks for the review, @chensun. Addressed your comments.

connor-mccarthy · 2023-10-18T00:52:24Z

sdk/python/kfp/dsl/pipeline_channel.py

+        # the combination of OneOf and an f-string is not common, so prefer
+        # deferring implementation
+        raise NotImplementedError(
+            f'dsl.{OneOf.__name__} is not yet supported in f-strings.')


Nice catch. Updated.

connor-mccarthy · 2023-10-18T01:10:37Z

sdk/python/kfp/dsl/pipeline_channel.py

+    ) -> None:
+        # avoid circular imports
+        from kfp.dsl import tasks_group
+        if not isinstance(parent_group.groups[-1], tasks_group.Else):


Yes, exactly.

For posterity, I'll leave a few notes to document my thinking:

This is aligned with the current treatment of value_from_oneof field in PipelineSpec on Vertex Pipelines, which is a backend that supports such collection.

If we want different behavior... We can, of course, prescribe new behavior for any backend and drop this validation in the KFP SDK. Since the current implementation is forward compatible with looser validation, my preference is to proceed with the strong validation and listen for user demands to loosen. This helps us get this feature in the hands of users faster with effectively no added cost for users or maintainers.

If we like this behavior, but want to permit more flexibility... It's technically possible that a user might want to use only dsl.If and dsl.Elif and ensure the conditions are mutually exclusive themselves (no dsl.Else). Since this comes with the cost of shifting errors right (to runtime, instead of compilation time as in the current implementation) and because this seems like it wouldn't be the typical case, my preference is again to start conservatively with the current approach.

connor-mccarthy · 2023-10-18T01:14:40Z

sdk/python/kfp/dsl/pipeline_channel.py

+            if not isinstance(channel, expected_type):
+                raise TypeError(
+                    f'Task outputs passed to dsl.{OneOf.__name__} must be the same type. Got two channels with different types: {readable_name_map[expected_type]} at index 0 and {readable_name_map[type(channel)]} at index {i}.'
+                )


I simplified the logic a bit and the error message. Please feel free to re-open if there are more changes you suggest.

connor-mccarthy · 2023-10-18T01:16:09Z

sdk/python/kfp/dsl/tasks_group.py

@@ -150,6 +153,16 @@ def __init__(self) -> None:
            is_root=False,
        )

+    def _get_oneof_id(self) -> int:


Yes, makes sense.

chensun

/lgtm
/approve

Thanks!

google-oss-prow · 2023-10-18T02:34:45Z

[APPROVALNOTIFIER] This PR is APPROVED

This pull-request has been approved by: chensun

The full list of commands accepted by this bot can be found here.

The pull request process is described here

Needs approval from an approver in each of these files:

~~sdk/OWNERS~~ [chensun]

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

… `dsl.OneOf` (kubeflow#10067) * support dsl.OneOf * address review feedback * address review feedback

…beflow#10128)

google-oss-prow bot requested review from IronPan and zijianjoy October 6, 2023 22:54

google-oss-prow bot added the size/XXL label Oct 6, 2023

connor-mccarthy force-pushed the oneof-7 branch 9 times, most recently from 35df06d to e2ecf3e Compare October 8, 2023 16:47

connor-mccarthy marked this pull request as draft October 8, 2023 16:57

google-oss-prow bot added the do-not-merge/work-in-progress label Oct 8, 2023

support dsl.OneOf

dddfbac

connor-mccarthy force-pushed the oneof-7 branch from e2ecf3e to dddfbac Compare October 9, 2023 22:16

connor-mccarthy marked this pull request as ready for review October 9, 2023 22:34

google-oss-prow bot removed the do-not-merge/work-in-progress label Oct 9, 2023

google-oss-prow bot requested a review from james-jwu October 9, 2023 22:35

connor-mccarthy mentioned this pull request Oct 11, 2023

sdk: add dsl.OneOf docs kubeflow/website#3605

Merged

google-oss-prow bot assigned chensun Oct 12, 2023

chensun reviewed Oct 14, 2023

View reviewed changes

sdk/python/kfp/compiler/compiler_test.py Outdated Show resolved Hide resolved

sdk/python/kfp/compiler/compiler_test.py Outdated Show resolved Hide resolved

sdk/python/kfp/compiler/compiler_test.py Show resolved Hide resolved

address review feedback

f914893

connor-mccarthy requested a review from chensun October 16, 2023 15:38

chensun reviewed Oct 17, 2023

View reviewed changes

google-oss-prow bot added the lgtm label Oct 17, 2023

address review feedback

8021e90

google-oss-prow bot removed the lgtm label Oct 18, 2023

connor-mccarthy commented Oct 18, 2023

View reviewed changes

connor-mccarthy requested a review from chensun October 18, 2023 01:22

chensun approved these changes Oct 18, 2023

View reviewed changes

google-oss-prow bot added the lgtm label Oct 18, 2023

google-oss-prow bot added the approved label Oct 18, 2023

google-oss-prow bot merged commit 2d3171c into kubeflow:master Oct 18, 2023
19 checks passed

connor-mccarthy added a commit to connor-mccarthy/pipelines that referenced this pull request Oct 19, 2023

chore(sdk): add test for key error bug resolved in kubeflow#10067

4248105

connor-mccarthy mentioned this pull request Oct 19, 2023

chore(sdk): add test for key error bug resolved in #10067 #10128

Merged

1 task

connor-mccarthy added a commit to connor-mccarthy/pipelines that referenced this pull request Oct 19, 2023

chore(sdk): add test for key error bug resolved in kubeflow#10067

96ae8c5

connor-mccarthy added a commit to connor-mccarthy/pipelines that referenced this pull request Oct 19, 2023

chore(sdk): add test for key error bug resolved in kubeflow#10067

5863872

google-oss-prow bot pushed a commit that referenced this pull request Oct 19, 2023

chore(sdk): add test for key error bug resolved in #10067 (#10128)

c0adea9

stijntratsaertit pushed a commit to stijntratsaertit/kfp that referenced this pull request Feb 16, 2024

feat(sdk): support collecting outputs from conditional branches using…

4106993

… `dsl.OneOf` (kubeflow#10067) * support dsl.OneOf * address review feedback * address review feedback

stijntratsaertit pushed a commit to stijntratsaertit/kfp that referenced this pull request Feb 16, 2024

chore(sdk): add test for key error bug resolved in kubeflow#10067 (ku…

d794821

…beflow#10128)

rimolive mentioned this pull request Jun 10, 2024

[bug] Unable to compile pipeline after upgrading to KFP 2.4 #10722

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(sdk): support collecting outputs from conditional branches using `dsl.OneOf` #10067

feat(sdk): support collecting outputs from conditional branches using `dsl.OneOf` #10067

connor-mccarthy commented Oct 6, 2023 •

edited

Loading

connor-mccarthy commented Oct 12, 2023

chensun left a comment

chensun Oct 17, 2023

connor-mccarthy Oct 18, 2023

chensun Oct 17, 2023

connor-mccarthy Oct 18, 2023

chensun Oct 17, 2023

connor-mccarthy Oct 18, 2023

chensun Oct 17, 2023

chensun Oct 17, 2023

connor-mccarthy Oct 18, 2023

connor-mccarthy left a comment

connor-mccarthy Oct 18, 2023

connor-mccarthy Oct 18, 2023

connor-mccarthy Oct 18, 2023

connor-mccarthy Oct 18, 2023

chensun left a comment

google-oss-prow bot commented Oct 18, 2023

feat(sdk): support collecting outputs from conditional branches using dsl.OneOf #10067

feat(sdk): support collecting outputs from conditional branches using dsl.OneOf #10067

Conversation

connor-mccarthy commented Oct 6, 2023 • edited Loading

connor-mccarthy commented Oct 12, 2023

chensun left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

connor-mccarthy left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chensun left a comment

Choose a reason for hiding this comment

google-oss-prow bot commented Oct 18, 2023

feat(sdk): support collecting outputs from conditional branches using `dsl.OneOf` #10067

feat(sdk): support collecting outputs from conditional branches using `dsl.OneOf` #10067

connor-mccarthy commented Oct 6, 2023 •

edited

Loading