柔軟なKubernetes活用で分散機械学習や負荷テストも実現 LINEが開発した分散処理用RPCライブラリ

Kubernetesジョブ定義の例

大東哲平氏（以下、大東）：こんにちは。Machine Learning Infrastructureチームの大東と言います。このセッションでは、機械学習のために開発した、RPCライブラリを紹介します。

Machine Learning室では、Kubernetesを使い、GPUやCPUのノードを必要な数だけ確保して、相互に通信することで機械学習を行っています。機械学習のジョブは複数台で走り、相互に通信します。CPUノードではデータをストレージから取得し、適切な加工を行います。

そして、学習の要求に合わせてGPUノードへ送ります。KubernetesではSecretを使い、ストレージへの認証情報をPodに持たせられます。これにより、ソースコード管理から認証情報を分離できます。CPUやメモリのリソース制限や、各Podからログを集めて読むことも重要な機能です。まずは、簡単なKubernetesジョブ定義の例を見てみます。

これは、Pod内で"Hello, world!"と表示するために必要なジョブ定義です。いくつかの行を選んで解説していきます。同じ名前のジョブは投げられないので、ランダムサフィックスをつけて、何度も実行できるようにします。

Notebookからの実行でも、スケジューラからでも、勝手にリトライしてしまうのは使い勝手が悪いので、全ジョブ共通でリトライしないように設定します。

完了後のジョブを指定時間後に自動削除する機能です。もしこれを設定しないと、完了後のジョブがいつまでも残ってしまい、非常に見づらくなってしまいます。Kubernetesのバージョンや設定により使えないことがあるので、確認が必要です。

実行されるイメージとコマンドを書きます。

（スライドを指して）CPUやメモリを使用する上限はここで設定します。CPUはOSにより割り当て時間が制限され、それ以上に使えなくなります。1,000mでCPUの1スレッドです。この例でいうと、CPU使用率が20パーセントを超えないようになります。メモリは上限を超えて確保しようとするとOutOfMemoryErrorが発生して、ジョブは失敗となります。設定はこれだけです。

あとはこのファイルを適当な名前で保存し、kubectlで作成したあと、Podのログを表示すれば動作確認できます。万が一失敗した時には、イベントを読むと失敗原因がわかることが多いので、そのあたりも見られるといいですね。次は、実際にGPUにジョブを投げる例を、実行するBashスクリプトとともに見てみます。

GPUにジョブを投げる例

（スライドを指して）BashとYAMLを組み合わせて試すと、このようになります。${}で文字列を置換できるようになっています。ここではイメージ、コマンド、ジョブ名を置き換えられるようにしています。当然、イメージは自分で作ることになりますが、イメージとコマンドを指定するだけで運用できるケースは多くあります。

GPUを使うための設定です。GPUノードは数が少ないので、一般のジョブが実行されてしまわないように防御されています。また、GPUのドライバのバージョンや、プロダクト名を指定して走らせられます。多少の変更はあれど、ほぼすべてのジョブがこのような設定をつけることになります。

オブジェクトストレージの認証情報をSecretに入れておき、それをPodの環境変数として読み出せます。オブジェクトストレージからは、モデルやデータを読み出して計算処理を行います。Kubernetesジョブを使い始めた時はYAMLと文字列置換を使い実行するにつれて、Bashスクリプトがどんどん長くなりました。

本来であればジョブを作成してログを読むだけで十分ですが、Bashスクリプトがここまで複雑になるのには理由があります。

実行しているログを取得して表示するため、waitコマンドとlogsコマンドを組み合わせます。

例えば、waitコマンドはPodの状態が実行中になるのを待つことしかできません。そのため完了を検知できず、例えば起動してwaitする前に完了、あるいは起動前にエラーとなった場合、永遠に待ち続けてしまいます。

逆にlogsコマンドはPodが実行中になる前はエラーとなるので、開始を待つ必要があります。つまり、単純にwaitしてからlogsという方法は取れず、タイムアウトを設定して、ループしながら待つことになってしまうのです。

exitCodeはバージョンにより0だったり0.0だったりします。まだ他にも問題が多くありますが、このあたりで諦めたほうがよいでしょう。もしかしたら諦めるのが遅かったかもしれません。

先ほど挙げたいくつかの点を解決するコマンドを開発しました。Swimmyと呼びます。Podを1つ起動してコマンドを実行するだけであれば、Pythonコードを書く必要はありません。単に、イメージ指定とコマンドを1つ実行したいだけというケースが多かったため、一部のライブラリ開発者を除き、これがもっとも多い使い方です。

ジョブに対応しているバージョンのKubernetesなら実行可能であり、常駐Podなどの事前準備は一切不要で、どのクラスタでも走らせられます。

リソース制限も可能です。利用メモリの少なさ、起動時間の短さも念頭に置いて開発しているため、クラスタ管理やイメージのテストのために使うケースもあります。

環境変数のついたvalueを入れない場合は、ローカルの環境変数がそのまま転送されます。

Podのテンプレート機能が使えます。多くのPodで共通の情報はあらかじめテンプレートとして保存しておき、共通の設定として使うと利便性が高くなります。ローカルからのデータアクセスを禁止して、KubernetesのSecretを使用したりするケースでも利用されています。

テンプレートを複数指定した場合は、必要に応じてマージされます。もし同じ設定項目がある場合、あとで設定されたものにより上書きされます。

用途やSLA（Service Level Agreement）に応じてクラスタが複数あるため、ConfigMapにテンプレートを置くことで、クラスタ固有の差異を吸収できます。例えば、クラスタAのGPUとクラスタBのGPUで、設定やドライバが異なっていても、同じように使えます。それらの設定はジョブ間で共通であることが多いため、クラスタ内に置かれていることに妥当性があります。

任意のファイルを送れるので、汎用イメージを作り、実行したいコマンドを送りつけて実行することで、毎回イメージを作る手間を削減し、開発を効率化できます。

指定時間までに開始できない場合や、途中で万が一固まってしまった場合など、実行開始後に指定時間を超えたら自動的に失敗となる機能を組み合わせると、安心して実行できます。

たった1つの制約として、ドライバと同じバージョンのPythonがイメージに入っている必要があります。必要なライブラリはエージェントが自動的に集めるので、特に気にする必要はありません。自動的に集めるとはいえ、依存の少なさも重視しているため、イメージが重くなければ数秒で立ち上がります。

Podが複数になる場合、swimmy.cmdだけでカバーできるケースは少ないため、Pythonでコードを書いてもらうことになります。また、Pythonで書くことによる、より大きなメリットもあります。Jupyterなどのインタラクティブな環境でPodを立ち上げたまま、手元で書いたfileやfunctionを送って実行させ、開発を進められるのです。

プログラムが実行させたプロセスの動き

ここからは内部の話をします。このプログラムが実行させたプロセスを、ドライバと呼んでいます。ドライバはPodSpecを読み、Kubernetes用のジョブ情報に読み替え、APIサーバーへ作成を依頼します。先ほど説明したテンプレートやConfigMapも、もちろん使えます。コードで書いたほうが、設定可能な項目は多くなります。

APIサーバーは指定されたとおりジョブを作成し、その後、ジョブで指定された個数だけPodを作ります。

Podは起動時のイメージにSwimmyがあらかじめあれば、それを使います。なければ社内リポジトリからSwimmyと依存ライブラリを取得し、SwimmyAgentを起動します。エージェント側で使用しているライブラリはZMQとシリアライズと暗号化だけなので、Dependency Hellを回避しやすくなっています。

Hadoop対応など、便利な機能をつけたい気持ちは抑えています。過去にHadoopクラスタは複数存在し、アップグレードや移行などを行い、その度にクライアントの更新を行いました。特定のバージョンに紐づいてしまうと、機能以外の部分でアップグレードを行わざるを得なくなります。環境と機能を分離していくために、小さく保っていくのはよい手法です。

また、同様にKubernetesも用途に合わせて複数存在し、これも更新頻度が高くなります。例えばマルチテナント環境ではカスタマイズできないため、純粋なKubernetesの機能だけで構成しておくことにも大きなメリットがあります。

ドライバはジョブを作成する前にZMQでサーバーを1つ立ち上げ、Pod側から読めるURLを環境変数として埋め込んでおきます。立ち上がったエージェントは、ドライバに立てられたサーバーへ接続することで通信が確立します。ここからドライバはエージェントを自由に操作できるようになります。

機械学習ライブラリを開発する際に、ノード間の転送方式をいくつか検討しました。ZMQが十分に速く、また大量にデータを転送するにあたり、書き込みバッファなどの設定を細かくできたところが魅力的で使ってきました。Swimmyにおいてもその知見を活かして採用しています。

各RPCの結果が順不同で返ってくるため、その処理を効率よく行うために、asyncioとZMQを組み合わせています。

もしここでファイルが指定されていた場合、そのファイルをエージェントへ送る必要があるため、ドライバはHTTPサーバーを立てます。そして指定されたファイル、またはディレクトリを圧縮してエージェントから取得させます。取得されたファイルはPythonパスが通っているので、コード上で簡単にインポートできます。

重要な機能の1つとして、ヘルスチェックがあります。定期的にドライバ側から通信を行うことでドライバはエージェントが生きていることを知り、エージェントもまたドライバが生きていることがわかります。

ヘルスチェックが連続で失敗した場合、ドライバは異常終了します。また、エージェントも一定時間ヘルスチェックを受け取れない場合に異常終了します。削除ではなく異常終了をしてくれるため、どこまで実行されたか明確で、デバッグしやすくなります。

KubernetesにはactiveDeadlineSecondsという機能があります。これは指定した時間より長く起動している場合、Kubernetesが止めてくれる機能です。最初のヘルスチェックが成功したら、activeDeadlineSecondsを伸ばします。なぜこのような方法を取るかというと、ZMQによる接続確立前にドライバがエラーとなった場合、ジョブを削除する以外に止める方法が存在しないからです。

最初はコネクションタイムアウトと同じ時間に設定しておき、起動に失敗したら自動的に終了しておくようにします。

Podのログは非同期で取得して、エージェントのPod名とともにログとして出力できます。コードの例であれば標準エラーに出力されることになります。ログレベルはinfoなので、ロガーを適切に設定したほうがよいでしょう。異常系のイベントやその他の警告はwarningで出力されます。

ここでようやく任意のファンクションを送りつけられるようになります。（スライドを指して）ここではエージェントの1つに対し、ファンクションを送りつけています。async functionでも、どちらでも送れます。すべてのエージェントに対して同じファンクションを送ったり、リストを渡して全エージェントへ異なるパラメーターを渡すことなども可能です。

実行を完了すると、タイトルのとおり "Hello, Swimmy!" と表示されるはずです。もし実行に問題があれば、解決に必要なメッセージが表示されるはずです。

機械学習ライブラリの事例

ここからはいくつかの事例を紹介します。

社内で使われている機械学習ライブラリであるGheeは、Swimmy上で動いています。機械学習に利用する場合、Pod間のデータ転送では、MPIなどのより速くて多機能な通信ライブラリを使います。Swimmyでは、お手軽にPodが起動でき、内部の実装は専用ライブラリに任せるスタンスで開発して、その手軽さを保っています。

他にも機械学習APIでリクエストがあった時にSwimmyからPodを起動して計算処理を走らせるものや、多数のノードを使ったオフラインテストフレームワークなども開発されています。

数十台から構成されるRedisクラスタから数億キーをスキャンする際に、単に1クライアントから実行すると1時間かかりましたが、接続対象と同じ数だけのエージェントを準備して同時にアクセスすることで、5分で完了できました。10倍の速度向上が得られたことになります。

任意の負荷ツールと、自分で書いたシナリオをマルチノードで走らせて連携できます。例えば高い負荷を与えるために、専用ツールのためのPod群をいくつか立ち上げ、対象に負荷をかけます。そのタイミングで別のPod群から認証、更新、値の確認などを順次行うことで、正常に処理されているかを確認できます。

それらを1つのスクリプトで書けて、柔軟にテストを改善できます。今ではKubernetes上で何かを実行したい時に欠かせないライブラリとなっています。

Kubernetesでジョブを走らせる方法と、そのための便利なライブラリについて解説しました。用途に応じてKubernetesクラスタが分かれている場合など、特に威力を発揮します。また、Notebookとバッチ処理で同じコードを走らせることで、開発速度を上げられています。発表は以上です。ありがとうございました。