Fix issue where pod template override pod spec was missing (#3270)

* Fix issue where pod template override pod spec was conditionally included

Signed-off-by: Jason Parraga <sovietaced@gmail.com>

* Add unit test

Signed-off-by: Jason Parraga <sovietaced@gmail.com>

* Make unit test parameterized

Signed-off-by: Jason Parraga <sovietaced@gmail.com>

---------

Signed-off-by: Jason Parraga <sovietaced@gmail.com>

Author: Sovietaced

flytekit/tools/translator.py

@@ -461,8 +461,9 @@ def get_serializable_node(
     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))
+        if entity._pod_template is not None:
+            if 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
             )

tests/flytekit/unit/test_translator.py

@@ -2,7 +2,7 @@
 from collections import OrderedDict
 
 import flytekit.configuration
-from flytekit import ContainerTask, Resources
+from flytekit import ContainerTask, Resources, PodTemplate
 from flytekit.configuration import FastSerializationSettings, Image, ImageConfig
 from flytekit.core.base_task import kwtypes
 from flytekit.core.launch_plan import LaunchPlan, ReferenceLaunchPlan
@@ -12,7 +12,9 @@
 from flytekit.deck import Deck
 from flytekit.models.core import identifier as identifier_models
 from flytekit.models.task import Resources as resource_model
-from flytekit.tools.translator import get_serializable, Options
+from flytekit.tools.translator import get_serializable
+from kubernetes import client
+from kubernetes.client import V1PodSpec, V1Container
 import pytest
 
 default_img = Image(name="default", fqn="test", tag="tag")
@@ -183,3 +185,52 @@ def morning_greeter_caller(day_of_week: str) -> str:
     assert len(task_spec.template.interface.outputs) == 1
     assert len(task_spec.template.nodes) == 1
     assert len(task_spec.template.nodes[0].inputs) == 2
+
+@pytest.mark.parametrize(
+    "fast_registration_enabled",
+    [
+        pytest.param(
+            True, id="fast registration enabled"
+        ),
+        pytest.param(
+            False, id="fast registration disabled"
+        ),
+    ],
+)
+def test_task_with_pod_template_override(fast_registration_enabled: bool):
+
+    custom_pod_template = PodTemplate(pod_spec=V1PodSpec(
+        containers=[
+            V1Container(
+                name="primary",
+                env=[
+                    client.V1EnvVar(name="MY_KEY", value="MY_VALUE"),
+                ]
+            )
+        ]
+    ))
+
+    @task
+    def t(a: str) -> str:
+        return a
+
+    @workflow
+    def wf():
+        t("Hello World").with_overrides(pod_template=custom_pod_template)
+
+    settings = (
+        serialization_settings.new_builder()
+        .with_fast_serialization_settings(FastSerializationSettings(enabled=fast_registration_enabled))
+        .build()
+    )
+
+    task_spec = get_serializable(OrderedDict(), settings, wf)
+    assert len(task_spec.template.nodes) == 1
+    node = task_spec.template.nodes[0]
+    assert node.metadata.name == "t"
+    assert node.task_node.overrides.pod_template is not None
+    pod_template_override = node.task_node.overrides.pod_template
+    assert pod_template_override.pod_spec # validate not empty
+    assert len(pod_template_override.pod_spec['containers']) == 1
+    container = pod_template_override.pod_spec['containers'][0]
+    assert container['env'] == [{'name': 'MY_KEY', 'value': 'MY_VALUE'}]