Add an example of how to use BatchSize (#980)

* Update flyte_pickle.py Signed-off-by: Yicheng-Lu-llll <[email protected]> * Update flyte_pickle.py Signed-off-by: Yicheng-Lu-llll <[email protected]> * update flytekit version Signed-off-by: Yicheng-Lu-llll <[email protected]> * Update requirements.txt Signed-off-by: Yicheng-Lu-llll <[email protected]> * nit Signed-off-by: Yicheng-Lu-llll <[email protected]> --------- Signed-off-by: Yicheng-Lu-llll <[email protected]> Co-authored-by: Kevin Su <[email protected]>
flyteorg · May 19, 2023 · 16eae65 · 16eae65
1 parent 99c9427
commit 16eae65
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 2 deletions.
diff --git a/cookbook/core/requirements.txt b/cookbook/core/requirements.txt
@@ -74,9 +74,9 @@ docstring-parser==0.15
     # via flytekit
 executing==1.2.0
     # via stack-data
-flyteidl==1.3.11
+flyteidl==1.3.18
     # via flytekit
-flytekit==1.4.1
+flytekit==1.5.0
     # via
     #   -r ../common/requirements-common.in
     #   flytekitplugins-deck-standard

diff --git a/cookbook/core/type_system/flyte_pickle.py b/cookbook/core/type_system/flyte_pickle.py
@@ -57,3 +57,31 @@ def welcome(name: str) -> People:
     the custom object (People) will be marshalled to and from python pickle.
     """
     welcome(name="Foo")
+
+
+# %%
+# By default, if the list subtype is unrecognized, a single pickle file is generated.
+# To also improve serialization and deserialization performance for cases with millions of items or large list items,
+# users can specify a batch size, processing each batch as a separate pickle file.
+# Example below shows how users can set batch size.
+from flytekit.types.pickle.pickle import BatchSize
+from typing import List
+from typing_extensions import Annotated
+
+@task
+def greet_all(names: List[str]) -> Annotated[List[People],BatchSize(2)]:
+    return [People(name) for name in names]
+
+
+@workflow
+def welcome_all(names: List[str]) -> Annotated[List[People],BatchSize(2)]:
+    return greet_all(names=names)
+
+
+if __name__ == "__main__":
+    """
+    In this example, two pickle files will be generated:
+    - One containing two People objects
+    - One containing one People object
+    """
+    welcome_all(names=["f","o","o"])