Skip to content

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

Draft
KR-Ravindra wants to merge 1 commit into
flyteorg:masterfrom
KR-Ravindra:fix/dict-transformer-deterministic-msgpack
Draft

fix(core): sort dict keys before msgpack encoding in DictTransformer#3467
KR-Ravindra wants to merge 1 commit into
flyteorg:masterfrom
KR-Ravindra:fix/dict-transformer-deterministic-msgpack

Conversation

@KR-Ravindra

Copy link
Copy Markdown

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 and lists). 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 and lists are rewritten. 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.

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

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

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>
@KR-Ravindra

Copy link
Copy Markdown
Author

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] DictTransformer has non-deterministic serialization causing unexpected cache misses

1 participant