Fix for concat map dataset
State: closed | Pull request: yes

# What does this PR do ?
Fix for concat map dataset #5117

Only caveat I know of is that it will not reshuffle each epoch when using round robin. I'm not sure if anyone actually uses round robin in practice, though.

# Usage

```python

from collections import defaultdict

import torch
from pytorch_lightning import seed_everything

from nemo.collections.common.data import ConcatMapDataset

seed_everything(42)


class IntegerDataset(torch.utils.data.Dataset):
    def __init__(self, n: int, dataset_id: int):
        self._ints = list(range(n))
        self._id = torch.tensor(dataset_id)

    def __len__(self):
        return len(self._ints)

    def __getitem__(self, idx):
        worker_info = torch.utils.data.get_worker_info()
        worker_id = worker_info.id
        return self._id, worker_id, self._ints[idx]


for sampling_technique in ["temperature", "random", "round-robin"]:
    probs = [0.1, 0.2, 0.3]
    concat_dataset = ConcatMapDataset(
        datasets=[
            IntegerDataset(5, 0),
            IntegerDataset(10, 1),
            IntegerDataset(20, 2),
        ],
        seed=123456,
        sampling_temperature=2,
        sampling_probabilities=probs,
        sampling_technique=sampling_technique
    )

    dataloader = torch.utils.data.DataLoader(
        dataset=concat_dataset,
        batch_size=1,
        num_workers=2,
        shuffle=False  # Should be False to preserve round robin
    )

    print(f"Item-by-item results for technique '{sampling_technique}':")
    collected = defaultdict(list)
    for batch in dataloader:
        batch_dataset_id, batch_worker_id, batch_value = batch
        print(
            f"Dataset {batch_dataset_id.item()} worker {batch_worker_id.item()} fetched {batch_value.item()}"
        )
        collected[batch_dataset_id.item()].append(batch_value.item())

    print("*" * 80)
    print(f"Summary of returned results for each dataset for technique '{sampling_technique}'")
    if sampling_technique == "random":
        print(f"\t{probs=}")
    for dataset_id, values in sorted(collected.items(), key=lambda x: x[0]):
        dataset_raw_length = len(concat_dataset.datasets[dataset_id])
        print(f"\tDataset #{dataset_id} (raw length {dataset_raw_length}): {values}")
    print("*" * 80)
    print()
```

Output:

```
Item-by-item results for technique 'temperature':
Global seed set to 42
Dataset 2 worker 0 fetched 0
Dataset 2 worker 1 fetched 12
Dataset 1 worker 0 fetched 3
Dataset 1 worker 1 fetched 1
Dataset 0 worker 0 fetched 0
Dataset 2 worker 1 fetched 18
Dataset 1 worker 0 fetched 5
Dataset 0 worker 1 fetched 4
Dataset 0 worker 0 fetched 3
Dataset 2 worker 1 fetched 1
Dataset 0 worker 0 fetched 2
Dataset 0 worker 1 fetched 1
Dataset 2 worker 0 fetched 3
Dataset 1 worker 1 fetched 6
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 2
Dataset 2 worker 0 fetched 5
Dataset 0 worker 1 fetched 3
Dataset 1 worker 0 fetched 9
Dataset 0 worker 1 fetched 0
Dataset 2 worker 0 fetched 9
Dataset 2 worker 1 fetched 8
Dataset 2 worker 0 fetched 13
Dataset 2 worker 1 fetched 7
Dataset 1 worker 0 fetched 4
Dataset 2 worker 1 fetched 19
Dataset 2 worker 0 fetched 4
Dataset 2 worker 1 fetched 16
Dataset 2 worker 0 fetched 10
Dataset 2 worker 1 fetched 17
Dataset 0 worker 0 fetched 2
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 7
Dataset 1 worker 1 fetched 0
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 8
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 5
Dataset 2 worker 0 fetched 14
Dataset 2 worker 1 fetched 2
Dataset 1 worker 0 fetched 0
Dataset 2 worker 1 fetched 6
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 15
Dataset 0 worker 0 fetched 3
Dataset 2 worker 1 fetched 11
********************************************************************************
Summary of returned results for each dataset for technique 'temperature'
	Dataset #0; raw length 5; num drawn = 13; values = [0, 4, 3, 2, 1, 1, 3, 0, 2, 4, 1, 2, 3]
	Dataset #1; raw length 10; num drawn = 13; values = [3, 1, 5, 6, 2, 9, 4, 7, 0, 8, 5, 0, 7]
	Dataset #2; raw length 20; num drawn = 20; values = [0, 12, 18, 1, 3, 5, 9, 8, 13, 7, 19, 4, 16, 10, 17, 14, 2, 6, 15, 11]
********************************************************************************

Item-by-item results for technique 'random':
Dataset 2 worker 0 fetched 0
Dataset 2 worker 1 fetched 12
Dataset 2 worker 0 fetched 18
Dataset 1 worker 1 fetched 3
Dataset 0 worker 0 fetched 0
Dataset 2 worker 1 fetched 1
Dataset 1 worker 0 fetched 1
Dataset 0 worker 1 fetched 4
Dataset 0 worker 0 fetched 3
Dataset 2 worker 1 fetched 3
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 5
Dataset 1 worker 0 fetched 6
Dataset 2 worker 1 fetched 5
Dataset 1 worker 0 fetched 2
Dataset 2 worker 1 fetched 9
Dataset 2 worker 0 fetched 8
Dataset 2 worker 1 fetched 13
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 9
Dataset 1 worker 0 fetched 4
Dataset 0 worker 1 fetched 0
Dataset 2 worker 0 fetched 7
Dataset 2 worker 1 fetched 19
Dataset 2 worker 0 fetched 4
Dataset 2 worker 1 fetched 16
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 10
Dataset 2 worker 0 fetched 17
Dataset 2 worker 1 fetched 14
Dataset 2 worker 0 fetched 2
Dataset 2 worker 1 fetched 6
Dataset 0 worker 0 fetched 3
Dataset 0 worker 1 fetched 2
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 0
Dataset 1 worker 0 fetched 8
Dataset 1 worker 1 fetched 8
Dataset 0 worker 0 fetched 4
Dataset 1 worker 1 fetched 7
Dataset 2 worker 0 fetched 15
Dataset 2 worker 1 fetched 11
********************************************************************************
Summary of returned results for each dataset for technique 'random'
	probs=[0.1, 0.2, 0.3]
	Dataset #0; raw length 5; num drawn = 10; values = [0, 4, 3, 2, 1, 0, 3, 2, 1, 4]
	Dataset #1; raw length 10; num drawn = 12; values = [3, 1, 5, 6, 2, 9, 4, 7, 0, 8, 8, 7]
	Dataset #2; raw length 20; num drawn = 20; values = [0, 12, 18, 1, 3, 5, 9, 8, 13, 7, 19, 4, 16, 10, 17, 14, 2, 6, 15, 11]
********************************************************************************

Item-by-item results for technique 'round-robin':
Dataset 0 worker 0 fetched 0
Dataset 1 worker 1 fetched 3
Dataset 2 worker 0 fetched 0
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 1
Dataset 2 worker 1 fetched 12
Dataset 0 worker 0 fetched 3
Dataset 1 worker 1 fetched 5
Dataset 2 worker 0 fetched 18
Dataset 0 worker 1 fetched 2
Dataset 1 worker 0 fetched 6
Dataset 2 worker 1 fetched 1
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 2
Dataset 2 worker 0 fetched 3
Dataset 0 worker 1 fetched 0
Dataset 1 worker 0 fetched 9
Dataset 2 worker 1 fetched 5
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 4
Dataset 2 worker 0 fetched 9
Dataset 0 worker 1 fetched 1
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 8
Dataset 0 worker 0 fetched 3
Dataset 1 worker 1 fetched 0
Dataset 2 worker 0 fetched 13
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 8
Dataset 2 worker 1 fetched 7
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 6
Dataset 2 worker 0 fetched 19
Dataset 0 worker 1 fetched 3
Dataset 1 worker 0 fetched 3
Dataset 2 worker 1 fetched 4
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 5
Dataset 2 worker 0 fetched 16
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 9
Dataset 2 worker 1 fetched 10
Dataset 0 worker 0 fetched 0
Dataset 1 worker 1 fetched 8
Dataset 2 worker 0 fetched 17
Dataset 0 worker 1 fetched 3
Dataset 1 worker 0 fetched 7
Dataset 2 worker 1 fetched 14
Dataset 0 worker 0 fetched 2
Dataset 1 worker 1 fetched 1
Dataset 2 worker 0 fetched 2
Dataset 0 worker 1 fetched 0
Dataset 1 worker 0 fetched 2
Dataset 2 worker 1 fetched 6
Dataset 0 worker 0 fetched 1
Dataset 1 worker 1 fetched 4
Dataset 2 worker 0 fetched 15
Dataset 0 worker 1 fetched 4
Dataset 1 worker 0 fetched 0
Dataset 2 worker 1 fetched 11
********************************************************************************
Summary of returned results for each dataset for technique 'round-robin'
	Dataset #0; raw length 5; num drawn = 20; values = [0, 4, 3, 2, 1, 0, 2, 1, 3, 4, 1, 3, 2, 4, 0, 3, 2, 0, 1, 4]
	Dataset #1; raw length 10; num drawn = 20; values = [3, 1, 5, 6, 2, 9, 4, 7, 0, 8, 6, 3, 5, 9, 8, 7, 1, 2, 4, 0]
	Dataset #2; raw length 20; num drawn = 20; values = [0, 12, 18, 1, 3, 5, 9, 8, 13, 7, 19, 4, 16, 10, 17, 14, 2, 6, 15, 11]
********************************************************************************

```

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation