Skip to content

Carry pod_template override through map_task serialization#3428

Open
1fanwang wants to merge 1 commit intoflyteorg:masterfrom
1fanwang:fix-7076-map-task-pod-template
Open

Carry pod_template override through map_task serialization#3428
1fanwang wants to merge 1 commit intoflyteorg:masterfrom
1fanwang:fix-7076-map-task-pod-template

Conversation

@1fanwang
Copy link
Copy Markdown

Tracking issue

Fixes flyteorg/flyte#7076

Why are the changes needed?

map_task(...).with_overrides(pod_template=PodTemplate(...)) silently drops the override at serialization. The override is accepted at workflow build time and stored on the Node (node._pod_template), but get_serializable_array_node_map_task in flytekit/tools/translator.py never reads that field — it builds TaskNodeOverrides with only resources, extended_resources, and container_image. The user thinks the override is applied, but the registered task spec has an empty overrides {} block.

This is the same class of bug as #6463, which #3270 fixed for the regular-task path of get_serializable_node. The array-node path has the same gap in a different function.

What changes were proposed in this pull request?

flytekit/tools/translator.pyget_serializable_array_node_map_task:

  • Read node._pod_template, build the override pod spec via _serialize_pod_spec, and include pod_template in the emitted TaskNodeOverrides.
  • Use entity.get_container(settings) (not _get_container) so prepare_target() substitutes the map-task command (pyflyte-map-execute) into the container args before the pod spec is built. ArrayNodeMapTask extends PythonTask directly (not PythonAutoContainerTask), so _get_container isn't on it anyway.
  • For fast-serialization, prefix container args post-build (matching the existing pattern in get_serializable_task at translator.py:170) instead of swapping command_fn via _fast_serialize_command_fn. The latter wraps task.get_default_command(settings), which on ArrayNodeMapTask returns the inherited pyflyte-execute — wrong shape for a map task.

tests/flytekit/unit/test_translator.py:

How was this patch tested?

  • New parametrized test passes (both cases).
  • Full tests/flytekit/unit/test_translator.py (11 tests) passes.
  • Adjacent suites that exercise the same surfaces all pass with no regressions:
    • tests/flytekit/unit/core/test_array_node_map_task.py (46 tests)
    • tests/flytekit/unit/core/test_map_task.py (23 tests)
    • tests/flytekit/unit/core/test_node_creation.py (26 tests)

Check all the applicable boxes

  • All new and existing tests passed.
  • All commits are signed-off.
  • I updated the documentation accordingly. (No user-facing behavior change beyond the bug fix; the documented behavior of with_overrides(pod_template=...) for map_task is what this PR makes work.)

map_task(...).with_overrides(pod_template=PodTemplate(...)) was being
silently dropped at serialization. The override is accepted at workflow
build time and stored on the Node, but get_serializable_array_node_map_task
never read node._pod_template — TaskNodeOverrides was emitted with only
resources/extended_resources/container_image, so the registered task spec
had an empty overrides {} block.

Same class of bug as #6463 / PR flyteorg#3270 fixed for regular tasks, but in
the array-node serialization path.

The map-task path needs two implementation tweaks vs. the regular-task
mirror:
  * Use entity.get_container() (not _get_container) so prepare_target()
    substitutes the map-task command (pyflyte-map-execute) into the
    container args before the override pod spec is built.
  * For fast-serialization, prefix container args post-build (matching
    get_serializable_task) instead of swapping command_fn via
    _fast_serialize_command_fn — the latter wraps the inherited
    pyflyte-execute default, which is wrong for map_task.

Adds a parametrized regression test covering both fast-registration
enabled and disabled.

Fixes flyteorg/flyte#7076

Signed-off-by: 1fanwang <1fannnw@gmail.com>
@1fanwang 1fanwang force-pushed the fix-7076-map-task-pod-template branch from d9e2fd3 to d2edd5d Compare April 27, 2026 09:02
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] with_overrides(pod_template=...) silently ignored for map_task

1 participant