Warmup input pipeline in the background.

clu/data/__init__.py

@@ -12,10 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Common utilities for loading data for JAX users.
+"""DatasetIterator is an interface for input pipelines."""
+# pylint: disable=g-multiple-import
+# pylint: disable=unused-import
 
-This will long term replace clu.deterministic_data.
-"""
-
-from clu.data.dataset_iterator import DatasetIterator
-from clu.data.dataset_iterator import TfDatasetIterator
+from clu.data.dataset_iterator import (
+    Array,
+    ArraySpec,
+    DatasetIterator,
+    Element,
+    ElementSpec,
+    TfDatasetIterator,
+    PeekableDatasetIterator,
+    PyTree,
+)

clu/data/dataset_iterator.py

@@ -29,12 +29,14 @@
 
 import abc
 import collections.abc
+import concurrent.futures
 import dataclasses
 import os
 import typing
 from typing import Any, Mapping, Optional, Sequence, Tuple, TypeVar, Union
 
 from absl import logging
+from clu import asynclib
 from etils import epath
 import jax.numpy as jnp  # Just for type checking.
 import numpy as np
@@ -227,8 +229,12 @@ class PeekableDatasetIterator(DatasetIterator):
   def __init__(self, it: DatasetIterator):
     self._it = it
     self._peek: Optional[Element] = None
+    self._peek_future = None
 
   def __next__(self) -> Element:
+    if self._peek_future:
+      self._peek_future.result()
+      self._peek_future = None
     if self._peek is None:
       return next(self._it)
     peek = self._peek
@@ -243,10 +249,36 @@ def element_spec(self) -> ElementSpec:
     return self._it.element_spec
 
   def peek(self) -> Element:
+    """Returns the next element without consuming it.
+
+    This will get the next element from the underlying iterator. The element
+    is stored and return on the next call of __next__().
+
+    Returns:
+      The next element.
+    """
     if self._peek is None:
       self._peek = next(self)
     return self._peek
 
+  def peek_async(self) -> concurrent.futures.Future[Element]:
+    """Same as peek() but returns the Future of the element.
+
+    Users can call this to warm up the iterator.
+
+    Returns:
+      Future with the next element. The element is also kept and returned on the
+      next call of __next__().
+    """
+    pool = asynclib.Pool()
+
+    @pool
+    def async_peek():
+      return self.peek()
+
+    self._peek_future = async_peek()
+    return self._peek_future
+
   def save(self, filename: epath.Path):
     self._it.save(filename)
 

clu/data/dataset_iterator_test.py

@@ -86,6 +86,14 @@ def test_peekable_dataset_iterator(self):
     self.assertEqual(next(it), {INDEX: [0, 2], "prime": [2, 5]})
     self.assertEqual(next(it), {INDEX: [4, 5], "prime": [11, 13]})
 
+  def test_peekable_dataset_iterator_async(self):
+    it = self._create_iterator(0)
+    it = dataset_iterator.PeekableDatasetIterator(it)
+    it.peek_async()
+    self.assertEqual(next(it), {INDEX: [0, 2], "prime": [2, 5]})
+    self.assertEqual(next(it), {INDEX: [0, 2], "prime": [2, 5]})
+    self.assertEqual(next(it), {INDEX: [4, 5], "prime": [11, 13]})
+
 
 if __name__ == "__main__":
   tf.test.main()