From 8ff18ef68320a630330d151d14269f1e344cbba5 Mon Sep 17 00:00:00 2001 From: sumana sree Date: Sun, 20 Oct 2024 23:57:03 +0530 Subject: [PATCH] added example page to the pydantic_plugin Signed-off-by: sumana sree --- examples/pydantic_plugin/README.md | 2 +- .../pydantic_integration_example.py | 45 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) create mode 100644 examples/pydantic_plugin/pydantic_plugin/pydantic_integration_example.py diff --git a/examples/pydantic_plugin/README.md b/examples/pydantic_plugin/README.md index c92ac09be..45bdd8119 100644 --- a/examples/pydantic_plugin/README.md +++ b/examples/pydantic_plugin/README.md @@ -18,7 +18,7 @@ pip install flytekitplugins-pydantic ``` ## Example usage -For a usage example, see {doc}`Pydantic example usage `. +For a usage example, see {doc}`Pydantic example usage `. ```{toctree} :maxdepth: -1 diff --git a/examples/pydantic_plugin/pydantic_plugin/pydantic_integration_example.py b/examples/pydantic_plugin/pydantic_plugin/pydantic_integration_example.py new file mode 100644 index 000000000..0c52f58c8 --- /dev/null +++ b/examples/pydantic_plugin/pydantic_plugin/pydantic_integration_example.py @@ -0,0 +1,45 @@ +# %% [markdown] +# (pydantic_integration_example)= +# +# # Pydantic Integration Example +# +# Pydantic is a data validation and settings management library for Python, enabling the creation of data models with type annotations. +# +# Flyte leverages Pydantic for robust input validation and serialization, ensuring that task inputs are correctly structured. + +# %% +from pydantic.v1 import BaseModel +from flytekit import task, workflow +from flytekit.types.file import FlyteFile +from typing import List + +# %% [markdown] +# Let's first define a Pydantic model for training configuration. +# %% +class TrainConfig(BaseModel): + lr: float = 1e-3 # Learning rate + batch_size: int = 32 # Batch size for training + files: List[FlyteFile] # List of file inputs for training + +# %% [markdown] +# Next, we use the Pydantic model in a Flyte task to train a model. +# %% +@task +def train(cfg: TrainConfig): + print(f"Training with learning rate: {cfg.lr} and batch size: {cfg.batch_size}") + for file in cfg.files: + print(f"Processing file: {file}") + +# %% [markdown] +# Now we define a Flyte workflow that utilizes the training task. +# %% +@workflow +def training_workflow(lr: float = 1e-3, batch_size: int = 32, files: List[FlyteFile] = []): + cfg = TrainConfig(lr=lr, batch_size=batch_size, files=files) + train(cfg=cfg) + +# %% [markdown] +# Finally, we execute the workflow with sample parameters. +# %% +if __name__ == "__main__": + training_workflow(lr=1e-3, batch_size=32, files=[FlyteFile(path="path/to/your/file")])