#439 [flyteidl] Support attribute access on promises Pull request opened by ByronHsu TL;DR Support attribute access on output promises for List, Dict, Dataclass, and the mix of them. After the change, users are able to pass in an attribute of output in the workflow, for example:

@task
def t1() -> Dict[str, str]:
  return {"a": "b"}

@task
def t2(x: str):
  return

@workflow
def wf():
  o = t1()
  t2(x=o["a"])

Type ☐ Bug Fix ☑︎ Feature ☐ Plugin Are all requirements met? ☐ Code completed ☐ Smoke tested ☑︎ Unit tests added ☐ Code documentation added ☐ Any pending items have an associated Issue Complete description 1. Add attribute path on promise • flyteidl: Add attribute path on output reference (promise) • flytekit: Append attribute to attribute path in promises by overriding

__getitem__

and

__getattr__

in class promise • flytekit: Modify flyteidlIdentity to serialize attribute path to protobuf 2. Compile Validation at flytepropeller Skip the type check if the promise contains attribute path because if the type is dataclass, we have no way to infer the resolved type. 3. Execution • Local Execution: In flytekit, Resolve the attribute path in translate_input_to_literals. • Remote Execution: In flytepropeller, Resolve the attribute path in Extract_Output. Although in different places, the logic is the same. It iterates the attribute path 1. If the current val is LiteralMap (dictionary), access the value by key 2. If the current val is LiteralCollection (List), access the value by index 3. If the current val is LiteralScalar (single value), break the loop If the remaining val is a LiteralScalar and it contains Dataclass, then iterate the remaining path to extract the value from dataclass. The resolved result of dataclass can be dataclass, list, or a scalar. 1. Dataclass: construct it as a protobuf struct, and wrap literal scalar generic on top. 2. List: construct it as a LiteralCollection 3. Scalar: construct it as a Literal, and wrap literal scalar primitive on top. Testing In the following workflow: •

basic_workflow

contains trivial examples to access output attributes •

failed_workflow

contains examples that causes exception (e.g. out of range) •

advanced_workflow

contains examples with more complex attribute access

from dataclasses import dataclass
from typing import Dict, List, NamedTuple

from dataclasses_json import dataclass_json

from flytekit import WorkflowFailurePolicy, task, workflow


@dataclass_json
@dataclass
class foo:
    a: str


@task
def t1() -> (List[str], Dict[str, str], foo):
    return ["a", "b"], {"a": "b"}, foo(a="b")


@task
def t2(a: str) -> str:
    print("a", a)
    return a


@task
def t3() -> (Dict[str, List[str]], List[Dict[str, str]], Dict[str, foo]):
    return {"a": ["b"]}, [{"a": "b"}], {"a": foo(a="b")}


@task
def t4(a: List[str]):
    print("a", a)


@task
def t5(a: Dict[str, str]):
    print("a", a)


@workflow
def basic_workflow():
    l, d, f = t1()
    t2(a=l[0])
    t2(a=d["a"])
    t2(a=f.a)


@workflow(
    # The workflow doesn't fail when one of the nodes fails but other nodes are still executable
    failure_policy=WorkflowFailurePolicy.FAIL_AFTER_EXECUTABLE_NODES_COMPLETE
)
def failed_workflow():
    # This workflow is supposed to fail due to exceptions
    l, d, f = t1()
    t2(a=l[100])
    t2(a=d["b"])
    t2(a=f.b)


@workflow
def advanced_workflow():
    dl, ld, dd = t3()
    t2(a=dl["a"][0])
    t2(a=ld[0]["a"])
    t2(a=dd["a"].a)

    t4(a=dl["a"])
    t5(a=ld[0])

Local Execution

$ pyflyte run test_workflow.py basic_workflow   

a a
a b
a b

$ pyflyte run test_workflow.py failed_workflow

Failed with Exception Code: USER:PromiseAttributeResolveError
Underlying Exception: Failed to resolve attribute path [100] in promise Resolved(o0=<FlyteLiteral collection { literals { scalar { primitive { string_value: "a" } } } literals { scalar { primitive { string_value: "b" } } } }>), index 100 out of range 2
Encountered error while executing workflow 'test_workflow.failed_workflow':
  Error encountered while executing 'failed_workflow':
  Failed to resolve attribute path [100] in promise Resolved(o0=<FlyteLiteral collection { literals { scalar { primitive { string_value: "a" } } } literals { scalar { primitive { string_value: "b" } } } }>), index 100 out of range 2

$ pyflyte run test_workflow.py advanced_workflow
a b
a b
a b
a ['b']
a {'a': 'b'}

Remote Execution • Basic • Failed • Advanced Tracking Issue flyteorg/flyte#3864 flyteorg/flyteidl ✅ All checks have passed 13/13 successful checks