Function async_save_tf_savables

axlearn/common/checkpointer.py:203–229

Asynchronously saves TF savables from `value_map` into `dir`. When this call returns, `value_map` can be safely mutated, but saving to `dir` will not complete unless the returned future is set.

(
    value_map: Nested[Any], *, executor: futures.ThreadPoolExecutor, dir: str
)

# pylint: disable=redefined-builtin
def async_save_tf_savables(
    value_map: Nested[Any], *, executor: futures.ThreadPoolExecutor, dir: str
) -> futures.Future:
    """Asynchronously saves TF savables from `value_map` into `dir`.

    When this call returns, `value_map` can be safely mutated, but saving to `dir` will not
    complete unless the returned future is set.
    """
    # pylint: disable-next=consider-using-with
    f = tempfile.TemporaryDirectory()
    for path, value in utils.flatten_items(value_map):
        if isinstance(value, tf.data.Iterator):
            tf_checkpoint = tf.train.Checkpoint(value)
            tf_checkpoint.write(os.path.join(f.name, f"tf_{jax.process_index()}", path))
        elif isinstance(value, ElasticDatasetIterator):
            primary_tf_checkpoint = tf.train.Checkpoint(value.primary_iterator)
            primary_tf_checkpoint.write(os.path.join(f.name, f"tf_{jax.process_index()}", path))

            if value.elastic_iterator is not None:
                if value.is_primary_for_checkpoint:
                    for pid in value.elastic_process_ids:
                        elastic_tf_checkpoint = tf.train.Checkpoint(value.elastic_iterator)
                        elastic_tf_checkpoint.write(os.path.join(f.name, f"tf_{pid}", path))
        else:
            raise NotImplementedError(f"Unknown dataset iterator: {value}")

    return executor.submit(_upload_dir, f, dst_dir=dir)


# pylint: disable-next=redefined-builtin