Carry pod_template override through map_task serialization#3428
Open
1fanwang wants to merge 1 commit intoflyteorg:masterfrom
Open
Carry pod_template override through map_task serialization#34281fanwang wants to merge 1 commit intoflyteorg:masterfrom
1fanwang wants to merge 1 commit intoflyteorg:masterfrom
Conversation
88a1cd2 to
d9e2fd3
Compare
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>
d9e2fd3 to
d2edd5d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 theNode(node._pod_template), butget_serializable_array_node_map_taskinflytekit/tools/translator.pynever reads that field — it buildsTaskNodeOverrideswith onlyresources,extended_resources, andcontainer_image. The user thinks the override is applied, but the registered task spec has an emptyoverrides {}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.py—get_serializable_array_node_map_task:node._pod_template, build the override pod spec via_serialize_pod_spec, and includepod_templatein the emittedTaskNodeOverrides.entity.get_container(settings)(not_get_container) soprepare_target()substitutes the map-task command (pyflyte-map-execute) into the container args before the pod spec is built.ArrayNodeMapTaskextendsPythonTaskdirectly (notPythonAutoContainerTask), so_get_containerisn't on it anyway.get_serializable_taskattranslator.py:170) instead of swappingcommand_fnvia_fast_serialize_command_fn. The latter wrapstask.get_default_command(settings), which onArrayNodeMapTaskreturns the inheritedpyflyte-execute— wrong shape for a map task.tests/flytekit/unit/test_translator.py:test_map_task_with_pod_template_overridecovering bothfast_registration_enabled=TrueandFalse. Mirrors the structure of the existingtest_task_with_pod_template_overridefrom Fix issue where pod template override pod spec was missing #3270.How was this patch tested?
tests/flytekit/unit/test_translator.py(11 tests) passes.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
with_overrides(pod_template=...)formap_taskis what this PR makes work.)