[BUG] Serialization bug with pydantic.BaseModel and multiple StructuredDataset
- Dominant language
- Go
- Stars
- 7.5k
- Forks
- 886
- Avg merge
- 1d 14h
- Merged PRs (30d)
- 120
Description
### Flyte & Flytekit version
Flyte 1.15.1
Flytekit 1.16.3
### Describe the bug
Flyte fails to serialize a pydantic.BaseModel with multiple StructuredDataset fields, but dataclass works.
```python
import flytekit as fl
from pydantic import BaseModel
from pydantic.dataclasses import dataclass
# This becomes an empty dict (does not work) after serialization by Flyte when used as output in a task.
class BaseModelTwo(BaseModel):
foo: fl.StructuredDataset
bar: fl.StructuredDataset
# This is serialized (works) by Flyte when used as output in a task.
@dataclass
class DataClassTwo:
foo: fl.StructuredDataset
bar: fl.StructuredDataset
```
### Expected behavior
Both `pydantic.Basemodel` and `dataclass` should work to annotate the output.
### Additional context to reproduce
One can test with these tasks.
```python
import flytekit as fl
import numpy as np
import pandas as pd
# Does not work (empty dict as output)
@fl.task
def return_basemodel_two() -> BaseModelTwo:
return BaseModelTwo(
foo=fl.StructuredDataset(dataframe=pd.DataFrame(np.random.rand(3, 3))),
bar=fl.StructuredDataset(dataframe=pd.DataFrame(np.random.rand(3, 3))),
)
# Works (returns populated dataclass as expected)
@fl.task
def return_dataclass_two() -> DataClassTwo:
return DataClassTwo(
foo=fl.StructuredDataset(dataframe=pd.DataFrame(np.random.rand(3, 3))),
bar=fl.StructuredDataset(dataframe=pd.DataFrame(np.random.rand(3, 3))),
)
```
### Screenshots
_No response_
### Are you sure this issue hasn't been raised already?
- [x] Yes
### Have you read the Code of Conduct?
- [x] Yes
Contributor guide
Research direction
Reproduce the issue with Flyte 1.15.1 and Flytekit 1.16.3 using the return_basemodel_two and return_dataclass_two tasks shown in the report. Compare serialization of BaseModelTwo and DataClassTwo with their two StructuredDataset fields; done means the BaseModel output is populated rather than an empty dict and the dataclass behavior remains working.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100