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

[ebran]

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.

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.

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?

• Yes

Have you read the Code of Conduct?

• Yes

[welcome]

Thank you for opening your first issue here! 🛠

[runllm]

This is a known issue: Flytekit currently has serialization problems when using pydantic.BaseModel with multiple StructuredDataset fields, resulting in empty dicts or failed outputs, while dataclasses work as expected. The root cause is that Flytekit's type transformer and serialization logic do not fully support multiple StructuredDataset fields inside a pydantic.BaseModel, but do support them in dataclasses, as confirmed by user reports and test cases in the Flyte repository. There is no official fix yet, and the recommended workaround is to use dataclasses instead of pydantic.BaseModel for such cases. See the discussion and examples in flyte/issues/5925 and test_pydantic_basemodel_in_dataclass.py.

Would you like a step-by-step explanation of why this happens and possible workarounds?

Sources:

• test_pydantic_basemodel_in_dataclass.py
• flyte/issues/5925

---

This is an AI-generated response and your feedback is appreciated! Please leave a 👍 if this is helpful and 👎 if it is not.

Mention @RunLLM to ask follow-up questions. See runllm.com for more!