PR1: Add TensorBlobManager for Efficient Tensor Storage (#156)

Summary:
## 📝 Summary

This PR introduces `TensorBlobManager`, a comprehensive system for efficient content-addressed storage of tensor data with automatic compression, deduplication, and quota management. This enables tritonparse to save tensor inputs/outputs during kernel tracing without consuming excessive disk space.

## 🎯 Motivation

When tracing Triton kernel launches, we often need to save tensor data for later analysis or reproduction. However, naive tensor storage faces several challenges:
- **Disk space**: Large tensors can quickly fill disk
- **Duplicates**: Same tensors may be traced multiple times
- **Performance**: Reading/writing large files is slow
- **Safety**: Need to prevent runaway disk usage

This PR addresses all these concerns with a production-ready blob storage system.

## 🚀 Key Features

### 1. **Content-Addressed Storage**
- Uses BLAKE2b hashing for content addressing
- Ensures data integrity through hash verification
- Two-level directory structure (`xx/hash.bin.gz`) to avoid filesystem limits
- Automatic deduplication: identical tensors stored only once

### 2. **Smart Compression**
- Automatic gzip compression for large blobs (>1MB threshold)
- Small tensors stored uncompressed to avoid overhead
- Configurable compression level (default: 4 for balanced speed/ratio)
- Atomic writes using temporary files + rename for safety

### 3. **Resource Management**
- Storage quota enforcement (default: 100GB)
- Automatic disabling when quota exceeded
- Per-tensor size limit (default: 10GB) to prevent OOM
- Graceful degradation: logs warnings but doesn't crash

### 4. **Observability**
- Real-time statistics logging every 100 blobs
- Tracks: saved count, total count, dedup hits, compression ratio
- Final statistics on storage disable
- Debug logging for troubleshooting

## 📊 Changes Overview

### Modified Files
- `tritonparse/structured_logging.py` (+306/-3)

## 🔧 Implementation Details

### New Class: `TensorBlobManager`

```python
class TensorBlobManager:
    """Manager for storing tensor data as content-addressed blobs."""

    def __init__(self, root_dir=None, storage_quota=None)
    def set_root_dir(self, root_dir: str)
    def save_tensor_blob(self, tensor) -> Dict[str, Any]
```

**Key Methods**:
- `save_tensor_blob()`: Main entry point, returns metadata dict with hash, path, sizes
- `_compute_hash()`: BLAKE2b hashing for content addressing
- `_get_blob_path()`: Two-level directory structure generation
- `_log_statistics()`: Progress tracking and reporting
- `_disable_storage()`: Graceful shutdown on quota/error

### Configuration (Environment Variables)

| Variable | Default | Description |
|----------|---------|-------------|
| `TRITONPARSE_SAVE_TENSOR_BLOBS` | `"0"` | Enable/disable blob storage |
| `TRITONPARSE_TENSOR_SIZE_LIMIT` | `10GB` | Max single tensor size |
| `TRITONPARSE_TENSOR_STORAGE_QUOTA` | `100GB` | Total storage quota (compressed) |
| `TRITONPARSE_COMPRESSION_THRESHOLD` | `1MB` | Compress blobs >= this size |
| `TRITONPARSE_COMPRESSION_LEVEL` | `4` | Gzip compression level (0-9) |
| `TRITONPARSE_STATS_LOG_FREQUENCY` | `100` | Log stats every N blobs |

### Integration Points

1. **Global Instance**: `TENSOR_BLOB_MANAGER` singleton initialized in `init_logs()`
2. **Tensor Logging**: Integrated into `_log_torch_tensor_info()` function
3. **API**: `init()` function accepts `enable_tensor_blob_storage` and `tensor_storage_quota` parameters
4. **Cleanup**: `clear_logging_config()` resets the manager

## 📁 Storage Structure

```
trace_output_dir/
└── saved_tensors/
    ├── 00/
    │   ├── 00a1b2c3...def.bin      # Small tensor (uncompressed)
    │   └── 00f9e8d7...abc.bin.gz   # Large tensor (compressed)
    ├── 01/
    │   └── 01234567...890.bin.gz
    └── ff/
        └── ffabcdef...123.bin.gz
```

**Naming Convention**: `{first_2_hex_chars}/{full_hash}{.bin|.bin.gz}`

## 🔒 Safety Features

### Error Handling
- **Disk Full**: Automatically disables storage, logs error
- **Large Tensors**: Skips with warning, continues tracing
- **Quota Exceeded**: Disables storage before write, shows statistics
- **Missing PyTorch**: Returns error dict, doesn't crash

### Atomic Operations
- Uses `tempfile.NamedTemporaryFile` + `Path.rename()` for atomic writes
- No partial files left on crash
- Thread-safe hash cache lookup

### Data Integrity
- Hash verification on filename
- Compression/decompression round-trip tested
- Graceful handling of corrupted files

## 📈 Performance Characteristics

**Time Complexity**:
- Save (new blob): O(n) where n = tensor size
- Save (duplicate): O(1) hash cache lookup
- Compression: O(n) for blobs >1MB

**Space Efficiency**:
- Zeros: ~1000x compression
- Random data: ~1.1x compression
- Typical kernels: 10-50x effective savings with dedup

**Benchmarks** (from testing):
- 2KB tensor: <1ms (uncompressed)
- 20MB tensor: ~50ms (compressed)
- 400MB tensor: ~2s (compressed)
- Dedup hit: <1ms (cache lookup)

## 🧪 Testing Strategy

This PR focuses on core implementation. Testing is on follow-up PR.

## 📚 API Example

```python
from tritonparse.structured_logging import init

# Enable blob storage with custom quota
init(
    trace_folder="/tmp/triton_trace",
    enable_trace_launch=True,
    enable_tensor_blob_storage=True,  # NEW
    tensor_storage_quota=50 * 1024**3,  # 50GB (NEW)
)
```

Tensors are automatically saved during kernel launches when tracing is enabled.

## ⚠️ Breaking Changes

None. This is purely additive functionality:
- Default: blob storage is **disabled**
- No changes to existing behavior
- Opt-in via environment variable or API parameter

## ✅ Checklist

- [x] Core `TensorBlobManager` class implemented
- [x] Environment variables and configuration
- [x] Integration with tensor logging pipeline
- [x] API parameters for `init()` function
- [x] Cleanup in `clear_logging_config()`
- [x] Error handling and safety features
- [x] Statistics logging
- [x] Documentation (docstrings)
- [ ] Unit tests (deferred to later PR )

## 🎉 Impact

This PR enables efficient tensor storage for tritonparse, making it practical to:
- Save tensor data for large-scale tracing runs
- Build reproducers with actual tensor values
- Debug numerical issues in Triton kernels
- Analyze kernel input/output distributions

With compression and deduplication, we can trace workloads that would otherwise consume terabytes of disk space.

Pull Request resolved: #156

Reviewed By: sfzhu93

Differential Revision: D84021513

Pulled By: FindHao

fbshipit-source-id: ec30f19f6fe29a8b5238ffffbeb7094bc2f457a2

Author: FindHao