Ray AIR Deep Dive

Ray 2.0, released in August 2022, introduces Ray AI Runtime (AIR), a unified layer for preprocessing, training, tuning, and serving ML models. AIR integrates Ray Data, Ray Train, Ray Tune, and Ray Serve into a cohesive API.

Pipeline Example

python 
import ray
from ray import train
from ray.train.torch import TorchTrainer
from ray.train import ScalingConfig

def train_func():
    import torch
    model = torch.nn.Linear(10, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    for epoch in range(10):
        loss = model(torch.randn(32, 10)).sum()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
        train.report({'loss': loss.item()})

trainer = TorchTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True),
)
result = trainer.fit()
print(f'Final loss: {result.metrics["loss"]:.4f}')

Sources