MosaicML Streaming

Warnings

• breaking Python 3.9 support was deprecated in `v0.12.0`. Users on Python 3.9 must upgrade their Python version to 3.10 or higher (3.12+ is fully supported) to use `mosaicml-streaming` versions 0.12.0 and above.
  Fix: Upgrade Python environment to 3.10 or later.
• gotcha Proper authentication/credentials are critical for streaming from cloud object stores (S3, GCS, Azure Blob, HF Hub). Incorrectly configured credentials are a common source of errors.
  Fix: Ensure environment variables (e.g., `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AZURE_STORAGE_ACCOUNT_NAME`, `AZURE_STORAGE_ACCOUNT_KEY`, `HF_TOKEN`) or cloud provider CLI configurations are correctly set for the target remote storage. Consult cloud provider documentation for specific setup.
• gotcha Earlier versions (<0.8.1) could experience dataloader hangs between epochs, significantly impacting training time. This issue was resolved in v0.8.1.
  Fix: Upgrade `mosaicml-streaming` to version 0.8.1 or newer to benefit from the fix.
• gotcha Prior to v0.10.0, the library created a new cloud client for each download, potentially leading to 'too many open sockets' errors or excessive cloud authentication requests. Version 0.10.0 introduced client reuse.
  Fix: Upgrade `mosaicml-streaming` to version 0.10.0 or newer to utilize reusable cloud download clients and improve stability.

Install

• pip install mosaicml-streaming Install stable version

Imports

• StreamingDataset

  from streaming import StreamingDataset

• MDSWriter

  from streaming import MDSWriter

  Used for creating MosaicML Streaming (MDS) datasets.

Quickstart

This quickstart demonstrates how to initialize `StreamingDataset` and integrate it with `torch.utils.data.DataLoader`. It sets up a minimal local MDS dataset for immediate testing. For cloud usage, ensure the `remote` parameter points to your cloud object storage path and that necessary cloud provider credentials are correctly configured in your environment.

import os
import torch
from streaming import StreamingDataset
from torch.utils.data import DataLoader
import json

# Define local paths for quickstart demonstration
# In a real scenario, 'remote' would point to your cloud MDS dataset
# (e.g., "s3://my-bucket/data" or "gs://my-bucket/data").
# Ensure cloud credentials are set in environment variables for cloud remotes.
local_remote_path = "quickstart_mds_data"
local_cache_path = "quickstart_mds_cache"

# --- Create a dummy MDS dataset for local testing if it doesn't exist ---
# For actual use, you'd generate MDS datasets with `streaming.MDSWriter`
# or point to existing ones in cloud storage.
if not os.path.exists(local_remote_path):
    print(f"Creating dummy MDS data in '{local_remote_path}'...")
    os.makedirs(local_remote_path)
    # A minimal `index.json` is required by StreamingDataset
    index_data = {
        "version": 2,
        "shards": [
            {"shard": 0, "samples": 2, "hash": "dummy_hash", "size": 100, 
             "raw_data_size": 80, "zip_data_size": 20, "compression": None, 
             "format": None}
        ]
    }
    with open(os.path.join(local_remote_path, 'index.json'), 'w') as f:
        json.dump(index_data, f)
    # A minimal shard file is also expected, though its content won't be processed
    # in this simplified example without actual schema.
    with open(os.path.join(local_remote_path, '00000.mds'), 'wb') as f:
        f.write(b'dummy_data_content_for_shard_0')
    print("Dummy MDS data created.")
else:
    print(f"Using existing dummy MDS data in '{local_remote_path}'.")

os.makedirs(local_cache_path, exist_ok=True)
# --- End of dummy MDS creation ---

# 1. Initialize the StreamingDataset
dataset = StreamingDataset(
    local=local_cache_path,  # Local cache directory for downloaded shards
    remote=local_remote_path, # Path to your MDS dataset (local or cloud)
    shuffle=True,
    batch_size=1, # Adjust batch size for internal buffering
    # Other parameters like `predownload` can be tuned for performance
)

# 2. Create a PyTorch DataLoader
dataloader = DataLoader(
    dataset=dataset,
    batch_size=1, # DataLoader batch size
    num_workers=0, # Use 0 workers for simple local testing to avoid multiprocessing issues
)

# 3. Iterate over the data
print(f"Dataset has {len(dataset)} samples.")
for i, batch in enumerate(dataloader):
    # In this dummy setup, 'batch' will be raw bytes as no actual data schema is defined.
    # With a real MDS dataset, this would be structured data (e.g., dicts, tensors).
    print(f"Batch {i}: {batch}")
    if i >= 1: # Process a few batches for demonstration
        break

# Note: For production use, remember to configure cloud credentials
# (e.g., via environment variables like AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY,
# or cloud provider CLI configs) if 'remote' points to cloud storage.