Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix pydantic basemodel default input #3013

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Fix pydantic basemodel default input #3013

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
90450d1
Fix pydantic default input
Future-Outlier Dec 18, 2024
d46d292
add pydantic integration test
Future-Outlier Dec 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 16 additions & 2 deletions flytekit/clis/sdk_in_container/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
labels_callback,
)
from flytekit.interaction.string_literals import literal_string_repr
from flytekit.lazy_import.lazy_module import is_imported
from flytekit.loggers import logger
from flytekit.models import security
from flytekit.models.common import RawOutputDataConfig
Expand Down Expand Up @@ -475,8 +476,21 @@ def to_click_option(
If no custom logic exists, fall back to json.dumps.
"""
with FlyteContextManager.with_context(flyte_ctx.new_builder()):
encoder = JSONEncoder(python_type)
default_val = encoder.encode(default_val)
if is_imported("pydantic"):
try:
from pydantic import BaseModel as BaseModelV2
from pydantic.v1 import BaseModel as BaseModelV1

if issubclass(python_type, BaseModelV2):
default_val = default_val.model_dump_json()
elif issubclass(python_type, BaseModelV1):
default_val = default_val.json()
except ImportError:
# Pydantic BaseModel v1
default_val = default_val.json()
Comment on lines +479 to +490
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it enough to duck type here?

if hasattr(default_val, "model_dump_json"):
    # pydantic v2
    default_val = default_val.model_dump_json()
elif hasattr(default_val, "json"):
    # pydantic v1
    default_val = default_val.json()
else:
    encoder = JSONEncoder(python_type)
    default_val = encoder.encode(default_val)

else:
encoder = JSONEncoder(python_type)
default_val = encoder.encode(default_val)
if literal_var.type.metadata:
description_extra = f": {json.dumps(literal_var.type.metadata)}"

Expand Down
8 changes: 8 additions & 0 deletions tests/flytekit/integration/remote/test_remote.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,14 @@ def test_remote_eager_run():
# child_workflow.parent_wf asynchronously register a parent wf1 with child lp from another wf2.
run("eager_example.py", "simple_eager_workflow", "--x", "3")

def test_pydantic_default_input_with_map_task():
execution_id = run("pydantic_wf.py", "wf")
remote = FlyteRemote(Config.auto(config_file=CONFIG), PROJECT, DOMAIN)
execution = remote.fetch_execution(name=execution_id)
execution = remote.wait(execution=execution, timeout=datetime.timedelta(minutes=5))
print("Execution Error:", execution.error)
assert execution.closure.phase == WorkflowExecutionPhase.SUCCEEDED, f"Execution failed with phase: {execution.closure.phase}"


def test_generic_idl_flytetypes():
os.environ["FLYTE_USE_OLD_DC_FORMAT"] = "true"
Expand Down
20 changes: 20 additions & 0 deletions tests/flytekit/integration/remote/workflows/basic/pydantic_wf.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
from pydantic import BaseModel

from flytekit import map_task
from typing import List
from flytekit import task, workflow


class MyBaseModel(BaseModel):
my_floats: List[float] = [1.0, 2.0, 5.0, 10.0]

@task
def print_float(my_float: float):
print(f"my_float: {my_float}")

@workflow
def wf(bm: MyBaseModel = MyBaseModel()):
map_task(print_float)(my_float=bm.my_floats)

if __name__ == "__main__":
wf()
Loading