google / google/grain

Amount of time spent accessing thread locks for ArrayRecord datasources

Open
#808 1 comment 1 reaction 0 assignees View on GitHub
type:performance
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**

![Image](https://github.com/user-attachments/assets/cc0880ab-90e0-4e2d-8283-6f028743430d)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.