🎄 本記事は ZOZO Advent Calendar 2024 シリーズ 2 の 20 日目です。
推薦基盤ブロックの19日目の記事は荒木さんによる強化学習で作る最強のCCレモンAI~修行編~でした。

はじめに

こんにちは！データシステム部推薦基盤ブロックのスミスと申します。
今回は、Rayというイケてるフレームワークについて勉強してみたので、共有してみたいと思います。

Rayとは

Rayは、Pythonベースの機械学習（ML）や人工知能（AI）アプリケーションを、手軽にスケーリング（並列化・分散化）するためのオープンソース統合フレームワークです。

Rayのフレームワークは、以下の3つのレイヤーで成り立っています。

Ray AI Libraries - MLエンジニアやデータサイエンティスト向けの、特定領域に特化したPythonライブラリ群。これらを使えば、MLアプリケーションを簡単かつスケーラブルに構築できる
Ray Core - 一般的なPythonプログラムやML処理を、わずかな変更で分散実行可能にするオープンソースの分散計算ライブラリ。これにより、コードを大規模環境に容易に拡張できる
Ray Clusters - Rayのクラスター環境。アプリケーションが要求するリソースに応じて自動的にスケールアップ・スケールダウン（ノード数を増減）することも可能で、柔軟かつ効率的なリソース利用が実現できる

（公式ドキュメントより抜粋）

今回は、Ray AI Librariesの中のRay Trainについて使用してみました。

Ray Trainとは

大規模なデータセットや高コストな処理を扱う際、通常は並列処理や分散処理の設計や実装が複雑になりがちです。しかし、Ray trainではこの部分を抽象化し、簡潔なコードで高性能な分散処理パイプラインを構築できるようにしてくれます。これにより、シングルマシン上のサンプルコードを、容易にマルチノード・マルチGPU環境へとスケールアップすることが可能になります。

さらに、PyTorch、TensorFlow、Hugging Faceのような有名な機械学習フレームワークとも統合可能で、既存のML/DLワークフローに自然に組み込める柔軟性が特徴です。

セットアップ

以下を実行してください。

pip install ray xgboost scikit-learn pandas

Rayによる簡単なスケールアップの例

機械学習のトレーニングコードを、ほぼそのままRayに移行することで、分散学習（データ並列）を簡単に実現できます。ここでは、まずRayを使用しない従来のトレーニング例を示し、その後、わずか数行の変更でRayによる分散学習に対応させる流れを紹介します。以下の資料を元に実装を行なっています。

従来のシンプルトレーニング（Ray未使用）

以下のコードは、California HousingデータセットでXGBoostを用いた単純な回帰タスクを学習する、従来の（Ray未使用）例です。
ここでは単一スレッド(nthread=1)でのみ学習を行っています。

# 従来のシンプルなトレーニング例（Ray未使用）
import pandas as pd
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
import xgboost as xgb
import time

# データ準備
housing = fetch_california_housing()
X = pd.DataFrame(housing.data, columns=housing.feature_names)
y = housing.target

X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.3, random_state=42)

# XGBoostデータセット化
dtrain = xgb.DMatrix(X_train, label=y_train)
dvalid = xgb.DMatrix(X_valid, label=y_valid)

params = {
    "objective": "reg:squarederror",
    "eval_metric": "rmse",
    "nthread": 1
}

bst = xgb.train(
    params,
    dtrain,
    num_boost_round=30,
    evals=[(dtrain, "train"), (dvalid, "valid")]
)

print("Training time (no Ray, single thread): {:.2f} sec".format(end - start))

このコードは単純でわかりやすいですが、大規模データを扱う場合、より高速な学習や大規模並列化をしたくなるかもしれません。しかし、従来なら分散学習に対応するためには、別途コードを大幅に書き換えたり、複雑な並列処理フレームワークを組み込む必要がありました。

Rayを使用した場合

Rayを使うと、以下のように同様の処理を簡単に分散化できます。
ポイントは、XGBoostTrainerやScalingConfigを使い、num_workersを指定するだけでデータ並列学習が可能になることです。

# Rayを使用したトレーニング例
import pandas as pd
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
import time
import ray
from ray.train import ScalingConfig
from ray.train.xgboost import XGBoostTrainer

# データ準備
# scikit-learnからCalifornia Housingデータセットを取得
housing = fetch_california_housing()
X = pd.DataFrame(housing.data, columns=housing.feature_names)
y = housing.target
df = X.copy()
df["target"] = y

# ローカルにCSV保存（この例では "california_housing.csv" ）
csv_path = "california_housing.csv"
df.to_csv(csv_path, index=False)

# Rayを初期化
ray.init()

# Ray DataでCSVデータセットを読み込み
dataset = ray.data.read_csv(csv_path)

train_dataset, valid_dataset = dataset.train_test_split(test_size=0.3)

# 分散学習用の設定
scaling_config = ScalingConfig(
    num_workers=1,
    use_gpu=False
)

# RayのXGBoostTrainerを用いた学習
trainer = XGBoostTrainer(
    scaling_config=scaling_config,
    label_column="target",
    num_boost_round=100,
    params={
        "objective": "reg:squarederror",
        "eval_metric": ["rmse"]
    },
    datasets={"train": train_dataset, "valid": valid_dataset},
)

result = trainer.fit()
print("Metrics:", result.metrics)

ここまでで、従来コードからの主な変更は以下の点です。

Rayの初期化 (ray.init())
XGBoostTrainerとScalingConfigクラスの使用
データ読み込みをRay Data経由に変更（ray.data.read_csv）

データ並列化をさらに簡単に

この段階で、num_workers=1をnum_workers=4などに増やすだけで、同じコードがデータ並列での分散学習にスケールアップします。

scaling_config = ScalingConfig(
    num_workers=4,
    use_gpu=False
)

これだけで、Rayが内部でデータを分割し、複数ワーカーで並列学習を行います。
従来なら複雑なコード変更が必要だった部分を、たった一行変更するだけで対応可能です。

そのほかにも

Rayは、以下のような機能を提供しています。

Ray Serve: オンライン推論APIを構築するためのモデルサービングライブラリ
Ray Tune: 機械学習モデルのハイパラチューニングを効率的に行うためのPythonライブラリ
Ray Cluster: Kubernetesなどのクラスター上で、Rayアプリケーションを簡潔にデプロイでき、複数のノードを用いたクラスタ構成を容易にするもの。

興味が出てきた方は、ぜひ、使ってみてください！