Feat: Add type support for pydantic BaseModels

[ArthurBook]

TL;DR

Lets you use pydantic basemodels as inputs and outputs in tasks.
Also extends the basemodels to have any flyte supported file as field.

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

Added serializer/deserializers for flytepath and flytedirectory s.t. pydantic can jsonify these objects too.

Type Example

from pydantic import BaseModel

class Config(BaseModel):
    lr: float = 1e-3
    batch_size: int = 32
    files: List[FlyteFile]
    directories: List[FlyteDirectory]

@task
def train(cfg: Config):
    ...

Tracking Issue

flyteorg/flyte#2686

Follow-up issue

NA

[welcome]

Thank you for opening this pull request! 🙌

These tips will help get your PR across the finish line:

• Most of the repos have a PR template; if not, fill it out to the best of your knowledge.
• Sign off your commits (Reference: DCO Guide).

[codecov]

Codecov Report

Patch coverage has no change and project coverage change: -52.29% ⚠️

Comparison is base (47be974) 71.01% compared to head (f58a55b) 18.73%.
Report is 97 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #1660       +/-   ##
===========================================
- Coverage   71.01%   18.73%   -52.29%     
===========================================
  Files         336      332        -4     
  Lines       30726    31398      +672     
  Branches     5567     3082     -2485     
===========================================
- Hits        21821     5882    -15939     
- Misses       8361    25432    +17071     
+ Partials      544       84      -460     

Files Changed	Coverage Δ
flytekit/clis/sdk_in_container/run.py	0.00% <0.00%> (-85.57%)	⬇️

... and 348 files with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[pingsutw]

Thanks @ArthurBook for working on this.
To serialize pydantic.baseModel, we have to serialize the object to json. Do I understand correctly?
If so, I have some thoughts

1. Add to_json to typeTransformer?
2. Make lv become Union type (union[json, Literal]) or add another method in transformer to_pydantic_model?
3. Create a generic FlyteDirSerializerDeserializer, and use transformer.to_json to serialize all the python value.
4. I think it's worth to move pydantic transformer to flytekit/core/type

wdyt?

[pingsutw]

@ArthurBook sorry for the late reply, we had some discussion internally. let's keep pydantic transformer in the plugin first. overall LGTM, leave a comment below.

[wild-endeavor]

@ArthurBook @fg91 should we just add the pydantic version gate and then merge this in? tested on a flytectl demo environment right? also lint is breaking, could you take a look please?

[wild-endeavor]

@fg91 sorry i think i lost the slack thread. did you find issues in testing?

[fg91]

@fg91 sorry i think i lost the slack thread. did you find issues in testing?

Yes, I did, I boiled it down to this minimal working example:

from flytekit import task, workflow, dynamic

from pydantic import BaseModel


class Config(BaseModel):
    path: str


@task
def train(cfg: Config):
    print(cfg)


@dynamic(cache=True, cache_version="0.3")
def sub_wf(cfg: Config):
    train(cfg=cfg)


@workflow
def wf():
    sub_wf(cfg=Config(path="bar"))


if __name__ == "__main__":
    wf()

Upon local execution with python test.py, the error is:

│ /home/fabiogratz/miniconda3/envs/dev/lib/python3.9/site-packages/flytekit/core/workflow.py:271 in __call__                  │
│                                                                                                                             │
│ ❱ 271 │   │   │   raise exc                                                                                                 │
│                                                                                                                             │
│ /home/fabiogratz/miniconda3/envs/dev/lib/python3.9/site-packages/flytekit/core/workflow.py:268 in __call__                  │
│                                                                                                                             │
│ ❱ 268 │   │   │   return flyte_entity_call_handler(self, *args, **input_kwargs)                                             │
│                                                                                                                             │
│ /home/fabiogratz/miniconda3/envs/dev/lib/python3.9/site-packages/flytekit/core/promise.py:1022 in flyte_entity_call_handler │
│                                                                                                                             │
│ ❱ 1022 │   │   │   result = cast(LocallyExecutable, entity).local_execute(child_ctx, **kwargs)                              │
│                                                                                                                             │
│ /home/fabiogratz/miniconda3/envs/dev/lib/python3.9/site-packages/flytekit/core/workflow.py:290 in local_execute             │
│                                                                                                                             │
│ ❱ 290 │   │   function_outputs = self.execute(**kwargs_literals)                                                            │
│                                                                                                                             │
│ /home/fabiogratz/miniconda3/envs/dev/lib/python3.9/site-packages/flytekit/core/workflow.py:729 in execute                   │
│                                                                                                                             │
│ ❱ 729 │   │   return exception_scopes.user_entry_point(self._workflow_function)(**kwargs)                                   │
│                                                                                                                             │
│ /home/fabiogratz/miniconda3/envs/dev/lib/python3.9/site-packages/flytekit/exceptions/scopes.py:203 in user_entry_point      │
│                                                                                                                             │
│ ❱ 203 │   │   │   │   raise type(exc)(f"Error encountered while executing '{fn_name}':\n  {exc                              │
╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
TypeError: Encountered error while executing workflow 'wf':
  Error encountered while executing 'wf':
  Parameter to MergeFrom() must be instance of same class: expected <class 'flyteidl.core.literals_pb2.Literal'> got <class 
'flyteidl.core.literals_pb2.LiteralMap'>.

When setting cache=False it works. But having forgotten to bump the cache version appears not to be the problem. When activating the cache and setting a cache version that 100% was never used before, it fails with this error on the first try.

[ArthurBook]

Hi all, haven't forgotten about this - just busy with a product launch + travel.
I'll aim to get to this ASAP but need to prioritize work stuff first.

Thanks @fg91 for testing this out and finding that!

[ArthurBook]

Hi all, found the long longed for moment to look this over.
Its a great catch @fg91. The issue is that nested literalmaps cause hickups:

> import from flyteidl.core import literals_pb2
> literals_pb2.LiteralMap(literals={'boo': literals_pb2.LiteralMap(literals={})})

>>> Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File ".../.virtualenvs/flytekit/lib/python3.8/site-packages/google/protobuf/message.py", line 126, in CopyFrom
    self.MergeFrom(other_msg)
TypeError: Parameter to MergeFrom() must be instance of same class: expected <class 'flyteidl.core.literals_pb2.Literal'> got <class 'flyteidl.core.literals_pb2.LiteralMap'>.

In the previous code we had a nested literalmap in the serialized basemodel here. I've pushed a change that works around this. Also added your example as tests so we avoid it in the future. Works smooth now at least on my end - let me know what you think!

[wild-endeavor]

is this a common occurrence though?

literals_pb2.LiteralMap(literals={'boo': literals_pb2.LiteralMap(literals={})})

wouldn't it be

literals_pb2.LiteralMap(literals={'boo': literals_pb2.Literal(map=literals_pb2.LiteralMap(literals={})}))

not sure if this is relevant at all.

[ArthurBook]

That would work too! Let me know if you have a preference to serializing to that rather than a flat literalmap like here?

[wild-endeavor]

can't commit to your branch but added a unit test that at least helps me understand what's going on. - 01c065e and updated some formatting. can you pull this in and then we should merge.

Thanks!

old review

[welcome]

Congrats on merging your first pull request! 🎉