add annotation option for serialization

[elibixby]

TL;DR

Allow users to specify TypeTransformers via Annotation

This is a followup to conversation in #1565

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

How did you fix the bug, make the feature etc. Link to any design docs etc

Tracking Issue

https://github.com/flyteorg/flyte/issues/

Follow-up issue

[elibixby]

@wild-endeavor @pingsutw

Based on our usage at Cradle I think this is the highest priority improvement to come out of #1565 (which is a little overstuffed :) )

A pattern we find ourselves reaching for pretty frequently for is:

class SerializationWrapper:

    def __class_getitem__(cls, wrapped_type):
        TypeEngine.register(Transformer(wrapped_type), wrapped_type)
        return wrapped_type

However this is limited if we want to specify alternate serialization for types, as TypeEngine enforces a 1:1 mapping from python types to TypeTransformer.

This small change unlocks that requirement, by storing the type serialization information in the annotation of the type.

[codecov]

Codecov Report

Merging #1615 (1254e50) into master (baf0b89) will increase coverage by 1.18%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #1615      +/-   ##
==========================================
+ Coverage   70.10%   71.28%   +1.18%     
==========================================
  Files         328      334       +6     
  Lines       30156    30495     +339     
  Branches     5452     5512      +60     
==========================================
+ Hits        21141    21739     +598     
+ Misses       8485     8205     -280     
- Partials      530      551      +21     

Impacted Files	Coverage Δ
flytekit/core/type_engine.py	61.28% <100.00%> (-0.23%)	⬇️
tests/flytekit/unit/core/test_type_engine.py	98.65% <100.00%> (+0.02%)	⬆️

... and 48 files with indirect coverage changes

[wild-endeavor]

So the idea is that we have

class MyType(object):
    ...

class MyTypeTransformer(TypeTransformer[MyType]):
    ...

class SecondaryMyTypeTransformer(TypeTransformer[MyType]):
    ...

MM = Annotated[MyType, SecondaryMyTypeTransformer]
m1 = MyType()
m2 = MM()

then m1 would get the first transformer and m2 would get the second transformer right? I think in spirit I'm mostly okay with this because it's like a manual bypass for the type engine. The user is overriding the type engine and declaring what transformer should be used.

I don't think it should be saved though, or added to guess_python_types. Honestly the guessing bit has led to a lot of frustration on the user side especially as we've added more complex types like StructuredDataset. We eventually gave up and just added the ability to specify a type to the LiteralsResolver object.

Would like to get @eapolinario's thoughts on this though.

[pingsutw]

qq: when do you use two transformer for certain type? why can't just override default flytekit transformer?

[elibixby]

qq: when do you use two transformer for certain type? why can't just override default flytekit transformer?

I can think there are two main cases:

1. You have a generic type which has many possible sources
2. A base type where you don't want to override the implementation but it comes from another source

In general, these boil down to the same thing: explicitly separating the source of the data (as a generic, or type factory), and the type of the data when it is decoded. I can see this being used in FlyteFile, StructuredDataset and essentially any other non-primitive type, to create very useful and intuitive behavior.

For example:

@typing.protocol
class FromParquet[T]:

     @classmethod
     def from_parquet(cls, file: typing.BytesIO, columns: List[str]) -> T:
         ...

class StructuredDataset:

    TextIO = Annotated[typing.TextIO, FlyteFileTypeTransformer(format='parquet')]
    Path = Annotated[os.PathLike, BlobMessagePathTransformer()]

    def __class_getitem__(self, decoded_as: Type[FromParquet[T]], kwtypes: OrderedDict[str, type]) -> Type[FromParquet[T]]:
         return Annotated[decoded_as, StructuredDatasetTypeTransformer(decoded_as=decoded_as, schema=kwtypes)]
 
 
schema = kwtypes(foo=int, bar=str)
 
@task
def my_task_dataframe(dataset: StructuredDataset[pd.DataFrame, schema]):
    assert isinstance(dataset, pd.DataFrame)
    
@task
def my_task_dataset(dataset: StructuredDataset[torch.data.Dataset, schema]):
   ...
   
@task
def my_task_file(dataset_file: StructuredDataset.TextIO):
   ...
   
@task
def my_task_filepath(dataset_path: StructuredDataset.Path):
   ...
 

In particular, the last two are not possible under the current system, because not every TextIO or str would be applicable, (you'd only want files or file-paths that represent parquet)

At cradle we want to be able to chain together workflows while keeping provenance information. We've decided the easiest way to do this is to have an OutputRef type which uses FlyteRemote to fetch the output of a specified execution from the Flyte server. The ideal syntax here is OutputRef[MyType], which inside the function is treated as MyType. But this doesn't work because by definition MyType will already have a registered type transformer (it is the output of another workflow after all).

We also would like to be able to pass references to tasks/launch plans/workflows into workflow functions for incorporation into the dag, as part of a dynamic call. The ideal approach here would be similar RemoteRef[Callable[...]], but registering generics dynamically is currently not possible.

[welcome]

Congrats on merging your first pull request! 🎉