wehub-resource-sync
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
from typing import TYPE_CHECKING, Optional
|
|
|
|
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
|
from ray.data.block import BlockMetadata
|
|
from ray.data.datasource.datasource import Datasource, ReadTask
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
|
|
from ray.data.context import DataContext
|
|
|
|
|
|
TORCH_DATASOURCE_READER_BATCH_SIZE = 32
|
|
|
|
|
|
class TorchDatasource(Datasource):
|
|
"""Torch datasource, for reading from `Torch
|
|
datasets <https://pytorch.org/docs/stable/data.html/>`_.
|
|
This datasource implements a streaming read using a single read task.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dataset: "torch.utils.data.Dataset",
|
|
):
|
|
self._dataset = dataset
|
|
|
|
def get_read_tasks(
|
|
self,
|
|
parallelism: int,
|
|
per_task_row_limit: Optional[int] = None,
|
|
data_context: Optional["DataContext"] = None,
|
|
):
|
|
assert parallelism == 1
|
|
|
|
meta = BlockMetadata(
|
|
# Note: avoid len(self._dataset) because it will trigger
|
|
# iterating through IterableDataset, which can cause OOM.
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
input_files=None,
|
|
exec_stats=None,
|
|
)
|
|
read_task = ReadTask(
|
|
lambda subset=self._dataset: _read_subset(
|
|
subset,
|
|
),
|
|
metadata=meta,
|
|
per_task_row_limit=per_task_row_limit,
|
|
)
|
|
|
|
return [read_task]
|
|
|
|
def estimate_inmemory_data_size(self):
|
|
return None
|
|
|
|
|
|
def _read_subset(subset: "torch.utils.data.Subset"):
|
|
batch = []
|
|
|
|
# Get items from dataset based on its type
|
|
if hasattr(subset, "__iter__"):
|
|
# IterableDataset: Use the iterator directly
|
|
items = subset
|
|
else:
|
|
# Map-style dataset: Respect __len__
|
|
items = (subset[i] for i in range(len(subset)))
|
|
|
|
# Process items in batches
|
|
for item in items:
|
|
batch.append(item)
|
|
if len(batch) == TORCH_DATASOURCE_READER_BATCH_SIZE:
|
|
builder = DelegatingBlockBuilder()
|
|
builder.add_batch({"item": batch})
|
|
yield builder.build()
|
|
batch.clear()
|
|
|
|
# Handle any remaining items
|
|
if len(batch) > 0:
|
|
builder = DelegatingBlockBuilder()
|
|
builder.add_batch({"item": batch})
|
|
yield builder.build()
|