feat(k8s_pod_plugin): Add ImageSpec to pod tasks (#1698)

Signed-off-by: Kevin Su <[email protected]>

.github/workflows/checks.yml

@@ -90,7 +90,7 @@ jobs:
           uv venv
           source .venv/bin/activate
           if [ -f requirements.in ]; then uv pip install -r requirements.in; fi
-          uv pip install "flytekit>=1.12.2" flytekitplugins-envd
+          uv pip install "flytekit>=1.12.2" flytekitplugins-envd "numpy<2.0.0"
           pip freeze
       - name: Set up Docker Buildx
         uses: docker/setup-buildx-action@v3
@@ -292,7 +292,7 @@ jobs:
           pip install uv
           uv venv
           source .venv/bin/activate
-          uv pip install "flytekit>=1.12.2" flytekitplugins-deck-standard torch tabulate
+          uv pip install "flytekit>=1.12.2" flytekitplugins-deck-standard flytekitplugins-envd torch tabulate
           pip freeze
       - name: Checkout flytesnacks
         uses: actions/checkout@v3

examples/k8s_pod_plugin/k8s_pod_plugin/pod.py

@@ -19,7 +19,7 @@
 from pathlib import Path
 from typing import List
 
-from flytekit import Resources, TaskMetadata, dynamic, map_task, task, workflow
+from flytekit import ImageSpec, Resources, TaskMetadata, dynamic, map_task, task, workflow
 from flytekitplugins.pod import Pod
 from kubernetes.client.models import (
     V1Container,
@@ -30,6 +30,8 @@
     V1VolumeMount,
 )
 
+image_spec = ImageSpec(registry="ghcr.io/flyteorg", packages=["flytekitplugins-pod"])
+
 
 # %% [markdown]
 # ## Add additional properties to the task container
@@ -48,6 +50,7 @@
     requests=Resources(
         mem="1G",
     ),
+    container_image=image_spec,
 )
 def pod_task() -> str:
     return "Hello from pod task!"
@@ -128,6 +131,7 @@ def pod_workflow() -> str:
     requests=Resources(
         mem="1G",
     ),
+    container_image=image_spec,
 )
 def multiple_containers_pod_task() -> str:
     # The code defined in this task will get injected into the primary container.
@@ -174,13 +178,14 @@ def multiple_containers_pod_workflow() -> str:
                 )
             ],
         ),
-    )
+    ),
+    container_image=image_spec,
 )
 def map_pod_task(int_val: int) -> str:
     return str(int_val)
 
 
-@task
+@task(container_image=image_spec)
 def coalesce(list_of_strings: List[str]) -> str:
     coalesced = ", ".join(list_of_strings)
     return coalesced
@@ -198,7 +203,7 @@ def map_pod_workflow(list_of_ints: List[int]) -> str:
 #
 # To use a pod task in a dynamic workflow, simply pass the pod task config to the annotated dynamic workflow.
 # %%
-@task
+@task(container_image=image_spec)
 def stringify(val: int) -> str:
     return f"{val} served courtesy of a dynamic pod task!"
 
@@ -216,7 +221,8 @@ def stringify(val: int) -> str:
                 )
             ],
         ),
-    )
+    ),
+    container_image=image_spec,
 )
 def dynamic_pod_task(val: int) -> str:
     return stringify(val=val)