Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

mp_prefetch silently drops elements that cannot be pickled

Open
  • Is something wrong?
#1,415 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

@rioyu123 is already working on this.

Since Oct 10, 2026.

  • #1428 by @rioyu123 — open

Assessment

Difficulty
3/5
Estimated time
1-2 days
Newbie friendliness
78/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python
Domain
backend

Research direction

Start in grain/_src/python/ipc/queue.py at add_element_to_queue, where multiprocessing.Queue.put returns before the feeder thread pickles; then read _put_dataset_elements_in_buffer in process_prefetch.py to see how the existing error tuple is produced and later raised by __next__. Pickle the element in the worker before putting it (see _check_picklable / _serialize_dataset for prior art) and route a failure through that error tuple. Done when the repro's unpicklable element makes iteration raise instead of silently dropping the item, covered by a test running mp_prefetch with one worker.

Written by the indexing model from the issue text.

Description

type:bug

mp_prefetch silently drops any element that cannot be pickled. Iteration ends normally with that element missing: no exception reaches the caller, and the only trace is a traceback printed by a background thread.

Repro
import grain
import numpy as np
from absl import flags


def element(i):
    # Element 2 carries a value that pickle cannot serialise.
    extra = (lambda: i) if i == 2 else i
    return np.full(4, i), {"extra": extra}


if __name__ == "__main__":
    flags.FLAGS.mark_as_parsed()  # see google/grain#1411
    ds = grain.MapDataset.range(5).map(element).to_iter_dataset()
    ds = ds.mp_prefetch(grain.multiprocessing.MultiprocessingOptions(num_workers=1))
    print("served:", [int(x[0]) for x, _ in ds])

uv run --isolated --no-project --python 3.12 --with grain==0.2.18 python repro.py

Output (exit code 0), the same with or without --with jax:

Traceback (most recent call last):
  File ".../multiprocessing/queues.py", line 264, in _feed
    obj = _ForkingPickler.dumps(obj)
  File ".../multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
AttributeError: Can't get local object 'element.<locals>.<lambda>'
served: [0, 1, 3, 4]

Expected: an exception at the element that cannot be sent (or at least not a silent gap). With every element unpicklable, the iterator yields nothing and still exits normally. The same happens with 2 workers.

Cause

The worker puts each element with grain_queue.add_element_to_queue (grain/_src/python/ipc/queue.py), which calls multiprocessing.Queue.put and returns True. multiprocessing.Queue pickles in its feeder thread after put has returned, so a pickling error there is only printed and the item is dropped. The worker loop (process_prefetch.py, _put_dataset_elements_in_buffer) carries on, and the parent receives the remaining elements and then StopIteration.

Grain already checks that the dataset pickles before starting workers (_check_picklable / _serialize_dataset), but nothing checks an element. The put path is the same at v0.2.18 and current main.

One possible fix: pickle the element in the worker before put (or send it as bytes) and route a failure through the existing error tuple, so __next__ raises it the way it raises other worker errors.

Environment: grain 0.2.18, Python 3.12, Linux x86_64.

Dominant language
Python
Stars
792
Forks
89
Avg merge
1d 9h
Merged PRs (30d)
15

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from google/grain

All issues in google/grain

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.