(self, checkpointer_cls)
| 550 | |
| 551 | @parameterized.parameters([Checkpointer, OrbaxCheckpointer]) |
| 552 | def test_input_iterator(self, checkpointer_cls): |
| 553 | self.skipTest("TODO(mark-b-lee): figure out why it fails on CI.") |
| 554 | mesh_shape = (1, 1) |
| 555 | if not test_utils.is_supported_mesh_shape(mesh_shape): |
| 556 | return |
| 557 | with _mesh(mesh_shape): |
| 558 | cfg = _checkpointer_config(checkpointer_cls) |
| 559 | ckpt: Checkpointer = cfg.instantiate(parent=None) |
| 560 | input_iter = iter(tf.data.Dataset.from_tensor_slices([1, 2, 3])) |
| 561 | # Move the input_iter. |
| 562 | self.assertEqual(next(input_iter), 1) |
| 563 | state0 = dict( |
| 564 | x=jnp.zeros([], dtype=jnp.int32), |
| 565 | input_iter=input_iter, |
| 566 | ) |
| 567 | |
| 568 | self.assertEqual([], os.listdir(cfg.dir)) |
| 569 | |
| 570 | ckpt.save(step=100, state=state0) |
| 571 | ckpt.wait_until_finished() |
| 572 | |
| 573 | # Check that input iterators are saved under a per-worker path. |
| 574 | # E.g., /path/to/<step>/[state/]tf_0/input_iter.index. |
| 575 | state_dir = ckpt.ckpt_dir(100) |
| 576 | if "state" in os.listdir(state_dir): |
| 577 | state_dir = os.path.join(state_dir, "state") |
| 578 | self.assertIn("tf_0", os.listdir(state_dir)) |
| 579 | |
| 580 | state0_specs = dict( |
| 581 | x=utils.TensorSpec(shape=[], dtype=jnp.int32), |
| 582 | # The same iterator, but with the position at 0. |
| 583 | input_iter=iter(tf.data.Dataset.from_tensor_slices([1, 2, 3])), |
| 584 | ) |
| 585 | |
| 586 | def tensors_only(tree): |
| 587 | return ( |
| 588 | utils.prune_tree( |
| 589 | tree, should_prune=lambda _, v: not isinstance(v, utils.Tensor) |
| 590 | ), |
| 591 | ) |
| 592 | |
| 593 | step, restored_state = ckpt.restore(step=None, state=state0_specs) |
| 594 | self.assertEqual(100, step) |
| 595 | # The iterators will be different (despite pointing to the same values). |
| 596 | self.assertNestedEqual(tensors_only(state0), tensors_only(restored_state)) |
| 597 | # The restored_state contains the input_iter pointing to the next value. |
| 598 | self.assertEqual(next(restored_state["input_iter"]), 2) |
| 599 | self.assertEqual(next(restored_state["input_iter"]), 3) |
| 600 | ckpt.stop() |
| 601 | |
| 602 | @parameterized.parameters([Checkpointer, OrbaxCheckpointer]) |
| 603 | def test_python_savable(self, checkpointer_cls): |
nothing calls this directly
no test coverage detected