Ray Data

# Ray Data - Scalable ML Data Processing Distributed data processing library for ML and AI workloads. ## When to use Ray Data **Use Ray Data when:** - Processing large datasets (>100GB) for ML training - Need distributed data preprocessing across cluster - Building batch inference pipelines - Loading multi-modal data (images, audio, video) - Scaling data processing from laptop to cluster **Key features**: - **Streaming execution**: Process data larger than memory - **GPU support**: Accelerate transforms with GPUs - **Framework integration**: PyTorch, TensorFlow, HuggingFace - **Multi-modal**: Images, Parquet, CSV, JSON, audio, video **Use alternatives instead**: - **Pandas**: Small data ( 100) ``` ### Group by and aggregate ```python # Group by column ds = ds.groupby("category").count() # Custom aggregation ds = ds.groupby("category").map_groups(lambda group: {"sum": group["value"].sum()}) ``` ## GPU-accelerated transforms ```python # Use GPU for preprocessing def preprocess_images_gpu(batch): import torch images = torch.tensor(batch["image"]).cuda() # GPU preprocessing processed = images * 255 return {"processed": processed.cpu().numpy()} ds = ds.map_batches( preprocess_images_gpu, batch_size=64, num_gpus=1 # Request GPU ) ``` ## Writing data ```python # Write to Parquet ds.write_parquet("s3://bucket/output/") # Write to CSV ds.write_csv("output/") # Write to JSON ds.write_json("output/") ``` ## Performance optimization ### Repartition ```python # Control parallelism ds = ds.repartition(100) # 100 blocks for 100-core cluster ``` ### Batch size tuning ```python # Larger batches = faster vectorized ops ds.map_batches(process_fn, batch_size=10000) # vs batch_size=100 ``` ### Streaming execution ```python # Process data larger than memory ds = ray.data.read_parquet("s3://huge-dataset/") for batch in ds.iter_batches(batch_size=1000): process(batch) # Streamed, not loaded to memory ``` ## Common patterns ### Batch inference ```python import ray # Load model def load_model(): # Load once per worker return MyModel() # Inference function class BatchInference: def __init__(self): self.model = load_model() def __call__(self, batch): predictions = self.model(batch["input"]) return {"prediction": predictions} # Run distributed inference ds = ray.data.read_parquet("s3://data/") predictions = ds.map_batches(BatchInference, batch_size=32, num_gpus=1) predictions.write_parquet("s3://output/") ``` ### Data preprocessing pipeline ```python # Multi-step pipeline ds = ( ray.data.read_parquet("s3://raw/") .map_batches(clean_data) .map_batches(tokenize) .map_batches(augment) .write_parquet("s3://processed/") ) ``` ## Integration with ML frameworks ### PyTorch ```python # Convert to PyTorch torch_ds = ds.to_torch(label_column="label", batch_size=32) for batch in torch_ds: # batch is dict with tensors inputs, labels = batch["features"], batch["label"] ``` ### TensorFlow ```python # Convert to TensorFlow tf_ds = ds.to_tf(feature_columns=["image"], label_column="label", batch_size=32) for features, labels in tf_ds: # Train model pass ``` ## Supported data formats | Format | Read | Write | Use Case | |--------|------|-------|----------| | Parquet | ✅ | ✅ | ML data (recommended) | | CSV | ✅ | ✅ | Tabular data | | JSON | ✅ | ✅ | Semi-structured | | Images | ✅ | ❌ | Computer vision | | NumPy | ✅ | ✅ | Arrays | | Pandas | ✅ | ❌ | DataFrames | ## Performance benchmarks **Scaling** (processing 100GB data): - 1 node (16 cores): ~30 minutes - 4 nodes (64 cores): ~8 minutes - 16 nodes (256 cores): ~2 minutes **GPU acceleration** (image preprocessing): - CPU only: 1,000 images/sec - 1 GPU: 5,000 images/sec - 4 GPUs: 18,000 images/sec ## Use cases **Production deployments**: - **Pinterest**: Last-mile data processing for model training - **ByteDance**: Scaling offline inference with multi-modal LLMs - **Spotify**: ML platform for batch inference ## References - **[Transformations Guide](references/transformations.md)** - Map, filter, groupby operations - **[Integration Guide](references/integration.md)** - Ray Train, PyTorch, TensorFlow ## Resources - **Docs**: https://docs.ray.io/en/latest/data/data.html - **GitHub**: https://github.com/ray-project/ray ⭐ 36,000+ - **Version**: Ray 2.40.0+ - **Examples**: https://docs.ray.io/en/latest/data/examples/overview.html