Skip to content

Pickle nonimportable Enums so Dataset.map(num_proc>1) and hashing work - #8610

Open
behroozazarkhalili wants to merge 3 commits into
huggingface:mainfrom
behroozazarkhalili:fix/2643-enum-dill-map
Open

Pickle nonimportable Enums so Dataset.map(num_proc>1) and hashing work#8610
behroozazarkhalili wants to merge 3 commits into
huggingface:mainfrom
behroozazarkhalili:fix/2643-enum-dill-map

Conversation

@behroozazarkhalili

Copy link
Copy Markdown
Contributor

Resolves #2643

Root cause

dill serializes a class it cannot import by its module path through _create_type, which rebuilds the class from type(name, bases, dict). That path cannot rebuild an Enum: the members, _value2member_map_, aliases, Flag boundary and slot descriptors are produced by EnumMeta.__new__, so an Enum defined in __main__ or inside a function fails to pickle. Dataset.map(num_proc>1) sends the shard function through multiprocess, whose ForkingPickler uses the same dispatch, so workers died unpickling; the fingerprint Hasher hit the same error and fell back to a random fingerprint.

Fix

  • New Enum reducer in 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 at dump rather than producing a wrong class in the worker.
  • Members and class-referencing hooks (__init__, __new__, _generate_next_value_, zero-argument super()) are restored after the class is memoized, so self-referencing Enums no longer hit dill's recursion guard.
  • Lazily populated caches (_inverted_, negative-value lookups) are excluded from the serialized state, so the hash of an Enum does not depend on which lookups happened before.
  • Explicit contracts win: class-, base- and instance-level __reduce_ex__/__reduce__, copyreg and dill.register reducers take precedence over the Enum reducer.
  • Faithfulness checks use _name_/_value_ internals, treat NotImplemented and raising __eq__ as unknown rather than failure, and compare array values with array_equal, so nothing that upstream could pickle becomes unpicklable. Slot discovery runs once per class; dumps are linear in member count.
  • The reducer is registered on datasets' Pickler and on multiprocess.reduction.ForkingPickler. Pool binds the shared ForkingPickler with no per-pool injection point; test_map_enum_requires_worker_dispatch shows a two-worker map fails without it. The effect on other multiprocess users is documented at the registration site.

Tests

41 new tests in tests/test_fingerprint.py: nonimportable Enum/IntEnum/StrEnum/Flag/IntFlag round-trips through dumps and ForkingPickler, real two-worker Dataset.map with a local Enum (fork and spawn contexts), 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 (date mixin, values changed after construction).

Behrooz Azarkhalili Aghmiyouni 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.
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.

Enum used in map functions will raise a RecursionError with dill.

1 participant