Make tfds.data_source pickable.

PiperOrigin-RevId: 636470622

Author: marcenacp

tensorflow_datasets/core/data_sources/array_record.py

@@ -20,13 +20,8 @@
 """
 
 import dataclasses
-from typing import Any, Optional
 
-from tensorflow_datasets.core import dataset_info as dataset_info_lib
-from tensorflow_datasets.core import decode
-from tensorflow_datasets.core import splits as splits_lib
 from tensorflow_datasets.core.data_sources import base
-from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import array_record_data_source
 
 
@@ -42,18 +37,9 @@ class ArrayRecordDataSource(base.BaseDataSource):
   source.
   """
 
-  dataset_info: dataset_info_lib.DatasetInfo
-  split: splits_lib.Split = None
-  decoders: Optional[type_utils.TreeDict[decode.partial_decode.DecoderArg]] = (
-      None
-  )
-  # In order to lazy load array_record, we don't load
-  # `array_record_data_source.ArrayRecordDataSource` here.
-  data_source: Any = dataclasses.field(init=False)
-  length: int = dataclasses.field(init=False)
-
   def __post_init__(self):
-    file_instructions = base.file_instructions(self.dataset_info, self.split)
+    dataset_info = self.dataset_builder.info
+    file_instructions = base.file_instructions(dataset_info, self.split)
     self.data_source = array_record_data_source.ArrayRecordDataSource(
         file_instructions
     )

tensorflow_datasets/core/data_sources/base.py

@@ -17,12 +17,14 @@
 
 from collections.abc import MappingView, Sequence
 import dataclasses
+import functools
 import typing
 from typing import Any, Generic, Iterable, Protocol, SupportsIndex, TypeVar
 
 from tensorflow_datasets.core import dataset_info as dataset_info_lib
 from tensorflow_datasets.core import decode
 from tensorflow_datasets.core import splits as splits_lib
+from tensorflow_datasets.core.features import top_level_feature
 from tensorflow_datasets.core.utils import shard_utils
 from tensorflow_datasets.core.utils import type_utils
 from tensorflow_datasets.core.utils.lazy_imports_utils import tree
@@ -54,6 +56,14 @@ def file_instructions(
   return split_dict[split].file_instructions
 
 
+class _DatasetBuilder(Protocol):
+  """Protocol for the DatasetBuilder to avoid cyclic imports."""
+
+  @property
+  def info(self) -> dataset_info_lib.DatasetInfo:
+    ...
+
+
 @dataclasses.dataclass
 class BaseDataSource(MappingView, Sequence):
   """Base DataSource to override all dunder methods with the deserialization.
@@ -64,22 +74,28 @@ class BaseDataSource(MappingView, Sequence):
   deserialization/decoding.
 
   Attributes:
-    dataset_info: The DatasetInfo of the
+    dataset_builder: The dataset builder.
     split: The split to load in the data source.
     decoders: Optional decoders for decoding.
     data_source: The underlying data source to initialize in the __post_init__.
   """
 
-  dataset_info: dataset_info_lib.DatasetInfo
+  dataset_builder: _DatasetBuilder
   split: splits_lib.Split | None = None
   decoders: type_utils.TreeDict[decode.partial_decode.DecoderArg] | None = None
   data_source: DataSource[Any] = dataclasses.field(init=False)
 
+  @functools.cached_property
+  def _features(self) -> top_level_feature.TopLevelFeature:
+    """Caches features because we log the use of dataset_builder.info."""
+    features = self.dataset_builder.info.features
+    if not features:
+      raise ValueError('No feature defined in the dataset buidler.')
+    return features
+
   def __getitem__(self, key: SupportsIndex) -> Any:
     record = self.data_source[key.__index__()]
-    return self.dataset_info.features.deserialize_example_np(
-        record, decoders=self.decoders
-    )
+    return self._features.deserialize_example_np(record, decoders=self.decoders)
 
   def __getitems__(self, keys: Sequence[int]) -> Sequence[Any]:
     """Retrieves items by batch.
@@ -98,24 +114,24 @@ def __getitems__(self, keys: Sequence[int]) -> Sequence[Any]:
     if not keys:
       return []
     records = self.data_source.__getitems__(keys)
-    features = self.dataset_info.features
     if len(keys) != len(records):
       raise IndexError(
           f'Requested {len(keys)} records but got'
           f' {len(records)} records.'
           f'{keys=}, {records=}'
       )
     return [
-        features.deserialize_example_np(record, decoders=self.decoders)
+        self._features.deserialize_example_np(record, decoders=self.decoders)
         for record in records
     ]
 
   def __repr__(self) -> str:
     decoders_repr = (
         tree.map_structure(type, self.decoders) if self.decoders else None
     )
+    name = self.dataset_builder.info.name
     return (
-        f'{self.__class__.__name__}(name={self.dataset_info.name}, '
+        f'{self.__class__.__name__}(name={name}, '
         f'split={self.split!r}, '
         f'decoders={decoders_repr})'
     )

tensorflow_datasets/core/data_sources/base_test.py

@@ -15,13 +15,15 @@
 
 """Tests for all data sources."""
 
+import pickle
 from unittest import mock
 
+import cloudpickle
 from etils import epath
 import pytest
 import tensorflow_datasets as tfds
 from tensorflow_datasets import testing
-from tensorflow_datasets.core import dataset_builder
+from tensorflow_datasets.core import dataset_builder as dataset_builder_lib
 from tensorflow_datasets.core import dataset_info as dataset_info_lib
 from tensorflow_datasets.core import decode
 from tensorflow_datasets.core import file_adapters
@@ -77,7 +79,7 @@ def mocked_parquet_dataset():
 )
 def test_read_write(
     tmp_path: epath.Path,
-    builder_cls: dataset_builder.DatasetBuilder,
+    builder_cls: dataset_builder_lib.DatasetBuilder,
     file_format: file_adapters.FileFormat,
 ):
   builder = builder_cls(data_dir=tmp_path, file_format=file_format)
@@ -106,28 +108,34 @@ def test_read_write(
 ]
 
 
-def create_dataset_info(file_format: file_adapters.FileFormat):
+def create_dataset_builder(file_format: file_adapters.FileFormat):
   with mock.patch.object(splits_lib, 'SplitInfo') as split_mock:
     split_mock.return_value.name = 'train'
     split_mock.return_value.file_instructions = _FILE_INSTRUCTIONS
     dataset_info = mock.create_autospec(dataset_info_lib.DatasetInfo)
     dataset_info.file_format = file_format
     dataset_info.splits = {'train': split_mock()}
     dataset_info.name = 'dataset_name'
-    return dataset_info
+
+    dataset_builder = mock.create_autospec(dataset_builder_lib.DatasetBuilder)
+    dataset_builder.info = dataset_info
+
+    return dataset_builder
 
 
 @pytest.mark.parametrize(
     'data_source_cls',
     _DATA_SOURCE_CLS,
 )
 def test_missing_split_raises_error(data_source_cls):
-  dataset_info = create_dataset_info(file_adapters.FileFormat.ARRAY_RECORD)
+  dataset_builder = create_dataset_builder(
+      file_adapters.FileFormat.ARRAY_RECORD
+  )
   with pytest.raises(
       ValueError,
       match="Unknown split 'doesnotexist'.",
   ):
-    data_source_cls(dataset_info, split='doesnotexist')
+    data_source_cls(dataset_builder, split='doesnotexist')
 
 
 @pytest.mark.usefixtures(*_FIXTURES)
@@ -136,8 +144,10 @@ def test_missing_split_raises_error(data_source_cls):
     _DATA_SOURCE_CLS,
 )
 def test_repr_returns_meaningful_string_without_decoders(data_source_cls):
-  dataset_info = create_dataset_info(file_adapters.FileFormat.ARRAY_RECORD)
-  source = data_source_cls(dataset_info, split='train')
+  dataset_builder = create_dataset_builder(
+      file_adapters.FileFormat.ARRAY_RECORD
+  )
+  source = data_source_cls(dataset_builder, split='train')
   name = data_source_cls.__name__
   assert (
       repr(source) == f"{name}(name=dataset_name, split='train', decoders=None)"
@@ -150,9 +160,11 @@ def test_repr_returns_meaningful_string_without_decoders(data_source_cls):
     _DATA_SOURCE_CLS,
 )
 def test_repr_returns_meaningful_string_with_decoders(data_source_cls):
-  dataset_info = create_dataset_info(file_adapters.FileFormat.ARRAY_RECORD)
+  dataset_builder = create_dataset_builder(
+      file_adapters.FileFormat.ARRAY_RECORD
+  )
   source = data_source_cls(
-      dataset_info,
+      dataset_builder,
       split='train',
       decoders={'my_feature': decode.SkipDecoding()},
   )
@@ -181,3 +193,18 @@ def test_data_source_is_sliceable():
     file_instructions = mock_array_record_data_source.call_args_list[1].args[0]
     assert file_instructions[0].skip == 0
     assert file_instructions[0].take == 30000
+
+
+# PyGrain requires that data sources are picklable.
+@pytest.mark.parametrize(
+    'file_format',
+    file_adapters.FileFormat.with_random_access(),
+)
+@pytest.mark.parametrize('pickle_module', [pickle, cloudpickle])
+def test_data_source_is_picklable_after_use(file_format, pickle_module):
+  with tfds.testing.tmp_dir() as data_dir:
+    builder = tfds.testing.DummyDataset(data_dir=data_dir)
+    builder.download_and_prepare(file_format=file_format)
+    data_source = builder.as_data_source(split='train')
+    assert data_source[0] == {'id': 0}
+    assert pickle_module.loads(pickle_module.dumps(data_source))[0] == {'id': 0}

tensorflow_datasets/core/data_sources/parquet.py

@@ -57,7 +57,8 @@ class ParquetDataSource(base.BaseDataSource):
   """ParquetDataSource to read from a ParquetDataset."""
 
   def __post_init__(self):
-    file_instructions = base.file_instructions(self.dataset_info, self.split)
+    dataset_info = self.dataset_builder.info
+    file_instructions = base.file_instructions(dataset_info, self.split)
     filenames = [
         file_instruction.filename for file_instruction in file_instructions
     ]

tensorflow_datasets/core/dataset_builder.py

@@ -774,13 +774,13 @@ def build_single_data_source(
       file_format = self.info.file_format
       if file_format == file_adapters.FileFormat.ARRAY_RECORD:
         return array_record.ArrayRecordDataSource(
-            self.info,
+            self,
             split=split,
             decoders=decoders,
         )
       elif file_format == file_adapters.FileFormat.PARQUET:
         return parquet.ParquetDataSource(
-            self.info,
+            self,
             split=split,
             decoders=decoders,
         )

tensorflow_datasets/testing/mocking.py

@@ -83,13 +83,25 @@ class PickableDataSourceMock(mock.MagicMock):
   """Makes MagicMock pickable in order to work with multiprocessing in Grain."""
 
   def __getstate__(self):
-    return {'num_examples': len(self), 'generator': self._generator}
+    return {
+        'num_examples': len(self),
+        'generator': self._generator,
+        'serialize_example': self._serialize_example,
+    }
 
   def __setstate__(self, state):
-    num_examples, generator = state['num_examples'], state['generator']
+    num_examples, generator, serialize_example = (
+        state['num_examples'],
+        state['generator'],
+        state['serialize_example'],
+    )
     self.__len__.return_value = num_examples
-    self.__getitem__ = functools.partial(_getitem, generator=generator)
-    self.__getitems__ = functools.partial(_getitems, generator=generator)
+    self.__getitem__ = functools.partial(
+        _getitem, generator=generator, serialize_example=serialize_example
+    )
+    self.__getitems__ = functools.partial(
+        _getitems, generator=generator, serialize_example=serialize_example
+    )
 
   def __reduce__(self):
     return (PickableDataSourceMock, (), self.__getstate__())
@@ -99,50 +111,33 @@ def _getitem(
     self,
     record_key: int,
     generator: RandomFakeGenerator,
-    serialized: bool = False,
+    serialize_example=None,
 ) -> Any:
   """Function to overwrite __getitem__ in data sources."""
+  del self
   example = generator[record_key]
-  if serialized:
+  if serialize_example:
     # Return serialized raw bytes
-    return self.dataset_info.features.serialize_example(example)
+    return serialize_example(example)
   return example
 
 
 def _getitems(
     self,
     record_keys: Sequence[int],
     generator: RandomFakeGenerator,
-    serialized: bool = False,
+    serialize_example=None,
 ) -> Sequence[Any]:
   """Function to overwrite __getitems__ in data sources."""
   items = [
-      _getitem(self, record_key, generator, serialized=serialized)
+      _getitem(self, record_key, generator, serialize_example=serialize_example)
       for record_key in record_keys
   ]
-  if serialized:
+  if serialize_example:
     return np.array(items)
   return items
 
 
-def _deserialize_example_np(serialized_example, *, decoders=None):
-  """Function to overwrite dataset_info.features.deserialize_example_np.
-
-  Warning: this has to be defined in the outer scope in order for the function
-  to be pickable.
-
-  Args:
-    serialized_example: the example to deserialize.
-    decoders: optional decoders.
-
-  Returns:
-    The serialized example, because deserialization is taken care by
-      RandomFakeGenerator.
-  """
-  del decoders
-  return serialized_example
-
-
 class MockPolicy(enum.Enum):
   """Strategy to use with `tfds.testing.mock_data` to mock the dataset.
 
@@ -385,21 +380,27 @@ def mock_as_data_source(self, split, decoders=None, **kwargs):
         # Force ARRAY_RECORD as the default file_format.
         return_value=file_adapters.FileFormat.ARRAY_RECORD,
     ):
-      self.info.features.deserialize_example_np = _deserialize_example_np
+      # Make mock_data_source pickable with a given len:
       mock_data_source.return_value.__len__.return_value = num_examples
+      # Make mock_data_source pickable with a given generator:
       mock_data_source.return_value._generator = (  # pylint:disable=protected-access
           generator
       )
+      # Make mock_data_source pickable with a given serialize_example:
+      mock_data_source.return_value._serialize_example = (  # pylint:disable=protected-access
+          self.info.features.serialize_example
+      )
+      serialize_example = self.info.features.serialize_example
       mock_data_source.return_value.__getitem__ = functools.partial(
-          _getitem, generator=generator
+          _getitem, generator=generator, serialize_example=serialize_example
       )
       mock_data_source.return_value.__getitems__ = functools.partial(
-          _getitems, generator=generator
+          _getitems, generator=generator, serialize_example=serialize_example
       )
 
       def build_single_data_source(split):
         single_data_source = array_record.ArrayRecordDataSource(
-            dataset_info=self.info, split=split, decoders=decoders
+            dataset_builder=self, split=split, decoders=decoders
         )
         return single_data_source
 

tensorflow_datasets/testing/mocking_test.py

@@ -392,3 +392,12 @@ def test_as_data_source_fn():
     assert imagenet[0] == 'foo'
     assert imagenet[1] == 'bar'
     assert imagenet[2] == 'baz'
+
+
+# PyGrain requires that data sources are picklable.
+def test_mocked_data_source_is_pickable():
+  with tfds.testing.mock_data(num_examples=2):
+    data_source = tfds.data_source('imagenet2012', split='train')
+    pickled_and_unpickled_data_source = pickle.loads(pickle.dumps(data_source))
+    assert len(pickled_and_unpickled_data_source) == 2
+    assert isinstance(pickled_and_unpickled_data_source[0]['image'], np.ndarray)