Dataloader: access state from each worker

Sam_Shleifer1: <p>I have a <code>torch.utils.data.IterableDataset</code> dataset that loops over files and generates batches. Neither num files nor how many batches in each file are known ahead of time, hence the need for IterableDataset.</p>
<p>at the beginning of  <code>dataset.__iter__</code>,</p>
<p>I set self.worker_id, self.num_workers torch.utils.data.get_worker_info().id<br>
and use this information to split the files between workers, so that they don’t all just generate identical batches.</p>
<p>Each worker/dataset instance stores a file_idx and sample_idx so that they can be resumed if the job dies.</p>
<p>I want the dataloader to collect dictionaries of</p>
<pre><code class="lang-auto">{worker_id: dataset_instance.get_state()}
</code></pre>
<p>so I started passing worker states through Dataset.<strong>iter</strong> and collate_fn</p>
<p>but there is no guarantee that I have ever seen the state of every worker.</p>
<p>In my test, for example, it seems like every batch element of every batch is generated by worker 0, so the dataloader only persists the state of worker 0.</p>
<pre><code class="lang-auto">class StatefulDataLoader(torch.utils.data.DataLoader):
    def __init__(self, *args,  **kwargs):
        super().__init__(*args, **kwargs)
        self._batch_num = 0
        self.global_worker_states = {}

    def __iter__(self):

        iterator = super().__iter__()
        return self._wrap_iterator(iterator)

    def _wrap_iterator(self, iterator):
        for batch, worker_state in iterator:
            self.global_worker_states.update(worker_state)
            print(f'Batch: {self._batch_num}: GlobalWorkerStates: {self.global_worker_states}')
            yield batch, self.global_worker_states
            self._batch_num += 1
</code></pre>
<p>Is there a mechanism to either:<br>
a) guarantee that batches are composed of entries from every worker<br>
b) call <code>dataset.get_state()</code> on every worker from the dataloader<br>
c) A more general pattern to save the state of every worker instance</p>
<p>My high level goal of using extra workers is to prefetch batches and hide IO latency from the training process.</p>

andrewkho: <p>Hi Sam, please have a look at Stateful DataLoader in the torchdata repo. It’s a drop-in replacement for torch.utils.data.DataLoader and handles what you’re requesting <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<aside class="onebox githubfolder" data-onebox-src="https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h3><a href="https://github.com/pytorch/data/tree/main/torchdata/stateful_dataloader" target="_blank" rel="noopener nofollow ugc">data/torchdata/stateful_dataloader at main · pytorch/data</a></h3>


  <p><span class="label1">A PyTorch repo for data loading and utilities to be shared by the PyTorch domain libraries. - pytorch/data</span></p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>We are currently preparing the next release of torchdata 0.8.0 to follow pytorch 2.4.0 and it will include this new StatefulDataLoader. Until then, you can install from the Nightly or 0.8.0rcX if they are available.</p>