Fix forked DataLoader workers hanging after the parent iterated the dataset

[avgvi]

Problem

Starting fork DataLoader workers after the parent process has iterated the dataset hangs. A forked child inherits the parent's ThreadPoolExecutor, whose idle-thread bookkeeping claims workers that do not exist in the child, so the download and ready threads submitted in __iter__ never run and the yield loop waits forever. If the parent is in the middle of an epoch, the child also calls exit() on the inherited iterator and waits for threads it does not have.

0.13.0 with filelock 3.31 or newer failed at once with RuntimeError: ... inherited across fork. #26 recreates the locks per process, which removed that error and exposed the hang; it is the known issue listed in the 0.14.0 release notes.

A dataset that has iterated also no longer pickled (TypeError: cannot pickle '_queue.SimpleQueue' object), so spawn and forkserver workers after parent use failed as well, in 0.13.0 too.

Fix

• __iter__ keys the thread pool, event and iterator on os.getpid(): a new process gets fresh ones and drops the inherited iterator.
• __getstate__ leaves the pool, event and iterator out of the pickled state, like the cache filelock.

Tests

• test_dataloader_after_parent_use replaces the fork-only test_dataloader_fork_after_parent_use: fork, spawn and forkserver, each after dataset[0], a full epoch and an epoch in progress, with DataLoader(timeout=60) so a regression fails instead of hanging.
• test_streaming, test_eviction and test_barrier pass on Linux 3.12 (226 tests); the nine start-method cases also pass on macOS.

Found by the differential test of 0.14.1 against 0.13.0.

[read-the-docs-community]

Documentation build overview

📚 streaming-community | 🛠️ Build #35082792 | 📁 Comparing e030b0d against latest (f19ad4a)

  🔍 Preview build  

1 file changed

± _modules/streaming/base/dataset.html