Release all the StreamingDataset resources during job termination

[karan6181]

Description of changes:

• Release all the StreamingDataset resources during job termination and fixed all the memory leaks
• Before:

$ pytest tests/

========================================================================= short test summary info 
=========================================================================
SKIPPED [6] tests/test_distributed.py:66: Fails due to new shared Filelock. See https://mosaicml.atlassian.net/browse/CO- 
1403
=================================================== 488 passed, 6 skipped, 5 deselected, 1 warning in 
129.29s (0:02:09) ===================================================
/opt/homebrew/Cellar/python@3.10/3.10.7/Frameworks/Python.framework/Versions/3.10/lib/python3.10/multiprocessing/reso urce_tracker.py:224: UserWarning: resource_tracker: 
There appear to be 378 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '

• After:

$ pytest tests/


========================================================================= short test summary info 
=========================================================================
SKIPPED [6] tests/test_distributed.py:66: Fails due to new shared Filelock. See https://mosaicml.atlassian.net/browse/CO- 
1403
================================================== 488 passed, 6 skipped, 5 deselected, 2 warnings in 
121.43s (0:02:01) ===================================================

• Ran the gpt-125m LLM model training
• Before

/usr/lib/python3.10/multiprocessing/resource_tracker.py:224: UserWarning: resource_tracker: There appear to be 6 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
/usr/lib/python3.10/multiprocessing/resource_tracker.py:237: UserWarning: resource_tracker: '/88fcf1_next_epoch': [Errno 2] No such file or directory: '/88fcf1_next_epoch'
  warnings.warn('resource_tracker: %r: %s' % (name, e))
/usr/lib/python3.10/multiprocessing/resource_tracker.py:237: UserWarning: resource_tracker: '/6f2a6f_shard_states': [Errno 2] No such file or directory: '/6f2a6f_shard_states'
  warnings.warn('resource_tracker: %r: %s' % (name, e))
/usr/lib/python3.10/multiprocessing/resource_tracker.py:237: UserWarning: resource_tracker: '/6f2a6f_barrier': [Errno 2] No such file or directory: '/6f2a6f_barrier'
  warnings.warn('resource_tracker: %r: %s' % (name, e))
/usr/lib/python3.10/multiprocessing/resource_tracker.py:237: UserWarning: resource_tracker: '/88fcf1_barrier': [Errno 2] No such file or directory: '/88fcf1_barrier'
  warnings.warn('resource_tracker: %r: %s' % (name, e))
/usr/lib/python3.10/multiprocessing/resource_tracker.py:237: UserWarning: resource_tracker: '/6f2a6f_next_epoch': [Errno 2] No such file or directory: '/6f2a6f_next_epoch'
  warnings.warn('resource_tracker: %r: %s' % (name, e))
/usr/lib/python3.10/multiprocessing/resource_tracker.py:237: UserWarning: resource_tracker: '/88fcf1_shard_states': [Errno 2] No such file or directory: '/88fcf1_shard_states'
  warnings.warn('resource_tracker: %r: %s' % (name, e))
Waiting up to 30 seconds for all training processes to terminate. Press Ctrl-C to exit immediately.

• After

wandb: Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)
wandb: Find logs at:
Waiting up to 30 seconds for all training processes to terminate. Press Ctrl-C to exit immediately.

Issue #, if available:

CO-1619

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.

[knighton]

Impressive