Map pod task example (flyteorg#321)

eapolinario · Jul 6, 2021 · 73cfdfc · 73cfdfc
1 parent c26db87
commit 73cfdfc
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 11 deletions.
diff --git a/.gitignore b/.gitignore
@@ -23,4 +23,5 @@ _build/
 *-out.html
 *-out.ipynb
 .python-version
-cookbook/release-snacks
+cookbook/release-snacks
+.kube/
diff --git a/Makefile b/Makefile
@@ -24,7 +24,7 @@ docker exec -it \
 	-e MAKEFLAGS \
 	-e REGISTRY \
 	-e VERSION \
-	-w /usr/src \
+        -w /root \
 	$(FLYTE_SANDBOX_NAME) \
 	$(1)
 endef

diff --git a/cookbook/integrations/kubernetes/pod/pod.py b/cookbook/integrations/kubernetes/pod/pod.py
@@ -26,6 +26,7 @@
 # containers. The secondary writes a file that the primary waits on before completing.
 import os
 import time
+from typing import List
 
 from flytekit import task, workflow
 from flytekitplugins.pod import Pod
@@ -47,17 +48,26 @@ def generate_pod_spec_for_task():
     primary_container = V1Container(name="primary")
 
     # Note: for non-primary containers we must specify an image.
-    secondary_container = V1Container(name="secondary", image="alpine",)
+    secondary_container = V1Container(
+        name="secondary",
+        image="alpine",
+    )
     secondary_container.command = ["/bin/sh"]
-    secondary_container.args = ["-c", "echo hi pod world > {}".format(_SHARED_DATA_PATH)]
+    secondary_container.args = [
+        "-c",
+        "echo hi pod world > {}".format(_SHARED_DATA_PATH),
+    ]
 
     resources = V1ResourceRequirements(
         requests={"cpu": "1", "memory": "100Mi"}, limits={"cpu": "1", "memory": "100Mi"}
     )
     primary_container.resources = resources
     secondary_container.resources = resources
 
-    shared_volume_mount = V1VolumeMount(name="shared-data", mount_path="/data",)
+    shared_volume_mount = V1VolumeMount(
+        name="shared-data",
+        mount_path="/data",
+    )
     secondary_container.volume_mounts = [shared_volume_mount]
     primary_container.volume_mounts = [shared_volume_mount]
 
@@ -96,5 +106,53 @@ def PodWorkflow() -> str:
     return s
 
 
+# %%
+# To use a pod task as part of a map task in your workflow, use the :py:func:`flytekit:flytekit.core.map_task` function
+# and pass in the pod task definition. This will run your pod task across a collection of inputs.
+
+from flytekit import map_task, TaskMetadata
+
+
+@task(
+    task_config=Pod(
+        pod_spec=V1PodSpec(
+            containers=[
+                V1Container(
+                    name="primary",
+                    resources=V1ResourceRequirements(
+                        requests={"cpu": ".5", "memory": "500Mi"},
+                        limits={"cpu": ".5", "memory": "500Mi"},
+                    ),
+                )
+            ],
+            init_containers=[
+                V1Container(
+                    name="init",
+                    command=["/bin/sh"],
+                    args=["-c", 'echo "I\'m a customizable init container"'],
+                )
+            ],
+        ),
+        primary_container_name="primary",
+    )
+)
+def my_pod_task(stringify: int) -> str:
+    return str(stringify)
+
+
+@task
+def coalesce(b: List[str]) -> str:
+    coalesced = ", ".join(b)
+    return coalesced
+
+
+@workflow
+def my_map_workflow(a: List[int]) -> str:
+    mapped_out = map_task(my_pod_task, metadata=TaskMetadata(retries=1))(stringify=a)
+    coalesced = coalesce(b=mapped_out)
+    return coalesced
+
+
 if __name__ == "__main__":
     pass
+
diff --git a/cookbook/integrations/kubernetes/pod/requirements.in b/cookbook/integrations/kubernetes/pod/requirements.in
@@ -1,2 +1,2 @@
 -r ../../../common/requirements-common.in
-flytekitplugins-pod>=0.16.0
+flytekitplugins-pod>=0.20.0
diff --git a/cookbook/integrations/kubernetes/pod/requirements.txt b/cookbook/integrations/kubernetes/pod/requirements.txt
@@ -30,7 +30,7 @@ dirhash==0.2.1
     # via flytekit
 docker-image-py==0.1.10
     # via flytekit
-flyteidl==0.19.9
+flyteidl==0.19.11
     # via flytekit
 flytekit==0.20.0
     # via
@@ -44,7 +44,7 @@ grpcio==1.38.1
     # via flytekit
 idna==2.10
     # via requests
-importlib-metadata==4.6.0
+importlib-metadata==4.6.1
     # via keyring
 keyring==23.0.1
     # via flytekit
@@ -74,7 +74,7 @@ numpy==1.21.0
     #   pyarrow
 oauthlib==3.1.1
     # via requests-oauthlib
-pandas==1.2.5
+pandas==1.3.0
     # via flytekit
 pathspec==0.8.1
     # via scantree
@@ -113,7 +113,7 @@ pytz==2018.4
     #   pandas
 pyyaml==5.4.1
     # via kubernetes
-regex==2021.7.1
+regex==2021.7.6
     # via docker-image-py
 requests==2.25.1
     # via
@@ -168,7 +168,7 @@ wrapt==1.12.1
     # via
     #   deprecated
     #   flytekit
-zipp==3.4.1
+zipp==3.5.0
     # via importlib-metadata
 
 # The following packages are considered to be unsafe in a requirements file: