diff --git a/flytekit/core/type_engine.py b/flytekit/core/type_engine.py index 982d9d031c..9d6412243b 100644 --- a/flytekit/core/type_engine.py +++ b/flytekit/core/type_engine.py @@ -72,6 +72,27 @@ 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, 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. + """ + 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] + if isinstance(obj, tuple): + return tuple(_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 +2311,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..2510800b1c 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,43 @@ 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 + + # 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 + # 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(): @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..c8be06b46d 100644 --- a/tests/flytekit/unit/core/test_type_hints.py +++ b/tests/flytekit/unit/core/test_type_hints.py @@ -1653,9 +1653,13 @@ 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(