fix(core): sort dict keys before msgpack encoding in DictTransformer

[KR-Ravindra]

Tracking issue

Closes flyteorg/flyte#6776

Problem

DictTransformer serializes untyped dicts (and typed dicts with non-str keys) to a msgpack binary scalar. msgpack preserves insertion order, so two dicts that are equal in Python but were built in a different key order produce different bytes. flytepropeller hashes the raw literal bytes to build the cache key, so logically identical inputs cause spurious cache misses. This shows up in practice with large, generated, nested config dicts.

Root cause

flytekit/core/type_engine.py:2292-2293 (DictTransformer.dict_to_binary_literal):

encoder = MessagePackEncoder(python_type)
msgpack_bytes = encoder.encode(v)

v is encoded as-is, so the byte layout depends on the caller's insertion order.

Fix

This is option 2 from the discussion in the issue (sort before encoding, no new dependency):

• Add a small helper _sort_dict_keys in flytekit/core/type_engine.py that returns a copy of the value with dict keys sorted recursively (through nested dicts, lists and tuples). Keys are ordered by (type name, key) so mixed-type keys such as int and str can still be ordered; if keys cannot be compared at all, the original order is kept so nothing that encodes today starts failing.
• Call it in DictTransformer.dict_to_binary_literal right before encoder.encode(...).

Only dicts, lists and tuples are rewritten. Because the sorted copy is what gets encoded, to_python_value on such a literal returns the dict with keys in sorted order rather than the caller's insertion order. Dataclasses, FlyteFile, FlyteDirectory and other values nested inside the dict are passed through untouched, so mashumaro's SerializableType handling is unchanged. mashumaro preserves iteration order for dict, Dict[str, Any], Dict[int, str] and nested typed dicts, so sorting the input is enough to make the output canonical.

Scope: this covers DictTransformer only. A Dict[...] field inside a dataclass goes through DataclassTransformer and is not changed here.

test_guess_dict3 in tests/flytekit/unit/core/test_type_hints.py compared the output bytes against msgpack.dumps() of the insertion-ordered dict, which no longer holds by design; it now decodes the payload and compares the value (and still checks the msgpack tag).

How tested

New unit test test_dict_to_binary_literal_is_independent_of_key_order in tests/flytekit/unit/core/test_type_engine.py: two permutations of a nested dict (dicts inside lists inside dicts) must produce byte-identical Literal.scalar.binary.value, round-trip back to the original value, and the same for Dict[int, str] and for a dict with mixed str/int/None keys. A tuple-nested dict is checked for byte equality only, since msgpack decodes arrays as lists.

Before the fix:

>       assert lv1.scalar.binary.value == lv2.scalar.binary.value
E       AssertionError: assert equals failed
E          -b'\x83\xa1a\x01\xa1b\x91\x82\xa   +b'\x83\xa1c\x82\xa1x\x82\xa2k1\
E          -1y\x01\xa1x\x02\xa1c\x82\xa1y\x   +x02\xa2k2\x01\xa1y\x01\xa1b\x91
E          -01\xa1x\x82\xa2k2\x01\xa2k1\x02   +\x82\xa1x\x02\xa1y\x01\xa1a\x01
FAILED tests/flytekit/unit/core/test_type_engine.py::test_dict_to_binary_literal_is_independent_of_key_order

After the fix:

1 passed, 226 deselected in 0.13s

tests/flytekit/unit/core with the change: 1386 passed, 10 skipped; the only failures on this machine (test_schema_in_dataclass, test_union_in_dataclass, test_schema_in_dataclassjsonmixin) fail identically on master with DataFrames of type pandas.DataFrame are not supported currently and are unrelated.

ruff check / ruff format --check (v0.8.3, the pre-commit pin) report no new findings on the changed files; the findings they do report in the two test files exist on master before this change.

Cost check for the concern raised in the issue about large nested dicts: on a ~200k-leaf nested dict (2.1 MB msgpack) the sort takes about 78 ms versus about 15 ms for the encode itself, i.e. a few times the encode cost and well below the network round trips involved in a task launch.

Setup process

uv venv .venv --python 3.12 && source .venv/bin/activate
SETUPTOOLS_SCM_PRETEND_VERSION=1.999.0dev0 uv pip install -r dev-requirements.in
python -m pytest tests/flytekit/unit/core/test_type_engine.py -k key_order

Screenshots

N/A

Check all the applicable boxes

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

Related PRs

flyteorg/flyte#7075 (closed, Go-side normalization) explored fixing this in propeller instead.

Links

• [BUG] DictTransformer has non-deterministic serialization causing unexpected cache misses flyte#6776
• [BUG] Spurious cache misses on dicts that are nested (_calculate_cache_key is nondeterministic) flyte#2924 (earlier report of the same cache-key behaviour)

This change was prepared with an AI agent operated by KR-Ravindra, who reviewed and tested it.

[KR-Ravindra]

Round 1 self-review.

1. tests/flytekit/unit/core/test_type_hints.py:1658 is 122 chars, over the 120 line-length in pyproject.toml; ruff format wants the expected dict wrapped. Reformat that assert.
2. _sort_dict_keys recurses through list but not tuple. msgpack encodes tuples as arrays, so {"a": ({"y": 1, "x": 2},)} is still order-dependent. Add tuple to the container branch (returning a tuple) and a case for it in the new test.
3. Description: after this change to_python_value returns keys in sorted order rather than insertion order. Equality is unaffected, but callers iterating over the decoded dict see a different order; say so explicitly.

Confirmed the new test fails on master and passes with the change, and the 34 dict/binary tests in test_type_engine.py still pass. Staying draft until 1 and 2 are pushed.

[KR-Ravindra]

Round 2 self-review. Changes requested: all three round-1 items are still open on 71fe3f0.

1. Wrap the 122-char assert at tests/flytekit/unit/core/test_type_hints.py:1658 under the 120 line-length (ruff format).
2. Recurse through tuple as well as list in _sort_dict_keys (return a tuple), with a test case for a tuple-nested dict.
3. State in the description that to_python_value now returns keys in sorted rather than insertion order.

Staying draft until 1 and 2 are pushed.

[KR-Ravindra]

Round 3 self-review. Changes requested on 9153bf5.

1. The new tuple case in test_dict_to_binary_literal_is_independent_of_key_order fails: msgpack decodes arrays as lists, so to_python_value returns {'a': [{'x': 2, 'y': 1}], ...} and comparing it to d3 (which holds a tuple) is False.

   E         {'a': [{'x': 2, 'y': 1}]} != {'a': ({'y': 1, 'x': 2},)}
   tests/flytekit/unit/core/test_type_engine.py:617: AssertionError
   

   Test the property instead: encode {"a": ({"y": 1, "x": 2},)} and {"a": ({"x": 2, "y": 1},)} and assert their scalar.binary.value bytes are equal; drop the == d3 round-trip or compare against the list form.
2. The description still says the helper recurses "through nested dicts and lists" and does not mention that to_python_value now returns keys in sorted rather than insertion order (round-2 item 3).
3. Minor: ruff format renders the wrapped assert in test_type_hints.py as assert msgpack.loads(binary_idl_obj.value, strict_map_key=False) == { with the dict split one key per line.

Staying draft until 1 is fixed and the test passes locally.

[KR-Ravindra]

Round 4 self-review on c1322ae.

1. Round-3 item 1: the tuple case now asserts byte equality of two key-order permutations of {"a": ({...},), ...} and no longer round-trips through to_python_value. pytest tests/flytekit/unit/core/test_type_engine.py -k key_order passes (1 passed) and test_type_hints.py -k guess_dict passes (4 passed) on this head.
2. Round-3 item 3: c1322ae reformats the wrapped assert in test_guess_dict3 the way ruff format (0.8.3, the pre-commit pin) renders it. ruff format --diff on the two test files now reports the same 81 pre-existing hunks as master, so the branch adds no findings.
3. Round-3 item 2: description updated (helper recurses through dicts, lists and tuples; to_python_value returns keys in sorted order).
4. No other open or merged PR for [BUG] DictTransformer has non-deterministic serialization causing unexpected cache misses flyte#6776 in this repo; fix: normalize msgpack Binary scalars for deterministic cache hashing flyte#7075 was the closed propeller-side attempt.

Marking ready. The workflows still need a maintainer's approval to run.