Support with_overrides setting metadata for map_task subnode instead of parent node

[pvditt]

Tracking issue

fixes: https://linear.app/unionai/issue/COR-2498/with-overrides-sets-metadata-for-parent-instead-of-subnode-for-map

Why are the changes needed?

with_overrides doesn't work for map_tasks

What changes were proposed in this pull request?

create a new field in array node map task to explicitl

How was this patch tested?

• added unit tests
• tested workflows locally that overrode timeouts, cache, container image, etc

Setup process

from flytekit import task, workflow, map_task, reference_launch_plan, reference_task, ImageSpec, LaunchPlan, Resources, TaskMetadata
from datetime import timedelta


@task(
    cache=True,
    cache_version="v1.0",
    timeout=timedelta(seconds=20),
    container_image="pvditt/flytekit:15",
    interruptible=False,
    retries=12,
)
def basic_cache(num: str) -> str:
    return f"Hi: {num[0]} ugh-1"


@workflow()
def map_metadata_wf() -> list[str]:
    a = ["1", "2", "3", "4", "5"]
    return map_task(basic_cache)(num=a).with_overrides(
        cache_version="v1.1", 
        timeout=timedelta(seconds=10), 
        container_image="pvditt/flytekit:20",
        interruptible=True,
        retries=10,
    )

Screenshots

Check all the applicable boxes

• I updated the documentation accordingly.
• All new and existing tests passed.
• All commits are signed-off.

Related PRs

Docs link

Summary by Bito

This PR addresses a critical bug in map_tasks where metadata was incorrectly being set on parent nodes instead of subnodes. The implementation introduces a dedicated sub_node_metadata field and refactors metadata override logic to ensure proper application of settings like timeout, cache, and container image to mapped subtasks.

Unit tests added: False

Estimated effort to review (1-5, lower is better): 2

[flyte-bot]

Code Review Agent Run Status

• Limitations and other issues: ❌ Failure - The AI Code Review Agent skipped reviewing this change because it is configured to exclude certain pull requests based on the source/target branch or the pull request status. You can change the settings here, or contact the agent instance creator at [email protected].

[pvditt]

noticed a separate issue: flyteorg/flyte#6153 - will get a quick fix for this after this is merged. Also need to follow up with fixing this for ArrayNode (mapping over ref tasks that I need to upstream the BE changes for)

[flyte-bot]

Code Review Agent Run #a0a338

Actionable Suggestions - 2

• flytekit/core/array_node_map_task.py - 1

  • Consider using constructor for name setting · Line 131-132
• flytekit/tools/translator.py - 1

  • Consider backward compatibility for method rename · Line 627-627

Additional Suggestions - 1

• flytekit/core/node.py - 1

  • Consider simplifying retry strategy initialization · Line 155-156

Review Details

• Files reviewed - 4 · Commit Range: 7a54e96..fe58977

  • flytekit/core/array_node_map_task.py
  • flytekit/core/node.py
  • flytekit/tools/translator.py
  • tests/flytekit/unit/core/test_array_node_map_task.py
• Files skipped - 0
• Tools

  • Whispers (Secret Scanner) - ✔︎ Successful
  • Detect-secrets (Secret Scanner) - ✔︎ Successful
  • MyPy (Static Code Analysis) - ✔︎ Successful
  • Astral Ruff (Static Code Analysis) - ✔︎ Successful

---

AI Code Review powered by

[flyte-bot]

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Bug Fix - Fix map_task metadata override functionality	- array_node_map_task.py - Refactored metadata handling for array node map tasks - node.py - Added _override_node_metadata method to handle metadata overrides - translator.py - Updated metadata retrieval method call - test_array_node_map_task.py - Added comprehensive tests for map task metadata overrides

[flyte-bot]

Consider using constructor for name setting

Consider using the constructor parameters to set the name property when creating NodeMetadata instead of modifying the protected _name attribute directly. This would follow better encapsulation practices.

Code suggestion

Check the AI-generated fix before applying

Suggested change

	self.sub_node_metadata: NodeMetadata = super().construct_node_metadata()
	self.sub_node_metadata._name = self.name
	self.sub_node_metadata: NodeMetadata = NodeMetadata(name=self.name, timeout=self.metadata.timeout, retries=self.metadata.retry_strategy, interruptible=self.metadata.interruptible)

Code Review Run #a0a338

---

Is this a valid issue, or was it incorrectly flagged by the Agent?

• it was incorrectly flagged

[flyte-bot]

Consider backward compatibility for method rename

Consider if renaming the method from construct_sub_node_metadata() to get_sub_node_metadata() maintains backward compatibility. This change could potentially break existing code that relies on the old method name.

Code suggestion

Check the AI-generated fix before applying

Suggested change

	metadata=entity.get_sub_node_metadata(),
	metadata=entity.construct_sub_node_metadata() if hasattr(entity, 'construct_sub_node_metadata') else entity.get_sub_node_metadata(),

Code Review Run #a0a338

---

Is this a valid issue, or was it incorrectly flagged by the Agent?

• it was incorrectly flagged

[thomasjpfan]

Minor nit, otherwise LGTM

[thomasjpfan]

The pattern in flytekit is to use a property + a private attribute:

@property
def sub_node_metadata(self) -> NodeMetadata:
    return self._sub_node_metadata