Improve shard efficiency of sampling for fractional stream repeats.

[knighton]

Description of changes:

Adds a StreamingDataset argument sampling_granularity: int = 1, which determines how many samples are drawn from the same shard at a time when sampling a stream.

At 1, each sample is drawn independently for the partial repeat if any.

At 100, up to a hundred samples are pulled at a time from the same shard, limited by the number of un-picked samples left in the shard, for the partial repeat if any.

If shard usage efficiency is especially a concern, you would set sampling_granularity to the size of the biggest shard (500? 50000?) in order to fully exhaust one shard a time.

Issue #, if available:

Merge Checklist:

Put an x without space in the boxes that apply. If you are unsure about any checklist, please don't hesitate to ask. We are here to help! This is simply a reminder of what we are going to look for before merging your pull request.

General

• I have read the contributor guidelines
• This is a documentation change or typo fix. If so, skip the rest of this checklist.
• I certify that the changes I am introducing will be backward compatible, and I have discussed concerns about this, if any, with the MosaicML team.
• I have updated any necessary documentation, including README and API docs (if appropriate).

Tests

• I ran pre-commit on my change. (check out the pre-commit section of prerequisites)
• I have added tests that prove my fix is effective or that my feature works (if appropriate).
• I ran the tests locally to make sure it pass. (check out testing)
• I have added unit and/or integration tests as appropriate to ensure backward compatibility of the changes.

[snarayan21]

just some clarification stuff like comments and duplicate variable names, other than that lgtm