Ray Hello World AI Runtime örnekleri

Important

Bu özellik Genel Önizleme aşamasındadır.

Bu sayfada, AI Runtime'daki aşağıdaki Ray kütüphanelerinin her biri için basit bir çalışma örneği bulunmaktadır:

• Ray Core
• Tışın Treni
• Tışın Verileri
• Ray Tune

Prerequisites

En son Databricks CLI kuruldu ve doğrulandı. Bkz. Databricks CLI'yı yükleme veya güncelleştirme ve Databricks CLI için Kimlik Doğrulaması.

Ray kümesi önyüklemesi

Bir iş yükü gönderdiğinizde, command tüm düğümlerde aynı anda çalışır. Ray’ı birden çok düğümde kullanmak için önyükleme betiği, her düğümün rolüne karar vermek üzere NODE_RANK kullanır. 0 Ray kafasını başlatır ve diğer her şey işçi olarak katılır.

Bu sayfadaki her örnek, bu kurulumu yönetmek için ortak bir ray_bootstrap.sh kullanır. Bu dosyanın sadece bir kopyasına ve çalıştırdığınız örnek betiklere ihtiyacınız var.

#!/bin/bash
# NODE_RANK=0 is the Ray head: it starts the cluster and runs the entrypoint
# script, then tears the cluster down. Every other rank joins as a worker and
# stays until the head goes away.
#
# The entrypoint to run on the head is passed via RAY_ENTRYPOINT, a path
# relative to CODE_SOURCE_PATH (e.g. "ray_train.py").
set -e

if [ -z "${RAY_ENTRYPOINT:-}" ]; then
    echo "RAY_ENTRYPOINT is not set; expected a script path relative to CODE_SOURCE_PATH." >&2
    exit 1
fi

RAY_HEAD_PORT=6379
GPUS_PER_NODE=${LOCAL_WORLD_SIZE:-1}

if [ "${NODE_RANK:-0}" = "0" ]; then
    echo "NODE_RANK=0: Starting Ray head node with $GPUS_PER_NODE GPU(s)..."
    ray start --head \
        --port=$RAY_HEAD_PORT \
        --num-gpus=$GPUS_PER_NODE \
        --dashboard-host=0.0.0.0

    # Always stop the cluster on exit, even if the entrypoint fails.
    trap 'ray stop' EXIT

    echo "Ray head node started. Running $RAY_ENTRYPOINT..."
    python "$CODE_SOURCE_PATH/$RAY_ENTRYPOINT"
else
    echo "NODE_RANK=$NODE_RANK: Connecting to Ray head at $MASTER_ADDR:$RAY_HEAD_PORT..."
    # Retry loop to wait for head to be ready. Note: omit --block, since it runs
    # forever and the head's `ray stop` only tears down local processes, leaving
    # the worker stuck. Without --block, `ray start` returns once this node joins
    # and we control our own exit below.
    joined=""
    for i in $(seq 1 12); do
        if ray start --address="$MASTER_ADDR:$RAY_HEAD_PORT" --num-gpus=$GPUS_PER_NODE 2>/dev/null; then
            joined=1
            break
        fi
        echo "Attempt $i failed, retrying in 5s..."
        sleep 5
    done
    if [ -z "$joined" ]; then
        echo "Worker failed to join the Ray head after all retries; aborting." >&2
        exit 1
    fi

    # `ray health-check` exits non-zero once the head runs `ray stop`, letting
    # this worker exit so the whole job can terminate. The counter backstops
    # against a hang.
    echo "Worker joined; waiting for the head to finish its work..."
    for _ in $(seq 1 360); do
        if ! ray health-check --address "$MASTER_ADDR:$RAY_HEAD_PORT" 2>/dev/null; then
            break
        fi
        sleep 5
    done
    echo "Head is no longer healthy; stopping local Ray and exiting."
    ray stop
fi

Her YAML örneği, RAY_ENTRYPOINT ayarlayıp ray_bootstrap.sh çağrısı yaparak bootstrap işlemini başlatır:

command: |
  cd $CODE_SOURCE_PATH
  RAY_ENTRYPOINT=ray_train.py bash ray_bootstrap.sh

LOCAL_WORLD_SIZE AI Runtime tarafından her düğümdeki GPU sayısına ayarlanır, böylece GPUS_PER_NODE istediğiniz GPU tipiyle otomatik ölçeklenir. MASTER_ADDR AI Runtime tarafından baş düğümün IP adresine ayarlanır; çalışanlar bu adresi Tışın kümesini bulmak ve birleştirmek için kullanırlar.

Işın Çekirdeği

Örnek, @ray.remote(num_gpus=1) kullanarak kümedeki her GPU’da işlerin nasıl zamanlanacağını gösterir; bu, Ray’e her görevi ayrı bir GPU’ya yerleştirmesini söyler. Her görev, hangi düğüme ve fiziksel GPU'ya yerleştirildiğini bildirerek, görevlerin tek bir düğüm üzerinde toplanmak yerine düğümler arasında dağıtıldığını doğrular.

Workload YAML

ray_core.yaml her biri 1 A10 GPU olan 2 düğüm talep eder (GPU_1xA10), küme toplamda 2 GPU verir:

experiment_name: ray-core-example

environment:
  version: '5'
  dependencies:
    - ray[default]

code_source:
  type: snapshot
  snapshot:
    root_path: .

compute:
  num_accelerators: 2
  accelerator_type: GPU_1xA10

command: |
  cd $CODE_SOURCE_PATH
  RAY_ENTRYPOINT=ray_core.py bash ray_bootstrap.sh

max_retries: 0
timeout_minutes: 15
env_variables:
  NCCL_DEBUG: INFO

Script

ray_core.py GPU başına bir görev yapar. Ray, her görev içinde CUDA_VISIBLE_DEVICES’ı atanan tek GPU’ya ayarladığından, current_device() her zaman 0 döndürür. Betik, gerçek fiziksel atamayı bildirmek için ray.get_gpu_ids() ve CUDA_VISIBLE_DEVICES kullanır:

@ray.remote(num_gpus=1)
def hello_from_gpu():
    node_rank = os.environ.get("NODE_RANK", "?")
    ray_gpu_ids = ray.get_gpu_ids()
    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    gpu_name = subprocess.run(
        ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
        capture_output=True, text=True, check=True,
    ).stdout.strip()
    return f"Hello from node {node_rank} | Ray GPU id {ray_gpu_ids} | CUDA_VISIBLE_DEVICES={visible} | {gpu_name}"

total_gpus = int(ray.cluster_resources().get("GPU", 0))
futures = [hello_from_gpu.remote() for _ in range(total_gpus)]
results = ray.get(futures)

Tam senaryo bu sayfanın sonunda Full scripts bölümünde yer almaktadır.

Çalıştırmayı gönderme

databricks air run -f ray_core.yaml --watch

Tışın Treni

Örnek, küçük bir MLP'yi sentetik veri üzerinde eğitiyor. prepare_model modeli işçinin GPU'suna taşıyor ve DDP ile sarıyor. prepare_data_loader, her çalışanın verinin farklı bir dilimini görmesini sağlamak için bir DistributedSampler ekler ve ray.train.report, epok başına metrikleri sürücüye iletir.

Workload YAML

ray_train.yaml Her biri 1 A10 GPU olan 2 düğüm talep ediyor. ray[train] Ray Train ekstralarını kurar:

experiment_name: ray-train-example

environment:
  version: '5'
  dependencies:
    - ray[train]
    - torch

code_source:
  type: snapshot
  snapshot:
    root_path: .

compute:
  num_accelerators: 2
  accelerator_type: GPU_1xA10

command: |
  cd $CODE_SOURCE_PATH
  RAY_ENTRYPOINT=ray_train.py bash ray_bootstrap.sh

max_retries: 0
timeout_minutes: 15
env_variables:
  NCCL_DEBUG: INFO

Eğitim senaryosu

ray_train.py çalışan başına bir eğitim döngüsü tanımlar ve kümedeki tüm GPU'ları kullanmak üzere yapılandırır TorchTrainer :

def train_loop_per_worker(config):
    model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 10))
    model = prepare_model(model)  # DDP wrap + move to this worker's GPU

    x = torch.randn(1024, 128)
    y = torch.randint(0, 10, (1024,))
    loader = DataLoader(TensorDataset(x, y), batch_size=64, shuffle=True)
    loader = prepare_data_loader(loader)  # adds DistributedSampler

    for epoch in range(config["epochs"]):
        ...
        ray.train.report({"epoch": epoch, "loss": epoch_loss / len(loader)})

trainer = TorchTrainer(
    train_loop_per_worker,
    train_loop_config={"lr": 1e-3, "epochs": 5},
    scaling_config=ScalingConfig(num_workers=total_gpus, use_gpu=True),
)
result = trainer.fit()

Tam senaryo bu sayfanın sonunda Full scripts bölümünde yer almaktadır.

Çalıştırmayı gönderme

databricks air run -f ray_train.yaml --watch

Tışın Verileri

Örnek, sentetik bir iş hattı oluşturur: satır başına çalışan bir map türetilmiş özellikler ekler, bir filter yalnızca çift numaralı satırları tutar ve bir map_batches vektörleştirilmiş bir NumPy dönüşümü uygular. Sonda count() ve sum() çağrılması yürütmeyi tetikler.

Workload YAML

ray_data.yaml 2 düğüm talep ediyor. Heterojen CPU/GPU kümeleri AI Runtime üzerindeki Ray Data’da henüz desteklenmediğinden, bu örnek iş hattını CPU’larda çalıştırır. GPU_1xA10 düğüm tipi, küme boyutunu belirler:

experiment_name: ray-data-example

environment:
  version: '5'
  dependencies:
    - ray[data]

code_source:
  type: snapshot
  snapshot:
    root_path: .

compute:
  num_accelerators: 2
  accelerator_type: GPU_1xA10

command: |
  cd $CODE_SOURCE_PATH
  RAY_ENTRYPOINT=ray_data.py bash ray_bootstrap.sh

max_retries: 0
timeout_minutes: 15

İşleme betiği

ray_data.py Üç aşamalı bir boru hattı tanımlar ve toplu sonuçları bastırır:

ds = ray.data.range(10_000)

def add_features(row):
    n = row["id"]
    return {"id": n, "squared": n * n, "is_even": n % 2 == 0}

def scale_batch(batch):
    batch["scaled"] = batch["squared"] * 0.001
    return batch

# Ray executes these stages in parallel across the cluster.
ds = ds.map(add_features)
ds = ds.filter(lambda row: row["is_even"])
ds = ds.map_batches(scale_batch, batch_format="numpy")

print(f"Pipeline produced {ds.count()} rows")
print(f"Sum of scaled feature: {ds.sum('scaled'):.2f}")

Tam senaryo bu sayfanın sonunda Full scripts bölümünde yer almaktadır.

Çalıştırmayı gönderme

databricks air run -f ray_data.yaml --watch

Ray Tune

Örnek, 4 GPU arasında aynı anda 4 deneme olmak üzere 8 deneme çalıştırıyor. Her deneme, öğrenme hızı, gizli boyut ve parti büyüklüğünün örneklenmiş bir kombinasyonuyla sentetik verilerde küçük bir MLP eğitir.

Workload YAML

ray_tune.yaml her biri 1 A10 GPU bulunan 4 düğüm talep eder, böylece eşzamanlı 4 deneme için 4 GPU sağlar:

experiment_name: ray-tune-example

environment:
  version: '5'
  dependencies:
    - ray[tune]
    - torch

code_source:
  type: snapshot
  snapshot:
    root_path: .

compute:
  num_accelerators: 4
  accelerator_type: GPU_1xA10

command: |
  cd $CODE_SOURCE_PATH
  RAY_ENTRYPOINT=ray_tune.py bash ray_bootstrap.sh

max_retries: 0
timeout_minutes: 30

Ayarlama betiği

ray_tune.py arama alanını yapılandırır ve ASHA ile 8 deneme başlatır, bu da erken başarısız denemeleri durdurur:

tuner = tune.Tuner(
    tune.with_resources(train_fn, resources={"gpu": 1}),
    param_space={
        "lr": tune.loguniform(1e-4, 1e-1),
        "hidden_size": tune.choice([64, 128, 256]),
        "batch_size": tune.choice([32, 64, 128]),
    },
    tune_config=tune.TuneConfig(
        metric="loss",
        mode="min",
        scheduler=ASHAScheduler(max_t=20, grace_period=3, reduction_factor=2),
        num_samples=8,
    ),
)
results = tuner.fit()
best = results.get_best_result("loss", "min")
print(f"Best config: {best.config}")

tune.with_resources(train_fn, resources={"gpu": 1}) Her deneme için bir GPU ayırıyor. 4 GPU ile Ray Tune aynı anda 4 deneme çalıştırıyor ve denemeler bittikçe bir sonraki partiye başlıyor. Tam senaryo bu sayfanın sonunda Full scripts bölümünde yer almaktadır.

Çalıştırmayı gönderme

databricks air run -f ray_tune.yaml --watch

Bir koşu kontrol et

Gönderdikten sonra durum ve yayın kayıtlarını kontrol edebilirsiniz:

databricks air get <run-id>
databricks air logs <run-id>

databricks air logs varsayılan olarak node 0'dan yayın yapar, Ray sürücüsü de burada çalışır. Bir işçi düğümündeki kayıtları görüntülemek için --node 1, --node 2 ve benzeri seçenekleri belirtin.

Sonraki Adımlar

• Ray Train ile dağıtılmış eğitim
• Ray Data ve vLLM ile toplu çıkarım
• Databricks CLI'yi yükleme veya güncelleştirme
• İş yükü YAML başvurusu

Tam senaryolar

ray_core.py

"""Ray Core remote-task example on AI Runtime.

Dispatches one @ray.remote task per GPU across the cluster. Each task prints
which node and physical GPU it was assigned to, confirming tasks reached every
node. Run after ray_bootstrap.sh has started the cluster.
"""

import os
import subprocess
import time

import ray

ray.init(address="auto")

num_nodes = int(os.environ.get("NUM_NODES", 1))
gpus_per_node = int(os.environ.get("LOCAL_WORLD_SIZE", 1))
expected_gpus = num_nodes * gpus_per_node

for _ in range(30):
    if len(ray.nodes()) >= num_nodes and ray.cluster_resources().get("GPU", 0) >= expected_gpus:
        break
    time.sleep(2)

total_gpus = int(ray.cluster_resources().get("GPU", 0))
if total_gpus < expected_gpus:
    raise SystemExit(
        f"Expected {expected_gpus} GPU(s) but Ray only sees {total_gpus}; " "check GPU discovery on all nodes."
    )

print(f"Ray cluster ready: {len(ray.nodes())} node(s), {total_gpus} GPU(s)")
print(f"Cluster resources: {ray.cluster_resources()}\n")


@ray.remote(num_gpus=1)
def hello_from_gpu():
    node_rank = os.environ.get("NODE_RANK", "?")
    # Ray sets CUDA_VISIBLE_DEVICES to the single assigned GPU, so
    # current_device() always returns 0. Report the physical GPU via
    # nvidia-smi and the Ray GPU ID instead.
    ray_gpu_ids = ray.get_gpu_ids()
    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    gpu_name = subprocess.run(
        ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
        capture_output=True,
        text=True,
        check=True,
    ).stdout.strip()
    return f"Hello from node {node_rank} | Ray GPU id {ray_gpu_ids} | CUDA_VISIBLE_DEVICES={visible} | {gpu_name}"


print(f"Launching {total_gpus} task(s), one per GPU across the cluster...")
futures = [hello_from_gpu.remote() for _ in range(total_gpus)]
results = ray.get(futures)

for r in results:
    print(r)

ray.shutdown()

ray_train.py

"""Ray Train distributed training example on AI Runtime.

Trains a small MLP on synthetic data with one training worker per GPU using
Ray Train's TorchTrainer. Ray Train places the workers across the cluster
(one per GPU) and wires up torch.distributed; the per-worker train loop just
uses `ray.train.torch` helpers to move the model/data to the right device.
"""

import os

import ray
import torch
import torch.nn as nn
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer, prepare_data_loader, prepare_model
from torch.utils.data import DataLoader, TensorDataset

# Connect to the cluster started by ray_bootstrap.sh.
ray.init(address="auto")

num_nodes = int(os.environ.get("NUM_NODES", 1))
total_gpus = int(ray.cluster_resources().get("GPU", 0))
if total_gpus < 1:
    raise SystemExit("No GPUs registered with Ray; check GPU discovery on the cluster.")
print(f"Cluster ready: {num_nodes} node(s), {total_gpus} GPU(s) available")
print(f"Launching a Ray Train run with {total_gpus} worker(s), one per GPU\n")


def train_loop_per_worker(config):
    """Runs on each Ray Train worker; one worker is pinned to one GPU."""
    # prepare_model wraps the model in DDP and moves it to this worker's GPU.
    model = nn.Sequential(nn.Linear(128, 256), nn.ReLU(), nn.Linear(256, 10))
    model = prepare_model(model)

    x = torch.randn(1024, 128)
    y = torch.randint(0, 10, (1024,))
    loader = DataLoader(TensorDataset(x, y), batch_size=64, shuffle=True)
    # prepare_data_loader shards the data across workers and moves batches to the GPU.
    loader = prepare_data_loader(loader)

    optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
    loss_fn = nn.CrossEntropyLoss()

    for epoch in range(config["epochs"]):
        model.train()
        epoch_loss = 0.0
        for inputs, labels in loader:
            optimizer.zero_grad()
            loss = loss_fn(model(inputs), labels)
            loss.backward()
            optimizer.step()
            epoch_loss += loss.item()
        # ray.train.report surfaces metrics back to the driver.
        ray.train.report({"epoch": epoch, "loss": epoch_loss / len(loader)})


trainer = TorchTrainer(
    train_loop_per_worker,
    train_loop_config={"lr": 1e-3, "epochs": 5},
    scaling_config=ScalingConfig(num_workers=total_gpus, use_gpu=True),
)

result = trainer.fit()
# result.metrics holds the last reported dict (may be None if nothing was
# reported on the final iteration); fall back to a plain message.
print(f"\nTraining finished. Final metrics: {result.metrics or 'see per-worker logs above'}")

ray.shutdown()

ray_data.py

"""Ray Data distributed preprocessing example on AI Runtime.

Builds a Ray Dataset and runs a distributed map / map_batches / filter
pipeline across CPU actors spread over the cluster. On AI Runtime, Ray Data
runs on CPU actors (heterogeneous CPU/GPU clusters are not supported yet), so
this example deliberately keeps the transforms on CPU. The common shape is Ray
Data preprocessing feeding into a Ray Train run.
"""

import os

import ray

# Connect to the cluster started by ray_bootstrap.sh.
ray.init(address="auto")

num_nodes = int(os.environ.get("NUM_NODES", 1))
num_cpus = int(ray.cluster_resources().get("CPU", 0))
print(f"Cluster ready: {num_nodes} node(s), {num_cpus} CPU(s) available")

# A simple synthetic dataset; range() produces a distributed Ray Dataset.
ds = ray.data.range(10_000)


def add_features(row):
    """Per-row transform, runs distributed across CPU tasks."""
    n = row["id"]
    return {"id": n, "squared": n * n, "is_even": n % 2 == 0}


def scale_batch(batch):
    """Vectorized per-batch transform (numpy), more efficient than per-row."""
    batch["scaled"] = batch["squared"] * 0.001
    return batch


# Distributed pipeline: map -> filter -> map_batches, then aggregate.
ds = ds.map(add_features)
ds = ds.filter(lambda row: row["is_even"])
ds = ds.map_batches(scale_batch, batch_format="numpy")

count = ds.count()
total = ds.sum("scaled")
print(f"\nPipeline produced {count} rows (even numbers only)")
print(f"Sum of scaled feature: {total:.2f}")
print("\nSample of 5 processed rows:")
for row in ds.take(5):
    print(f"  {row}")

ray.shutdown()

ray_tune.py

"""Ray Tune hyperparameter search example on AI Runtime.

Runs 8 trials across all available GPUs in the cluster (one GPU per trial).
Uses ASHA scheduler to prune unpromising trials early.
"""

import os
import ray
import torch
import torch.nn as nn
from ray import tune
from ray.tune.schedulers import ASHAScheduler

ray.init(address="auto")

num_nodes = int(os.environ.get("NUM_NODES", 1))
total_gpus = int(ray.cluster_resources().get("GPU", 0))
if total_gpus < 1:
    raise SystemExit("No GPUs registered with Ray; check GPU discovery on the cluster.")
print(f"Cluster ready: {num_nodes} node(s), {total_gpus} GPU(s) available")
print(f"Running 8 trials with up to {total_gpus} in parallel\n")


def train_fn(config):
    """Single trial: trains a small MLP on synthetic data for one GPU."""
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = nn.Sequential(
        nn.Linear(128, config["hidden_size"]),
        nn.ReLU(),
        nn.Linear(config["hidden_size"], 10),
    ).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=config["lr"])
    loss_fn = nn.CrossEntropyLoss()

    for epoch in range(20):
        x = torch.randn(config["batch_size"], 128, device=device)
        y = torch.randint(0, 10, (config["batch_size"],), device=device)

        optimizer.zero_grad()
        loss = loss_fn(model(x), y)
        loss.backward()
        optimizer.step()

        tune.report({"loss": loss.item(), "epoch": epoch})


tuner = tune.Tuner(
    tune.with_resources(train_fn, resources={"gpu": 1}),
    param_space={
        "lr": tune.loguniform(1e-4, 1e-1),
        "hidden_size": tune.choice([64, 128, 256]),
        "batch_size": tune.choice([32, 64, 128]),
    },
    tune_config=tune.TuneConfig(
        metric="loss",
        mode="min",
        scheduler=ASHAScheduler(max_t=20, grace_period=3, reduction_factor=2),
        num_samples=8,
    ),
)

results = tuner.fit()
best = results.get_best_result("loss", "min")
print(f"\nBest config: {best.config}")
print(f"Best loss:   {best.metrics['loss']:.4f}")

ray.shutdown()