Use registry when creating Stream in StreamingDataset

[es94129]

Description of changes:

Add registry support to allow custom implementations of components like Stream, (in the future) batching methods and shuffling algorithms, similar to LLM foundry's registry.

Changes made in this PR:

• Copy the registry_utils.py file and corresponding tests from LLM foundry
• Set up a streams_registry that contains the Stream class
• Enable StreamingDataset to use the stream instance constructed from streams_registry

Existing Stream and StreamingDataset functionality remains unchanged, with added support for custom Stream implementations. See test_registry.py and LLM foundry's README of registry for more usage examples. Minimal example:

from streaming.base.registry_utils import construct_from_registry
from streaming.base.stream import streams_registry

streams_registry.register('stream', func=MyStream)

kwargs = {
    'remote': remote,
    'local': local,
    **kwargs,
}

stream = construct_from_registry(
    'stream',
    streams_registry,
    partial_function=False,
    kwargs=kwargs,
)

assert local == stream.local

Issue #, if available:

N/A

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). -- added under the section "Using Streams"

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. -- covered by existing tests

[es94129]

Hoping to request a review from @XiaohanZhangCMU too, but Github wouldn't let me do that

[snarayan21]

code and tests all lgtm, thanks for adding this! Could you please add a few sentences about registries in the streaming docs, likely in this section: https://docs.mosaicml.com/projects/streaming/en/stable/dataset_configuration/mixing_data_sources.html#using-streams

Can be very similar to what you have in the PR description. And a link to LLM Foundry's READMEs on registry usage would be useful too.

Thanks!

[snarayan21]

thanks for adding tests and updating the docs! Just minor changes.

[XiaohanZhangCMU]

Overall looking good! Left a few minor comments.

[snarayan21]

minor nit, other than that LGTM