Repository: horovod/horovod

Issue:
[Elastic Horovod]  It will loss some indices of processed samples in hvd.elastic.state when some nodes dropped
**Environment:**
1. Framework: PyTorch
2. Framework version: 1.7.0+cu101
3. Horovod version: 0.22.1

**Checklist:**
1. Did you search issues to find if somebody asked this question before?  
    yes
2. If your question is about hang, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/running.rst)?
    yes
3. If your question is about docker, did you read [this doc](https://github.com/horovod/horovod/blob/master/docs/docker.rst)? 
    yes
4. Did you check if you question is answered in the [troubleshooting guide](https://github.com/horovod/horovod/blob/master/docs/troubleshooting.rst)?
    yes

**Bug report:**
horovod/torch/elastic/state.py
```python
class SamplerStateHandler(StateHandler):
    def __init__(self, sampler):
        super().__init__(sampler)
        self._saved_sampler_state = copy.deepcopy(self.value.state_dict())

    def save(self):
        self._saved_sampler_state = copy.deepcopy(self.value.state_dict())

    def restore(self):
        self.value.load_state_dict(self._saved_sampler_state)

    def sync(self):
        # Get the set of processed indices from all workers
        world_processed_indices = _union(allgather_object(self.value.processed_indices))

        # Replace local processed indices with global indices
        state_dict = self.value.state_dict()
        state_dict['processed_indices'] = world_processed_indices

        # Broadcast and load the state to make sure we're all in sync
        self.value.load_state_dict(broadcast_object(state_dict))
```
when `state.commit()` is called, function `save()` above only save the local state of `ElasticSampler` locally. If one node is dropped by some reason, the indices of processed samples on this node is lost. So after restart and sync, those samples whill be processed again, which is not what we want.

**Steps to reproduce**
1. Add some log to `SamplerStateHandler.sync()`
```python
class SamplerStateHandler(StateHandler):
    ......
    def sync(self):
        # Get the set of processed indices from all workers
        world_processed_indices = _union(allgather_object(self.value.processed_indices))
        print(f"world_processed_indices: {world_processed_indices }")
        ......
```
2. Use the code below to reproduce. Note not do shuffle for the convenience of observation.
```python
#! /usr/bin/env python3
# -*- coding: utf-8 -*-


import time
import torch
import horovod.torch as hvd


BATCH_SIZE_PER_GPU = 2


class MyDataset(torch.utils.data.Dataset):
    def __init__(self, n):
        self.n = n
    
    def __getitem__(self, index):
        index = index % self.n
        return index

    def __len__(self):
        return self.n


@hvd.elastic.run
def train(state, data_loader, a):
    rank = hvd.rank()
    print(f"train rank={rank}")
    total_epoch = 100
    for epoch in range(state.epoch, total_epoch):
        print(f"epoch={epoch}")

        print("Epoch {} / {}, Start training".format(epoch, total_epoch))
        print(f"train... rank={rank}")
        print(f"start enumerate train_loader... rank={rank}")
        batch_offset = state.batch
        for i, d in enumerate(data_loader):
            state.batch = batch_idx = batch_offset + i
            if state.batch % 5 == 0:
                t1 = time.time()
                state.commit()
                print(f"time: {time.time() - t1}")
            state.check_host_updates()
            b = hvd.allreduce(a)
            print(f"b: {b}")
            state.train_sampler.record_batch(i, BATCH_SIZE_PER_GPU)
    
            # if rank == 0:
            msg = 'Epoch: [{0}][{1}/{2}]\t'.format(
                state.epoch, state.batch, len(data_loader))

            print(msg)
            time.sleep(0.5)
        state.epoch += 1
        state.batch = 0
        data_loader.sampler.set_epoch(epoch)
        state.commit()


def main():
    hvd.init()
    torch.manual_seed(219)
    torch.cuda.set_device(hvd.local_rank())

    dataset = MyDataset(2000)
    sampler = hvd.elastic.ElasticSampler(dataset, shuffle=False)
    data_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=BATCH_SIZE_PER_GPU,
        shuffle=False,
        num_workers=2,
        sampler=sampler,
        worker_init_fn=None,
        drop_last=True,
    )
    a = torch.Tensor([1,2,3,4])
    state = hvd.elastic.TorchState(epoch=0,
                                   train_sampler=sampler,
                                   batch=0)

    train(state, data_loader, a)


if __name__ == "__main__":
    main()

```
3. Use elastic horovod to run the above code on some nodes, for example, on 3 nodes.
4. Kill the processes on one node after a while.
5. Observe the log we added to `SamplerStateHandler.sync()`.

**Sloutions**
1. sloution 1
Save the global state. To get the global state of ElasticSampler on every node, function `save` should call function `sync` first.
```python
class SamplerStateHandler(StateHandler):
    ......
    def save(self):
        self.sync()
        self._saved_sampler_state = copy.deepcopy(self.value.state_dict())
    ......
```
But this will  cause `state.commit()` to take a long time.
2. sloution 2
Maybe we can save the number of processed samples instead of save all the processed indices. The number of processed samples can be calculated locally by `batch_size` and `num_replicas`.



Discussion:

Resolving patch:
diff --git a/horovod/torch/elastic/sampler.py b/horovod/torch/elastic/sampler.py
--- a/horovod/torch/elastic/sampler.py
+++ b/horovod/torch/elastic/sampler.py
@@ -32,7 +32,7 @@ class ElasticSampler(torch.utils.data.Sampler):
     In order to use this object successfully it is recommended that the user:
 
     1. Include this object in the `TorchState`.
-    2. Call `record_batch` or `record_indices` after processing a set of samples.
+    2. Call `record_batch` after processing a set of samples.
     3. Call `set_epoch` at the end of each epoch to clear the processed indices.
 
     Args:
@@ -54,6 +54,7 @@ def __init__(self, dataset, shuffle=True, seed=0):
         self.remaining_indices = []
         self.num_samples = 0
         self.total_size = 0
+        self.processed_num = 0
 
         self.reset()
 
@@ -71,33 +72,22 @@ def set_epoch(self, epoch):
             epoch: Epoch number.
         """
         self.epoch = epoch
-        self.processed_indices = set()
+        self.processed_num = 0
         self.reset()
 
     def record_batch(self, batch_idx, batch_size):
-        """Record indices at batch `batch_idx` with length `batch_size` as processed."""
-        indices = set(self.get_indices(batch_idx, batch_size))
-        self.record_indices(indices)
-
-    def record_indices(self, indices):
-        """Record set `indices` as processed."""
-        self.processed_indices.update(indices)
-
-    def get_indices(self, batch_idx, batch_size):
-        """Return list of indices at batch `batch_idx` with length `batch_size`."""
-        start_idx = batch_idx * batch_size
-        end_idx = min(start_idx + batch_size, len(self.indices))
-        return self.indices[start_idx:end_idx]
+        """Record the number of processed samples."""
+        self.processed_num += (batch_size * self.num_replicas)
 
     def load_state_dict(self, state_dict):
         self.epoch = state_dict['epoch']
-        self.processed_indices = state_dict['processed_indices']
+        self.processed_num = state_dict["processed_num"]
         self.reset()
 
     def state_dict(self):
         return dict(
             epoch=self.epoch,
-            processed_indices=self.processed_indices
+            processed_num=self.processed_num
         )
 
     def reset(self):
@@ -105,18 +95,18 @@ def reset(self):
         self.rank = rank()
 
         # Exclude any samples we have already processed this epoch
-        self.remaining_indices = [idx for idx in range(len(self.dataset))
-                                  if idx not in self.processed_indices]
+        all_indices = [idx for idx in range(len(self.dataset))]
+        if self.shuffle:
+            # Shuffle indices across workers deterministically in place
+            seed = self.seed + self.epoch
+            random.Random(seed).shuffle(all_indices)
+        self.remaining_indices = all_indices[self.processed_num:]
 
         self.num_samples = int(math.ceil(len(self.remaining_indices) * 1.0 / self.num_replicas))
         self.total_size = self.num_samples * self.num_replicas
 
     def __iter__(self):
         self.indices = self.remaining_indices[:]
-        if self.shuffle:
-            # Shuffle indices across workers deterministically in place
-            seed = self.seed + self.epoch
-            random.Random(seed).shuffle(self.indices)
 
         # add extra samples to make it evenly divisible
         self.indices += self.indices[:(self.total_size - len(self.indices))]
diff --git a/horovod/torch/elastic/state.py b/horovod/torch/elastic/state.py
--- a/horovod/torch/elastic/state.py
+++ b/horovod/torch/elastic/state.py
@@ -128,12 +128,7 @@ def restore(self):
         self.value.load_state_dict(self._saved_sampler_state)
 
     def sync(self):
-        # Get the set of processed indices from all workers
-        world_processed_indices = _union(allgather_object(self.value.processed_indices))
-
-        # Replace local processed indices with global indices
         state_dict = self.value.state_dict()
-        state_dict['processed_indices'] = world_processed_indices
 
         # Broadcast and load the state to make sure we're all in sync
         self.value.load_state_dict(broadcast_object(state_dict))


Test patch:
diff --git a/test/single/test_torch_elastic.py b/test/single/test_torch_elastic.py
--- a/test/single/test_torch_elastic.py
+++ b/test/single/test_torch_elastic.py
@@ -112,7 +112,7 @@ def __len__(self):
         state.sync()
 
         assert state.sampler.epoch == 0
-        assert len(state.sampler.processed_indices) == 0
+        assert state.sampler.processed_num == 0
 
         # Normal usage, no errors
         epochs = 2
@@ -120,12 +120,8 @@ def __len__(self):
         for epoch in range(epochs):
             sampler.set_epoch(epoch)
             for batch_idx, batch in enumerate(data_loader):
-                batch_indices = sampler.get_indices(batch_idx, batch_size)
-                batch_data = [dataset[idx] for idx in batch_indices]
-                assert batch_data == batch.numpy().tolist()
-
                 sampler.record_batch(batch_idx, batch_size)
-                assert len(sampler.processed_indices) == batch_size * (batch_idx + 1)
+                assert sampler.processed_num == batch_size * (batch_idx + 1)
 
                 total_batches += 1
         assert total_batches == (samples_per_worker / batch_size) * epochs
@@ -133,47 +129,44 @@ def __len__(self):
         # Do not reset epoch: processed samples are retained and data loader repeats
         total_batches = 0
         for _ in enumerate(data_loader):
-            assert len(sampler.processed_indices) == len(sampler)
+            assert sampler.processed_num == len(sampler)
             total_batches += 1
         assert total_batches == samples_per_worker / batch_size
 
         # Elastic: partial epoch + commit
         sampler.set_epoch(2)
-        assert len(sampler.processed_indices) == 0
+        assert sampler.processed_num == 0
 
         sampler.record_batch(0, batch_size)
         sampler.record_batch(1, batch_size)
-        assert len(sampler.processed_indices) == 2 * batch_size
+        assert sampler.processed_num == 2 * batch_size
 
-        committed_indices = copy.copy(sampler.processed_indices)
+        committed_num = copy.copy(sampler.processed_num)
         state.commit()
 
         # Elastic: partial epoch + restore
         sampler.record_batch(2, batch_size)
         sampler.record_batch(3, batch_size)
-        assert len(sampler.processed_indices) == 4 * batch_size
+        assert sampler.processed_num == 4 * batch_size
 
         state.restore()
 
-        assert len(sampler.processed_indices) == 2 * batch_size
-        assert sampler.processed_indices == committed_indices
+        assert sampler.processed_num == 2 * batch_size
+        assert sampler.processed_num == committed_num
 
         # Elastic: sync across workers and verify non-overlap of processed samples
         sampler.record_batch(2, batch_size)
-        assert len(sampler.processed_indices) == 3 * batch_size
+        assert sampler.processed_num == 3 * batch_size
 
         state.commit()
         state.sync()
 
-        assert len(sampler.processed_indices) == 3 * batch_size * hvd.size()
+        assert sampler.processed_num == 3 * batch_size * hvd.size()
 
         # After the sync, the remaining indices should be updated and repartitioned
         total_batches = 0
         assert len(sampler) == batch_size
         for batch_idx, batch in enumerate(data_loader):
-            batch_indices = sampler.get_indices(batch_idx, batch_size)
-            overlap_indices = set(batch_indices) & sampler.processed_indices
-            assert overlap_indices == set()
             total_batches += 1
         assert total_batches == 1