diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml index dab23b073..b9382e576 100644 --- a/.github/workflows/checks.yml +++ b/.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 diff --git a/examples/k8s_pod_plugin/k8s_pod_plugin/pod.py b/examples/k8s_pod_plugin/k8s_pod_plugin/pod.py index 46eb2ed51..3e7fb2003 100644 --- a/examples/k8s_pod_plugin/k8s_pod_plugin/pod.py +++ b/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)