Support overriding task pod_template via with_overrides (flyteorg#2981)

Signed-off-by: Nelson Chen <[email protected]>
Signed-off-by: Kevin Su <[email protected]>
Co-authored-by: Kevin Su <[email protected]>
Signed-off-by: Atharva <[email protected]>

Author: 2 people

flytekit/core/node.py

@@ -6,6 +6,7 @@
 
 from flyteidl.core import tasks_pb2
 
+from flytekit.core.pod_template import PodTemplate
 from flytekit.core.resources import Resources, convert_resources_to_resource_model
 from flytekit.core.utils import _dnsify
 from flytekit.extras.accelerators import BaseAccelerator
@@ -67,6 +68,7 @@ def __init__(
         self._resources: typing.Optional[_resources_model] = None
         self._extended_resources: typing.Optional[tasks_pb2.ExtendedResources] = None
         self._container_image: typing.Optional[str] = None
+        self._pod_template: typing.Optional[PodTemplate] = None
 
     def runs_before(self, other: Node):
         """
@@ -191,6 +193,7 @@ def with_overrides(
         cache: Optional[bool] = None,
         cache_version: Optional[str] = None,
         cache_serialize: Optional[bool] = None,
+        pod_template: Optional[PodTemplate] = None,
         *args,
         **kwargs,
     ):
@@ -241,6 +244,10 @@ def with_overrides(
 
         self._override_node_metadata(name, timeout, retries, interruptible, cache, cache_version, cache_serialize)
 
+        if pod_template is not None:
+            assert_not_promise(pod_template, "podtemplate")
+            self._pod_template = pod_template
+
         return self
 
 

flytekit/models/core/workflow.py

@@ -1,18 +1,21 @@
 import datetime
+import json
 import typing
 
 from flyteidl.core import tasks_pb2
 from flyteidl.core import workflow_pb2 as _core_workflow
+from google.protobuf import json_format, struct_pb2
 from google.protobuf.wrappers_pb2 import BoolValue
 
+from flytekit.core.pod_template import PodTemplate
 from flytekit.models import common as _common
 from flytekit.models import interface as _interface
 from flytekit.models import types as type_models
 from flytekit.models.core import condition as _condition
 from flytekit.models.core import identifier as _identifier
 from flytekit.models.literals import Binding as _Binding
 from flytekit.models.literals import RetryStrategy as _RetryStrategy
-from flytekit.models.task import Resources
+from flytekit.models.task import K8sObjectMetadata, Resources
 
 
 class IfBlock(_common.FlyteIdlEntity):
@@ -615,10 +618,12 @@ def __init__(
         resources: typing.Optional[Resources],
         extended_resources: typing.Optional[tasks_pb2.ExtendedResources],
         container_image: typing.Optional[str] = None,
+        pod_template: typing.Optional[PodTemplate] = None,
     ):
         self._resources = resources
         self._extended_resources = extended_resources
         self._container_image = container_image
+        self._pod_template = pod_template
 
     @property
     def resources(self) -> Resources:
@@ -632,11 +637,27 @@ def extended_resources(self) -> tasks_pb2.ExtendedResources:
     def container_image(self) -> typing.Optional[str]:
         return self._container_image
 
+    @property
+    def pod_template(self) -> typing.Optional[PodTemplate]:
+        return self._pod_template
+
     def to_flyte_idl(self):
         return _core_workflow.TaskNodeOverrides(
             resources=self.resources.to_flyte_idl() if self.resources is not None else None,
             extended_resources=self.extended_resources,
             container_image=self.container_image,
+            pod_template=tasks_pb2.K8sPod(
+                metadata=K8sObjectMetadata(
+                    labels=self.pod_template.labels if self.pod_template else None,
+                    annotations=self.pod_template.annotations if self.pod_template else None,
+                ).to_flyte_idl()
+                if self.pod_template is not None
+                else None,
+                pod_spec=json_format.Parse(json.dumps(self.pod_template.pod_spec), struct_pb2.Struct())
+                if self.pod_template
+                else None,
+                primary_container_name=self.pod_template.primary_container_name if self.pod_template else None,
+            ),
         )
 
     @classmethod

flytekit/models/task.py

@@ -1009,13 +1009,15 @@ def __init__(
         metadata: K8sObjectMetadata = None,
         pod_spec: typing.Dict[str, typing.Any] = None,
         data_config: typing.Optional[DataLoadingConfig] = None,
+        primary_container_name: typing.Optional[str] = None,
     ):
         """
         This defines a kubernetes pod target.  It will build the pod target during task execution
         """
         self._metadata = metadata
         self._pod_spec = pod_spec
         self._data_config = data_config
+        self._primary_container_name = primary_container_name
 
     @property
     def metadata(self) -> K8sObjectMetadata:
@@ -1029,6 +1031,10 @@ def pod_spec(self) -> typing.Dict[str, typing.Any]:
     def data_config(self) -> typing.Optional[DataLoadingConfig]:
         return self._data_config
 
+    @property
+    def primary_container_name(self) -> typing.Optional[str]:
+        return self._primary_container_name
+
     def to_flyte_idl(self) -> _core_task.K8sPod:
         return _core_task.K8sPod(
             metadata=self._metadata.to_flyte_idl() if self.metadata else None,

flytekit/tools/translator.py

@@ -5,7 +5,7 @@
 
 from flyteidl.admin import schedule_pb2
 
-from flytekit import ImageSpec, PythonFunctionTask, SourceCode
+from flytekit import ImageSpec, PodTemplate, PythonFunctionTask, SourceCode
 from flytekit.configuration import Image, ImageConfig, SerializationSettings
 from flytekit.core import constants as _common_constants
 from flytekit.core import context_manager
@@ -25,7 +25,7 @@
 from flytekit.core.python_function_task import EagerAsyncPythonFunctionTask
 from flytekit.core.reference_entity import ReferenceEntity, ReferenceSpec, ReferenceTemplate
 from flytekit.core.task import ReferenceTask
-from flytekit.core.utils import ClassDecorator, _dnsify
+from flytekit.core.utils import ClassDecorator, _dnsify, _serialize_pod_spec
 from flytekit.core.workflow import ReferenceWorkflow, WorkflowBase
 from flytekit.models import common as _common_models
 from flytekit.models import interface as interface_models
@@ -453,6 +453,13 @@ def get_serializable_node(
         # if entity._aliases:
         #     node_model._output_aliases = entity._aliases
     elif isinstance(entity.flyte_entity, PythonTask):
+        # handle pod template overrides
+        override_pod_spec = {}
+        if entity._pod_template is not None and settings.should_fast_serialize():
+            entity.flyte_entity.set_command_fn(_fast_serialize_command_fn(settings, entity.flyte_entity))
+            override_pod_spec = _serialize_pod_spec(
+                entity._pod_template, entity.flyte_entity._get_container(settings), settings
+            )
         task_spec = get_serializable(entity_mapping, settings, entity.flyte_entity, options=options)
         node_model = workflow_model.Node(
             id=_dnsify(entity.id),
@@ -466,6 +473,16 @@ def get_serializable_node(
                     resources=entity._resources,
                     extended_resources=entity._extended_resources,
                     container_image=entity._container_image,
+                    pod_template=PodTemplate(
+                        pod_spec=override_pod_spec,
+                        labels=entity._pod_template.labels if entity._pod_template.labels else None,
+                        annotations=entity._pod_template.annotations if entity._pod_template.annotations else None,
+                        primary_container_name=entity._pod_template.primary_container_name
+                        if entity._pod_template.primary_container_name
+                        else None,
+                    )
+                    if entity._pod_template
+                    else None,
                 ),
             ),
         )

tests/flytekit/unit/core/test_array_node_map_task.py

@@ -9,7 +9,7 @@
 import pytest
 from flyteidl.core import workflow_pb2 as _core_workflow
 
-from flytekit import dynamic, map_task, task, workflow, eager, PythonFunctionTask, Resources
+from flytekit import dynamic, map_task, task, workflow, eager, PythonFunctionTask
 from flytekit.configuration import FastSerializationSettings, Image, ImageConfig, SerializationSettings
 from flytekit.core import context_manager
 from flytekit.core.array_node_map_task import ArrayNodeMapTask, ArrayNodeMapTaskResolver
@@ -21,7 +21,6 @@
     LiteralMap,
     LiteralOffloadedMetadata,
 )
-from flytekit.models.task import Resources as _resources_models
 from flytekit.tools.translator import get_serializable
 from flytekit.types.directory import FlyteDirectory
 
@@ -350,59 +349,16 @@ def my_wf1() -> typing.List[typing.Optional[int]]:
         assert my_wf1() == [1, None, 3, 4]
 
 
-@task
-def my_mappable_task(a: int) -> typing.Optional[str]:
-    return str(a)
-
-
-@task(
-    container_image="original-image",
-    timeout=timedelta(seconds=10),
-    interruptible=False,
-    retries=10,
-    cache=True,
-    cache_version="original-version",
-    requests=Resources(cpu=1)
-)
-def my_mappable_task_1(a: int) -> typing.Optional[str]:
-    return str(a)
-
-
-@pytest.mark.parametrize(
-    "task_func",
-    [my_mappable_task, my_mappable_task_1]
-)
-def test_map_task_override(serialization_settings, task_func):
-    array_node_map_task = map_task(task_func)
+def test_map_task_override(serialization_settings):
+    @task
+    def my_mappable_task(a: int) -> typing.Optional[str]:
+        return str(a)
 
     @workflow
     def wf(x: typing.List[int]):
-        array_node_map_task(a=x).with_overrides(
-            container_image="new-image",
-            timeout=timedelta(seconds=20),
-            interruptible=True,
-            retries=5,
-            cache=True,
-            cache_version="new-version",
-            requests=Resources(cpu=2)
-        )
-
-    assert wf.nodes[0]._container_image == "new-image"
-
-    od = OrderedDict()
-    wf_spec = get_serializable(od, serialization_settings, wf)
+        map_task(my_mappable_task)(a=x).with_overrides(container_image="random:image")
 
-    array_node = wf_spec.template.nodes[0]
-    assert array_node.metadata.timeout == timedelta()
-    sub_node_spec = array_node.array_node.node
-    assert sub_node_spec.metadata.timeout == timedelta(seconds=20)
-    assert sub_node_spec.metadata.interruptible
-    assert sub_node_spec.metadata.retries.retries == 5
-    assert sub_node_spec.metadata.cacheable
-    assert sub_node_spec.metadata.cache_version == "new-version"
-    assert sub_node_spec.target.overrides.resources.requests == [
-        _resources_models.ResourceEntry(_resources_models.ResourceName.CPU, "2")
-    ]
+    assert wf.nodes[0]._container_image == "random:image"
 
 
 def test_serialization_metadata(serialization_settings):

tests/flytekit/unit/core/test_map_task.py

@@ -3,9 +3,10 @@
 from collections import OrderedDict
 
 import pytest
+from kubernetes.client import V1PodSpec, V1Container, V1EnvVar
 
 import flytekit.configuration
-from flytekit import LaunchPlan, Resources
+from flytekit import LaunchPlan, Resources, PodTemplate
 from flytekit.configuration import Image, ImageConfig
 from flytekit.core.legacy_map_task import MapPythonTask, MapTaskResolver, map_task
 from flytekit.core.task import TaskMetadata, task
@@ -354,6 +355,39 @@ def wf(x: typing.List[int]):
 
     assert wf.nodes[0]._container_image == "random:image"
 
+def test_map_task_pod_template_override(serialization_settings):
+    @task
+    def my_mappable_task(a: int) -> typing.Optional[str]:
+        return str(a)
+
+    @workflow
+    def wf(x: typing.List[int]):
+        map_task(my_mappable_task)(a=x).with_overrides(pod_template=PodTemplate(
+        primary_container_name="primary1",
+        labels={"lKeyA": "lValA", "lKeyB": "lValB"},
+        annotations={"aKeyA": "aValA", "aKeyB": "aValB"},
+        pod_spec=V1PodSpec(
+            containers=[
+                V1Container(
+                    name="primary1",
+                    image="random:image",
+                    env=[V1EnvVar(name="eKeyC", value="eValC"), V1EnvVar(name="eKeyD", value="eValD")],
+                ),
+                V1Container(
+                    name="primary2",
+                    image="random:image2",
+                    env=[V1EnvVar(name="eKeyC", value="eValC"), V1EnvVar(name="eKeyD", value="eValD")],
+                ),
+            ],
+        )
+    ))
+
+
+    assert wf.nodes[0]._pod_template.primary_container_name == "primary1"
+    assert wf.nodes[0]._pod_template.pod_spec.containers[0].image == "random:image"
+    assert wf.nodes[0]._pod_template.labels == {"lKeyA": "lValA", "lKeyB": "lValB"}
+    assert wf.nodes[0]._pod_template.annotations["aKeyA"] == "aValA"
+
 
 def test_bounded_inputs_vars_order(serialization_settings):
     mt = map_task(functools.partial(t3, c=1.0, b="hello", a=1))

tests/flytekit/unit/core/test_node_creation.py

@@ -4,9 +4,10 @@
 from dataclasses import dataclass
 
 import pytest
+from kubernetes.client import V1PodSpec, V1Container, V1EnvVar
 
 import flytekit.configuration
-from flytekit import Resources, map_task
+from flytekit import Resources, map_task, PodTemplate
 from flytekit.configuration import Image, ImageConfig
 from flytekit.core.dynamic_workflow_task import dynamic
 from flytekit.core.node_creation import create_node
@@ -470,6 +471,39 @@ def wf() -> str:
 
     assert wf.nodes[0]._container_image == "hello/world"
 
+def test_pod_template_override():
+    @task
+    def bar():
+        print("hello")
+
+    @workflow
+    def wf() -> str:
+        bar().with_overrides(pod_template=PodTemplate(
+        primary_container_name="primary1",
+        labels={"lKeyA": "lValA", "lKeyB": "lValB"},
+        annotations={"aKeyA": "aValA", "aKeyB": "aValB"},
+        pod_spec=V1PodSpec(
+            containers=[
+                V1Container(
+                    name="primary1",
+                    image="random:image",
+                    env=[V1EnvVar(name="eKeyC", value="eValC"), V1EnvVar(name="eKeyD", value="eValD")],
+                ),
+                V1Container(
+                    name="primary2",
+                    image="random:image2",
+                    env=[V1EnvVar(name="eKeyC", value="eValC"), V1EnvVar(name="eKeyD", value="eValD")],
+                ),
+            ],
+        )
+        ))
+        return "hi"
+
+    assert wf.nodes[0]._pod_template.primary_container_name == "primary1"
+    assert wf.nodes[0]._pod_template.pod_spec.containers[0].image == "random:image"
+    assert wf.nodes[0]._pod_template.labels == {"lKeyA": "lValA", "lKeyB": "lValB"}
+    assert wf.nodes[0]._pod_template.annotations["aKeyA"] == "aValA"
+
 
 def test_override_accelerator():
     @task(accelerator=T4)