Fix dict annotations by wild-endeavor · Pull Request #3123 · flyteorg/flytekit

wild-endeavor · 2025-02-10T22:17:20Z

Tracking issue

Internal ticket only.

Why are the changes needed?

Customer reported issue with internal model registry. The issue though is that annotations don't work well with dictionary return types.

What changes were proposed in this pull request?

Most of the changes here are in the dict transformer.

The extract_types_or_metadata function is inconsistent, returning different things in different scenarios - sometimes it would return annotations, sometimes it would return the key/value type of a dictionary. This has been changed to always return the latter, or an empty tuple () in the case of an untyped dictionary. In all cases now, the caller interprets the result as a tuple of key type,value type. Renaming function to just extract_types.
the is_pickle function returned both a boolean, as well as types, but the latter isn't really necessary. Changing this to just return the pickle bool.
dict_types is a function that was almost the same as extract_types_or_metadata, removing in favor of that one and manually extracting the annotation args where needed.
Remove a duplicate test

How was this patch tested?

Setup process

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

Enhanced DictTransformer implementation by refactoring type handling and simplifying type extraction functionality. The extract_types_or_metadata function has been renamed to extract_types, with improved type casting. Reorganized test structure by eliminating duplicate tests and properly relocating dictionary type extraction tests. Changes improve consistency in dictionary type handling and focus on type safety improvements.

Unit tests added: True

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

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

flyte-bot · 2025-02-10T22:17:34Z

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 eduardo@union.ai.

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

codecov · 2025-02-10T22:55:35Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.75%. Comparing base (1eb6743) to head (127d737).

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #3123       +/-   ##
===========================================
+ Coverage   79.60%   92.75%   +13.15%     
===========================================
  Files         203        7      -196     
  Lines       21625      497    -21128     
  Branches     2788        0     -2788     
===========================================
- Hits        17214      461    -16753     
+ Misses       3629       36     -3593     
+ Partials      782        0      -782

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

eapolinario

this looks good. Left a few nits, but won't block on any of them.

eapolinario · 2025-02-10T23:37:18Z

                _, v_type = transformer_override.extract_types_or_metadata(t_value_type)  # type: ignore
            else:
-                _, v_type = DictTransformer.extract_types_or_metadata(t_value_type)  # type: ignore
+                _, v_type = DictTransformer.extract_types(t_value_type)  # type: ignore


can we remove the # type: ignore now?

eapolinario · 2025-02-10T23:38:57Z

    @staticmethod
-    def is_pickle(python_type: Type[dict]) -> typing.Tuple[bool, Type]:
-        base_type, *metadata = DictTransformer.extract_types_or_metadata(python_type)
+    def is_pickle(python_type: Type[dict]) -> bool:


the change of return type in is_pickle is the best change in this PR.

eapolinario · 2025-02-10T23:44:41Z

-@pytest.mark.parametrize(
-    "t,expected",
-    [
-        (None, (None, None)),
-        (typing.Dict, ()),
-        (typing.Dict[str, str], (str, str)),
-        (
-                Annotated[typing.Dict[str, str], kwtypes(allow_pickle=True)],
-                (typing.Dict[str, str], kwtypes(allow_pickle=True)),
-        ),
-        (typing.Dict[Annotated[str, "a-tag"], int], (Annotated[str, "a-tag"], int)),
-    ],
-)
-def test_dict_get(t, expected):
-    assert DictTransformer.extract_types_or_metadata(t) == expected
-
-


nit: can you remove the other one instead?

flyte-bot · 2025-02-10T23:45:59Z

Code Review Agent Run #d12db3

Actionable Suggestions - 5

tests/flytekit/unit/core/test_generice_idl_type_engine.py - 1
- Consider preserving kwtypes annotation for testing · Line 2869-2869
flytekit/core/type_engine.py - 4
- Consider preserving metadata extraction capability · Line 881-881
- Consider updating all function references · Line 2021-2021
- Potential tuple unpacking issue in is_pickle · Line 2146-2146
- Potential IndexError in type check condition · Line 2174-2174

Review Details

Files reviewed - 5 · Commit Range: aa72916..825264c
- flytekit/core/promise.py
- flytekit/core/type_engine.py
- tests/flytekit/unit/core/test_annotated_bindings.py
- tests/flytekit/unit/core/test_generice_idl_type_engine.py
- tests/flytekit/unit/core/test_type_engine.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-02-10T23:51:43Z

Changelist by Bito

This pull request implements the following key changes.

Key Change	Files Impacted
Feature Improvement - Dictionary Type Handling Enhancement	- `promise.py` - Updated dictionary type extraction with improved type casting - `type_engine.py` - Refactored DictTransformer with simplified type extraction and pickle handling - `test_annotated_bindings.py` - Added new tests for annotated dynamic dictionary types - `test_generice_idl_type_engine.py` - Removed redundant dictionary type extraction tests - `test_type_engine.py` - Updated dictionary type extraction test cases

flyte-bot · 2025-02-10T23:51:46Z


        if isinstance(val, dict):
-            ktype, vtype = DictTransformer.extract_types_or_metadata(t)
+            ktype, vtype = DictTransformer.extract_types(t)


Consider preserving metadata extraction capability

Consider keeping the extract_types_or_metadata method call instead of changing to extract_types. The original method may handle metadata that could be important for dictionary type transformations.

Code suggestion

Check the AI-generated fix before applying

Suggested change

ktype, vtype = DictTransformer.extract_types(t)

ktype, vtype = DictTransformer.extract_types_or_metadata(t)

Code Review Run #d12db3

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

it was incorrectly flagged

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

flyte-bot · 2025-02-11T01:15:59Z

Code Review Agent Run #bebaf5

Actionable Suggestions - 0

Review Details

Files reviewed - 3 · Commit Range: 825264c..aaa1984
- flytekit/core/promise.py
- tests/flytekit/unit/core/test_generice_idl_type_engine.py
- tests/flytekit/unit/core/test_type_engine.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

Future-Outlier · 2025-02-11T02:06:16Z

+            if _origin is dict and _args is not None:
                return _args  # type: ignore
+            elif _origin is Annotated:
+                return DictTransformer.extract_types(_args[0])
+            else:
+                raise ValueError(f"Trying to extract dictionary type information from a non-dict type {t}")


Just use debugger to run the unit test of this PR, this make sense to me.

checking other usecase in annotation, put tuple in it

oh I see, tuple type is restricted now

the if else statement in this function looks super weird, trying to fix it

Future-Outlier · 2025-02-11T02:28:00Z

+
+

Let's add test when Tuple IDL is supported!
flyteorg/flyte#5699

def test_pydantic_annotated_tuple(): from pydantic import BaseModel from flytekit.core.artifact import Artifact from flytekit import task, dynamic from flytekit.configuration import SerializationSettings, ImageConfig, Image from flytekit.core.context_manager import FlyteContextManager, ExecutionState, FlyteContext from typing import Annotated, Union, Tuple # Local Pydantic model for demonstration class Config2(BaseModel): param: int = 42 @task def get_config() -> Config2: return Config2() # We'll reuse the Prophet class already defined above in the same file. # If needed, ensure it's defined at a higher scope or imported in your actual code. def dt_function(): # Create Pydantic object cfg = get_config() # Create a dummy Prophet dict best_models = {"model_1": Prophet(param=0.1)} # Return a tuple of (Config2, dict[str, Prophet]) return cfg, best_models a = Artifact(name="my_tuple_model") # A dynamic workflow returning an Annotated tuple @dynamic def dt_tuple() -> Annotated[Tuple[Config2, dict[str, Prophet]], a]: return dt_function() # Serialization settings similar to other tests ss = SerializationSettings( project="test_proj", domain="test_domain", version="abc", image_config=ImageConfig(Image(name="name", fqn="image", tag="name")), env={}, ) # Compile the dynamic workflow into a spec and ensure it behaves as expected with FlyteContextManager.with_context( FlyteContextManager.current_context().with_serialization_settings(ss) ) as ctx: new_exc_state = ctx.execution_state.with_params(mode=ExecutionState.Mode.TASK_EXECUTION) with FlyteContextManager.with_context(ctx.with_execution_state(new_exc_state)): dynamic_job_spec_tuple = dt_tuple.compile_into_workflow(ctx, dt_tuple._task_function) # We expect exactly one node in the spec (this dynamic call itself), # and a single output binding containing our tuple assert len(dynamic_job_spec_tuple.nodes) == 1 assert len(dynamic_job_spec_tuple.outputs) == 1 # Here, we can check that the output is stored as a collection of two items bd = dynamic_job_spec_tuple.outputs[0].binding assert bd.collection is not None assert len(bd.collection.bindings) == 2

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com> Signed-off-by: Chih Tsung Lu <lu001224@gmail.com>

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com> Signed-off-by: Atharva <atharvakulkarni172003@gmail.com>

fix dict annotations

aa72916

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

lint

127d737

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

rename

825264c

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

wild-endeavor changed the title ~~fix dict annotations~~ Fix dict annotations Feb 10, 2025

wild-endeavor marked this pull request as ready for review February 10, 2025 23:14

wild-endeavor requested review from Future-Outlier, cosmicBboy, eapolinario, kumare3, pingsutw, samhita-alla and thomasjpfan as code owners February 10, 2025 23:14

eapolinario previously approved these changes Feb 10, 2025

View reviewed changes

flyte-bot reviewed Feb 10, 2025

View reviewed changes

Comment thread tests/flytekit/unit/core/test_generice_idl_type_engine.py Outdated

flyte-bot reviewed Feb 10, 2025

View reviewed changes

Comment thread flytekit/core/type_engine.py

flyte-bot reviewed Feb 10, 2025

View reviewed changes

Comment thread flytekit/core/type_engine.py

flyte-bot reviewed Feb 10, 2025

View reviewed changes

Comment thread flytekit/core/type_engine.py

pr comments

aaa1984

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com>

wild-endeavor dismissed eapolinario’s stale review via aaa1984 February 11, 2025 00:11

eapolinario approved these changes Feb 11, 2025

View reviewed changes

wild-endeavor merged commit d732d6f into master Feb 11, 2025

Future-Outlier reviewed Feb 11, 2025

View reviewed changes

Future-Outlier mentioned this pull request Feb 11, 2025

Refactor extract_types in Dict Transformer #3124

Merged

3 tasks

ChihTsungLu pushed a commit to ChihTsungLu/flytekit that referenced this pull request Feb 17, 2025

Fix dict annotations (flyteorg#3123)

5009ab2

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com> Signed-off-by: Chih Tsung Lu <lu001224@gmail.com>

Atharva1723 pushed a commit to Atharva1723/flytekit that referenced this pull request Oct 5, 2025

Fix dict annotations (flyteorg#3123)

460e16d

Signed-off-by: Yee Hing Tong <wild-endeavor@users.noreply.github.com> Signed-off-by: Atharva <atharvakulkarni172003@gmail.com>

	ktype, vtype = DictTransformer.extract_types(t)
	ktype, vtype = DictTransformer.extract_types_or_metadata(t)

Uh oh!

Conversation

wild-endeavor commented Feb 10, 2025 • edited by flyte-bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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

Uh oh!

flyte-bot commented Feb 10, 2025

Code Review Agent Run Status

Uh oh!

codecov Bot commented Feb 10, 2025

Codecov Report

Uh oh!

eapolinario left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

flyte-bot commented Feb 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #d12db3

Uh oh!

flyte-bot commented Feb 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changelist by Bito

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

flyte-bot commented Feb 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Code Review Agent Run #bebaf5

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wild-endeavor commented Feb 10, 2025 •

edited by flyte-bot

Loading

flyte-bot commented Feb 10, 2025 •

edited

Loading

flyte-bot commented Feb 10, 2025 •

edited

Loading

flyte-bot commented Feb 11, 2025 •

edited

Loading