Lightning-AI / Lightning-AI/litData
If data samples contain a Python list with a variable number of elements, type inference will fail
- Dominant language
- Python
- Stars
- 614
- Forks
- 106
- Avg merge
- 15h 8m
- Merged PRs (30d)
- 22
Description
## 🐛 Bug
When optimizing a dataset, [`BinaryWriter.serialize()`](https://github.com/Lightning-AI/litdata/blob/main/src/litdata/streaming/writer.py#L167) will first flatten the sample dictionary and infer the type of each element. Then it will assume that all data samples have the same types and the same number of elements. If a list contains a variable number of elements, the type information will be incorrect for subsequent data samples. This is not documented and can cause some confusion.
### To Reproduce
When storing the data in NumPy arrays, there's no problem, because each array is considered as one element in the flattened type list. So this works:
```python
import numpy as np
from PIL import Image
import litdata as ld
def random_images(index):
fake_images = Image.fromarray(np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8))
num_labels = np.random.randint(10)
fake_labels = np.random.randint(10, size=num_labels)
# use any key:value pairs
data = {"index": index, "image": fake_images, "class": fake_labels}
return data
if __name__ == "__main__":
# the optimize function outputs data in an optimized format (chunked, binerized, etc...)
ld.optimize(
fn=random_images, # the function applied to each input
inputs=list(range(1000)), # the inputs to the function (here it's a list of numbers)
output_dir="my_optimized_dataset", # optimized data is stored here
num_workers=4, # The number of workers on the same machine
chunk_bytes="64MB" # size of each chunk
)
```
When storing the data in a Python list, the type of each list element is inferred separately. If a list contains a variable number of elements, the type information of one sample is not useful for other samples. Now, if we add an element that has a different type (say, a string) after a variable-length list, the function will crash:
```python
def random_images(index):
fake_images = Image.fromarray(np.random.randint(0, 256, (32, 32, 3), dtype=np.uint8))
num_labels = np.random.randint(10)
fake_labels = np.random.randint(10, size=num_labels).tolist()
name = "name"
# use any key:value pairs
data = {"index": index, "image": fake_images, "class": fake_labels, "name": name}
return data
```
The error:
```
Rank 0 inferred the following `['int', 'pil', 'int', 'int', 'int', 'str']` data format. | 0/1000 [00:00
ld.optimize(
File ".../lib/python3.12/site-packages/litdata/processing/functions.py", line 432, in optimize
data_processor.run(
File ".../lib/python3.12/site-packages/litdata/processing/data_processor.py", line 1055, in run
self._exit_on_error(error)
File ".../lib/python3.12/site-packages/litdata/processing/data_processor.py", line 1119, in _exit_on_error
raise RuntimeError(f"We found the following error {error}.")
RuntimeError: We found the following error Traceback (most recent call last):
File ".../lib/python3.12/site-packages/litdata/processing/data_processor.py", line 665, in _handle_data_chunk_recipe
chunk_filepath = self.cache._add_item(self._index_counter, item_data_or_generator)
File ".../lib/python3.12/site-packages/litdata/streaming/cache.py", line 134, in _add_item
return self._writer.add_item(index, data)
File ".../lib/python3.12/site-packages/litdata/streaming/writer.py", line 305, in add_item
data, dim = self.serialize(items)
File ".../lib/python3.12/site-packages/litdata/streaming/writer.py", line 179, in serialize
self._serialize_with_data_format(flattened, sizes, data, self._data_format)
File ".../lib/python3.12/site-packages/litdata/streaming/writer.py", line 210, in _serialize_with_data_format
serialized_item, _ = serializer.serialize(element)
File ".../lib/python3.12/site-packages/litdata/streaming/serializers.py", line 347, in serialize
return obj.encode("utf-8"), None
AttributeError: 'int' object has no attribute 'encode'
```
### Expected behavior
It should be documented that every data sample must have the same elements and every list must be the same size.
I wonder if caching the data types is necessary. Afterall, the `optimize()` call doesn't have to be that fast. If the data types were instead inferred for every sample separately, it would be possible to use variable-length lists.
Would it make sense to at least check in `BinaryWriter.serialize()` that `self._data_format` has the same number of elements as the data sample?
### Environment
- PyTorch Version (e.g., 1.0): 2.3.1
- OS (e.g., Linux): Linux
- How you installed PyTorch (`conda`, `pip`, source): pip
- Python version: 3.12.3
Contributor guide
Research direction
Start with src/litdata/streaming/writer.py, especially BinaryWriter.serialize(), and reproduce the variable-length Python-list case from the issue. Compare the inferred format with later samples and review the expected behavior discussion; done should clearly address the documented fixed-element and fixed-list-size requirement or the requested validation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- data-engineering
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100