google / google/grain

ThreadPrefetch `close()` cannot cancel a producer blocked in `parent.__next__()`

Open
#1,372 0 comments 0 reactions 0 assignees View on GitHub
type:bug
Dominant language
Python
Stars
779
Forks
86
Avg merge
2d 6h
Merged PRs (30d)
6

Description

This is separate from #1196. It reproduces on CPython 3.12 today and is not fixed by the `sys.is_finalizing()` join guard.

### Symptom

`ThreadPrefetchDatasetIterator.close()` hangs when its producer is blocked inside `parent.__next__()`.

### Mechanism (mainline before the fix)

1. `close` sets `_closed`, calls `_stop_prefetch`, then `parent.close()`.
2. `_stop_prefetch` sets the stop event, drains the buffer, and **joins the producer with no timeout**.
3. The producer only checks the stop event at the top of its loop. Inside `parent.__next__()` the stop event does nothing.
4. `parent.close()` runs only after the join, so a parent that unblocks only when closed never runs. Deadlock.

### Evidence

Source yields one element, then blocks in `__next__` until `_closed` is set. Consumer takes one element, raises, keeps the iterator, calls `close()`.

| Pipeline | Hang rate |
| --- | --- |
| Single-level `ThreadPrefetch` | 10 / 10 |
| Nested `ThreadPrefetch -> map -> ThreadPrefetch` (device_put shape) | 10 / 10 |

Not nesting-only. Nesting is how real pipelines surface it; the defect is join-before-cancel on every ThreadPrefetch node. The full-buffer exception-path `put` did not hang alone (`_clear_buffer` mitigates it).

### Regression test

`test_close_does_not_hang_when_producer_blocked_in_parent`, single-level and nested. Subprocess-isolated with a hard join timeout. Fails at pristine base on behaviour; passes with two-phase close.

### Fix direction

- `request_stop()`: non-blocking cancel (stop event, buffer wake, parent propagation, no join).
- Explicit `close()`: `request_stop`, then `parent.close()`, then join (still skips join when `sys.is_finalizing()`).
- `__del__` calls `request_stop` only, never join.

### Question

Should non-blocking `request_stop()` live on base `DatasetIterator`, or stay local to ThreadPrefetch? Explicit `close()` would remain the blocking path that joins.

### Second bug (not fixed here)

Parent `StopIteration` is a `BaseException`, so the producer `except Exception` never puts end-of-stream on the buffer and a consumer blocked in `get` is not woken. Not required for the blocked-parent tests; left as follow-up.

Contributor guide

Open the contributing guide

Research direction

Start with ThreadPrefetchDatasetIterator.close(), _stop_prefetch(), and the producer path around parent.__next__(). Run the named test_close_does_not_hang_when_producer_blocked_in_parent in both single-level and nested configurations; done means the subprocess completes without hanging and explicit close performs the intended parent propagation and join behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.