Large number of chunks causes `OSError: [Errno 24] Too many open files`

[fdalvi]

🐛 Bug

If an optimized dataset has too many chunks (can be replicated using small chunk size or reasonable chunk size with lots of data), while reading (e.g. using StreamingDataLoader), a Too many open files error is encountered at some point. This is because each chunk is mmap'ed as it is loaded, and the mmap'ed handle is never released, leading to too many open files across the system:

https://github.com/Lightning-AI/litdata/blob/7efd76197e0191c0a07c9f0583ce9aec4b80a49e/src/litdata/streaming/item_loader.py#L320-L322

One could ofcourse use ulimit to increase it, but at some point the number of chunks can really be too high.

It is easily resolved by keeping the number of mmap'ed files under a certain limit (for example by occasionally deleting some entries in the above dict, either randomly or smartly using FIFO or so).

I wanted to understand if this fix is alright, I am happy to send a PR for this if it makes sense.

To Reproduce

Steps to reproduce the behavior:

1. Optimize a dataset that results in a large number of chunks
2. Use StreamingDataset and StreamingDataLoader to load the data
3. At some point, OSError: [Errno 24] Too many open files will be encountered

Minimal code sample below to reproduce the issue.

Code sample

import glob
import numpy as np
import random

from pathlib import Path

from litdata import optimize
from litdata.streaming import StreamingDataLoader, StreamingDataset, TokensLoader

# Fake tokenizer
def tokenize_fn(filepath):
    yield np.array([random.randint(0, 10000) for _ in range(random.randint(100, 1000))])

def main():
    Path("fake_file.txt").touch()
    outputs = optimize(
        fn=tokenize_fn,
        inputs=["fake_file.txt" for i in range(10000)], # increase number of files if error is not encountered on a specific machine
        output_dir="./optimized/",
        chunk_bytes="10KB",
        num_workers=1,
        item_loader=TokensLoader(block_size=1024),
    )
    


    train_dataset = StreamingDataset(
        input_dir="./optimized/",
        item_loader=TokensLoader(block_size=1024),
        shuffle=True,
        drop_last=False,
    )

    train_dataloader = StreamingDataLoader(
        train_dataset, batch_size=1, pin_memory=False, num_workers=1, drop_last=False
    )

    total_tokens = 0
    total_batches = 0
    for sample in train_dataloader:
        total_batches += 1
        total_tokens += np.prod(sample.shape)

        print(total_batches, total_tokens)
            
    print("Batches:", total_batches)
    print("Tokens:", total_tokens)

if __name__ == "__main__":
    main()

Expected behavior

Environment

• LitData version: 0.2.26
• OS (e.g., Linux): Linux
• How you installed: pip
• Python version: 3.12

[github-actions]

Hi! thanks for your contribution!, great first issue!

[bhimrazy]

Hi @fdalvi, Thank you for reporting the issue.

I was trying to reproduce the issue using the provided script to understand the bug, but so far, I’ve had no luck on both macOS and Linux (default GitHub Codespace).

Could you please take another look at the provided script and verify its details?

[fdalvi]

Hi @bhimrazy

Thanks for taking a look! Can you please increase the inputs array by a factor of 10 (so 100000 'files' instead of 10000)? I think this should trigger the issue.

In the meantime I'll try to find another machine to test this on as well.

Best,
Fahim

[fdalvi]

Hi @bhimrazy

I was able to replicate the same issue on an macOS, so I checked and both my systems have a ulimit -n value of around 1024. My assumption is that in your case, these values are much higher. You can temporarily lower the value of this to confirm the bug by adding this to the top of the script:

import resource
resource.setrlimit(resource.RLIMIT_NOFILE, (1024, 1024))

I am not sure of the practical implications of increasing this value by arbitrary values to be honest, and it might need escalated privileges on shared clusters. Given that this may not be in a user's control, would it be useful to read this value (resource.getrlimit(resource.RLIMIT_NOFILE)[0]) and manage the mmap's appropriately?

Its also useful to think if we really need to have handles to each and every chunk of a huge dataset.

Best Regards,
Fahim

[bhimrazy]

Thank you, @fdalvi. I was able to replicate the issue.

Hi @tchaton, Do you have any suggestions here ?

[tchaton]

We need to close the memap in the TokensLoader. There used to be some issues but worth try it again.

[fdalvi]

The way I have currently resolved this issue locally is when the number of items inside _mmaps reaches a certain limit, I remove half its entries (randomly at the moment). However, since dictionary order is now guaranteed to be insertion order, we can just remove the first half/close those handles.

The limit can be something we supply in the TokensLoader initialization or even automatically derived from the system limit. Do you think this is a good way to go?

[tchaton]

Yes, this is a good option. A slightly better one would be to call the item_loader.close(...) there: https://github.com/Lightning-AI/litdata/blob/main/src/litdata/streaming/reader.py#L307 and add a new method to the item loader to close the handle: https://github.com/Lightning-AI/litdata/blob/main/src/litdata/streaming/item_loader.py#L361