☁️

RayクラスターをGKE Autopilotを使って構築する

2023/07/15に公開

Kubernetes

GKE

ray

tech

背景

最近ネットサーフィンしていたら以下の記事を見つけた。

端的に言えば、みなさんご存知 ChatGPT は分散システムのRayを使って構築されているらしい。

https://www.ray.io

前々からRayの存在自体は知っていたのだが、当時はクラスタを構築する暇がなかったので、スルーしていたのだけれど、せっかくCKAに合格したしもう少し具体的なk8sの利用法を模索したいなということで、まず、GKE Autopilotを使ってRayのクラスタを構築してみようと思う。

Rayとは

Qiitaのこの記事を見てもらえれば基本的な部分を把握できると思う。

要は、Python、Java、C++で利用できる分散処理用ライブラリ&実行環境のことだ。
使い方としてはかなり簡単で、Pythonの場合、以下のようにライブラリをインポートし、初期化、分散処理させる関数にデコレータをつけてあげればOKというもの。

import ray
ray.init()

@ray.remote
def f(x):
    return x * x


futures = [f.remote(i) for i in range(4)]
print(ray.get(futures))

クラスタを用意しなくてもRayを利用して遊ぶことはできるのだが、せっかくなので、学習兼おもちゃとしてクラウド上に構築してみようというのが今回の目的。

GKE クラスタの用意

普通に gcloud を使って、クラスタ名 ray-playground を東京リージョンに構築する。

$ gcloud container clusters create-auto ray-playground --region asia-northeast1

クラスタが出来るまで待ち、kubectl get nodes でローカルから制御可能か確認する。

Rayクラスタの構築

ありがたいことに、簡単にk8sクラスタ上にRayクラスタを構築できるよう helm のチャートを用意してくれている。それを使って環境を整備しよう。

まずは、 KubeRay オペレータをクラスタ上にデプロイする。

$ helm repo add kuberay https://ray-project.github.io/kuberay-helm/
$ helm install kuberay-operator kuberay/kuberay-operator --version 0.5.2

kubectl get pods を叩いて、正常に kuberay-operator がデプロイされていることを確認したら次に進む。

$ kubectl get pods               
NAME                                          READY   STATUS    RESTARTS   AGE
kuberay-operator-657c9454c6-xchjq             1/1     Running   0          57m

いよいよ RayCluster をデプロイする。

$ helm install raycluster kuberay/ray-cluster --version 0.5.2

RayCluster は CRなので、以下のコマンドを叩けるようになる。

$ kubectl get rayclusters.ray.io

実行すると以下のような出力が得られるはずだ。

NAME                 DESIRED WORKERS   AVAILABLE WORKERS   STATUS   AGE
raycluster-kuberay   1                 1                   ready    60m

KubeRay オペレータが RayCluster を検知すると、head と worker というクラスタを構成する Podを生成してくれる。

以下のコマンドを叩けば生成されたPodを見ることができる。

$ kubectl get pods --selector=ray.io/cluster=raycluster-kuberay

各Podの起動に時間がかかるが、正常に起動すれば以下のような表示が得られるはずだ。

NAME                                          READY   STATUS    RESTARTS   AGE
raycluster-kuberay-head-tmrvw                 1/1     Running   0          64m
raycluster-kuberay-worker-workergroup-cd4s4   1/1     Running   0          64m

Rayクラスタでアプリを実行する

Rayクラスタの構築が完了したので、いよいよアプリをクラスタ上で実行してみようと思う。

ローカルにRay CLIをインストールする

まずはアプリをクラスタに渡すためのCLIを準備する。
Python 3.7以上をインストールしたPCを用意し、まずは仮想環境を作成しよう。

$ python -m venv .venv
$ source .venv/bin/activate

次に、以下のコマンドでRayをインストールする。

$ pip install -U "ray[default]"

これで仮想環境内で Ray CLIを叩けるようになったはずだ。

RayクラスタにJobを渡す

Rayクラスタは先ほど作成した head PodにJobを投げることで処理をしてくれる構造になっている。
KubeRay オペレータは、head Pod用のサービスを立てているので、以下のコマンドでどのポートを使って通信すればいいのか確認してみる。

kubectl get service raycluster-kuberay-head-svc

実行結果はこうなるはずだ。

NAME                          TYPE        CLUSTER-IP     EXTERNAL-IP   PORT(S)                                         AGE
raycluster-kuberay-head-svc   ClusterIP   10.114.1.218   <none>        10001/TCP,8265/TCP,8080/TCP,6379/TCP,8000/TCP   79m

これでどのポートにアクセスすれば良いのか分かったので、port-forward を使ってRayクラスタに指示を送れるようにする。

以下のコマンドで port-forward できるようにする。

$ kubectl port-forward --address 0.0.0.0 service/raycluster-kuberay-head-svc 8265:8265

ブラウザで localhost:8265 にアクセスすれば、Rayクラスタのダッシュボードが見られるはずだ。

次に、Rayクラスタ上で実行させる簡単なコードを記述する。
まず、ディレクトリ構造は以下の通り。

project
├── LICENSE
├── README.md
└── main.py

main.py の中身は以下の通り。

main.py

import ray
ray.init()

@ray.remote
def f(x):
    return x * x


futures = [f.remote(i) for i in range(4)]
print(ray.get(futures))

main.py が入ったディレクトリ全体を Rayクラスタに渡して実行結果を得てみよう。
以下のコマンドを実行する。

$ ray job submit --address=http://127.0.0.1:8265 --working-dir $PWD -- python main.py

実行すると以下のような出力がされ最終的に結果を得ることができる。

2023-07-15 16:16:31,613 INFO dashboard_sdk.py:385 -- Package gcs://_ray_pkg_dfb5338fb5514dcc.zip already exists, skipping upload.

-------------------------------------------------------
Job 'raysubmit_cSwdVzUHEdydpeVA' submitted successfully
-------------------------------------------------------

Next steps
  Query the logs of the job:
    ray job logs raysubmit_cSwdVzUHEdydpeVA
  Query the status of the job:
    ray job status raysubmit_cSwdVzUHEdydpeVA
  Request the job to be stopped:
    ray job stop raysubmit_cSwdVzUHEdydpeVA

Tailing logs until the job exits (disable with --no-wait):
2023-07-15 00:16:35,799 INFO worker.py:1315 -- Using address 10.113.128.148:6379 set in the environment variable RAY_ADDRESS
2023-07-15 00:16:35,799 INFO worker.py:1432 -- Connecting to existing Ray cluster at address: 10.113.128.148:6379...
2023-07-15 00:16:35,814 INFO worker.py:1622 -- Connected to Ray cluster. View the dashboard at http://10.113.128.148:8265 
[0, 1, 4, 9]

------------------------------------------
Job 'raysubmit_cSwdVzUHEdydpeVA' succeeded
------------------------------------------

以上、駆け足だったけどひとまずGKE上にRayクラスタを建てて簡単なプログラムを実行させることが出来た。
次は、PyTorchを分散学習させてみようと思う。あと、できるのであれば、ラズパイとかJetson Nanoとかを使ってEdgeコンピューティングをRayを使って実現できるか試してみたいな。