grain.experimental.ElasticIterator

grain.experimental.ElasticIterator#

class grain.experimental.ElasticIterator(ds, global_batch_size, shard_options, *, read_options=ReadOptions(num_threads=16, prefetch_buffer_size=500), multiprocessing_options=None)#

Iterator supporting recovery from a checkpoint after changes in sharding.

The input dataset is expected to be unbatched and unsharded. In order to provide elasticity guarantee this iterator includes both, batching and sharding. The iterator supports elastic re-configuration by having each shard produce the same exact checkpoint (while producing different data) as long as they are advanced the same number of steps.

State of any shard can be used to restore the state of all of the shards after changes in sharding and global batch size.

This iterator explicitly disallows many-to-one transformations without a fixed ratio, like filter and generic IterDataset transformations.

Parameters:
__init__(ds, global_batch_size, shard_options, *, read_options=ReadOptions(num_threads=16, prefetch_buffer_size=500), multiprocessing_options=None)#
Parameters:

Methods

__init__(ds, global_batch_size, shard_options, *)

get_state()

Returns the current state of the iterator.

load(directory)

Loads the iterator state from a directory.

save(directory)

Saves the iterator state to a directory.

set_state(state)

Sets the current state of the iterator.

start_prefetch()

Asynchronously starts processing and buffering elements.
Contents