Support overriding task annotations and labels via with_overrides #3061

arbaobao · 2025-01-15T09:06:03Z

Tracking issue

Why are the changes needed?

Improvement to allow setting labels and annotations dynamically at run time for things like cost allocation.

What changes were proposed in this pull request?

Adds labels and annotations to with_overrides().

How was this patch tested?

Excute a workflow and using with_override(labels={"lKeyA": "lValA"},annotations={"lKeyB": "lValB"}).

Setup process

python

@task
def say_hello() -> str:
    return "Hello, World!"

@workflow
def hello_world_wf() -> str:
    res = say_hello().with_overrides(labels={"lKeyA": "lValA"},annotations={"lKeyB": "lValB"})
    return res

It can be tested by using kubectl describe pods {pod_name} -n flytesnacks-development

Screenshots

Check all the applicable boxes

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

Related PRs

flyte 6171

Docs link

Summary by Bito

Implementation of task annotations and labels override functionality via with_overrides() method. Added label and annotation fields to Node class and TaskNodeOverrides with serialization support. Enables dynamic runtime configuration of labels/annotations for cost allocation purposes. Includes comprehensive test coverage across various task types.

Unit tests added: True

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

flyte-bot · 2025-01-15T09:06:17Z

Code Review Agent Run #2c0fc5

Actionable Suggestions - 2

flytekit/core/node.py - 1
- Consider consolidating duplicate validation logic · Line 228-234
flytekit/tools/translator.py - 1
- Consider validating node override fields · Line 469-470

Review Details

Files reviewed - 6 · Commit Range: 44e8c1b..44e8c1b
- flytekit/core/node.py
- flytekit/models/core/workflow.py
- flytekit/tools/translator.py
- tests/flytekit/unit/core/test_array_node_map_task.py
- tests/flytekit/unit/core/test_map_task.py
- tests/flytekit/unit/core/test_node_creation.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-15T09:11:29Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Task Node Override Enhancement	- `node.py` - Added support for labels and annotations in Node class - `workflow.py` - Implemented label and annotation fields in TaskNodeOverrides - `translator.py` - Added serialization support for labels and annotations
Testing - Test Coverage for Label and Annotation Overrides	- `test_array_node_map_task.py` - Added tests for map task overrides with labels and annotations - `test_map_task.py` - Added tests for map task label and annotation overrides - `test_node_creation.py` - Added dedicated tests for label and annotation overrides

flyte-bot · 2025-01-15T09:11:31Z

flytekit/core/node.py

+        if annotations is not None:
+            assert_not_promise(annotations, "annotations")
+            self._annotations = annotations
+
+        if labels is not None:
+            assert_not_promise(labels, "labels")
+            self._labels = labels


Consider consolidating duplicate validation logic

Consider consolidating the duplicate validation logic for annotations and labels into a helper function to reduce code duplication. The same pattern of checking for None and calling assert_not_promise() is repeated.

Code suggestion

Check the AI-generated fix before applying

Suggested change

if annotations is not None:

assert_not_promise(annotations, "annotations")

self._annotations = annotations

if labels is not None:

assert_not_promise(labels, "labels")

self._labels = labels

def _validate_dict_param(value: Optional[Dict[str, str]], param_name: str) -> None:

if value is not None:

assert_not_promise(value, param_name)

setattr(self, f"_{param_name}", value)

_validate_dict_param(annotations, "annotations")

_validate_dict_param(labels, "labels")

Code Review Run #2c0fc5

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

it was incorrectly flagged

flyte-bot · 2025-01-15T09:11:32Z

flytekit/tools/translator.py

+                    annotations=entity._annotations,
+                    labels=entity._labels,


Consider validating node override fields

Consider adding validation for _annotations and _labels before passing them to TaskNodeOverrides. These fields could potentially be None which may cause issues downstream.

Code suggestion

Check the AI-generated fix before applying

Suggested change

annotations=entity._annotations,

labels=entity._labels,

annotations=entity._annotations if entity._annotations is not None else {},

labels=entity._labels if entity._labels is not None else {},

Code Review Run #2c0fc5

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

it was incorrectly flagged

Signed-off-by: Nelson Chen <asd3431090@gmail.com>

flyte-bot · 2025-01-15T16:43:40Z

Code Review Agent Run #db4a2c

Actionable Suggestions - 0

Review Details

Files reviewed - 6 · Commit Range: e485d6a..7df32f2
- flytekit/core/node.py
- flytekit/models/core/workflow.py
- flytekit/tools/translator.py
- tests/flytekit/unit/core/test_array_node_map_task.py
- tests/flytekit/unit/core/test_map_task.py
- tests/flytekit/unit/core/test_node_creation.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

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

arbaobao mentioned this pull request Jan 15, 2025

Support overriding task annotations and labels via with_overrides flyteorg/flyte#6171

Open

3 tasks

flyte-bot reviewed Jan 15, 2025

View reviewed changes

arbaobao added 2 commits January 15, 2025 23:44

add with override for annotation and label

e485d6a

Signed-off-by: Nelson Chen <asd3431090@gmail.com>

lint

7df32f2

Signed-off-by: Nelson Chen <asd3431090@gmail.com>

arbaobao force-pushed the override_label branch from 84c3cad to 7df32f2 Compare January 15, 2025 15:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support overriding task annotations and labels via with_overrides #3061

Support overriding task annotations and labels via with_overrides #3061

arbaobao commented Jan 15, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 15, 2025 •

edited

Loading

Code Review Agent Run #2c0fc5

flyte-bot commented Jan 15, 2025 •

edited

Loading

Changelist by Bito

flyte-bot Jan 15, 2025

flyte-bot Jan 15, 2025

flyte-bot commented Jan 15, 2025 •

edited

Loading

Code Review Agent Run #db4a2c

Support overriding task annotations and labels via with_overrides #3061

Are you sure you want to change the base?

Support overriding task annotations and labels via with_overrides #3061

Conversation

arbaobao commented Jan 15, 2025 • 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 15, 2025 • edited Loading

Code Review Agent Run #2c0fc5

flyte-bot commented Jan 15, 2025 • edited Loading

Changelist by Bito

flyte-bot Jan 15, 2025

Choose a reason for hiding this comment

flyte-bot Jan 15, 2025

Choose a reason for hiding this comment

flyte-bot commented Jan 15, 2025 • edited Loading

Code Review Agent Run #db4a2c

arbaobao commented Jan 15, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Jan 15, 2025 •

edited

Loading

flyte-bot commented Jan 15, 2025 •

edited

Loading

flyte-bot commented Jan 15, 2025 •

edited

Loading