Skip to content

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

Description

@MohammadSadeghSalehi

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    type:bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions