Make `epoch_sample_ids` cachable

🚀 Feature Request

It would be awesome to enable caching of epoch_sample_ids.

Motivation

Caching would remove a lot of redundant work that currently is re-executed at each run. It takes 20 minutes for my dataset's sample IDs to be created. This wastes a lot of budget for large-scale runs.

In my case, I'll specifically focus on the implementation in https://github.com/mosaicml/streaming/blob/2e9db78db6dd4108b697cfde92a95cd0de80539c/streaming/base/batching/random.py. Specifically, what takes long is dataset.resample_streams (with sampling_method="balanced") and get_shuffle (with shuffle_algo="py1e" in my case).

[Optional] Implementation

I've looked into this a bit, but get_shuffle's indirect dependence through get_partitions on sample_in_epoch (drop_first in the called functions) seems to make this very difficult. Maybe someone with more knowledge of the codebase can chime in on this, though. I would personally be happy with a simple hacky solution for now. :)

For now I've implemented a stupid NumPy file hash for dataset.resample_streams, which already saves around 40–50% of the time.

Start in streaming/base/batching/random.py and trace epoch_sample_ids through dataset.resample_streams, get_shuffle, and get_partitions, including sample_in_epoch and drop_first. Compare the current balanced sampling and py1e shuffle paths with the reported NumPy file-hash workaround. Done means repeated runs can reuse epoch_sample_ids without repeating the expensive work.