インタラクティブ シェルを使用した Vertex AI トレーニング ジョブのデバッグ
Google Cloud Japan Team
※この投稿は米国時間 2021 年 10 月 8 日に、Google Cloud blog に投稿されたものの抄訳です。
機械学習モデルのトレーニングを成功させることは、困難で多大な時間を要する作業になる場合があります。通常のソフトウェア開発とは違い、トレーニングの結果はトレーニング コードと入力データの両方に左右されます。そのため、ローカルマシンで実行している場合でも、トレーニング ジョブのデバッグが複雑なプロセスとなる可能性があります。リモートのインフラストラクチャでコードを実行すると、その作業がさらに困難になるかもしれません。
ローカルでプログラムのデバッグに使用される標準ツールは、マネージド環境に対応していないため、マネージドのクラウド環境で実行されるコードのデバッグは、面倒でエラーが発生しやすいプロセスとなりえます。また、トレーニング ジョブが動かなくなり、ログや指標を出力することなく進行を停止することもあります。そこでジョブに対しインタラクティブにアクセスできれば、デバッグのプロセス全体を大幅に簡易化できる可能性があります。
今回の記事では、Vertex AI カスタム トレーニング ジョブのユーザーの方々がご利用可能な新ツールであるインタラクティブ シェルを紹介します。この機能により、コードを実行している VM へのシェルのような直接アクセスが可能になります。また、任意のコマンドを実行して、ログやモニタリング指標で解決できない問題をプロファイリングまたはデバッグできます。また、トレーニング コードと同じ認証情報を使用してコマンドを実行し、権限の問題や、ローカルで再現できない問題などについて調査できます。インタラクティブ シェルへのアクセスは、通常のカスタム トレーニング ジョブと同じ一連の IAM 権限を使用して認証され、Vertex AI トレーニング環境への安全なインターフェースを提供します。
例: TensorFlow の分散トレーニング
Vertex AI でインタラクティブ シェルをトレーニング プログラムのデバッグに役立てる一例を取り上げます。このケースでは、デッドロックして進行を停止するジョブを Vertex AI トレーニングに故意に送信します。インタラクティブ シェルで py-spy を使用して、問題の根本的な原因を把握します。
Vertex AI は、トレーニング ジョブのスケールアップに役立つマネージド ML プラットフォームです。スケールアップすることで、追加のコンピューティング リソースをフルに活用できるようになります。TensorFlow トレーナーを複数のノードまたはアクセラレータにまたがって実行するにあたり、分散計算を実行するための TensorFlow モジュールである TensorFlow の分散戦略 API を使用できます。それぞれが 1 つ以上の GPU を備えた複数のワーカーを使用するには tf.distribute.MultiWorkerMirroredStrategy を使用します。これは、all-reduce アルゴリズムを使い、複数のデバイス間で勾配更新を同期します。
コードの設定
ここでは、Vertex AI Multi-Worker Training Codelab の例を使用します。この Codelab では、Imagenet で事前にトレーニング済みの ResNet50 モデルと Tensorflow の Cassava データセットを利用して画像分類モデルをトレーニングします。tf.distribute.MultiWorkerMirroredStrategy を使用して、複数のノードでトレーニング ジョブを実行します。
この Codelab では、トレーニング コード用のカスタム コンテナを作成し、GCP プロジェクトの Google Container Registry（GCR）に push します。
ジョブの送信
tf.distribute.MultiWorkerMirroredStrategy は、同期データ並列型アルゴリズムなので、すべてのワーカーに同じ数の GPU があてがわれている必要があります。MultiWorkerMirroredStrategy のドキュメントには、「すべてのワーカーは同じ数のデバイスを使用する必要があります。そうでない場合の動作は未定義です。」とあります。
2 つのワーカープールに異なる数の GPU を割り当てたうえでトレーニング ジョブを送信することで、デッドロック動作の例をトリガーできます。これにより、このジョブが無期限に停止することがわかります。ジョブが停止している間はログが出力されず、使用率の指標にも使用状況が表示されないため、インタラクティブ シェルを使用して調べます。これによりジョブが停止した時点のコールスタックを取得し、詳しい分析に役立てることができます。
Vertex AI トレーニング ジョブの送信には Cloud Console、REST API、Python SDK または gcloud を使用できますが、ジョブ リクエストの際に customJobSpec.enableWebAccess API フィールドを true に設定するだけでインタラクティブ シェルを利用できるようになります。ここでは Cloud Console を使用して、インタラクティブ シェルを有効にしジョブを送信します。
Cloud Console を使用する場合:
[トレーニング方法] で、[管理されているデータセットなし] と [カスタム トレーニング] を選択します。
[モデルの詳細] で、[詳細設定] のプルダウンを開き、[トレーニングのデバッグを有効にする] を選択します。これでトレーニング ジョブのインタラクティブ シェルが有効になります。
3. [トレーニング コンテナ] で、[カスタム コンテナ] を選択し、次に Codelab で GCR に push したコンテナ（gcr.io/$PROJECT_ID/multiworker:cassava）を選択します。
4. [コンピューティングと料金] で、2 つのワーカープールを作成します。ワーカープール 0 には、1 GPU を有する単一のチーフノードを置き、ワーカープール 1 には、それぞれ 2 つの GPU を有する 2 つのワーカーを設定します。これは Codelab で使用されている構成から外れており、ワーカープールごとに GPU の数が異なるため、デッドロック動作がトリガーされます。
[トレーニングを開始] ボタンをクリックし、ジョブがプロビジョニングされるのを待ちます。ジョブの実行が開始されてからログと指標を確認すると、デッドロックが発生したことを確認できます。CPU 使用率の指標は 0% で停止し、ログには何も出力されていません。
インタラクティブ シェルへのアクセス
インタラクティブ シェルはジョブと連動して作成されるため、ジョブが実行中の状態にあるときにのみ使用できます。ジョブが完了するか、キャンセルされると、そのシェルにアクセスできなくなります。今回のサンプル ジョブの場合、ジョブがプロビジョニングされ開始されるまで、約 5～10 分かかります。
ジョブが実行されると、ジョブの詳細ページにインタラクティブ シェルへのリンクが表示されます。ジョブのノードごとに、ウェブ ターミナルのリンクが 1 つ作成されます。
ウェブ ターミナルのリンクの 1 つをクリックすると、新しいタブが開き、トレーニング ジョブのライブ VM でシェル セッションが実行されます。
デバッグに py-spy を使用
py-spy は Python プログラムのサンプリング プロファイラであり、トレーニング アプリケーションの問題をコードを変更せずに調査できます。また、別プロセスでの Python プログラム プロファイリングにも対応しています。これは、インタラクティブ シェルが別のプロセスとして実行されるので便利です。
pip install py-spy を実行して、py-spy をインストールします。これは、コンテナのビルド時またはランタイム時に行えます。インタラクティブ シェルを有効にしたときにコンテナを変更しなかったため、ランタイム時に py-spy をインストールして、停止したジョブを調べます。
py-spy には、Python プロセスのデバッグとプロファイリングに役立つコマンドがいくつかあります。ジョブがデッドロックされているため、py-spy dump を使用して、各 Python スレッドの現行のコールスタックを出力します。
このコマンドをノード workerpool0-0（チーフノード）で実行すると、次の結果が得られます。
上記のコールスタックに役立つエントリがいくつかあります。以下に列挙します。
create_mirrored_variable (tensorflow/python/distribute/distribute_utils.py:306)
_create_variable (tensorflow/python/distribute/mirrored_strategy.py:538)
ResNet50 (tensorflow/python/keras/applications/resnet.py:458)
create_model (task.py:41)
これらのエントリから、コードの create_model 関数の create_mirrored_variable メソッドで停止したことがわかります。これは、TensorFlow により MultiWorkerMirroredStrategy が初期化され、所定のデバイス全体でモデルの変数が複製されたタイミングです。MultiWorkerMirroredStrategy ドキュメントで指定されているように、各マシンの GPU の数に不一致がある場合、その動作は未定義です。この場合、この複製ステップは永久に停止します。
py-spy には、トレーニング ジョブのデバッグに役立つコマンドが他にもあります。たとえば、py-spy top や py-spy record などです。 py-spy top では、プログラムの実行時のリアルタイム更新表示が可能になります。py-spy record では、Python のプロセスが定期的にサンプリングされ、各関数で要した時間を示すフレームグラフが作成されます。フレームグラフは、トレーニング ノードでローカルに書き込まれます。そのため、gsutil を使用して Cloud Storage にコピーし、詳しく分析できます。
クリーンアップ
ジョブが停止した場合でも、ジョブの実行中に使用したインフラストラクチャに対して課金されるため、手動でキャンセルして、超過料金を回避する必要があります。ジョブをキャンセルするには gcloud を使用することができます。
指定のタイムアウト値に達したときに、ジョブを自動的にキャンセルするように Vertex AI を構成することで、超過料金の発生を全般的に回避できます。CustomJobSpec.Scheduling.timeout フィールドを目的の値に設定すると、ジョブが自動的にキャンセルされます。
次のステップ
この投稿では、Vertex AI のインタラクティブ シェルを使用して、ライブ カスタム トレーニング ジョブをデバッグする方法を紹介しました。ランタイム時に py-spy プロファイラをインストールして、実行中のプロセスのライブ コールスタックを取得するために使用しました。TensorFlow 分散戦略の初期化コードでの問題の根本的な原因を特定するのに役立ちました。
インタラクティブ シェルを使用して、次のようなモニタリングやデバッグのユースケースを実現することもできます。
ノードの永続ディスクに書き込まれたローカル ファイルまたは一時ファイルの分析。
権限の問題についての調査。インタラクティブ シェルは、Vertex AI がトレーニング コードを実行する際と同じサービス アカウントを使用して認証されます。
実行中のトレーニング ジョブの、perf または nvprof（GPU ジョブの場合）を使用したプロファイリング。
詳細については、次のリソースをご確認ください。
インタラクティブ シェルを使用したトレーニングのモニタリングとデバッグに関するドキュメント。
Vertex AI に送信する前に、トレーニング コードをコンテナ化してローカルで実行することに関するドキュメント。
GitHub リポジトリにある Vertex AI を使用したサンプル。
- ソフトウェア エンジニア Rajiv Pasricha
