こんにちは、@moken20です。大学では量子機械学習の研究をしています。

かなり今更ですが、産総研AI橋渡しクラウド(ABCI)の環境構築で自分がつまずいた多段SSH接続やprivateリポジトリのcloneを中心にまとめておきたいと思います。ABCI 2.0 はもうすぐ運用終了してしまうようですが、次期ABCIも同様の方法でセットアップできることを期待して、誰かしらの力になれば幸いです。

より詳細なサーバーの構成やストレージ等の情報は公式ユーザーガイドを参照してください。

計算環境のセットアップ

説明やコマンドで登場する {USER} {GROUP} は、それぞれABCIのアカウント名（例: abc12345de）とグループ名（例: gxx12345）を示します。各自で置き換えるか、あらかじめ環境変数に設定しておくと良いです。なお、ABCIのサーバーで環境変数を設定したい場合は、~/.bash_profile に以下を追加してください。

$ export USER="abc12345de"
$ export GROUP="gxx12345"

SSH多段階接続

このステップはすべてローカルのコンピュータで実行します。

1. SSH公開鍵・秘密鍵ペアの生成

生成時に設定したパスフレーズは、SSH接続のたびに必要となるため、忘れないようにしてください。既存の鍵と使い分けたい場合は、別途 -f オプションなどでファイル名を指定する必要があります。ただし、以降はデフォルトのパスに鍵があることを前提として説明します。

$ cd ~/.ssh
$ ssh-keygen                 # ssh keyの生成

2. 公開鍵をABCIに登録

公開鍵認証を行うため、~/.ssh/id_rsa.pubに保存された公開鍵をABCIの利用者ポータルに登録してください。また、後にリモートサーバーでGitHubのprivateリポジトリをcloneしたい場合は、GitHubにも公開鍵を登録しのちのステップ3を実施してください。githubからcolneするつもりがない、またはリポジトリがpublicな場合は3をスキップしても問題ないです。
ABCI: 利用者ポータル
GitHub: SSH and GPG keys

$ pbcopy < ~/.ssh/id_rsa.pub  # クリップボードに公開鍵をコピー

(3. SSH-Agentに秘密鍵の登録 ※privateリポジトリをcloneしたい場合)

SSH Agent Forwarding 機能を利用して、ローカルに秘密鍵を置いたままリモートで SSH公開鍵認証 (git clone) を可能にします。詳しくは ssh-agent による公開鍵認証と SSH Agent Forwarding による多段ログインを参照してください。

$ ssh-add -K ~/.ssh/id_rsa  # ssh-agentに秘密鍵の登録

4. proxy jumpの設定

インタラクティブノードに接続するためには、アクセスサーバーを踏み台にする必要があります。本来はSSHトンネルの作成後にログインする必要がありますが、以下の ~/.ssh/config ファイルを作成することで、1つのコマンドでログインできるようになります。

$ vi ~/.ssh/config  # vimでconfigファイルを開く
$ i                 # インサートモードに切り替える
下記の内容をコピペし、Userの項目のみ自分のABCIアカウント名に書き換える
ESC                 # escapeキーでコマンドモードに戻る
$ :wq               # 保存して終了

Host abci
    HostName es
    User {USER}
    ProxyJump %r@as.abci.ai
    IdentityFile ~/.ssh/id_rsa
    HostKeyAlgorithms ssh-rsa
    ForwardAgent yes

Host abci-a
    HostName es-a
    User {USER}
    ProxyJump %r@as.abci.ai
    IdentityFile ~/.ssh/id_rsa
    ForwardAgent yes

Host as.abci.ai
    IdentityFile ~/.ssh/id_rsa
    ForwardAgent yes

5. VS CodeでSSH接続

VS Codeの拡張機能「Remote Development」を用いてインタラクティブノードに接続します。方法は、Visual Studio Code で Remote SSH するのステップ1とステップ5を実行すればよく、接続先として表示されるの3つの選択肢のうち、とりあえず abci というホスト名のサーバーに接続してください。

仮想環境の構築

このステップは接続したABCIのサーバーの中で実行します。

1. GitHubのリポジトリをクローン

すでにリモートサーバには git がインストールされているため、すぐにクローンしてしまって問題ないです。

# publicリポジトリの場合
$ git clone https://github.com/xxx/xxx.git

# privateリポジトリの場合
git clone git@github.com:xxx/xxx.git

privateなリポジトリを上記のコマンドでクローンできず Authentication error になってしまう場合は SSH Agent Forwarding がうまくできていない可能性が高いです。どうしてもエラーが解決できない場合は、個人のGitHubアカウントで PATを発行し以下コマンドで登録すれば HTTPS の方でクローンすることができます。しかしこの方法だとリポジトリにアクセスするための認証情報がリモートサーバーに平文で保存されるため、セキュリティリスクにはなってしまいます。

$ git config --global credential.helper store
$ export GITHUB_TOKEN=https://{GITHUB-USERNAME}:<PAT>@github.com/xxx  # usernameとpatは各自で
$ echo $GITHUB_TOKEN > $HOME/.git-credentials

$ git clone https://github.com/xxx/your_repository_name.git

2. 計算ノードに接続

インタラクティブノードと計算ノードはストレージを共有しているため、どちらでpythonの環境構築を行なっても問題ないらしいです。ただ、GPUを搭載しているサーバで行わないと、cudaに対応したpytorchがうまくインストールできない恐れがあるので、今回は計算ノードの中で作業します。

$ qrsh -g {GROUP} -l rt_G.small=1

3. Python環境の構築

接続した計算ノードの中で以下のコマンドを実行してください。
使用したいpytorchのバージョンがすでに決まっている場合、それと対応するcudaのバージョンをPyTorchとGPU/CUDA周りの環境構築のバージョン解決等を参考にしながら選んでください。また、仮想環境に関してはpipenvやpoetryはおそらく後のバッチジョブで使用できないため、venvの仮想環境を作成するとよいです。

$ module avail                      # loadできるモジュールの種類やバージョンを確認
$ module load python/x/x cuda/x/x   # 最低限GPUで学習・推論ができるモジュールの例(バージョンは要検討)
$ cd {your_repository_name}
$ python3 -m venv {newenvname}
$ source ~/venv/{newenvname}/bin/activate
$ pip install -r requirements.txt   # すでにrequirements.txtがある場合

上記を実行したのちにコマンドラインが (newenvname)[{USER}@gXXXX your_repository_name] $ となっていれば、環境構築が完了となります。さらにGPUが使用できる環境になっているかを確認したい場合は、以下を実行するとよいです。

$ python
>>> import torch
>>> print(torch.cuda.is_available())

4. 計算ノードからデタッチする

計算ノードに接続している間は料金が発生するため、環境の構築が確認できたら $ exit コマンドで終了してください。インタラクティブノードからもデタッチする場合は、同様に $ exit コマンドを打ってください。

pythonファイルの実行

1. 利用するGPUの種類を選択し、対応するインタラクティブノードに接続する

NVIDIAのA100かV100のどちらを計算資源として利用するか選択してください。 A100のほうが単体のメモリや演算性能は高いが料金も割高となります (利用料金表)。利用するGPUの種類によってSSH接続するインタラクティブノードが異なり、V100なら abci 、A100なら abci-a を選択してください。先のステップでは abci のノードでのみ環境構築をおこいましたが、abci-a ノードでもその環境は使用できます。

2. ジョブの実行タイプを選択し実行する

計算ノードで実行するジョブのタイプは以下の3つであり、用途に応じて使い分ける必要があります。コマンドオプション等の詳細は公式ドキュメントを参照するとよいです。

インタラクティブジョブ
プログラムのデバッグやjupyter notebookなどの対話的なアプリケーション、可視化ソフトウェアの実行に適したジョブです。
- 実行コマンド : qrsh -g {GROUP} -l resource_type=num [options]
- インタラクティブジョブ実行の例 :
```
$ qrsh -g {GROUP} -l rt_G.small=1 -l h_rt=1:00:00  # rt_G.smallのノードを1時間起動
$ module load python/x/x cuda/x/x  # モジュールのロードはジョブの度に必要
$ cd {your_repository_name}
$ source ~/venv/{newenvname}/bin/activate
$ python test.py
```

バッチジョブ
対話的な処理が必要のないアプリケーションに適したジョブです。インタラクティブジョブよりも長期間または並列度の高いジョブの実行が可能となります。バッチジョブでは計算ノードの中で作業ができないので、事前に実行するコマンドを.shファイルにまとめてそのパスを指定してください。
- 実行コマンド : $ qsub -g {GROUP} -l run.sh
- run.shファイルの例:
```
 #!/bin/bash
 #$ -l rt_G.small=1
 #$ -j y
 #$ -N test
 #$ -o logs/
 #$ -cwd

 # 仮想環境の初期化
 source /etc/profile.d/modules.sh
 # モジュールのロード
 module load python/x/x cuda/x/x
 # 仮想環境の有効化
 source ~/venv/{newenvname}/bin/activate
 # プログラムの実行
 python test.py
```

予約ジョブ
計算リソースを日単位で事前に予約して利用できるサービスです。バッチジョブより長時間のジョブ実行が可能となります。

予約コマンド : $ qrsub [options]

予約実行の例 :

# 2024年7月5日から1週間 (7日間) 計算ノード(V)4台を予約
$ qrsub -a 20240705 -d 7 -g {GROUP} -n 4 -N "Reserve_for_AI"
Your advance reservation 12345 has been granted

# 発行された予約IDを -ar オプションで指定しバッチジョブを実行
$ qsub -g {GROUP} -ar 12345 run.sh

notebook ファイルの実行

多段階のssh接続をしている関係で VSCode から計算環境の jupyterカーネルに接続できません。ipynbファイルを実行するには、ポートフォワーディングを利用してwebブラウザからリモートのjupyterサーバーに接続する必要があります。
参考サイト: Jupyter Notebook から ABCI の python仮想環境で tensorflow+keras を利用する

1. jupyterサーバーの起動

ローカルPCのターミナルからABCIサーバーに接続し、インタラクティブジョブで jupyter サーバーを起動します。以下の3つの目のhostnameコマンドで表示される、g1234.abci.localの'g1234'にあたる計算ノード名はステップ3で使用すためメモしておいてください。またjupyter関連のライブラリが仮想環境にインストールされていない場合は、6つめのコマンドの前にpipでインストールしておいてください。

$ ssh abci 
$ qrsh -g {GROUP} -l rt_G.small=1 -l h_rt=2:00:00
$ hostname
$ module load python/x/x cuda/x/x
$ source ~/venv/{newenvname}/bin/activate
$ jupyter notebook --ip=`hostname` --port=8888 --no-browser

2. アクセスサーバーへのポートフォワーディング

1とは別のターミナルを開き、ローカルPCからアクセスサーバーへの port forwardingを指定してください。

$ ssh -L 10022:es:22 -l {USER} as.abci.ai

3. jupyterサーバーへのポートフォワーディング

最後に3つ目のターミナルで以下のコマンドを実行してください。このコマンドによって、既に設定済みの localhostのポート10022からアクセスサーバーのポート22番のトンネルを介して転送さるようになります。また、ここで、ステップ1でメモした計算ノード名を{計算ノード名}に代入するよう注意してください。

$ ssh -N -L 8888:{計算ノード名}:8888 -l {USER} -p 10022 localhost

4. ローカルPCのWebブラウザからアクセス

Jpyter Notebookサーバーを起動させたときに表示された http//127.0.0.1:8888/?token=???? の形式のURLをローカルPCの Web ブラウザに入力すると無事アクセスすることができます。

さいごに

以上、ABCIのセットアップ方法を多段SSH接続の工程を中心に説明しました。わかりにくい点や誤り等ありましたら、コメントいただけると幸いです。