[RFC] Retool for iterator only

[andrewkho]

What would it look like to disallow generators?

Arrived at this PR after a couple of iterations:

BaseNodes that have .iterator() implemented as generators break down once we start to try multi-dataset sampling.

Example:

node = MultiDatasetMixer(
  datasets = {0: IterableWrapper(range(10)), 1: IterableWrapper(range(10))}, 
  weights = {0: 0.01, 1: 0.99},
  stop_criteria = "LOOP_UNTIL_ALL_ITERS_EXHAUSTED",
)
results = list(node)
state_dict_0 = node.state_dict()

# Start a new epoch
for idx, val in enumerate(node):
  if idx == 5:
    break

state_dict_1 = node.state_dict()

The behaviour of this node should be to mostly sample from dataset 1 (due to weight 0.99) and loop until we finally get 10 samples from dataset 0. Only once dataset 0 throws StopIteration, MultiDatasetMixer will also throw StopIteration.

What does state_dict_0 contain? It should contain the state of datasets 0 and 1, where 0 will be at the end of it's previous iteration.

Now what does state_dict_1 contain? It should contain the state of datasets 0 and 1, where 0 is almost certainly at or near the beginning of iteration. But it doesn't! It contains the state at the end of iteration. Loading this state_dict will return un-expected results.

Why is this the case? The problem is that Dataset 0's .iterator() method is defined by a generator. Even if you request another .iterator() from it, any initialization code does not get run until the first next() call occurs.

How can we deal with this and get the correct behaviour? I ran through a few scenarios and they all sound disgusting IMHO.

• Change IterableWrapper.iterator() to return an Iterator instead of a Generator. This doesn't work because of composability: all BaseNode subclasses would have to implement this. The only way this could work is if we "coloured" the nodes, so that certain nodes were only compatible with other nodes, or more realistically, to disallow generators entirely.
• add a has_next() or other data caching mechanism. Major downside: you now need to hold a batch from every source in memory, as well as deal with state either by holding a copy of the previous state_dict, or by storing the data itself in the state_dict.
• Ask generators implementations to handle re-initailization AFTER the loop is finished (before exiting .iterator()), but this is non-obvious, dependent on users doing the right thing, and also makes state management complicated to handle.

No Generators

Getting rid of generators has a major downside: it's a very Pythonic and way of writing and expressing iterators and we should be very wary of giving this up. Despite this massive downside, it's even more crucial for users to have a clear mental model of how Nodes should behave, and the fewer things to reason about, the easier this usually is.

Downsides of supporting generators:

• Two implementation paths to reason about: We need to support both Generators and Iterators, as Iterators.
• Explicitly Managing State: An Iterable that returns a Generator has its state managed handled implicitly by the runtime. Asking for a state_dict() from a generator isn't really possible, and instead we need to ask for state_dict() from the Iterable instance that defined the Generator. Users need to re-write their generators to use Instance-level variables
• No multi-iterator support: Having state managed in the Iterable instance member variables will break if someone ever creates multiple iterators from the same Iterable. This is a pretty uncommon scenario in data loading, and most folks would not be doing this, however it does force framework maintainers to make assumptions and reason about these scenarios, and how to catch/block them.
• Harder to reason about state: See above example, can we expect users and developers to reason about the current state of the iterator properly?

Iterators only

If we're going to disallow generators, we might as well get rid of the concept of Iterables as well, for the same reason as above (No multi-iterator support). Supporting both Iterables + Iterators gives us more surface area to cover:

• How would you trace the dag through inspection? From the Iterable or the iterator?
• From whom should you request the state, Iterable or Iterator?

The main downside of getting rid of Iterables is that Iterables give us a natural method of doing lazy initialization, and reseting iteration/getting a new iterator. With Iterables, you can delay expensive initialization until you call __iter__. With Iterators, you'd need to either make next lazy, or have the framework include some other call (we opt for the latter). With Iterables, getting a fresh iterator is as simple as calling __iter__. If there is tricky clean up or state to manage, it's easy to just throw away the old one and create a brand new Iterator. An Iterator only approach would require an explicit "reset" call or some other approach to handle multiple epochs.

We opt for the following API:

class BaseNode(Iterator[T]):
  def next(self): ... 
  def reset(self, initial_state: Optional[dict] = None): ...
  def get_state(self) -> dict: ...

Implementations need to define these three methods.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/data/1362

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 1fa022f with merge base 66a17ae ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build Linux Wheels / pytorch/data / upload / manywheel-py3_9-cpu (gh) (trunk failure)
  ##[error]Unable to find any artifacts for the associated workflow

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[andrewkho]

Updating based on feedback / ideation in #1363

[divyanshk]

@andrewkho is the problem you highlighted with generators also in the other samplers we already have , for eg in https://github.com/pytorch/data/blob/main/torchdata/stateful_dataloader/sampler.py#L63, https://github.com/pytorch/pytorch/blob/main/torch/utils/data/sampler.py#L288 ?

[divyanshk]

We don't force users to define reset?

[andrewkho]

We do, in a bit of roundabout way: if users don't implement .reset() (and call super().reset()) we throw an error during next() and state_dict() calls

[divyanshk]

Trying to understand the role of self.__initialized, when we initialize a node (and call super().init() in it), we set __initialized to false. but only when we reset we set it true

[andrewkho]

This is here to enable .reset() to contain initialization logic which users might want to be lazy initialized, because it's expensive and they eg want to load a state_dict before beginning iteration.

[andrewkho]

TODO: we should try and make single threaded mapper and ParallelMapper handle reset, similar to how persistent_workers currently works. However that's going to take a more significant refactor so leaving like this for now to minimize the changes

[andrewkho]

There might be a more-clever way to handle this, eg through a property or something, that won't make the type-hinters complain. Would prefer to do it in a separate diff as it's currently mostly sugar. Alternatively we can also assert self._it is not None to keep linters happy, but I don't like the idea of adding code to make linters happy

[divyanshk]

Lets gooo!!!

[andrewkho]

Linux builds failing are unrelated