📓

VertexAI WorkbenchのGPU付きインスタンスにVSCodeでSSH接続する

2022/07/12に公開

こんにちわ alivelimb です。
最近 GCP を触る機会が増えており、サクッと機械学習環境を整えたいケースがありました。
その際Vertex AI Workbenchを使うことで、環境構築が簡単に出来たので備忘録として残します。

はじめに

本記事では VSCode Remote SSH の設定手順については紹介しません。
詳しくは公式ドキュメントを適宜参照して下さい。

また、SSH キーペアについては既に作成していることを前提としています。
作成方法は公式ドキュメント等を適宜参照して下さい。

Vertex AI Workbench とは

ざっくり言うと  JupyterLab の GCP マネージドサービスです。
VM インスタンス立てて、Jupyter をインストールして、SSH/HTTP(S)アクセスできる様に設定する手間なく、簡単に JupyterLab 環境を用意することが出来ます。Pytorch や Tensorflow といった深層学習フレームワークを事前にインストールさせたり、GPU をアタッチしてすぐ使うことも出来ます。

また、Cloud Storage や BigQuery の接続が容易なためデータ分析〜モデル構築がスムーズに行えたり、Vertex AI Pipelinesを用いてパイプラインに組み込んだりすることが可能です。

GPU インスタンスを作成・設定する

以下の手順で GPU インスタンスを作成・設定していきます。

  1. GCP プロジェクトで GPU を割り当てできるか確認する
  2. GPU インスタンスを作成する
  3. SSH 公開鍵を登録する

GPU の割り当てを確認する

「IAMと管理」 > 「割り当て」を開きます。

割り当て画面

フィルタで指標: compute.googleapis.com/gpus_all_regionsと検索します。

フィルタ画面

上限が 0 であれば、右上の「割り当てを編集」から上限を 1 以上に設定します。
※承認されるまでしばらく待つ必要があります(検証した際は 5 分ほどで承認されました)。

GPU インスタンスを作成する

ユーザ管理のノートブックタブに合わせます。

画像

「新しいノートブック」から任意の言語・フレームワークを選択し、「With 1 Nvidia~」を選択します。
※検証環境では「With 1 Nvidia Tesla T4」となっていました。

ノートブック名などを適宜変更した後、「NVIDIA GPU ドライバを自動的にインストールする」にチェックをいれて作成を押下します

SSH 公開鍵を登録する

作成したノートブックの詳細を表示し、「VM の詳細を表示」を選択します。
※ノートブックが動作している VM インスタンスの詳細ページに遷移します。

ノートブック詳細

VM インスタンスに SSH の公開鍵を登録します

SSH公開鍵

インスタンスに SSH でアクセスする

以下のような SSH Config ファイルを作成し、VSCode で Remote SSH します。

Host user1
  HostName [インスタンスの外部IPアドレス]
  User [キーペアに登録したユーザ名]
  IdentityFile [インスタンスに登録した公開鍵に対応する秘密鍵へのファイルパス]

正しく設定出来ていれば、nvidia-smi, nvcc コマンド等が叩けると思います。

まとめ

本記事では VertexAI Workbench で GPU 付きインスタンスを作成し、VSCode で SSH 接続する方法を紹介しました。JupyterLab も便利ですが、使い慣れた VSCode も使いたいと言う方は是非一度試してみて下さい。

また、VSCode の拡張機能を活用し静的解析やフォーマッタを入れて、開発効率を向上させることも可能です。こちらについては記事にしているので、よろしければ参照して下さい。

Discussion