flyteorg / flyteorg/flyte

[BUG] Serialization bug with pydantic.BaseModel and multiple StructuredDataset

Open
#6,574 2 comments 0 reactions 0 assignees View on GitHub
bug untriaged
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.