fix(core): sort dict keys before msgpack encoding in DictTransformer - #3467
Draft
KR-Ravindra wants to merge 1 commit into
Draft
fix(core): sort dict keys before msgpack encoding in DictTransformer#3467KR-Ravindra wants to merge 1 commit into
KR-Ravindra wants to merge 1 commit into
Conversation
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>
Author
|
Round 1 self-review.
Confirmed the new test fails on master and passes with the change, and the 34 dict/binary tests in |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Tracking issue
Closes flyteorg/flyte#6776
Problem
DictTransformerserializes untyped dicts (and typed dicts with non-strkeys) 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):vis 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):
_sort_dict_keysinflytekit/core/type_engine.pythat 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 asintandstrcan still be ordered; if keys cannot be compared at all, the original order is kept so nothing that encodes today starts failing.DictTransformer.dict_to_binary_literalright beforeencoder.encode(...).Only dicts and lists are rewritten. Dataclasses,
FlyteFile,FlyteDirectoryand other values nested inside the dict are passed through untouched, so mashumaro'sSerializableTypehandling is unchanged. mashumaro preserves iteration order fordict,Dict[str, Any],Dict[int, str]and nested typed dicts, so sorting the input is enough to make the output canonical.Scope: this covers
DictTransformeronly. ADict[...]field inside a dataclass goes throughDataclassTransformerand is not changed here.test_guess_dict3intests/flytekit/unit/core/test_type_hints.pycompared the output bytes againstmsgpack.dumps()of the insertion-ordered dict, which no longer holds by design; it now decodes the payload and compares the value (and still checks themsgpacktag).How tested
New unit test
test_dict_to_binary_literal_is_independent_of_key_orderintests/flytekit/unit/core/test_type_engine.py: two permutations of a nested dict (dicts inside lists inside dicts) must produce byte-identicalLiteral.scalar.binary.value, round-trip back to the original value, and the same forDict[int, str]and for a dict with mixedstr/int/Nonekeys.Before the fix:
After the fix:
tests/flytekit/unit/corewith 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 onmasterwithDataFrames of type pandas.DataFrame are not supported currentlyand 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 onmasterbefore 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
Screenshots
N/A
Check all the applicable boxes
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.