AIStore for ASR datasets
State: closed | Pull request: yes
Labels: name: feature

name: core

name: ASR

name: CI

name: common

Signed-off-by: Ante Jukić <ajukic@nvidia.com>

# What does this PR do ?

This PR adds an option to use ASR datasets from an AIStore cluster.

## NOTE

Blocked by #5569.

**Collection**: [ASR]

# Changelog 
- Added AIS binary to Docker image
- Added env variables for configuring NeMo data store cache
- Added an option to defer dataloader initialization to `fit`, so data download can be synchronized across workers after a distributed environment is initialized
- For tarred datasets, added caching of manifests to store cache
- For non-tarred datasets, added caching of manifests and audio to store cache
- Utilities for handling data from AIStore cluster

# Usage

### AIS cluster config

AIStore endpoint needs to be configured before running the training command

```bash
export AIS_ENDPOINT=host:port
```

### Dataset setup

Manifests and tarred filepaths can point to AIS. For example, manifest can be set as

```bash
manifest_filepath = ‘ais://bucket/train_manifest.json'
```

and tarred filepaths as

```bash
tarred_audio_filepaths = ‘ais://bucket/audio__OP_0..511_CL_.tar'
```

To make sure data is download only on local or global rank zero, we use defer dataloader setup.
This can be achieved by setting `defer_setup = True` for the corresponding dataset or as

```bash
+model.train_ds.defer_setup=true \
+model.validation_ds.defer_setup=true \
+model.test_ds.defer_setup=true \
```

Note: A similar behavior could be achieved by using DataModule and controlling shared/non-shared caching using `datamodule.prepare_data_per_node`. However, that may require a more significant refactor.

### Cache setup

Local node-specific cache can be set up by setting the following env variables:

```bash
export NEMO_DATA_STORE_CACHE_DIR=/tmp
export NEMO_DATA_STORE_CACHE_SHARED=0
```
In this case, each local rank zero will cache data from AIS.

A shared cache location can be setup, for example, as

```bash
export NEMO_DATA_STORE_CACHE_DIR=/path/to/shared/dir
export NEMO_DATA_STORE_CACHE_SHARED=1
```
In this case, only the global rank zero will cache data from AIS.


# Before your PR is "Ready for review"
**Pre checks**:
- [x] 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**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Blocked by #5569. Currently rebased on top of that branch.