Amount of time spent accessing thread locks for ArrayRecord datasources
- Dominant language
- Python
- Stars
- 779
- Forks
- 86
- Avg merge
- 2d 6h
- Merged PRs (30d)
- 6
Description
Hey! I am trying to put together a fast dataloader with Grain for ArrayRecord files, but I am running into performance issues. It seems like using a batched `grain.DataLoader` or `grain.Dataset` is quite slow both with and without multiprocessing, and most of the time is spent acquiring thread locks. This slowdown effectively renders the dataloader unusable, I get only ~1 it/s with a batch size of 2048 even on fairly small datasets.
The delay grows with the shard size (the max I have tried is around 1GB per shard), although I think this is quite far from "a new frontier of IO efficiency" regardless of the size of my dataset.
I am wondering what I am missing here, it would be nice to have some instructions on how to speed things up.
Below is a demonstration of the performance bottleneck.
This output was produced on a single NVIDIA GH200 machine rented from Lambda, using `grain==0.2.6`.
**Dataset creation:**
```python
import os
import string
import random
from tqdm.auto import tqdm
import grain.python as grain
import array_record.python.array_record_module as array_record
out_dir = "/tmp/test_dataset"
base_fname = "shard"
size_limit = 1024
def random_str(length):
return ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(length))
data_iter = [random_str(500) for _ in range(50000)]
os.makedirs(out_dir, exist_ok=True)
shard = 0
current_bytes = 0
filename = f"{out_dir}/{base_fname}-{shard:05d}"
writer = array_record.ArrayRecordWriter(filename, "group_size:1")
for item in tqdm(data_iter):
record_bytes = item.encode("utf-8")
record_size = len(record_bytes)
if current_bytes + record_size > size_limit:
writer.close()
shard += 1
current_bytes = 0
filename = f"{out_dir}/{base_fname}-{shard:05d}"
writer = array_record.ArrayRecordWriter(filename, "group_size:1")
writer.write(record_bytes)
current_bytes += record_size
```
---
**Dataloader:**
```python
test_datasource = grain.ArrayRecordDataSource(
[f"{out_dir}/{record_path}" for record_path in os.listdir(out_dir)]
)
class DecodeString(grain.MapTransform):
def map(self, element):
return element.decode("utf-8")
sampler = grain.IndexSampler(
num_records=len(test_datasource),
num_epochs=1,
shard_options=grain.ShardOptions(shard_index=0, shard_count=1, drop_remainder=True),
shuffle=True,
seed=2002,
)
dataloader = grain.DataLoader(
data_source=test_datasource,
sampler=sampler,
shard_options=grain.NoSharding(),
operations=[
DecodeString(),
grain.Batch(batch_size=2048, drop_remainder=True),
],
worker_count=16,
worker_buffer_size=500,
enable_profiling=True,
)
print(f"Total shards: {len(os.listdir(out_dir))}")
max_len = 0
for step in tqdm(range(10)):
element = next(iter(dataloader))
x = max(max_len, element.shape[0])
```
---
**Example debugging output**

Contributor guide
Assessment
This issue has not been assessed yet.