Support with_overrides setting metadata for map_task subnode instead of parent node #2982

pvditt · 2024-12-05T20:22:55Z

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

Fixed a critical bug in map_tasks where metadata was incorrectly being set on parent nodes instead of subnodes. Implemented dedicated sub_node_metadata field and refactored override logic for proper application of settings (timeout, cache, container image) to mapped subtasks. Simplified metadata access by removing getter method and using direct property access while maintaining backward compatibility.

Unit tests added: False

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

Signed-off-by: Paul Dittamo <[email protected]>

flyte-bot · 2025-01-09T05:22:44Z

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].

Signed-off-by: Paul Dittamo <[email protected]>

pvditt · 2025-01-09T06:08:41Z

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 · 2025-01-09T06:42:45Z

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 · 2025-01-09T06:48:13Z

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` - Added sub_node_metadata field and refactored metadata handling - `node.py` - Implemented _override_node_metadata method for proper metadata inheritance - `translator.py` - Updated metadata reference to use sub_node_metadata - `test_array_node_map_task.py` - Added comprehensive tests for map task metadata overrides

flyte-bot · 2025-01-09T06:48:15Z

flytekit/core/array_node_map_task.py

+        self.sub_node_metadata: NodeMetadata = super().construct_node_metadata()
+        self.sub_node_metadata._name = self.name


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 · 2025-01-09T06:48:16Z

flytekit/tools/translator.py

@@ -624,7 +624,7 @@ def get_serializable_array_node_map_task(
    )
    node = workflow_model.Node(
        id=entity.name,
-        metadata=entity.construct_sub_node_metadata(),
+        metadata=entity.get_sub_node_metadata(),


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 · 2025-01-11T14:34:17Z

flytekit/core/array_node_map_task.py

+    def get_sub_node_metadata(self) -> NodeMetadata:
+        return self.sub_node_metadata


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

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

thanks for pointing this out. just updated

Signed-off-by: Paul Dittamo <[email protected]>

flyte-bot · 2025-01-13T08:02:46Z

Code Review Agent Run #e417ab

Actionable Suggestions - 0

Review Details

Files reviewed - 2 · Commit Range: fe58977..550ca24
- flytekit/core/array_node_map_task.py
- flytekit/tools/translator.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

pvditt added 3 commits December 5, 2024 12:21

test

7a54e96

Signed-off-by: Paul Dittamo <[email protected]>

Merge branch 'master' into bug/with-overrides-map-task

419ddaa

Signed-off-by: Paul Dittamo <[email protected]>

add support for with_overrides for map tasks

ba2f578

Signed-off-by: Paul Dittamo <[email protected]>

expand unit test

fe58977

Signed-off-by: Paul Dittamo <[email protected]>

pvditt marked this pull request as ready for review January 9, 2025 06:08

pvditt requested review from wild-endeavor, kumare3, eapolinario, pingsutw, cosmicBboy, samhita-alla, thomasjpfan and Future-Outlier as code owners January 9, 2025 06:08

flyte-bot reviewed Jan 9, 2025

View reviewed changes

thomasjpfan previously approved these changes Jan 11, 2025

View reviewed changes

cleanup

550ca24

Signed-off-by: Paul Dittamo <[email protected]>

pvditt dismissed thomasjpfan’s stale review via 550ca24 January 13, 2025 06:56

pvditt requested a review from thomasjpfan January 13, 2025 17:00

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support with_overrides setting metadata for map_task subnode instead of parent node #2982

Support with_overrides setting metadata for map_task subnode instead of parent node #2982

pvditt commented Dec 5, 2024 •

edited by flyte-bot

Loading

flyte-bot commented Jan 9, 2025

pvditt commented Jan 9, 2025 •

edited

Loading

flyte-bot commented Jan 9, 2025 •

edited

Loading

Code Review Agent Run #a0a338

flyte-bot commented Jan 9, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Jan 9, 2025

flyte-bot Jan 9, 2025

thomasjpfan left a comment

thomasjpfan Jan 11, 2025

pvditt Jan 13, 2025

flyte-bot commented Jan 13, 2025 •

edited

Loading

Code Review Agent Run #e417ab

		self.sub_node_metadata: NodeMetadata = super().construct_node_metadata()
		self.sub_node_metadata._name = self.name

	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)

	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(),

		def get_sub_node_metadata(self) -> NodeMetadata:
		return self.sub_node_metadata

Support with_overrides setting metadata for map_task subnode instead of parent node #2982

Are you sure you want to change the base?

Support with_overrides setting metadata for map_task subnode instead of parent node #2982

Conversation

pvditt commented Dec 5, 2024 • edited by flyte-bot Loading

Tracking issue

Why are the changes needed?

What changes were proposed in this pull request?

How was this patch tested?

Setup process

Screenshots

Check all the applicable boxes

Related PRs

Docs link

Summary by Bito

flyte-bot commented Jan 9, 2025

Code Review Agent Run Status

pvditt commented Jan 9, 2025 • edited Loading

flyte-bot commented Jan 9, 2025 • edited Loading

Code Review Agent Run #a0a338

flyte-bot commented Jan 9, 2025 • edited Loading

Changelist by Bito

flyte-bot Jan 9, 2025

Choose a reason for hiding this comment

flyte-bot Jan 9, 2025

Choose a reason for hiding this comment

thomasjpfan left a comment

Choose a reason for hiding this comment

thomasjpfan Jan 11, 2025

Choose a reason for hiding this comment

pvditt Jan 13, 2025

Choose a reason for hiding this comment

flyte-bot commented Jan 13, 2025 • edited Loading

Code Review Agent Run #e417ab

pvditt commented Dec 5, 2024 •

edited by flyte-bot

Loading

pvditt commented Jan 9, 2025 •

edited

Loading

flyte-bot commented Jan 9, 2025 •

edited

Loading

flyte-bot commented Jan 9, 2025 •

edited

Loading

flyte-bot commented Jan 13, 2025 •

edited

Loading