huggingface / huggingface/datasets
`IterableDataset.from_spark` does not support multiple workers in pytorch `Dataloader`
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### Describe the bug
Looks like `IterableDataset.from_spark` does not support multiple workers in pytorch `Dataloader` if I'm not missing anything.
Also, returns not consistent error messages, which probably depend on the nondeterministic order of worker executions
Some exampes I've encountered:
```
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 79, in __iter__
yield from self.generate_examples_fn()
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 49, in generate_fn
df_with_partition_id = df.select("*", pyspark.sql.functions.spark_partition_id().alias("part_id"))
File "/databricks/spark/python/pyspark/instrumentation_utils.py", line 54, in wrapper
logger.log_failure(
File "/databricks/spark/python/pyspark/databricks/usage_logger.py", line 70, in log_failure
self.logger.recordFunctionCallFailureEvent(
File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/java_gateway.py", line 1322, in __call__
return_value = get_return_value(
File "/databricks/spark/python/pyspark/errors/exceptions/captured.py", line 188, in deco
return f(*a, **kw)
File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/protocol.py", line 342, in get_return_value
return OUTPUT_CONVERTER[type](answer[2:], gateway_client)
KeyError: 'c'
```
```
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 79, in __iter__
yield from self.generate_examples_fn()
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 49, in generate_fn
df_with_partition_id = df.select("*", pyspark.sql.functions.spark_partition_id().alias("part_id"))
File "/databricks/spark/python/pyspark/sql/utils.py", line 162, in wrapped
return f(*args, **kwargs)
File "/databricks/spark/python/pyspark/sql/functions.py", line 4893, in spark_partition_id
return _invoke_function("spark_partition_id")
File "/databricks/spark/python/pyspark/sql/functions.py", line 98, in _invoke_function
return Column(jf(*args))
File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/java_gateway.py", line 1322, in __call__
return_value = get_return_value(
File "/databricks/spark/python/pyspark/errors/exceptions/captured.py", line 188, in deco
return f(*a, **kw)
File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/protocol.py", line 342, in get_return_value
return OUTPUT_CONVERTER[type](answer[2:], gateway_client)
KeyError: 'm'
```
```
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 79, in __iter__
yield from self.generate_examples_fn()
File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-68c05436-3512-41c4-88ca-5630012b70d1/lib/python3.10/site-packages/datasets/packaged_modules/spark/spark.py", line 49, in generate_fn
df_with_partition_id = df.select("*", pyspark.sql.functions.spark_partition_id().alias("part_id"))
File "/databricks/spark/python/pyspark/sql/utils.py", line 162, in wrapped
return f(*args, **kwargs)
File "/databricks/spark/python/pyspark/sql/functions.py", line 4893, in spark_partition_id
return _invoke_function("spark_partition_id")
File "/databricks/spark/python/pyspark/sql/functions.py", line 97, in _invoke_function
jf = _get_jvm_function(name, SparkContext._active_spark_context)
File "/databricks/spark/python/pyspark/sql/functions.py", line 88, in _get_jvm_function
return getattr(sc._jvm.functions, name)
File "/databricks/spark/python/lib/py4j-0.10.9.7-src.zip/py4j/java_gateway.py", line 1725, in __getattr__
raise Py4JError(message)
py4j.protocol.Py4JError: functions does not exist in the JVM
```
### Steps to reproduce the bug
```python
import pandas as pd
import numpy as np
batch_size = 16
pdf = pd.DataFrame({
key: np.random.rand(16*100) for key in ['feature', 'target']
})
test_df = spark.createDataFrame(pdf)
from datasets import IterableDataset
from torch.utils.data import DataLoader
ids = IterableDataset.from_spark(test_df)
for batch in DataLoader(ids, batch_size=16, num_workers=4):
for k, b in batch.items():
print(k, b.shape, sep='\t')
print('\n')
```
### Expected behavior
For `num_workers` equal to 0 or 1 works fine as expected:
```
feature torch.Size([16])
target torch.Size([16])
feature torch.Size([16])
target torch.Size([16])
....
```
Expected to support workers >1.
### Environment info
Databricks 13.3 LTS ML runtime - Spark 3.4.1
pyspark==3.4.1
py4j==0.10.9.7
datasets==2.13.1 and also tested with datasets==2.14.6
Contributor guide
Assessment
This issue has not been assessed yet.