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

Open
#8,595 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.