Add Grain iterator checkpoint/resume and fix num_batches

Enable deterministic mid-epoch resume for Grain datasets by saving and
restoring the DatasetIterator state through BackupAndRestore. Also fix
num_batches to return the actual count for finite MapDatasets so progress
bars work correctly.

Author: MarcosAsh

keras/src/backend/jax/trainer.py

@@ -418,6 +418,10 @@ def fit(
         self._symbolic_build(iterator=epoch_iterator)
         epoch_iterator.reset()
 
+        # Expose the iterator so callbacks (e.g. BackupAndRestore) can
+        # save / restore data-pipeline state for fault tolerance.
+        self._epoch_iterator = epoch_iterator
+
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
             callbacks = callbacks_module.CallbackList(
@@ -541,6 +545,7 @@ def fit(
             # are done.
             if getattr(self, "_eval_epoch_iterator", None) is not None:
                 del self._eval_epoch_iterator
+            self._epoch_iterator = None
             if training_finished:
                 callbacks.on_train_end(logs=training_logs)
             self._jax_state = None

keras/src/backend/tensorflow/trainer.py

@@ -372,6 +372,10 @@ def fit(
         self._maybe_symbolic_build(iterator=epoch_iterator)
         epoch_iterator.reset()
 
+        # Expose the iterator so callbacks (e.g. BackupAndRestore) can
+        # save / restore data-pipeline state for fault tolerance.
+        self._epoch_iterator = epoch_iterator
+
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
             callbacks = callbacks_module.CallbackList(
@@ -449,6 +453,7 @@ def fit(
         # If _eval_epoch_iterator exists, delete it after all epochs are done.
         if getattr(self, "_eval_epoch_iterator", None) is not None:
             del self._eval_epoch_iterator
+        self._epoch_iterator = None
         callbacks.on_train_end(logs=training_logs)
         return self.history
 

keras/src/backend/torch/trainer.py

@@ -236,6 +236,10 @@ def fit(
         self._symbolic_build(iterator=epoch_iterator)
         epoch_iterator.reset()
 
+        # Expose the iterator so callbacks (e.g. BackupAndRestore) can
+        # save / restore data-pipeline state for fault tolerance.
+        self._epoch_iterator = epoch_iterator
+
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
             callbacks = callbacks_module.CallbackList(
@@ -324,6 +328,7 @@ def fit(
         # If _eval_epoch_iterator exists, delete it after all epochs are done.
         if getattr(self, "_eval_epoch_iterator", None) is not None:
             del self._eval_epoch_iterator
+        self._epoch_iterator = None
         callbacks.on_train_end(logs=training_logs)
         return self.history
 

keras/src/callbacks/backup_and_restore.py

@@ -151,6 +151,14 @@ def _load_model(self):
             epoch = training_metadata["epoch"]
             self.model._initial_epoch = epoch
 
+            # Restore data-pipeline iterator state when available (e.g.
+            # Grain datasets support deterministic mid-epoch resume).
+            iterator_state = training_metadata.get("iterator_state")
+            if iterator_state is not None:
+                epoch_iterator = getattr(self.model, "_epoch_iterator", None)
+                if epoch_iterator is not None:
+                    epoch_iterator.set_iterator_state(iterator_state)
+
     def on_epoch_end(self, epoch, logs=None):
         self._current_epoch = epoch + 1
         self._last_batch_seen = 0
@@ -187,6 +195,13 @@ def _save_model(self):
                 "epoch": self._current_epoch,
                 "batch": self._last_batch_seen,
             }
+            # Persist data-pipeline iterator state when the adapter
+            # supports it (e.g. Grain).
+            epoch_iterator = getattr(self.model, "_epoch_iterator", None)
+            if epoch_iterator is not None:
+                iterator_state = epoch_iterator.get_iterator_state()
+                if iterator_state is not None:
+                    training_metadata["iterator_state"] = iterator_state
             f.write(json.dumps(training_metadata))
 
     def _should_save_on_batch(self, batch):

keras/src/trainers/data_adapters/data_adapter.py

@@ -88,6 +88,24 @@ def partial_batch_size(self):
         """
         raise NotImplementedError
 
+    def get_iterator_state(self):
+        """Return serializable state for the current data iterator.
+
+        Adapters that support deterministic checkpoint/resume (e.g. Grain)
+        override this to return a small dict that can reconstruct the
+        iterator position. The default returns ``None`` (not supported).
+        """
+        return None
+
+    def set_iterator_state(self, state):
+        """Restore the data iterator to a previously saved state.
+
+        Called before the next ``iter()`` call so the iterator resumes
+        from the saved position. Adapters that do not support
+        checkpointing ignore this (the default is a no-op).
+        """
+        pass
+
     def on_epoch_begin(self):
         """A hook called before each epoch."""
         pass

keras/src/trainers/data_adapters/grain_dataset_adapter.py

@@ -1,4 +1,5 @@
 import itertools
+import sys
 
 import numpy as np
 
@@ -9,17 +10,43 @@
 from keras.src.utils.module_utils import tensorflow as tf
 
 
+class _TrackableIterable:
+    """Wrapper that captures the live ``DatasetIterator`` on ``iter()``.
+
+    When the ``EpochIterator`` calls ``iter()`` on the object returned by
+    ``get_numpy_iterator()`` / ``get_jax_iterator()``, this wrapper
+    stores the resulting iterator on the adapter so that
+    ``get_iterator_state()`` can reach it.  If a pending state was
+    previously set via ``set_iterator_state()``, it is applied to the
+    fresh iterator immediately.
+    """
+
+    def __init__(self, dataset, adapter):
+        self._dataset = dataset
+        self._adapter = adapter
+
+    def __iter__(self):
+        it = iter(self._dataset)
+        self._adapter._live_iterator = it
+        if self._adapter._pending_iterator_state is not None:
+            if hasattr(it, "set_state"):
+                it.set_state(self._adapter._pending_iterator_state)
+            self._adapter._pending_iterator_state = None
+        return it
+
+
 class GrainDatasetAdapter(DataAdapter):
-    """Adapter that handles `grain.DataLoader`, `grain.MapDataset` and
-    `grain.IterDataset`.
+    """Adapter that handles ``grain.DataLoader``, ``grain.MapDataset`` and
+    ``grain.IterDataset``.
     """
 
     def __init__(self, dataset):
         """Initialize the GrainDatasetAdapter.
 
         Args:
             dataset: A Grain dataset instance. Must be one of
-                `grain.DataLoader`, `grain.MapDataset`, or `grain.IterDataset`.
+                ``grain.DataLoader``, ``grain.MapDataset``, or
+                ``grain.IterDataset``.
         """
 
         if not isinstance(
@@ -32,17 +59,19 @@ def __init__(self, dataset):
             )
 
         self._dataset = dataset
+        self._live_iterator = None
+        self._pending_iterator_state = None
 
         batch_size, output_signature = self._get_dataset_info(dataset)
         self._batch_size = batch_size
         self._output_signature = output_signature
         self._output_tf_signature = None
 
     def _get_dataset_info(self, dataset):
-        """Get the `batch_size` and `output_signature` from the dataset.
+        """Get the ``batch_size`` and ``output_signature`` from the dataset.
 
-        We use a small list of batches to infer the `batch_size` and
-        `output_signature`.
+        We use a small list of batches to infer the ``batch_size`` and
+        ``output_signature``.
         """
         batches = list(
             itertools.islice(
@@ -73,9 +102,9 @@ def convert_to_numpy(x):
             if isinstance(x, (np.ndarray, SharedMemoryArrayMetadata)):
                 return x
             else:
-                # Using `__array__` should handle `tf.Tensor`, `jax.np.ndarray`,
-                # `torch.Tensor`, as well as any other tensor-like object that
-                # has added numpy support.
+                # Using ``__array__`` should handle ``tf.Tensor``,
+                # ``jax.np.ndarray``, ``torch.Tensor``, as well as any
+                # other tensor-like object that has added numpy support.
                 if hasattr(x, "__array__"):
                     if data_adapter_utils.is_torch_tensor(x):
                         x = x.cpu()
@@ -90,20 +119,21 @@ def map(self, x):
 
         if isinstance(self._dataset, (grain.MapDataset, grain.IterDataset)):
             dataset = self._dataset.map(ConvertToNumpy())
+            return _TrackableIterable(dataset, self)
         else:
-            # Instantiate a new `DataLoader`.
+            # Instantiate a new ``DataLoader``.
             dataset = grain.DataLoader(
                 data_source=self._dataset._data_source,
                 sampler=self._dataset._sampler,
-                # Append `ConvertToNumpy`.
+                # Append ``ConvertToNumpy``.
                 operations=list(self._dataset._operations) + [ConvertToNumpy()],
                 worker_count=self._dataset._multiprocessing_options.num_workers,
                 worker_buffer_size=self._dataset._multiprocessing_options.per_worker_buffer_size,
                 shard_options=self._dataset._shard_options,
                 read_options=self._dataset._read_options,
                 enable_profiling=self._dataset._multiprocessing_options.enable_profiling,
             )
-        return dataset
+            return dataset
 
     def get_jax_iterator(self):
         def convert_to_jax_compatible(x):
@@ -121,12 +151,13 @@ def map(self, x):
 
         if isinstance(self._dataset, (grain.MapDataset, grain.IterDataset)):
             dataset = self._dataset.map(ConvertToJaxCompatible())
+            return _TrackableIterable(dataset, self)
         else:
-            # Instantiate a new `DataLoader`.
+            # Instantiate a new ``DataLoader``.
             dataset = grain.DataLoader(
                 data_source=self._dataset._data_source,
                 sampler=self._dataset._sampler,
-                # Append `ConvertToJaxCompatible`.
+                # Append ``ConvertToJaxCompatible``.
                 operations=list(self._dataset._operations)
                 + [ConvertToJaxCompatible()],
                 worker_count=self._dataset._multiprocessing_options.num_workers,
@@ -135,7 +166,7 @@ def map(self, x):
                 read_options=self._dataset._read_options,
                 enable_profiling=self._dataset._multiprocessing_options.enable_profiling,
             )
-        return dataset
+            return dataset
 
     def get_tf_dataset(self):
         def convert_to_tf(x):
@@ -151,7 +182,7 @@ class ConvertToTF(grain.transforms.Map):
             def map(self, x):
                 return tree.map_structure(convert_to_tf, x)
 
-        # `tf.data.Dataset.from_generator` does not support lists as output.
+        # ``tf.data.Dataset.from_generator`` does not support lists as output.
         # We convert lists to tuples.
         class ListToTuple(grain.transforms.Map):
             def map(self, x):
@@ -161,11 +192,11 @@ def map(self, x):
             dataset = self._dataset.map(ConvertToTF())
             dataset = dataset.map(ListToTuple())
         else:
-            # Instantiate a new `DataLoader`.
+            # Instantiate a new ``DataLoader``.
             dataset = grain.DataLoader(
                 data_source=self._dataset._data_source,
                 sampler=self._dataset._sampler,
-                # Append `ConvertToTF` and `ListToTuple`.
+                # Append ``ConvertToTF`` and ``ListToTuple``.
                 operations=list(self._dataset._operations)
                 + [ConvertToTF(), ListToTuple()],
                 worker_count=self._dataset._multiprocessing_options.num_workers,
@@ -196,13 +227,46 @@ def __init__(self, iterable):
             def __iter__(self):
                 return iter(self.iterable)
 
-        # `batch_size=None` indicates that we should not re-batch
+        if isinstance(self._dataset, (grain.MapDataset, grain.IterDataset)):
+            iterable = _TrackableIterable(self._dataset, self)
+        else:
+            iterable = self._dataset
+
+        # ``batch_size=None`` indicates that we should not re-batch
         return torch_data.DataLoader(
-            ConverterIterableDataset(self._dataset), batch_size=None
+            ConverterIterableDataset(iterable), batch_size=None
         )
 
+    # ------------------------------------------------------------------
+    # Iterator checkpoint / resume
+    # ------------------------------------------------------------------
+
+    def get_iterator_state(self):
+        if self._live_iterator is not None and hasattr(
+            self._live_iterator, "get_state"
+        ):
+            return self._live_iterator.get_state()
+        return None
+
+    def set_iterator_state(self, state):
+        if state is not None:
+            self._pending_iterator_state = state
+
+    # ------------------------------------------------------------------
+    # Metadata
+    # ------------------------------------------------------------------
+
     @property
     def num_batches(self):
+        if isinstance(self._dataset, grain.MapDataset):
+            try:
+                length = len(self._dataset)
+            except TypeError:
+                return None
+            # ``repeat(None)`` sets length to ``sys.maxsize``.
+            if length >= sys.maxsize:
+                return None
+            return length
         return None
 
     @property