https://flyte.org logo
#ask-the-community
Title
# ask-the-community
m

Máté Csákvári

02/08/2024, 9:21 AM
Hi! I found a bug in flytekit regarding the attribute access for dataclass type promises in case of int fields. The bug can be reproduced with flytekit==1.10.3 with the following example taken from the documentation by changing the type from str to int:
Copy code
@dataclass_json
@dataclass
class Fruit:
    name: int # was str

@task
def print_message(message: int):
    print(message)
    return

@task
def dataclass_task() -> Fruit:
    return Fruit(name=10)


@workflow
def dataclass_wf():
    fruit_instance = dataclass_task()
    print_message(message=fruit_instance.name)
Running this workflow results in:
Copy code
TypeTransformerFailedError: Error encountered while executing 'dataclass_wf':
  Failed to convert inputs of task 'access-test.print_message':
  Error converting input 'message' at position 0:
  Cannot convert literal <FlyteLiteral scalar { primitive { float_value: 10 } }> to <class 'int'>
This is happening because when a new promise is created from a dataclass instance by accessing one of its attributes, that promise has a pb struct representation with a float_value inside. That promise will be resolved in flytekit.core.promise:translate_inputs_to_literals to a float literal in flyte's type system. Later when that value is about to be converted to an integer using the TypeEngine with the registered SimpleTransformer for ints, it will fail to access the underlying data because its now stored inside lv.scalar.primitive.float_value instead of lv.scalar.primitive.integer:
Copy code
TypeEngine.register(
        SimpleTransformer(
            "int",
            int,
            _type_models.LiteralType(simple=_type_models.SimpleType.INTEGER),
            lambda x: Literal(scalar=Scalar(primitive=Primitive(integer=x))),
            lambda x: x.scalar.primitive.integer   # <- This doesn't exist. x.scalar.primite.float_value does.
        )
    )
If I replace that accessor for the transformer with:
Copy code
def _fix_int_access(lv: Literal) -> int:
    if lv.scalar.primitive.integer is None:
        return int(lv.scalar.primitive.float_value)
    return lv.scalar.primitive.integer
then it works fine. Also the unit test suite completes successfuly with this change. I see that this behaviour was encountered before when dataclasses are passed inside promises and these float values are converted back to integer after the python dataclass has been reconstructed. This is unfortunately not covered by that fix. I can also see a different fix for this right where this conversion first takes place in flytekit.core.promise:
Copy code
try:
  if type(v) is Promise:
    v = resolve_attr_path_in_promise(v)  # <-- HERE we can convert the Literal to contain an integer instead of a float_value
    result[k] = TypeEngine.to_literal(ctx, v, t, var.type)
except TypeTransformerFailedError as exc:
  raise TypeTransformerFailedError(f"Failed argument '{k}': {exc}") from exc
The straightforward fix here worked also, but messed up some of the unit tests which I haven't investigated further yet. I'm not sure which one would be the better solution, but I think fixing it right where the discrepancy first manifests between the type annotations and the underlying data in the flyte's type system seems better for me. I'm happy to create an issue and a PR but I first wanted to ask if this approach is the right way to go about this.
k

Kevin Su

02/08/2024, 10:27 AM
yes, mind creating a pr, we can take a look
cc @Byron Hsu
m

Máté Csákvári

02/08/2024, 11:29 AM
https://github.com/flyteorg/flytekit/pull/2180 Let me know if I missed some things in the process
b

Byron Hsu

02/09/2024, 6:25 AM
Hi @Máté Csákvári thanks for surfacing up the issue and the PR. I implemented attribute access log. I suspect the issue lies in the condition here in propeller. it should have gone into that line but instead to the float one. I will take a closer look cc @Chi-Sheng Liu(劉奇聖)
m

Máté Csákvári

02/09/2024, 7:56 AM
This issue came up at some point I believe when dataclass promises were implemented. This https://github.com/flyteorg/flytekit/blob/master/flytekit/core/type_engine.py#L722 fixes the same issue when all the int attributes of a dataclass are float when transforming back from protobuf. The actual place where float type gets into the typesystem I believe is here: https://github.com/flyteorg/flytekit/blob/master/flytekit/core/promise.py#L132. The type operator interprets the float value coming from protobuf as float type
b

Byron Hsu

02/09/2024, 8:31 AM
The solution makes sense to me at the first glance. Will take a closer look tomorrow
m

Máté Csákvári

02/09/2024, 8:32 AM
Thank you very much! Fixing it at that promise resolve level could also be possible I think, but when I tried it it affected some other parts of the code and I didn't go deeper into it.
b

Byron Hsu

02/09/2024, 8:34 AM
sticking with what flytekit currently has to resolve protobuf strict caveet is a safer way!
Very impressed by your thorough investigation. Thanks a lot!
m

Máté Csákvári

02/09/2024, 8:36 AM
No problem, I like this stuff 🙂 Also we really need that feature :)
b

Byron Hsu

02/09/2024, 8:54 AM
i just reviewed
m

Máté Csákvári

02/09/2024, 8:58 AM
thanks, i'll check it out soon
2 Views