Skip to content

Raise output pickling errors from process prefetch workers - #1428

Open
rioyu123 wants to merge 1 commit into
google:mainfrom
rioyu123:codex/grain-prefetch-serialization-errors
Open

rioyu123 wants to merge 1 commit into
google:mainfrom
rioyu123:codex/grain-prefetch-serialization-errors

Conversation

@rioyu123

@rioyu123 rioyu123 commented Oct 10, 2026 •

Copy link
Copy Markdown

Fixes #1415.

multiprocessing.Queue pickles items on its feeder thread. When an output cannot be pickled, that thread prints an error and drops the element, while the caller sees normal iteration over a truncated dataset.

The worker now pickles each output tuple with ForkingPickler before enqueueing it. If this fails, it releases the shared memory allocated for that output and sends the error after the preceding elements, following the existing transform-error path.

A small __reduce__ wrapper carries the pickled bytes. User objects are pickled once, and both the consumer and queue-cleanup code still receive the original tuples. There is no public API change.

This moves pickling onto the worker loop and adds copies of the serialized payload. Arrays passed through shared memory only add small metadata to that payload. In three interleaved local single-worker probes, median throughput differences were within run-to-run noise. These are not production benchmarks or a claim of zero overhead.

Tests:

  • Two real-worker regressions cover buffer sizes 1 and 4. They check that elements 0 and 1 arrive before the error from element 2. Both fail on main because no exception reaches the caller.
  • Unit tests check one serialization pass and that the failed output's shared-memory block is unlinked.
  • The full process_prefetch_test module passes on Linux CPython 3.12 and 3.14: 71 tests on each, using the native absltest runner.
  • The related queue and shared-memory-array modules pass on CPython 3.12: 6 and 11 tests respectively. Existing resource-tracker warnings are not addressed by this patch.

The full repository/Bazel CI was not run locally. This change does not include the separate shared-memory cleanup work in #1386, #1421 or #1422, and does not address exceptions that are themselves unpicklable.


📚 Documentation preview 📚: https://google-grain--1428.org.readthedocs.build/

Signed-off-by: Rio Yu <52408936+rioyu123@users.noreply.github.com>

This branch has not been deployed

No deployments
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.

mp_prefetch silently drops elements that cannot be pickled

1 participant