From 73cfdfc2c4a97987ab873fbce592cebc3fc96f11 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <katrina@nuclyde.io>
Date: Tue, 6 Jul 2021 11:55:15 -0700
Subject: [PATCH] Map pod task example (#321)

---
 .gitignore                                    |  3 +-
 Makefile                                      |  2 +-
 cookbook/integrations/kubernetes/pod/pod.py   | 64 ++++++++++++++++++-
 .../kubernetes/pod/requirements.in            |  2 +-
 .../kubernetes/pod/requirements.txt           | 10 +--
 5 files changed, 70 insertions(+), 11 deletions(-)

diff --git a/.gitignore b/.gitignore
index 55bf763acc..8877fafbac 100644
--- a/.gitignore
+++ b/.gitignore
@@ -23,4 +23,5 @@ _build/
 *-out.html
 *-out.ipynb
 .python-version
-cookbook/release-snacks
\ No newline at end of file
+cookbook/release-snacks
+.kube/
diff --git a/Makefile b/Makefile
index dd41a3b59d..cf69331d63 100644
--- 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
index 26f85be6d1..c6786a4eec 100644
--- 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,9 +48,15 @@ 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"}
@@ -57,7 +64,10 @@ def generate_pod_spec_for_task():
     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
index e1ed198888..16968efa94 100644
--- 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
index 21d234ae1a..60dfca763c 100644
--- 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: