Hi @aswinkks ,

The core issue is clear: cross-cloud file-level I/O (AWS S3 → Azure compute) with many small files is the worst-case scenario for training throughput. Here's a structured breakdown addressing each of your questions.

1. Recommended Architecture for S3 Data from Azure Databricks

Your current path S3 → UC Volume (external, read-only) → FUSE mount → Training adds multiple latency layers: cross-cloud network hop (~50-150ms per request), FUSE overhead, and per-file request amplification.

Recommended approaches (best → acceptable):

If you truly cannot replicate data to Azure, option B (MDS shards on S3 + Mosaic Streaming + local cache) or option C (UCVolumeDataset) are your best bets. The key insight is: you must amortize the cross-cloud transfer across epochs via local caching.

2 & 3. Mosaic Streaming Configuration for Cross-Cloud Image Datasets

The reason Mosaic Streaming may be slower in your tests is likely shard size too small and/or predownload too low, meaning the pipeline stalls waiting for cross-cloud downloads.

Recommended settings for cross-cloud image workloads:

import streaming

# --- Writing MDS shards ---
# Target 128-256 MB per shard for images (larger = fewer cross-cloud round trips)
# Default is 67 MB — too small for high-latency links
writer = streaming.MDSWriter(
    out="s3://your-bucket/mds-dataset",
    columns={"image": "jpeg", "label": "int"},
    size_limit=256 * 1024 * 1024,  # 256 MB per shard
)
# --- Reading / Training ---
dataset = streaming.StreamingDataset(
    remote="s3://your-bucket/mds-dataset",      # remote source
    local="/local_disk0/mds-cache",              # LOCAL SSD cache (critical!)
    shuffle=True,
    shuffle_block_size=262144,                   # 256K samples per shuffle block
    predownload=16,                              # download 16 batches ahead (increase for high latency!)
    cache_limit="100gb",                         # keep shards locally across epochs
    num_canonical_nodes=None,                    # auto
    batch_size=64,
)
dataloader = streaming.StreamingDataLoader(
    dataset,
    batch_size=64,
    num_workers=8,          # match to CPU cores available for I/O
    prefetch_factor=4,      # each worker prefetches 4 batches
    pin_memory=True,
    persistent_workers=True,  # avoid re-fork overhead between epochs
)

4. DataLoader Workers and Prefetching Impact

Yes — significantly, but only if paired with local caching. Without caching, more workers just amplify the cross-cloud request volume, potentially hitting S3 rate limits or saturating bandwidth without reducing latency.

With local caching:

• num_workers=6-8 + prefetch_factor=4 is the Databricks-recommended starting point (from the AI Runtime DataLoader defaults)
• This overlaps GPU compute with data fetch/decode, hiding I/O latency
• pin_memory=True + persistent_workers=True eliminate per-batch allocation and per-epoch fork overhead

If using Databricks Serverless GPU (AI Runtime 5+):

from serverless_gpu.data import UCVolumeDataset, DataLoader
# Automatic local caching + optimized prefetch (num_workers=6, prefetch_factor=4 by default)
path_dataset = UCVolumeDataset("/Volumes/catalog/schema/volume/images")
loader = DataLoader(path_dataset, batch_size=64)

This is Databricks' purpose-built solution for exactly your scenario — it caches each file to local NVMe on first access and serves from cache for all subsequent reads.

5. Local Caching vs. Mosaic Streaming for Multi-Epoch Image Training

For multi-epoch training, local caching wins after epoch 1. Here's the tradeoff:

Recommendation: If your dataset fits on local disk — use UCVolumeDataset (simplest, no format conversion). If it doesn't fit — use Mosaic Streaming with a generous cache_limit and LRU eviction.

The reason your tests showed UC Volume faster than Mosaic Streaming is likely because:

1. UC Volume FUSE has internal read-ahead/prefetch that partially hides latency for sequential access
2. Mosaic Streaming's shard download + decompression adds overhead that only pays off when shards are large and predownload is tuned high
3. If MDS shards were small (default 67 MB), you're making just as many cross-cloud requests but with added decompression cost

6. Diagnosing the Bottleneck

Run these in parallel during training to identify the chokepoint:

# === GPU Utilization (should be >90% if data pipeline is healthy) ===
# In a separate cell or terminal:
# !nvidia-smi dmon -s u -d 2  # GPU util every 2 sec

# === DataLoader timing (insert around your training loop) ===
import time
for epoch in range(num_epochs):
    data_time_total = 0
    compute_time_total = 0
    t0 = time.perf_counter()
    for batch in dataloader:
        data_time = time.perf_counter() - t0
        data_time_total += data_time
        
        # --- GPU forward/backward ---
        t1 = time.perf_counter()
        loss = model(batch)
        loss.backward()
        optimizer.step()
        compute_time = time.perf_counter() - t1
        compute_time_total += compute_time
        t0 = time.perf_counter()
    
    print(f"Epoch {epoch}: data_load={data_time_total:.1f}s, "
          f"compute={compute_time_total:.1f}s, "
          f"ratio={data_time_total/(data_time_total+compute_time_total)*100:.1f}% Waiting on data")

Summary: Recommended Action Plan

1. Quick win: Switch to UCVolumeDataset + serverless_gpu.data.DataLoader (if on Serverless GPU) or manually pre-copy a subset to /local_disk0/ before training
2. Medium-term: Convert images to MDS with size_limit=256MB, tune predownload=16, set cache_limit to fill local SSD
3. Long-term: If budget allows, stage the MDS shards to an Azure ADLS location (even if the raw images must stay in S3, derived MDS shards might be allowed?) — this eliminates the cross-cloud hop entirely

The fundamental law here: with cross-cloud access, you're paying ~100ms per request vs ~1ms for local storage. The only real fix is to minimize the number of remote requests (large shards) and cache aggressively (avoid re-fetching across epochs).

If my answer was helpful, please consider marking it as accepted solution!