NVIDIA AI Workbenchを使ってみる

概要

NVIDIA AI Workbenchを触ってみたため、構築手順と軽く使ってみるところまでを備忘として残します。

本記事の対象

• NVIDIA AI Workbenchを触ってみたい方
• NVIDIA GPUを使ったコンテナ環境での開発を便利にしたい方

NVIDIA AI Workbenchとは

NVIDIA AI WorkbenchはAIプロジェクト開発のための開発者ツールキットです。
用いることでAI Workbenchを通じてプロジェクトの管理やコンテナの実行、コードの作成などをスムーズに行えます。
また、インストールの過程でコンテナやGPUドライバなど開発で必要な環境も合わせてセットアップしてくれます。
(Windowsの場合WSL有効化もやってくれる)

本記事執筆時点では以下OSにて利用することができます。
※ Windows Server2022は試したところ、だめでした

• Windows 11 (build 22000 or later)
• Windows 10 (build 19044 or later)
• Ubuntu 24.04
• Ubuntu 22.04
• macOS Monterrey (12) or later

構築環境

今回はUbuntu ServerとWindowsクライアントマシンをそれぞれ1台ずつ用意しています。
Ubuntu Serverは物理GPUを搭載しており、実際にコード等を実行するコンテナが動く場所として用います。
WindowsクライアントはAI Workbenchの操作用マシンとして用います。

実際は両方ともAI Workbenchをインストールしておき、Windowsクライアント側のAI WorkbenchにRemote LocationとしてUbuntu Server側のAI Workbenchを登録して繋げる形です。
利用者がAI Workbenchを用いて開発などをする際はWindows側のAI Workbenchを操作して、Ubuntu Server上のコンテナを操作する感じですね。

※ Windows側のDockerは使わないですが、NVIDIA AI Workbenchのインストール過程で入ります

それぞれの簡易的なマシンの状態は以下です。

• Ubuntu Server
  • OS: Ubuntu Server 24.04
  • コンテナランタイム: Dockerインストール済
  • GPU: 有 (L40s)
  • GPUドライバ: 未
• Windowsクライアントマシン
  • OS: Windows10 Pro
  • コンテナランタイム: 未
  • WSL: 未
  • GPU: 無

流れ

流れとしては以下となります。

• NVIDIA AI Workbenchセットアップ
  1. Ubuntu Server側 (Remote Location) の準備
     1. NVIDIA AI Workbench インストール
     2. docker buildx インストール
  2. Windows側 (AI Workbench操作側) の準備
     1. SSH Keyの準備
     2. NVIDIA AI Workbench インストール
     3. Remote Locationの追加
• NVIDIA AI Workbenchを使ってみる
  1. Hello World
  2. コンテナカスタマイズ
  3. GPU割り当て

NVIDIA AI Workbenchセットアップ

Ubuntu Server側 (Remote Location) の準備

ここでの操作はすべてWindowsクライアント側からUbuntu ServerにSSHで接続して行います。

NVIDIA AI Workbench インストール

まずUbuntu ServerにNVIDIA AI Workbenchのインストールを行います。
公式ドキュメントに従い、以下コマンドを実行します。

mkdir -p $HOME/.nvwb/bin && \
curl -L https://workbench.download.nvidia.com/stable/workbench-cli/$(curl -L -s https://workbench.download.nvidia.com/stable/workbench-cli/LATEST)/nvwb-cli-$(uname)-$(uname -m) --output $HOME/.nvwb/bin/nvwb-cli && \
chmod +x $HOME/.nvwb/bin/nvwb-cli && \
sudo -E $HOME/.nvwb/bin/nvwb-cli install

実行するとインストールウィザードが表示されるため、案内に従いながら進めていきます。

まず初めにコンテナランタイムのインストールを聞かれます。
今回はDockerがインストール済みのため、そちらを選択して次に進みます。

続いてGPU Driverのインストールを聞かれます。
今回は未インストールのため、こちらでインストールを行ってもらうよう選択して次に進みます。

インストールプロセスが開始されるため、しばらく待ちます。
以下画像のようにこのタイミングでGPU Driverのインストール等も行ってくれます。

インストールが完了したら Reboot Nowを選択して再起動を待ちます。
この手順の中でGPU Driverのインストールを行わない場合は再起動は不要となります。

再起動完了後、以下コマンドが正常に動作することを確認しておきます。
GPU DriverとNVIDIA AI Workbenchが無事にインストールされていることが確認できます。

nvidia-smi

## 出力
| NVIDIA-SMI 550.90.07              Driver Version: 550.90.07      CUDA Version: 12.4 

nvwb version

## 出力
nvwb 0.27.7 (stable; linux; amd64; go1.21.11; Tue Jun 18 21:41:52 UTC 2024)

docker buildx インストール

NVIDIA AI WorkbenchではコンテナのBuildにbuildxを用います。
そのため、あらかじめインストールを行っておく必要があります。
以下コマンドにてインストールを実施しておきます。

sudo apt install docker-buildx

インストール後、実行できることを確認しておきます。

docker buildx version

## 出力
github.com/docker/buildx 0.12.1 0.12.1-0ubuntu2.1

以上でUbuntu Server側の準備は完了です。

Windows側 (AI Workbench操作側) の準備

続いてWindows側の準備を進めていきます。

SSH Keyの準備

後ほどUbuntu Serverを登録するために、公開鍵での認証が必要となります。
そのため、以下コマンドをWindowsのpowershellで実行し、鍵の作成とUbuntu側への公開鍵登録を行っておきます。

ssh-keygen -f "C:\Users\<ユーザー名>\.ssh\id_rsa" -t rsa -N '""'
type $env:USERPROFILE\.ssh\id_rsa.pub | ssh <ubuntuユーザー>@<ubuntu IP> "cat >> .ssh/authorized_keys"

※< >で囲っているところは環境に合わせて変更が必要

NVIDIA AI Workbench インストール

まずインストーラーを公式サイトからダウンロードしておきます。

インストーラーを実行すると以下のような画面が表示されるため、[Begin Installation]をクリックします。

するとWSLの有効化を聞かれるため、そのまま進んで有効化処理を待ちます。

有効化処理が完了すると再起動が案内されるため、再起動を行います。(WSL側の都合)

再起動完了後、NVIDIA AI Workbenchを起動するとWSLのDistributionのダウンロードが開始されます。
(WindowsアプリケーションにNVIDIA AI Workbenchが登録されています)

Distributionのダウンロードが完了すると、コンテナランタイムを聞かれます。
今回はDockerを選んで次に進みます。

すると、Docker Desktopをインストールするよう案内がされます。
ここは手動で表示されているリンクからインストーラーをダウンロードしてインストールを行っておきます。

Docker Desktopインストール後は一度サインアウト、サインインを行い再度AI Workbenchを起動します。

すると、再びコンテナランタイムを選択する画面が表示されるため、Dockerを選んで次に進みます。
(dockerの選択肢の左上にInstalledが記載されています)

パッケージのインストールが開始されるため、完了まで待ちます。

パッケージのインストールが完了するとGit Serverの選択肢が表示されます。
ここで登録しておくことで、AI Workbenchで作成したコードなどをそのままgitへコミットすることができたりできるようです。

今回は一旦下の[Skip and connect Git server later]を選択します。

ここまで実施するとインストール完了となります。

Remote Locationの追加

続いてUbuntu ServerをRemote Locationとして追加していきます。
インストール完了後は起動すると以下のようなLocation選択画面が表示されるようになります。
[Add Remote Location]を選択します。

Remote Locationの接続情報等を入力する画面が表示されます。
ここのSSH Ke Fileのところで準備しておいた秘密鍵を選択する形になります。
すべて入力を終えたらAdd Locationを選択します。

無事にLocationが追加されたら完了となります。

NVIDIA AI Workbenchを使ってみる

それでは最後にNVIDIA AI Workbenchを軽くコードの実行等を行っていきます。

Hello World

AI Workbenchを起動して、Locationを選択するとプロジェクトの選択画面が表示されます。
まだ1つも無い状態のため、[Start a New Project]で新しく作っていきます。

プロジェクト名などを入力して、次に進みます。
今回は[hello-world]にしています。

次は実行するコンテナイメージの選択する画面が表示されます。
今回はPython Basicを選択して[Create]を選択します。

するとプロジェクトが作成され、フォルダ構造等が表示されます。

なお最初はBuild中のため、画面下のOutputをクリックしてBuild完了まで待ちます。

Buildが完了したら左のペインから[Environment]を選択します。
ここでは含まれているPackageの確認や、ビルドのスクリプトの編集、コンテナにマウントするフォルダなどを設定したりすることができます。
一旦そのまま右上の[Start Environment]をクリックします。

[Start Environment]の処理完了後右上の[Open Jupyterlab]をクリックします。
するとブラウザでJupyter labが起動してきます。
今回はNotebookからPython3を選択します。

Hello Worldを表示するコードを実行して、問題なくPythonが動作することが確認できます

ちなみにこのタイミングでubuntu Server側にSSHで繋いでコンテナを確認すると実行されていることが確認できます。

docker ps

## 出力
CONTAINER ID   IMAGE                                       COMMAND                  CREATED          STATUS          PORTS                                       NAMES
86172c9477e5   project-hello-world                         "/entrypoint.sh tail…"   3 minutes ago    Up 3 minutes                                                project-hello-world
d2d3d2e3ebea   traefik:v2.10.7                             "/entrypoint.sh trae…"   38 minutes ago   Up 11 minutes   80/tcp, 127.0.0.1:10000->10000/tcp          workbench-proxy

コンテナカスタマイズ

続いてコンテナのカスタマイズを行います。
厳密には以下Python ライブラリの追加をここでは行ってみます。

torch
numpy
matplotlib
pandas
scikit-learn

Jupyterlabの画面で左のrequirements.txtを開き、先程のライブラリを追記して保存します。

AI Workbenchに戻ると、変更を検知してRebuildの案内が表示されています。
[Rebuild & ReStart]を選択します。

Restart完了後、Jupyterlabに戻りライブラリがインストールされていることを確認します。
以下のようにimportができて実行できることが確認できます。

ちなみにAI WorkbenchのEnvironmentのPackagesを見ると、追加したライブラリ群が認識されていることも確認できます。

GPU割り当て

それでは最後にGPUの割り当てを行っていきます。
デフォルトではGPUを使わない設定でコンテナを起動するため、手を加えます。

AI WorkbenchのEnvironmentの中のHardwareを確認します。
Number of Desired GPUsの項目で割り当てるGPU数を選ぶことができます。
今回は1を選択しておきます。
選択後はRestartの案内がでるため、そのまま[Restart]を選択します。

Restart完了後、Jupyterlabに戻りtorchのライブラリでGPUを確認できるため以下のように実行します。
すると無事GPUを認識できていることが確認できます。

まとめ

今回はNVIDIA AI Workbenchを使ってみました。
思ったより使いやすく、コンテナの実行環境を提供する際に便利だなと感じました。
今回は試しきれていませんが、VS Codeと連携したりGitとの連携等も別途触ろうと思います。

おまけ

記事上はきれいに順番を整理していますが、初め試行錯誤する中でAI WorkbenchのRemote Locationにて使われるworkbench-proxyのコンテナが落ちていたときがありました。
手順結構ごちゃった時になったため、基本気にする必要は無いのですが、もし落ちていたら再起動をしてあげてください。

docker restart workbench-proxy

参考にしたサイト