Method _check_iterator_saveable

axlearn/common/input_tf_data_test.py:497–510

(self, iterator: Iterable)

        self.assertEqual(num_batches, len(text_examples) // logical_feed_logical_batch_size)

    def _check_iterator_saveable(self, iterator: Iterable):
        # Check that we can save the data iterator.
        with tempfile.TemporaryDirectory() as td:
            save_dir = os.path.join(td, "ckpt")
            step = 100
            ckptr = (
                Checkpointer.default_config()
                .set(name="ckptr", dir=save_dir)
                .instantiate(parent=None)
            )
            with Mesh(jax.devices(), "data"):
                ckptr.save(step=step, state={"iterator": iterator}, evaler_summaries=None)
                ckptr.wait_until_finished()
            self.assertTrue(os.path.exists(os.path.join(save_dir, f"step_{step:08d}", "index")))

    @parameterized.product(
        num_physical_feeds=(2, 4),