Pickle nonimportable Enums so Dataset.map(num_proc>1) and hashing work - #8610
Open
behroozazarkhalili wants to merge 3 commits into
Open
Pickle nonimportable Enums so Dataset.map(num_proc>1) and hashing work#8610behroozazarkhalili wants to merge 3 commits into
behroozazarkhalili wants to merge 3 commits into
Conversation
added 3 commits
September 12, 2026 02:02
Enums defined in __main__ or a local scope could not be pickled by dill's default type path, so Dataset.map(num_proc>1) workers died unpickling them and the fingerprint hasher fell back to a random fingerprint (huggingface#2643). Add an Enum reducer that rebuilds such classes through their metaclass with their members, aliases (including Python 3.13+ unhashable value aliases), Flag boundary and slot state, and validates faithfulness at dump time by reconstructing and comparing member state. Members and class-referencing hooks (__init__/__new__/_generate_next_value_) are restored after the class is memoized, lazily populated caches (_inverted_, negative lookups) never enter the serialized state so hashes are lookup-independent, and explicit __reduce_ex__/copyreg/dill.register contracts, including instance-level ones, take precedence over the Enum reducer. Equality uses _name_/_value_ internals, handles NotImplemented and raising __eq__, and compares array values with array_equal, so nothing upstream could pickle becomes unpicklable. Slot discovery runs once per class; dumps stay linear in the member count. The reducer is registered on both datasets' Pickler and multiprocess's ForkingPickler: Pool binds the shared ForkingPickler with no per-pool injection point, which test_map_enum_requires_worker_dispatch demonstrates. Fixes huggingface#2643
…s stdlib-agnostic Fork CI for the Enum reducer was red on Python 3.10 and 3.14. On 3.9/3.10 two members declared with the same NaN object become two members but one _value2member_map_ key, so the reducer's alias collection dropped the shared value and dump-time validation rejected the class. Record which member body each shared value belongs to and reference it on rebuild instead of relying on the value map. The worker-dispatch test did not raise once tests/features/test_nifti.py was collected in the same session: that module imports the package under a second name, so _dill registered a second dispatch wrapper and removing one still left the Enum reducer active. The test now strips every wrapper and disables map-cache loading, so its negative assertion is exact. IntFlag inversion and negative lookups, and EnumMeta member creation cost, differ between 3.10 and 3.11+ in the standard library itself; the affected tests now compare against the running interpreter's stdlib behaviour or are gated to the versions whose semantics they assert.
…g oracle The deps-minimum CI lane installs dill 0.3.1.1, which has no PicklingWarning and cannot pickle local non-Enum classes defined in an importable module (its save_type only rebuilds nonimportable classes from __main__; local Value/Meta helpers fall through to save_global). Resolve the warning class with a fallback and skip only the local helper-class cases on dill < 0.3.5, keeping every importable case. The large-Enum scaling test compared process_time ratios and flaked under xdist on a loaded machine. It now counts calls to the per-class slot discovery, which is what became quadratic when done per member, so the assertion is independent of load and of EnumMeta construction cost. Reintroducing per-member discovery makes it fail; the real code passes.
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.
Resolves #2643
Root cause
dillserializes a class it cannot import by its module path through_create_type, which rebuilds the class fromtype(name, bases, dict). That path cannot rebuild anEnum: the members,_value2member_map_, aliases, Flag boundary and slot descriptors are produced byEnumMeta.__new__, so an Enum defined in__main__or inside a function fails to pickle.Dataset.map(num_proc>1)sends the shard function throughmultiprocess, whoseForkingPickleruses the same dispatch, so workers died unpickling; the fingerprintHasherhit the same error and fell back to a random fingerprint.Fix
datasets/utils/_dill.py: rebuilds nonimportable Enums through their own metaclass with members, aliases (including Python 3.13+ unhashable value aliases), Flag boundary and slot state, then validates at dump time by reconstructing and comparing member state so a class that cannot be rebuilt faithfully raises atdumprather than producing a wrong class in the worker.__init__,__new__,_generate_next_value_, zero-argumentsuper()) are restored after the class is memoized, so self-referencing Enums no longer hit dill's recursion guard._inverted_, negative-value lookups) are excluded from the serialized state, so the hash of an Enum does not depend on which lookups happened before.__reduce_ex__/__reduce__,copyreganddill.registerreducers take precedence over the Enum reducer._name_/_value_internals, treatNotImplementedand raising__eq__as unknown rather than failure, and compare array values witharray_equal, so nothing that upstream could pickle becomes unpicklable. Slot discovery runs once per class; dumps are linear in member count.datasets'Picklerand onmultiprocess.reduction.ForkingPickler.Poolbinds the sharedForkingPicklerwith no per-pool injection point;test_map_enum_requires_worker_dispatchshows a two-workermapfails without it. The effect on othermultiprocessusers is documented at the registration site.Tests
41 new tests in
tests/test_fingerprint.py: nonimportableEnum/IntEnum/StrEnum/Flag/IntFlaground-trips throughdumpsandForkingPickler, real two-workerDataset.mapwith a local Enum (fork andspawncontexts), aliases and NaN values, custom__new__/_missing_/__init__hooks, lookup-cache independence of the hash, instance and registered reducers, numpy-valued members, undefined equality, raising public accessors, constructor replay (documented behaviour), large-Enum scaling, and dump-time failures for classes that cannot be rebuilt faithfully (datemixin, values changed after construction).