fix: FCC behavior classes could not be pickled for distributed workers

[ikrommyd]

🤖 AI text below 🤖

The FCC tests pass when the file is run on its own but fail under pytest tests:

TypeError: ('Could not serialize object of type HighLevelGraph', '<ToPickle: HighLevelGraph ...>')

The dask_client fixture in tests/conftest.py is session-scoped and distributed.Client defaults to set_as_default=True, so as soon as any test that uses it has run, every later .compute() in the session goes through the distributed scheduler instead of the local threaded one. That means the schema's behavior dict has to be pickled, and for FCCSchema it can't be.

The vector.backends.awkward.<lambda> message in the traceback is a red herring — distributed.protocol.pickle.dumps tries plain pickle first and falls back to cloudpickle, which handles those lambdas fine. The error that actually escapes is from the fallback:

TypeError: cannot pickle '_DaskProperty' object
  when serializing dict item 'get_daughters'
  when serializing type state          <- a class being pickled *by value*
  when serializing dict item 'MCParticle'
  when serializing dict item 'behavior'

Shadowed behavior classes

methods/fcc.py defined MCParticle, ReconstructedParticle, ParticleID, Cluster and Track twice — once for behavior (pre-edm4hep1) and again for behavior_edm4hep1, with # noqa: F811. awkward.mixin_class injects the generated <Name>Record/<Name>Array classes as module attributes, so the second batch overwrote the first, and the classes still held by the pre-edm4hep1 behavior dict were no longer reachable as coffea.nanoevents.methods.fcc.MCParticleRecord.

cloudpickle only pickles a class by reference when getattr(module, cls.__qualname__) is cls. That check now failed, so it fell back to pickling the class by value, walked its __dict__, and hit a _DaskProperty.

The edm4hep1 classes are renamed with an _edm4hep1 suffix and keep their behavior keys via mixin_class(..., name=...), so nothing user-facing changes. FCCSchema was the only schema affected.

Note that by-value class pickling is worth avoiding on its own even when it succeeds: it bloats every graph and gives each worker a class object that is not the one in the module.

Unpicklable dask_property

Separately, property keeps fget/fset/fdel in C-level slots and provides no reduction, so CPython refuses to pickle it (pickle.dumps(property(f)) raises the same TypeError). Any behavior mixin using @dask_property and defined outside an importable module — __main__, a notebook — therefore cannot be sent to a distributed worker, whether or not it lives in coffea. _DaskProperty now defines __reduce__. _DaskMethod keeps all of its state in a normal instance __dict__ and already pickled fine, so it is unchanged.

coffea's _DaskProperty is a copy of dask-awkward's, which has the same bug; fixed upstream in dask-contrib/dask-awkward#629.

Tests

test_schema_behavior_survives_pickling checks, for every schema, that the behavior dict cloudpickles and that every class in it is reachable under its own qualified name. The second half is the important one: with the __reduce__ fix in place the shadowed FCC classes pickle successfully again, so only the qualname check still catches the regression.

test_dask_property_is_picklable round-trips a class defined in a function body (forcing the by-value path) and checks the eager getter, the .dask dispatch and no_dispatch=True all survive.