ray_on_aml を試してみた話
はじめに
ray_on_aml というのは AzureML 上で Ray を良しなに取り扱うための便利ライブラリです。
このパッケージを使用すると AzureML においてインタラクティブな分散処理ができるらしいと聞きつけたので試してみました。
前提知識
AzureML における分散処理
AzureML Studio の Notebook にアタッチして Jupyter 同様インタラクティブに処理を実行できる環境は以下 3 種類です。
- Compute Instance
- Synapse Apache Spark (preview)
- Azure ML Managed Spark (preview)
Compute Instance はシングルノードの計算リソースであるのに対し、Synapse Apache Spark と AzureML Managed Spark の 2 つは Spark ということで、Notebook にアタッチして分散処理が可能な計算リソースとなります。
Notebook ではなくジョブ実行専用の環境として Compute Cluster というものがあります。Compute Cluster はマルチノード構成の計算リソースで、 OpenMPI や NCCL といったノード間通信の仕組みも備えています。ジョブ定義に多少手を加えることで、複数ノードを使用した分散深層学習なども可能なように設計されています。
Ray
強化学習界隈などでは恐らく有名なのではないかと思いますが、Python における分散並列処理フレームワークです。
Pandas や NumPy 互換インターフェースを備えた分散データフレームツールである Dask のバックエンドとして使用したり、Ray Tune を活用してハイパラチューニングに使ったりと用途は極めて広いようです。
シングルノード内でのマルチプロセッシングを簡単に実装できるツールぐらいの認識しかしていなかったのですが、ユースケースによっては Spark を置き変えることも可能なぐらいのポテンシャルがありそうです。
ちなみにちょっと調べてみたところによれば、Ray on Spark (Databricks) とか Spark on Ray も可能なようです。
実験
ray_on_aml インタラクティブモードのアーキテクチャ
公式リポジトリに記載があった以下図が分かりやすかったので引用します。
Compute Instance をアタッチした Notebook がインターフェースとなり、Compute Instance と同じネットワークセグメント上に作成されて通信可能な Compute Cluster に対して Compute Instance から分散処理の指令を出すという構造の様子です。
Compute Instance が司令塔ということは、Compute Instance に対する接続は AzureML Studio の Notebook の他、Jupyter や VSCode でも問題なさそうです。
準備
リソース準備
AzureML Workspace の用意が終わっていることは前提条件とします。VNet を使用しますが、Private Endpoint は不要でパブリックのものでも構いません。(もちろん Private Endpoint で閉域化している環境でも問題ありません)
同じネットワークセグメント上に Compute Instance と Cluster を配置するという条件を満たすために、事前に VNet を作成しその中に展開するようにします。
Compute Instance の場合は「Next: Advanced Settings」をクリックすると VNet 内に展開するオプションが表示されます。
Compute Cluster の場合は 2 ページ目の真ん中あたりに「Advanced Settings」のメニュー以下に VNet オプションがあります。
依存ライブラリインストール
Compute Instance に新たにカーネルを作り、必要なパッケージをインストールします。
いつも忘れるのでメモがてらカーネルを追加するコマンドなども全部一緒に書いておきます。
conda create -n py310_ray_on_aml python=3.10
conda activate py310_ray_on_aml
pip install --upgrade ray==2.2.0 ray[air]==2.2.0 ray[data]==2.2.0 azure-ai-ml ray-on-aml
pip install mlflow azureml-mlflow
conda install notebook ipykernel
ipython kernel install --user --name=py310_ray_on_amlさらに、ray_on_aml の import 時に AzureML SDK v1 に対する依存の存在により正常に import できないエラーに遭遇したことから、AzureML SDK v1 も追加しておきます。既に v2 が入っているのでなんだか気持ち悪いですが、幸い別の名前が付いた別のパッケージとして整備されているので共存は可能です。
pip install azureml-core分散処理
ray_on_aml リポジトリのサンプルノートブックをベースに動かしていきます。
ノートブック環境は先に作った Compute Instance を想定します。VSCode 環境から試していますが、原理的には AzureML Studio でも Jupyter でも問題ないはずです。
まずはおなじみの AzureML SDK v2 による workspace への接続です。
from azure.ai.ml import MLClient
from azure.ai.ml import command, Input
from azure.identity import DefaultAzureCredential
subscription_id = "<subscription id>"
resource_group = "<resource group name>"
workspace = "<workspace name>"
# get a handle to the workspace
ml_client = MLClient(
DefaultAzureCredential(), subscription_id, resource_group, workspace
)続いて ray_on_aml のパッケージを import します。SDK v1 を入れてあれば、ここでエラーが出ることはないはずです。
from ray_on_aml.core import Ray_On_AML続いて分散処理を実際に実行するワーカーとなる Compute Cluster を指定して、Ray_On_AMLクラスからインスタンスを作成します。
ray_on_aml =Ray_On_AML(ml_client=ml_client, compute_cluster ="<cluster name>")インスタンスを作成したら、ワーカーノードを立ち上げます。
ray = ray_on_aml.getRay(ci_is_head=True, num_node=2,pip_packages=["ray[air]==2.2.0","torch==1.13.0","xgboost", "xgboost_ray","azureml-mlflow==1.48.0", "pyarrow==6.0.1"])ci_is_headは Compute Instance をヘッドノードとして使うかどうかを指定するフラグで、trueの場合は処理を実行している Compute Instance がヘッドノードになり、falseの場合は Compute Cluster のうち 1 ノードがヘッドノードとなり、残りがワーカーノードとなるようです。num_nodeは Compute Cluster のうち何台のノードを立ち上げるかを指定するパラメーターとなります。pip_packagesは後続の処理で実際に使うパッケージをワーカーノードにインストールするものです。
これらの記述を踏まえ、pip_packagesに基づいて AzureML でワーカーノードの環境を定義したコンテナのビルドが開始し、ジョブが開始されて Compute Cluster のうち指定されたノードが立ち上がります。このジョブはこちらから明示的に停止を行うまでは停止せず、ジョブによって確保されたノードは Compute Instance からの指示で様々な分散処理をインタラクティブに実行できるようになります。
初回もしくは依存関係を書き換えた場合、コンテナビルドが伴う関係上実行に数分 (ドキュメントによれば 7 分程度) かかります。2 回目以降はノードの立ち上げだけなので比較的高速です。
続いて、Ray を使った分散処理の init を行いますが、サンプルノートブック記載の記述ではエラーになったので、以下のように書き換えました。
client = ray.init(f"ray://{ray_on_aml.get_ip()}:10001")続いてサンプルノートブックに従って xgboost による分散学習へと進んでいきます。
シンプルに以下の記述で実行することができるようです。
from xgboost_ray import RayDMatrix, RayParams, train
from sklearn.datasets import load_breast_cancer
train_x, train_y = load_breast_cancer(return_X_y=True)
train_set = RayDMatrix(train_x, train_y)
evals_result = {}
bst = train(
{
"objective": "binary:logistic",
"eval_metric": ["logloss", "error"],
},
train_set,
evals_result=evals_result,
evals=[(train_set, "train")],
verbose_eval=False,
ray_params=RayParams(
num_actors=10, # Number of remote actors
cpus_per_actor=1))
bst.save_model("model.xgb")
print("Final training error: {:.4f}".format(
evals_result["train"]["error"][-1]))もう 1 つ、データのロードを自力で行うやり方もノートブックには書いてあります。
# Load data.
dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
# Split data into train and validation.
train_dataset, valid_dataset = dataset.train_test_split(test_size=0.3)
# Create a test dataset by dropping the target column.
test_dataset = valid_dataset.drop_columns(cols=["target"])pandas 互換っぽい API ですね。大規模データを相手にする場合、色々と応用がありそうです。
pandas 互換の分散処理というと PySpark に取り込まれた koalas や、dask が思い浮かびますが、dask については Ray でもサポートしているのでデータの処理がメインであればそちらを利用しても良さそうです。そちらの記述もサンプルノートブックにはあったので、興味ある人はそちらも試してみると良さそうです。
分散処理を試したところで、 Ray Dashboard を開きます。
VSCode で Compute Instance につなげた状態で Terminal に http://127.0.0.1:8265 と打って実行はせず、ctrl キーを押しながら打ち込んだアドレスをクリックするとうまく Compute Instance にポートフォワーディング的なことをしてつなげてくれるとドキュメントに書いてあります。
これでも良いのですが、ray_on_aml のドキュメントには書いていない裏技があるのでそちらをメモしておきます。
https://<compute instance name>-<port number>.<region>.instances.azureml.ms上記 URL にアクセスすることで、Compute Instance ローカルに立ち上げた様々な UI を開くことができます。
ポート番号を Ray Dashboard のポート番号である 8265 に置換し、その他の部分も自分の環境にあわせて置換した URL をブラウザで開くと、AAD 認証の後に Ray Dashboard を開くことができます。
最後に、放っておくといつまでも Compute Cluster のノードが起動し続けてお金がかかるので、ノードの停止処理をします。
ray_on_aml.shutdown()このコマンドを実行すると Compute Cluster のノードを確保していたジョブが停止し、しばらくして (自動停止を有効にしていれば) Compute Cluster が停止して課金が止まります。
おわりに
ray_on_aml を使用して AzureML 環境でインタラクティブに分散処理を実行する手順を試しました。
これまで AzureML で大規模なデータを処理する場合は Compute Cluster 上でジョブを実行するというやり方が定石でした。
しかしデータの前処理の段階や特徴量エンジニアリング、もしくはもう少し前のデータと向き合っている段階ではやはりインタラクティブな操作ができることが望ましかったのですが、唯一 AzureML でインタラクティブにデータを処理する方法だった Compute Instance を使う方法はシングルノードになるので、データの規模次第では AzureML を使うことは難しいと言わざるを得ませんでした。
今回 preview になった Spark を利用する方法に加え、ray_on_aml を使う方法も出てきたわけで、AzureML におけるインタラクティブな分散処理は一気に充実した印象です。