Make DatasetInfo pickable so that tfds.data_source works in multiprocessing externally.

PiperOrigin-RevId: 638272983

Author: marcenacp

tensorflow_datasets/core/data_sources/base_test.py

@@ -15,8 +15,10 @@
 
 """Tests for all data sources."""
 
+import pickle
 from unittest import mock
 
+import cloudpickle
 from etils import epath
 import pytest
 import tensorflow_datasets as tfds
@@ -181,3 +183,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/dataset_info.py

@@ -173,6 +173,7 @@ class DatasetInfo(object):
   """
 
   def __init__(
+      # LINT.IfChange(dataset_info_args)
       self,
       *,
       builder: Union[DatasetIdentity, Any],
@@ -186,6 +187,7 @@ def __init__(
       license: Optional[str] = None,  # pylint: disable=redefined-builtin
       redistribution_info: Optional[Dict[str, str]] = None,
       split_dict: Optional[splits_lib.SplitDict] = None,
+      # LINT.ThenChange(:setstate)
   ):
     # pyformat: disable
     """Constructs DatasetInfo.
@@ -864,6 +866,35 @@ def __repr__(self):
     lines.append(")")
     return "\n".join(lines)
 
+  def __getstate__(self):
+    return {
+        "builder": self._builder_or_identity,
+        "description": self.description,
+        "features": self.features,
+        "supervised_keys": self.supervised_keys,
+        "disable_shuffling": self.disable_shuffling,
+        "homepage": self.homepage,
+        "citation": self.citation,
+        "metadata": self.metadata,
+        "license": self.redistribution_info.license,
+        "split_dict": self.splits,
+    }
+  def __setstate__(self, state):
+    # LINT.IfChange(setstate)
+    self.__init__(
+        builder=state["builder"],
+        description=state["description"],
+        features=state["features"],
+        supervised_keys=state["supervised_keys"],
+        disable_shuffling=state["disable_shuffling"],
+        homepage=state["homepage"],
+        citation=state["citation"],
+        metadata=state["metadata"],
+        license=state["license"],
+        split_dict=state["split_dict"],
+    )
+    # LINT.ThenChange(:dataset_info_args)
+
 
 def _nest_to_proto(nest: Nest) -> dataset_info_pb2.SupervisedKeys.Nest:
   """Creates a `SupervisedKeys.Nest` from a limited `tf.nest` style structure.

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)