Try `obstore` for LitData Chunks

[bhimrazy]

🚀 Feature

Try obstore for LitData Chunks
In #652 , we noticed 2-3x improvements after using obstore.

Notes from @tchaton :

So, It would be good to try it out for LitData chunks too
Instead of boto3
Maybe it would get faster

[bhimrazy]

One thing to note when using obstore with multiprocessing is that it can get stuck without explicitly setting a start method. So, it’s better to check with mp methods.

It seems to work with spawn:

import multiprocessing as mp

mp.set_start_method("spawn", force=True)
print("Start method:", mp.get_start_method())

[bhimrazy]

With spawn, it did work, but it took at least 2 minutes to start streaming—probably to initialize all the workers.
code branch

⚡ feat/stream-litdata-chunks-with-obstore ~/litdata python benchmarks/litdata/stream_imagenet.py \
    --input_dir /teamspace/datasets/imagenet-1m-optimized-0.2.41-v2/train \
    --batch_size 256 \
    --epochs 2 \
    --num_workers 32
Start method: spawn
[INFO] Running streaming benchmark with arguments: Namespace(input_dir='/teamspace/datasets/imagenet-1m-optimized-0.2.41-v2/train', cache_dir='/cache/chunks', dtype='float32', batch_size=256, num_workers=32, drop_last=False, epochs=2, max_cache_size='200GB', use_pil=False, clear_cache=True)
Seed set to 42
[INFO] Clearing cache directory: /cache/chunks
[INFO] Initializing streaming dataset from: /teamspace/datasets/imagenet-1m-optimized-0.2.41-v2/train
[INFO] Starting benchmark for 2 epoch(s) with batch size 256 and 32 workers.
Epoch 1/2: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5005/5005 [07:17<00:00, 11.43it/s]
[RESULT] Epoch 1: Streamed 1281167 samples in 437.72s (2926.88 images/sec)
Epoch 2/2: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5005/5005 [06:52<00:00, 12.13it/s]
[RESULT] Epoch 2: Streamed 1281167 samples in 412.47s (3106.09 images/sec)
[INFO] Clearing cache directory after benchmark: /cache/chunks
[INFO] Finished streaming benchmark.