Ray Hello World AI 런타임 CLI 예시

Important

이 기능은 공개 미리보기 단계에 있습니다.

이 페이지에는 AI 런타임에 적용되는 다음 레이 라이브러리 각각에 대한 간단한 작동 예제가 있습니다:

• 레이 코어
• 레이 트레인
• 레이 데이터
• 레이 튠

사전 요구 사항

air CLI가 설치되고 인증되었습니다. AI 런타임 CLI 설치를 참조하세요.

레이 클러스터 부트스트랩

워크로드 command 를 제출하면 모든 노드에서 동시에 실행됩니다. 여러 노드에서 Ray를 사용할 때, 부트스트랩 스크립트 NODE_RANK 가 각 노드의 역할을 결정합니다. 0 레이 헤드가 시작하고, 나머지 모든 것은 일꾼으로 합류합니다.

이 페이지의 각 예제는 이 구성을 처리하기 위해 공유 ray_bootstrap.sh 를 사용합니다. 이 파일은 실행 중인 예제 스크립트와 함께 한 장만 있으면 됩니다.

#!/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

각 예제 YAML은 다음을 설정 RAY_ENTRYPOINT 하고 호출 ray_bootstrap.sh하여 부트스트랩을 호출합니다:

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

LOCAL_WORLD_SIZE AI 런타임에서 각 노드의 GPU 수로 설정되어 GPUS_PER_NODE , 요청한 GPU 유형에 따라 자동으로 스케일링됩니다. MASTER_ADDR AI 런타임이 헤드 노드의 IP 주소로 설정하며, 작업자들이 이를 이용해 Ray 클러스터를 찾고 가입합니다.

레이 코어

예시는 클러스터 내 모든 GPU에서 작업을 스케줄링하는 방법을 보여줍니다. 이 기능은 @ray.remote(num_gpus=1)Ray가 각 작업을 별도의 GPU에 배치하도록 지시합니다. 각 작업은 해당 작업이 어떤 노드와 물리적 GPU에 도달했는지 보고하여, 작업이 한 노드에 쌓여 있지 않고 노드 간에 분산되어 있음을 확인했습니다.

작업 부하 YAM

ray_core.yaml 각 노드 2개에 각각 1개의 A10 GPU를 요청합니다 (GPU_1xA10), 클러스터 총 2개의 GPU가 됩니다:

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

스크립트

ray_core.py GPU당 하나의 작업을 전송합니다. Ray는 각 작업 내에서 단일 할당된 GPU로 설정 CUDA_VISIBLE_DEVICES 하기 때문에 current_device() 항상 0을 반환합니다. 스크립트는 실제 신체 과제를 보고하기 위해 와 CUDA_VISIBLE_DEVICES 를 사용합니다ray.get_gpu_ids():

@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)

전체 대본은 이 페이지 끝에 있는 '전체 대본' 에 있습니다.

실행 제출

air run -f ray_core.yaml --watch

레이 트레인

예시는 소규모 MLP를 합성 데이터로 훈련시킵니다. prepare_model 모델을 워커의 GPU로 옮기고 DDP로 랩핑합니다. prepare_data_loader 각 작업자가 서로 다른 샤드를 볼 수 있도록 a DistributedSampler 를 추가하고, ray.train.report 에포크별 지표를 드라이버에게 다시 전달합니다.

작업 부하 YAM

ray_train.yaml 각 노드에 1개의 A10 GPU를 장착한 2개의 노드를 요청합니다. ray[train] Ray Train 부가 영상을 설치합니다:

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

훈련 스크립트

ray_train.py 작업자별 학습 루프를 정의하고 클러스터 내 모든 GPU를 사용하도록 구성 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()

전체 대본은 이 페이지 끝에 있는 '전체 대본' 에 있습니다.

실행 제출

air run -f ray_train.yaml --watch

광선 데이터

예시는 합성 파이프라인을 구축합니다: 각 행 map 은 파생 특징을 추가하고, a filter 는 짝수 행만 유지하며, a map_batches 는 벡터화된 NumPy 변환을 적용합니다. 그리고 끝에 호출하면 count()sum() 실행이 트리거됩니다.

작업 부하 YAM

ray_data.yaml 2개의 노드를 요청합니다. AI 런타임의 Ray Data에서는 아직 이기종 CPU/GPU 클러스터를 지원하지 않기 때문에, 이 예시는 파이프라인을 CPU에 유지합니다. 노드 유형은 GPU_1xA10 클러스터 크기를 결정합니다:

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

처리 스크립트

ray_data.py 3단계 파이프라인을 정의하고 집계 결과를 출력합니다:

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}")

전체 대본은 이 페이지 끝에 있는 '전체 대본' 에 있습니다.

실행 제출

air run -f ray_data.yaml --watch

레이 튠

예시에서는 4개의 GPU에서 4번씩 총 8번의 시도를 실행합니다. 각 실험은 학습률, 숨겨진 크기, 배치 크기의 샘플링 조합을 사용하여 합성 데이터를 대상으로 소규모 MLP를 훈련시킵니다.

작업 부하 YAM

ray_tune.yaml 각 노드 4개에 각각 1개의 A10 GPU를 요청하여 최대 4개의 동시 시험을 위해 4개의 GPU를 제공합니다:

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

튜닝 스크립트

ray_tune.py ASHA를 활용해 탐색 공간을 구성하고 8개의 시험을 시작하여 조기 저조 시험 저조 중단:

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}) 한 번에 GPU 1개를 예약합니다. 4개의 GPU를 사용할 때 Ray Tune은 한 번에 4개의 시행을 실행하고, 시행이 끝나는 다음 배치를 시작합니다. 전체 대본은 이 페이지 끝에 있는 '전체 대본' 에 있습니다.

실행 제출

air run -f ray_tune.yaml --watch

런을 점검하세요

제출 후에는 상태 및 스트리밍 로그를 확인할 수 있습니다:

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

air logs 기본적으로 노드 0에서 스트림을 전송하며, 여기서 레이 드라이버가 실행됩니다. 워커 노드에서 로그를 보려면, --node 1, --node 2, 이런 식으로 이동합니다.

다음 단계

• Ray Train을 사용하여 분산 학습
• Ray Data 및 vLLM을 사용하여 일괄 처리 유추
• AI 런타임 CLI 설치
• 워크로드 YAML 참조

전체 대본

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()