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):
| Strategy | Latency | Complexity |
| Pre-convert to large MDS/WebDataset shards, stage to Azure ADLS, read locally | Lowest | Medium (one-time ETL |
| Pre-convert to MDS shards on S3, use Mosaic Streaming with aggressive local caching | Medium | Medium |
| UCVolumeDataset with local caching (Databricks AI Runtime) | Medium-High (1st epoch), Low (subsequent) | Lowest |
| Direct UC Volume reads (no caching) | Highest | Lowest |
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:
| Approach | Epoch 1 | Epoch 2+ | Memory-efficient shuffle |
| Mosaic Streaming (MDS) | Moderate (sequential shard download) | Fast (if cache_limit retains shards) | Yes (built-in) |
| UCVolumeDataset + local cache | Slower (per-file cross-cloud fetch) | Very fast (local NVMe) | No (must implement) |
| Manual shutil.copytree upfront | Slowest start (copies everything) | Fastest (fully local) | No |
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:
- UC Volume FUSE has internal read-ahead/prefetch that partially hides latency for sequential access
- Mosaic Streaming's shard download + decompression adds overhead that only pays off when shards are large and predownload is tuned high
- 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
- Quick win: Switch to UCVolumeDataset + serverless_gpu.data.DataLoader (if on Serverless GPU) or manually pre-copy a subset to /local_disk0/ before training
- Medium-term: Convert images to MDS with size_limit=256MB, tune predownload=16, set cache_limit to fill local SSD
- 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!
