Quickstart Guide: dgen-py + StreamingCheckpointing

This guide helps you verify and test the two major optimizations introduced in this PR:

1. dgen-py Integration: 155x faster random tensor generation
2. StreamingCheckpointing: 192x memory reduction for checkpoints

Prerequisites

# Ensure virtual environment is activated
source .venv/bin/activate

# Verify dgen-py is installed
python -c "import dgen_py; print(f'dgen-py {dgen_py.__version__} installed')"

# If not installed:
uv pip install dgen-py

Quick Demo (5 minutes)

Run the comprehensive demo script:

# Simple test (1 GB, requires checkpoint directory)
export TEST_CHECKPOINT_DIR=/path/to/storage
./quickstart_demo.sh

# Larger test (24 GB - shows full memory reduction)
export TEST_SIZE_GB=24
export TEST_CHECKPOINT_DIR=/fast/nvme/storage
./quickstart_demo.sh

This script demonstrates:

• Part 1: File storage comparison (OLD vs NEW methods)
  • OLD: Pre-allocate full checkpoint in RAM
  • NEW: Stream with 192x less memory
• Part 2: Object storage with multi-library support
  • Tests s3dlio, minio, s3torchconnector (if credentials available)
  • Shows multi-endpoint load balancing (if configured)

Feature 1: dgen-py Integration

What It Does

Replaces Python-based random data generation (NumPy, PyTorch) with Rust-based dgen-py:

• 155x faster: 1.54 GB/s → 239 GB/s generation speed
• Drop-in replacement: No code changes to existing DLIO configs
• Zero-copy integration: Uses BytesView for memory efficiency

How to Verify

# Run checkpoint comparison test
./demo_checkpoint_methods.sh

Expected output:

[Original] Generation: 0.0042s @ 239.0 GB/s (dgen-py)
[Streaming] Generation throughput: 238.5 GB/s (dgen-py)

Compare this to NumPy baseline (~1.5 GB/s on same hardware).

Where It's Used

dgen-py is automatically used in:

• dlio_benchmark/utils/utility.py: gen_random_tensor() function
• dlio_benchmark/checkpointing/pytorch_checkpointing.py: get_tensor_core()
• dlio_benchmark/checkpointing/tf_checkpointing.py: TensorFlow tensor generation

Set DLIO_DATA_GEN=numpy environment variable to use NumPy instead (for comparison).

Feature 2: StreamingCheckpointing

What It Does

Implements producer-consumer pattern for checkpoint writing:

• 192x memory reduction: 24 GB → 128 MB for large checkpoints
• Overlapped I/O: Generation and writing happen in parallel
• Same performance: I/O throughput matches original method

How to Verify

# Compare memory usage between methods
./demo_checkpoint_methods.sh

# Expected output shows:
# - Original: ~24 GB memory for 24 GB checkpoint
# - Streaming: ~128 MB memory (64 buffers × 32 MB chunks ÷ 2)

Monitor memory with:

# In another terminal while test runs
watch -n 1 'ps aux | grep python | grep -v grep'

Architecture

Producer Thread              Shared Buffer Pool          Consumer Thread
───────────────              ──────────────────          ───────────────
                                                         
gen_random_tensor()  ──→  [Buffer 1: 32 MB]  ──→  write_chunk(buf1)
    (dgen-py)              [Buffer 2: 32 MB]  ──→  write_chunk(buf2)
    239 GB/s               [Buffer 3: 32 MB]  ──→  write_chunk(buf3)
                                  ...
                           [Buffer 64: 32 MB]

Total pool: 64 × 32 MB = 2 GB
Active memory: ~128 MB (only filled buffers)

Using in Your Code

from mlpstorage.checkpointing import StreamingCheckpointing

# Local file
checkpoint = StreamingCheckpointing(
    chunk_size=32 * 1024 * 1024,  # 32 MB chunks
    num_buffers=64,                # 2 GB pool
    use_dgen=True                  # Use dgen-py (default)
)
checkpoint.save('/tmp/checkpoint.pt', total_size_bytes=24 * (1024**3))

# Object storage (auto-detects library from URI)
checkpoint.save('s3://bucket/checkpoint.pt', total_size_bytes=24 * (1024**3))

Feature 3: Multi-Library Object Storage

Supported Backends

StreamingCheckpointing automatically detects and uses the appropriate library:

Library	URI Prefix	Use Case	Performance
s3dlio	s3://	Highest performance, Rust-based	Tested up to 7 GB/s per client
minio	s3://	Python SDK, widely compatible	Library/target dependent
s3torchconnector	s3://	AWS recommended for PyTorch	Library/target dependent
file	/path/to/	Local files with O_DIRECT	Local NVMe speeds

Performance Note: Tested results up to 7 GB/s per client, varies by library and storage target.

How to Test

# Set up credentials
cat > .env << EOF
AWS_ACCESS_KEY_ID=<your-access-key>
AWS_SECRET_ACCESS_KEY=<your-secret-key>
AWS_ENDPOINT_URL=<your-s3-endpoint>
AWS_REGION=us-east-1
EOF

# Test all 3 S3 libraries
python test_compare_backends.py --size-gb 1.0

Expected output:

Backend: s3dlio
  Elapsed: 1.234s
  Throughput: 810.5 MB/s
  
Backend: minio
  Elapsed: 1.456s
  Throughput: 686.3 MB/s
  
Backend: s3torchconnector
  Elapsed: 1.389s
  Throughput: 719.8 MB/s

Backend Selection

Explicit backend selection:

# Force specific backend
checkpoint = StreamingCheckpointing(
    backend='s3dlio',              # Explicitly use s3dlio
    part_size=32 * 1024 * 1024,    # 32 MB multipart
    max_in_flight=4                # Concurrent uploads
)

checkpoint = StreamingCheckpointing(
    backend='minio',
    part_size=32 * 1024 * 1024,
    num_parallel_uploads=4
)

checkpoint = StreamingCheckpointing(
    backend='s3torchconnector'     # Auto-managed multipart
)

Auto-detection based on URI:

# Detects s3:// prefix, uses default backend (s3dlio if available)
checkpoint.save('s3://bucket/key', total_size)

# Detects file path, uses local file backend with O_DIRECT
checkpoint.save('/nvme/checkpoint.pt', total_size)

Writer subprocess start method (fork vs forkserver)

The streaming writer runs in a subprocess. The multiprocessing start method used to create it is selectable, with a backend-aware default:

Target	'auto' default	Rationale
POSIX file (backend='file', or a scheme-less / file:// path)	fork	The parent never starts s3dlio's Tokio runtime, so fork is safe and ~40% faster — it keeps the parent's copy-on-write warm state and CPU/NUMA locality with the shared-memory buffers (#682).
Object storage (s3dlio, direct_fs, s3torchconnector, minio, or any s3:///az:///gs:///direct:// URI)	forkserver	The parent has a live s3dlio Tokio runtime (via ObjStoreLibStorage._preflight()) by the time the writer is created; fork() copies that mutex state and the writer futex-deadlocks (#642). forkserver gets a clean interpreter without spawn's __main__ re-import (which hangs MPI ranks).

Leave mp_start_method='auto' (the default) and each backend gets the right method automatically. To override:

# Constructor argument (highest precedence)
checkpoint = StreamingCheckpointing(
    backend='file',
    mp_start_method='forkserver',   # 'auto' | 'fork' | 'forkserver' | 'spawn'
)

# Or via environment variable (constructor arg wins if both are set)
export MLPS_CHECKPOINT_MP_START_METHOD=fork

Notes:

• Requesting mp_start_method='fork' on an object-storage target raises a ValueError (citing #642) rather than deadlocking.
• If the platform lacks the requested method (e.g. no fork/forkserver on Windows/macOS), it degrades gracefully: fork → forkserver → spawn.
• save() logs the resolved method, e.g. [Main] Writer start method: fork (policy='auto', backend=file), so each run records which method produced its result.

Feature 4: Multi-Endpoint Load Balancing

What It Does

Multi-endpoint support allows distributing I/O load across multiple storage endpoints:

• Round-robin: Distribute requests evenly across endpoints
• Least-connections: Route to endpoint with fewest active connections (s3dlio only)
• Automatic failover: Handle endpoint failures gracefully (s3dlio only)

Backend Support:

Backend	Native Multi-Endpoint	MPI Rank-Based	Load Balancing
s3dlio	✅ Yes	✅ Yes	Round-robin, Least-connections
minio	❌ No	✅ Yes	Round-robin (via MPI rank)
s3torchconnector	❌ No	✅ Yes	Round-robin (via MPI rank)

Key Differences:

• s3dlio: Uses native MultiEndpointStore with true load balancing across endpoints
• minio/s3torch: Each MPI rank selects one endpoint (round-robin), no per-request balancing

Use cases:

• Scale beyond single endpoint bandwidth
• Distribute load across multiple storage nodes
• High-availability configurations

Configuration Methods

Option 1: Comma-separated list

export S3_ENDPOINT_URIS='http://172.16.21.1:9000,http://172.16.21.2:9000,http://172.16.21.3:9000'
export S3_LOAD_BALANCE_STRATEGY=round_robin  # or least_connections

# Test with quickstart
./quickstart_demo.sh

Option 2: Template expansion

# Expands {1...8} to create 8 endpoint URIs
export S3_ENDPOINT_TEMPLATE='http://172.16.21.{1...8}:9000'
export S3_LOAD_BALANCE_STRATEGY=least_connections

./quickstart_demo.sh

Option 3: File with URIs

# Create file with one URI per line
cat > endpoints.txt << EOF
http://172.16.21.1:9000
http://172.16.21.2:9000
http://172.16.21.3:9000
http://172.16.21.4:9000
EOF

export S3_ENDPOINT_FILE=endpoints.txt
export S3_LOAD_BALANCE_STRATEGY=round_robin

./quickstart_demo.sh

MPI Distributed Mode

For distributed training with MPI, each rank automatically selects a different endpoint:

All backends (s3dlio, minio, s3torchconnector):

# Each of 8 ranks will use a different endpoint (round-robin)
export S3_ENDPOINT_URIS='http://172.16.21.1:9000,http://172.16.21.2:9000,http://172.16.21.3:9000,http://172.16.21.4:9000'

mpirun -np 8 python -m dlio_benchmark.main workload=unet3d_v100

# Rank 0 → endpoint 1
# Rank 1 → endpoint 2
# Rank 2 → endpoint 3
# Rank 3 → endpoint 4
# Rank 4 → endpoint 1 (wraps around)
# ... etc

How it works:

• s3dlio: Can use native MultiEndpointStore OR MPI rank selection (both work)
• minio: Uses MPI rank selection only (no native multi-endpoint)
• s3torchconnector: Uses MPI rank selection only (no native multi-endpoint)

For minio and s3torchconnector, each rank:

1. Detects its MPI rank via OMPI_COMM_WORLD_RANK or PMI_RANK
2. Selects endpoint using rank % num_endpoints
3. Uses that single endpoint for all requests (no per-request balancing)

For s3dlio, you have two options:

1. Native multi-endpoint: Set S3_ENDPOINT_URIS + S3_LOAD_BALANCE_STRATEGY

   • Each rank uses ALL endpoints with load balancing
   • Round-robin or least-connections per-request routing

2. MPI rank selection: Same as minio/s3torch

   • Each rank uses ONE endpoint
   • Simpler, but no per-request balancing

MPI environment variables automatically detected:

• Open MPI: OMPI_COMM_WORLD_RANK, OMPI_COMM_WORLD_SIZE
• MPICH: PMI_RANK, PMI_SIZE

See: https://docs.open-mpi.org/en/v5.0.x/tuning-apps/environment-var.html

Performance Impact

Multi-endpoint configuration can provide:

• Aggregate bandwidth: N endpoints × per-endpoint bandwidth
• Example: 4 endpoints × 2 GB/s = 8 GB/s aggregate
• Scalability: Add endpoints to scale beyond single node limits

Note: Actual performance depends on:

• Network topology (avoid oversubscription)
• Storage backend capabilities
• Workload characteristics (request size, pattern)

Integration with DLIO

Zero-Code Integration

Existing DLIO configs automatically benefit from dgen-py:

# Your existing DLIO workload
python -m dlio_benchmark.main workload=unet3d_v100

# dgen-py is automatically used for checkpoint generation
# No config changes needed!

Explicit StreamingCheckpointing

To use streaming checkpoints with DLIO:

# In your DLIO config YAML
checkpoint:
  checkpoint_folder: s3://bucket/checkpoints
  steps_between_checkpoints: 100
  checkpoint_mechanism: pytorch
  
  # StreamingCheckpointing configuration (optional)
  streaming:
    enabled: true
    chunk_size: 33554432      # 32 MB
    num_buffers: 64           # 2 GB pool
    use_dgen: true            # Use dgen-py
    backend: s3dlio           # Explicit backend (or auto-detect)

Performance Tuning

dgen-py Tuning

import dgen_py

# NUMA-aware generation (automatic in StreamingCheckpointing)
generator = dgen_py.Generator(
    size=total_bytes,
    dedup_ratio=1.0,        # No deduplication for checkpoints
    compress_ratio=1.0,     # No compression
    numa_mode="auto",       # Bind to NUMA nodes
    max_threads=None        # Use all cores
)

Generator thread count under MPI (#689): StreamingCheckpointing throttles dgen-py's max_threads to os.cpu_count() // <ranks sharing this node>, so that N ranks on one host don't each spin up a thread per core and oversubscribe it. The divisor is the local rank count (OMPI_COMM_WORLD_LOCAL_SIZE / MPI_LOCALNRANKS / MV2_COMM_WORLD_LOCAL_SIZE, first one present), not the global MPI world size — dividing a per-node CPU count by a cluster-wide rank count starves generation on multi-node runs (e.g. 192 CPUs / 64 global ranks = 3 threads, when only 4 ranks actually share that node and 192 / 4 = 48 is correct). If generation is showing up as the pipeline bottleneck (bottleneck: "Generation" in save()'s results, or a low throughput_ratio), check the reported local_ranks_per_node in the [Main] Initializing dgen-py (...) log line against how many ranks are actually launched per host.

StreamingCheckpointing Tuning

Chunk Size:

• Larger chunks: Better throughput, more memory
• Smaller chunks: Lower latency, less memory
• Recommended: 32 MB (aligns with dgen-py, S3 multipart)

Buffer Pool Size:

• More buffers: Better parallelism, more memory
• Fewer buffers: Lower memory, potential stalls
• Recommended: 64 buffers (2 GB pool, ~128 MB active)

S3-Specific:

# s3dlio tuning
checkpoint = StreamingCheckpointing(
    backend='s3dlio',
    part_size=32 * 1024 * 1024,   # Match chunk_size
    max_in_flight=8               # More for high-bandwidth links
)

# minio tuning
checkpoint = StreamingCheckpointing(
    backend='minio',
    part_size=32 * 1024 * 1024,
    num_parallel_uploads=8
)

Troubleshooting

dgen-py Import Error

ImportError: No module named 'dgen_py'

Solution: Install via pip:

uv pip install dgen-py

Low S3 Performance

If seeing <100 MB/s throughput:

1. Check network bandwidth: iperf3 -c <s3-endpoint>
2. Increase parallelism: max_in_flight=16 or higher
3. Try different backend: Some libraries work better with certain S3 implementations
4. Verify multipart is working: Check S3 server logs

Memory Usage Higher Than Expected

StreamingCheckpointing uses:

• Buffer pool: chunk_size × num_buffers (e.g., 32 MB × 64 = 2 GB)
• Active memory: ~50% of pool (only filled buffers)
• Per-backend overhead: ~10-50 MB

Total: ~1-2 GB for recommended configuration.

If seeing higher:

1. Reduce buffer pool: num_buffers=32 (1 GB pool)
2. Reduce chunk size: chunk_size=16*1024*1024 (16 MB)

Checkpoint Verification

Verify checkpoint integrity:

import torch

# Load checkpoint and verify
state = torch.load('/tmp/checkpoint.pt')
print(f"Checkpoint size: {os.path.getsize('/tmp/checkpoint.pt') / (1024**3):.2f} GB")
print(f"Keys: {state.keys()}")
print(f"Model params: {sum(p.numel() for p in state['model'].values())}")

Next Steps

• Performance benchmarks: See docs/PERFORMANCE.md
• Implementation details: See docs/IMPLEMENTATION_COMPARISON.md
• Test suite: See tests/checkpointing/compare_methods.py
• DLIO integration: See dlio_benchmark/utils/utility.py

Questions?

File an issue or check the test scripts:

• demo_checkpoint_methods.sh: Method comparison
• test_compare_backends.py: Multi-library S3 testing
• quickstart_demo.sh: Comprehensive demo (runs both above)