huggingface / huggingface/datasets
IterableDataset.decode(num_threads) busy-polls its ThreadPool futures, starving the decode threads — 0 rows on remote audio; callback-based bridging restores the intended 20x
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### Describe the bug
`decode(num_threads=N)` (added in #7450, "speeds up streaming up to 20 times") schedules the sync `decode_example` calls onto a `multiprocessing.pool.ThreadPool` and bridges them into the async map machinery with a busy-poll (`iterable_dataset.py`):
```python
async def _apply_async(pool, func, x):
future = pool.apply_async(func, (x,))
while not future.ready():
await asyncio.sleep(0) # <- hot loop
return future.get()
```
With the in-flight cap at `2 * num_threads`, the event loop spins up to 2N of these hot loops continuously. That is pure Python work holding the GIL, and it starves the very ThreadPool threads doing the decoding: in our stack dumps at `num_threads=32`, a worker spent **45+ seconds inside `ssl.create_default_context`** (a ~50 ms operation) because it could barely get scheduled. Net effect on remote (`hf://` path-referenced) audio: **zero rows yielded, ever** — we killed runs after 155 s with 0 rows, while plain sequential iteration of the same dataset runs fine at 2.3 ex/s.
A second, smaller cold-start problem stacks on top: `Audio.decode_example` lazily imports `datasets.features._torchcodec` (which pulls in torch) on first call, so all N threads immediately pile up on the importlib module lock.
### Reproduction
```python
import datasets.features._torchcodec # work around the import pile-up to isolate the busy-poll
from datasets import Audio, Dataset
files = [f"hf://datasets/datasets-examples/doc-audio-1/{i}.wav" for i in (1, 2, 3, 4)] * 100
ds = Dataset.from_dict({"audio": files}).cast_column("audio", Audio()).to_iterable_dataset()
for row in ds.decode(num_threads=32): # first row never arrives
...
```
### The fix (tested)
Bridge the ThreadPool future to asyncio with callbacks instead of polling — the loop then sleeps until a result arrives:
```python
async def _apply_async(pool, func, x):
loop = asyncio.get_running_loop()
fut = loop.create_future()
pool.apply_async(
func, (x,),
callback=lambda r: loop.call_soon_threadsafe(fut.set_result, r),
error_callback=lambda e: loop.call_soon_threadsafe(fut.set_exception, e),
)
return await fut
```
plus importing the lazy decoder deps eagerly before the pool starts (`datasets.features._torchcodec`; pre-importing `fsspec.implementations.chained` also helps, since fsspec runs that function-level import on every `open`).
Measured on the reproduction above (400 rows, `num_threads=32`, same machine):
| variant | throughput |
|---|---|
| sequential (`num_threads=0`) | 2.3 ex/s |
| `decode(num_threads=32)`, current code | 0 rows in 155 s (killed) |
| with the `_apply_async` fix above | 21.2 ex/s |
| with the fix + a huggingface_hub fix (see below) | **43.8 ex/s, t_first 1.2 s** |
43.8 ex/s is 19× over sequential — matching the original 20× claim of #7450, so this restores the feature to its designed behavior.
The remaining 2× (21 → 44) comes from a separate `huggingface_hub` bug, filed as huggingface/huggingface_hub#4861: `get_session()` lacks a re-check inside its lock, so N cold-starting threads each create and overwrite the global client, and the orphaned clients' sockets get closed under in-flight requests. Worth knowing about because `decode(num_threads)` is exactly the workload that triggers it, and datasets' open-retry loop silently absorbs the resulting connection errors.
### System info
datasets 4.7.1.dev0 (`_apply_async` and `decode()` are identical on current main), huggingface_hub 1.9.2, httpx 0.28.1, Python 3.14.0, Linux x86_64.
Contributor guide
Research direction
Start in iterable_dataset.py at _apply_async and inspect how decode(num_threads) submits work to the ThreadPool. Run the supplied remote-audio reproduction, then verify that callback-based future bridging yields rows and restores parallel throughput without relying on busy-polling; the eager decoder imports are also part of the reported fix.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data, performance
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100