From 71fe3f0cac2bc8f897215ef34997cf7338a3f594 Mon Sep 17 00:00:00 2001 From: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:29:12 +0000 Subject: [PATCH 1/4] fix(core): sort dict keys before msgpack encoding in DictTransformer DictTransformer.dict_to_binary_literal encoded the dict in insertion order, so equal dicts built in a different key order produced different msgpack bytes and therefore different propeller cache keys. Recursively sort dict keys (through nested dicts and lists) before encoding so the literal bytes are canonical. Dataclasses, FlyteFile and FlyteDirectory values inside the dict are passed through untouched. Signed-off-by: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> --- flytekit/core/type_engine.py | 21 +++++++++++++- tests/flytekit/unit/core/test_type_engine.py | 29 ++++++++++++++++++++ tests/flytekit/unit/core/test_type_hints.py | 6 ++-- 3 files changed, 52 insertions(+), 4 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 982d9d031c..9de8b5a4bf 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -72,6 +72,25 @@ def _default_msgpack_decoder(data: bytes) -> Any: return msgpack.unpackb(data, strict_map_key=False) +def _sort_dict_keys(obj: Any) -> Any: + """ + Recursively sort the keys of ``obj`` and of any dict nested inside dicts or lists, so that two dicts holding + the same items serialize to identical msgpack bytes regardless of insertion order. Propeller derives cache + keys from the raw literal bytes, so key order would otherwise cause spurious cache misses. + Keys are grouped by type name so that mixed-type keys (e.g. ``int`` and ``str``) can be ordered; if the keys + still cannot be compared, the original order is kept. + """ + if isinstance(obj, dict): + try: + keys = sorted(obj, key=lambda k: (type(k).__name__, k)) + except TypeError: + keys = list(obj) + return {k: _sort_dict_keys(obj[k]) for k in keys} + if isinstance(obj, list): + return [_sort_dict_keys(v) for v in obj] + return obj + + class BatchSize: """ This is used to annotate a FlyteDirectory when we want to download/upload the contents of the directory in batches. For example, @@ -2290,7 +2309,7 @@ async def dict_to_binary_literal( try: # Handle dictionaries with non-string keys (e.g., Dict[int, Type]) encoder = MessagePackEncoder(python_type) - msgpack_bytes = encoder.encode(v) + msgpack_bytes = encoder.encode(_sort_dict_keys(v)) return Literal(scalar=Scalar(binary=Binary(value=msgpack_bytes, tag=MESSAGEPACK))) except TypeError as e: if allow_pickle: diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index de3f7dfd02..2ccef9cb9c 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -29,6 +29,7 @@ from flytekit import dynamic, kwtypes, task, workflow from flytekit.core.annotation import FlyteAnnotation +from flytekit.core.constants import MESSAGEPACK from flytekit.core.context_manager import FlyteContext, FlyteContextManager from flytekit.core.data_persistence import flyte_tmp_dir from flytekit.core.hash import HashMethod @@ -583,6 +584,34 @@ def recursive_assert( assert d.to_python_value(ctx, lv, dict) == {"x": "hello"} +def test_dict_to_binary_literal_is_independent_of_key_order(): + ctx = FlyteContext.current_context() + + d1 = {"a": 1, "b": [{"y": 1, "x": 2}], "c": {"y": 1, "x": {"k2": 1, "k1": 2}}} + d2 = {"c": {"x": {"k1": 2, "k2": 1}, "y": 1}, "b": [{"x": 2, "y": 1}], "a": 1} + assert d1 == d2 + + lt = TypeEngine.to_literal_type(dict) + lv1 = TypeEngine.to_literal(ctx, d1, dict, lt) + lv2 = TypeEngine.to_literal(ctx, d2, dict, lt) + assert lv1.scalar.binary.tag == MESSAGEPACK + assert lv1.scalar.binary.value == lv2.scalar.binary.value + assert TypeEngine.to_python_value(ctx, lv1, dict) == d1 + + lt_int = TypeEngine.to_literal_type(Dict[int, str]) + lv1 = TypeEngine.to_literal(ctx, {2: "b", 1: "a"}, Dict[int, str], lt_int) + lv2 = TypeEngine.to_literal(ctx, {1: "a", 2: "b"}, Dict[int, str], lt_int) + assert lv1.scalar.binary.value == lv2.scalar.binary.value + assert TypeEngine.to_python_value(ctx, lv1, Dict[int, str]) == {1: "a", 2: "b"} + + mixed1 = {"b": 1, 1: "a", None: 2} + mixed2 = {None: 2, 1: "a", "b": 1} + lv1 = TypeEngine.to_literal(ctx, mixed1, dict, lt) + lv2 = TypeEngine.to_literal(ctx, mixed2, dict, lt) + assert lv1.scalar.binary.value == lv2.scalar.binary.value + assert TypeEngine.to_python_value(ctx, lv1, dict) == mixed1 + + def test_convert_marshmallow_json_schema_to_python_class(): @dataclass class Foo(DataClassJsonMixin): diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 26d8615eb2..4d053c4a7d 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1653,9 +1653,9 @@ def t2() -> dict: ctx = context_manager.FlyteContextManager.current_context() output_lm = t2.dispatch_execute(ctx, _literal_models.LiteralMap(literals={})) - msgpack_bytes = msgpack.dumps({"k1": "v1", "k2": 3, 4: {"one": [1, "two", [3]]}}) - binary_idl_obj = Binary(value=msgpack_bytes, tag=MESSAGEPACK) - assert output_lm.literals["o0"].scalar.binary == binary_idl_obj + binary_idl_obj = output_lm.literals["o0"].scalar.binary + assert binary_idl_obj.tag == MESSAGEPACK + assert msgpack.loads(binary_idl_obj.value, strict_map_key=False) == {"k1": "v1", "k2": 3, 4: {"one": [1, "two", [3]]}} @pytest.mark.skipif( From 9153bf5453065780d186b99bdbca5f81140917a0 Mon Sep 17 00:00:00 2001 From: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> Date: Fri, 11 Sep 2026 06:50:19 +0000 Subject: [PATCH 2/4] fix: add tuple support to _sort_dict_keys and fix line length - Add tuple support to _sort_dict_keys (returns tuple) - Add test case for tuple-nested dict - Fix 122-char assert to meet 120-char line limit in pyproject.toml Signed-off-by: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> --- flytekit/core/type_engine.py | 6 ++++-- tests/flytekit/unit/core/test_type_engine.py | 5 +++++ tests/flytekit/unit/core/test_type_hints.py | 4 +++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 9de8b5a4bf..9d6412243b 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -74,8 +74,8 @@ def _default_msgpack_decoder(data: bytes) -> Any: def _sort_dict_keys(obj: Any) -> Any: """ - Recursively sort the keys of ``obj`` and of any dict nested inside dicts or lists, so that two dicts holding - the same items serialize to identical msgpack bytes regardless of insertion order. Propeller derives cache + Recursively sort the keys of ``obj`` and of any dict nested inside dicts, lists, or tuples, so that two dicts + holding the same items serialize to identical msgpack bytes regardless of insertion order. Propeller derives cache keys from the raw literal bytes, so key order would otherwise cause spurious cache misses. Keys are grouped by type name so that mixed-type keys (e.g. ``int`` and ``str``) can be ordered; if the keys still cannot be compared, the original order is kept. @@ -88,6 +88,8 @@ def _sort_dict_keys(obj: Any) -> Any: return {k: _sort_dict_keys(obj[k]) for k in keys} if isinstance(obj, list): return [_sort_dict_keys(v) for v in obj] + if isinstance(obj, tuple): + return tuple(_sort_dict_keys(v) for v in obj) return obj diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 2ccef9cb9c..196407ccfe 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -611,6 +611,11 @@ def test_dict_to_binary_literal_is_independent_of_key_order(): assert lv1.scalar.binary.value == lv2.scalar.binary.value assert TypeEngine.to_python_value(ctx, lv1, dict) == mixed1 + d3 = {"a": ({"y": 1, "x": 2},), "b": [{"y": 1, "x": 2}]} + lv3 = TypeEngine.to_literal(ctx, d3, dict, lt) + assert lv3.scalar.binary.tag == MESSAGEPACK + assert TypeEngine.to_python_value(ctx, lv3, dict) == d3 + def test_convert_marshmallow_json_schema_to_python_class(): @dataclass diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 4d053c4a7d..57c07d7a13 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1655,7 +1655,9 @@ def t2() -> dict: output_lm = t2.dispatch_execute(ctx, _literal_models.LiteralMap(literals={})) binary_idl_obj = output_lm.literals["o0"].scalar.binary assert binary_idl_obj.tag == MESSAGEPACK - assert msgpack.loads(binary_idl_obj.value, strict_map_key=False) == {"k1": "v1", "k2": 3, 4: {"one": [1, "two", [3]]}} + assert msgpack.loads( + binary_idl_obj.value, strict_map_key=False + ) == {"k1": "v1", "k2": 3, 4: {"one": [1, "two", [3]]}} @pytest.mark.skipif( From d6feecfa79403ce361fe4adea1c568a0a13fee31 Mon Sep 17 00:00:00 2001 From: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> Date: Fri, 11 Sep 2026 07:20:33 +0000 Subject: [PATCH 3/4] fix: compare tuple dict bytes directly, not round-trip Signed-off-by: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> --- tests/flytekit/unit/core/test_type_engine.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_engine.py b/tests/flytekit/unit/core/test_type_engine.py index 196407ccfe..2510800b1c 100644 --- a/tests/flytekit/unit/core/test_type_engine.py +++ b/tests/flytekit/unit/core/test_type_engine.py @@ -611,10 +611,14 @@ def test_dict_to_binary_literal_is_independent_of_key_order(): assert lv1.scalar.binary.value == lv2.scalar.binary.value assert TypeEngine.to_python_value(ctx, lv1, dict) == mixed1 - d3 = {"a": ({"y": 1, "x": 2},), "b": [{"y": 1, "x": 2}]} - lv3 = TypeEngine.to_literal(ctx, d3, dict, lt) + # Test that tuples are handled deterministically: msgpack decodes tuples as lists, + # so we compare bytes directly rather than round-tripping through to_python_value. + d3a = {"a": ({"y": 1, "x": 2},), "b": [{"y": 1, "x": 2}]} + d3b = {"a": ({"x": 2, "y": 1},), "b": [{"y": 1, "x": 2}]} + lv3 = TypeEngine.to_literal(ctx, d3a, dict, lt) assert lv3.scalar.binary.tag == MESSAGEPACK - assert TypeEngine.to_python_value(ctx, lv3, dict) == d3 + # Two encodings with different key order must produce identical bytes + assert lv3.scalar.binary.value == TypeEngine.to_literal(ctx, d3b, dict, lt).scalar.binary.value def test_convert_marshmallow_json_schema_to_python_class(): From c1322ae94aef38e40f52dc90fe53ead303c45b03 Mon Sep 17 00:00:00 2001 From: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:50:16 -0700 Subject: [PATCH 4/4] test: format the decoded-payload assert in test_guess_dict3 with ruff Signed-off-by: KR Ravindra <42912207+KR-Ravindra@users.noreply.github.com> --- tests/flytekit/unit/core/test_type_hints.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/flytekit/unit/core/test_type_hints.py b/tests/flytekit/unit/core/test_type_hints.py index 57c07d7a13..c8be06b46d 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1655,9 +1655,11 @@ def t2() -> dict: output_lm = t2.dispatch_execute(ctx, _literal_models.LiteralMap(literals={})) binary_idl_obj = output_lm.literals["o0"].scalar.binary assert binary_idl_obj.tag == MESSAGEPACK - assert msgpack.loads( - binary_idl_obj.value, strict_map_key=False - ) == {"k1": "v1", "k2": 3, 4: {"one": [1, "two", [3]]}} + assert msgpack.loads(binary_idl_obj.value, strict_map_key=False) == { + "k1": "v1", + "k2": 3, + 4: {"one": [1, "two", [3]]}, + } @pytest.mark.skipif(