NVIDIA A100でMIGを切り替えて使う＋Slurmから便利に使う

# 概要
「A100のMIGの有効無効を切り替えるにはホストの再起動が必要で、H100ではそれがいらなくなった」的な認識だったのだが、改めてドキュメントを読んだらそんなことは書いてなかったぜ。
というわけで、調査をしてみたら、A100でも再起動などせずともGPUのリセットさえすれば切り替え可能であった。問題はGPUのリセットができるか、であった。というお話。
slurm管理下にある場合に便利にできないかな？というあたりも含めてメモしておく。

* 参考資料（MIGのユーザーガイド）

https://docs.nvidia.com/datacenter/tesla/mig-user-guide/index.html


**※ 途中にも書いたのだが、サービスの再起動をしている都合上、複数GPU環境において、特定のGPUだけを切り替えたい（他のGPUを利用中に特定のGPUだけ変更したい）という場合にはこの方法では駄目だと思われる点には注意が必要。あくまで今回の環境ではこの方法で良いというだけである。**



# そもそもなにがしたいの？
そもそもMIGを有効にしたり無効にしたりを頻繁に行いたいという需要自体あまりないと思うのだが、諸般の事情により`一時的にMIGの状態を変えてプログラムを実行したい`、という要求を満たしたい。1ジョブの性能をしっかり出したいならMIGを無効化したままで良いし、GPUを触れるジョブを多数流したいのであれば常にMIGを有効化するのが妥当。それがわかったうえで、都合により時々切り替えたいという人向けのニッチな話である。


# 実験環境（対象環境）
- NVIDIA A100 PCIe （1枚だけ）
- Ubuntu 22.04.4 LTS
- slurm経由で計算ノードとして利用
- ランレベルは multi-user.target
- ホスト名は fire （一部のスクリプト例で使っている）
ラックマウントのワークステーションをGPU実験機として利用している。ちなみにCPUはEPYC Milanだが今回の設定の話には関係ない。

# MIGの有効無効を切り替える
## 基本手順
`nvidia-smi -mig 0` で無効化。
`nvidia-smi -mig 1` で有効化。

## 何が問題なのか？
実際に切り替えようとすると以下のように別のプロセスがGPUを使ってるから駄目よと弾かれる。
```
$ sudo nvidia-smi -mig 1
Warning: MIG mode is in pending enable state for GPU 00000000:C1:00.0:In use by another client
00000000:C1:00.0 is currently being used by one or more other processes (e.g. CUDA application or a monitoring application such as another instance of nvidia-smi). Please first kill all processes using the device and retry the command or reboot the system to make MIG mode effective.
All done.
```

GPUのリセットをすれば良いということはユーザーガイドに書かれている。
しかしこれも同様にGPUが使われているとエラーしてしまう。
```
$ sudo nvidia-smi --gpu-reset
The following GPUs could not be reset:
  GPU 00000000:C1:00.0: In use by another client

1 device is currently being used by one or more other processes (e.g., Fabric Manager, CUDA application, graphics application such as an X server, or a monitoring application such as another instance of nvidia-smi). Please first kill all processes using this device and all compute applications running in the system.
```

ランレベルを`multi-user.target`にしているなどして画面出力は特にしていない状況でもこれが起きる。

~~＼詰んだ／~~ これをどうすれば良いのかというのが本題。


## 対処法：dcgmサービスを止める
色々と調べていたら、実はdcgmを止めてやればMIGの有効無効を切り替えられることがわかった。
```
systemctl stop nvidia-dcgm
nvidia-smi -mig 0
nvidia-smi -r
systemctl start nvidia-dcgm
```
とか
```
systemctl stop nvidia-dcgm
nvidia-smi -mig 1
nvidia-smi -r
systemctl start nvidia-dcgm
```
とかやれば良いだけ。`nvidia-smi`の`-r`は`--gpu-reset`の短縮形。ちなみに上記の手順ではGPUリセットを省略しても切り替わってくれているよう見えた。状況にもよるのか？

なお今回のように1ワークステーションを単独でてきとーに使っているだけではDCGMに明示的にお世話になることはなさそうな気がするのだが、しっかりした管理ソリューション的な何かを使っているときはこのサービスの再起動自体が何か悪さをする可能性があるかもしれない点には注意した方が良いかもしれない。
また、サービスの再起動（停止と起動）をしているということは、複数GPU環境において、特定のGPUだけを切り替えたい（他のGPUを利用中に特定のGPUだけ変更したい）という場合にはこの方法では駄目だと思うので注意が必要。

参考：DCGMの紹介ページ

https://developer.nvidia.com/dcgm

# Slurmから便利に使いたい
対象ワークステーションは数人で共有利用している環境にあり、Slurm環境下の計算ノードとして利用している。MIGを使いたい人は多くはない。どのように運用すると便利だろうか？

## ユーザのジョブスクリプト内で切り替える
ユーザがジョブスクリプト内で明示的に切り替えるにはどうすれば良いだろうか？

例えばこんな感じにすれば、特定のジョブのみMIG有効状態で使える。

```bash:job.sh
#!/bin/bash -x
#SBATCH -p fire
module load nvhpc
# enable MIG
sudo /usr/bin/systemctl stop nvidia-dcgm
sudo /usr/bin/nvidia-smi -mig 1
sudo /usr/bin/nvidia-smi -r
sudo /usr/bin/systemctl start nvidia-dcgm
sudo /usr/bin/nvidia-smi mig -cgi 1g.5gb,1g.5gb,1g.5gb,1g.5gb,1g.5gb,1g.5gb,1g.5gb -C
nvidia-smi -L | grep MIG | awk '{print $6}' | sed 's/)//' 2>&1 | tee mig.txt
mpirun -n 7 ./run2.sh ./a.out
# disable MIG
sudo /usr/bin/systemctl stop nvidia-dcgm
sudo /usr/bin/nvidia-smi -mig 0
sudo /usr/bin/nvidia-smi -r
sudo /usr/bin/systemctl start nvidia-dcgm
```

```bash:run2.sh
#!/bin/bash
ID=$(( ${OMPI_COMM_WORLD_RANK} + 1 ))
GPU=`head -n ${ID} ./mig.txt | tail -n 1`
export CUDA_VISIBLE_DEVICES=${GPU}
echo $ID $GPU
echo $@
$@
```

ジョブの冒頭で`systemctl`や`nvidia-smi`を叩いてMIGを有効化し、`nvidia-smi -L`の情報からGPUのUUIDを抽出、それを使ってMPIで並列実行する事例。
これなら特定のジョブだけMIG有効になるので他のジョブには迷惑をかけない。もちろん、これは`fire`というパーティションが対象ホストを1ジョブだけで占有する（他のジョブと同時には流れない）という構成である前提が必要である点には注意が必要。

なおジョブスクリプト内で`sudo`をせねばならないため、ホスト側にそれに対応した設定が必要。例えば`/etc/sudoers`の末尾に以下のような設定を書いてパスワードなし`sudo`でのコマンド実行を許す必要がある。大変お行儀が悪い気がするが、仕方なし。（もっと良い方法があるかも？）

```
user  ALL=(ALL) NOPASSWD: /usr/bin/nvidia-smi
user  ALL=(ALL) NOPASSWD: /usr/bin/systemctl
```


## prolog, epilogで切り替える
### epilogで破棄する
上述の例では、ジョブに問題が起きた場合などにMIGが有効なままで資源が解放されてしまい、次にジョブを実行する人が困ってしまう可能性ある。
対策としては、slurmのepilogでクリーンアップをすると良いだろう。

例えばこんな感じのスクリプトを書けばMIGが有効な時に破壊してくれる。MIGの状況確認コードなどが力技。本当はもっとスマートにできるのかもしれないが、とりあえずちゃんと動いているように見える。
（ついでにpersistent modeを破棄したりMPSを止めたりという処理も入れてみた。GPUリセットの際に破棄される気もするので改善の余地があるかもしれない。）
```bash:epilog.sh
#!/bin/bash
HOST=`hostname`
HOST2=${HOST%%.*}
if [ $HOST2 = fire ]; then
        echo "fire" 2>&1 | tee /tmp/log_epilog
        # if MIG is enabled, disable it
        ## if sub GPUs are existing, delete them
        if [ `nvidia-smi -L | grep MIG | wc -l` = '0' ]; then
                echo "MIG device(s) is not existing" 2>&1 | tee -a /tmp/log_epilog
        else
                echo "MIG device(s) is existing" 2>&1 | tee -a /tmp/log_epilog
                # delete MIG
                sudo /usr/bin/nvidia-smi mig -dci
                sudo /usr/bin/nvidia-smi mig -dgi
        fi
        ## disable MIG
        mig=`nvidia-smi|head -n 11|tail -n 1|grep Enabled`
        # echo $mig 2>&1 | tee -a /tmp/log_epilog
        if [[ -n $mig ]]; then
                echo "MIG is enabled, disable it" 2>&1 | tee -a /tmp/log_epilog
                sudo /usr/bin/systemctl stop nvidia-dcgm
                sudo /usr/bin/nvidia-smi -mig 0
                sudo /usr/bin/nvidia-smi -r
                sudo /usr/bin/systemctl start nvidia-dcgm
        else
                echo "MIG is disabled" 2>&1 | tee -a /tmp/log_epilog
        fi
        # disable PM
        sudo /usr/bin/nvidia-smi -pm 0
        # disable MPS
        echo quit | nvidia-cuda-mps-control || true
fi
```

判定に使っているホスト名については`SLURMD_NODENAME=fire`を使っても良いだろう。（むしろその方が確実か？）

`systemctl`や`nvidia-smi`をパスワードなしの`sudo`実行できるように`sudoers`の設定も忘れないこと。
```
slurm  ALL=(ALL) NOPASSWD: /usr/bin/nvidia-smi
slurm  ALL=(ALL) NOPASSWD: /usr/bin/systemctl
```


### MIG用のパーティションを用意する
どうせならジョブの設定でMIGの有無を分けた方が便利かもしれない。ジョブ毎の
`#SBATCH`オプションでどうにかなるかな？と思ったのだが、正直よくわからなかったので、パーティションを分けてみることにした。

これまで`slurm.conf`で用意していたパーティションはこんな感じ：
```
NodeName=fire CPUs=32 Boards=1 SocketsPerBoard=4 CoresPerSocket=4 ThreadsPerCore=2 RealMemory=515800 State=UNKNOWN
PartitionName=fire Nodes=fire Default=NO DefaultTime=10:00 MaxTime=INFINITE State=UP
```

そこに以下のような新しいパーティションを追加し……
```
PartitionName=fire_mig Nodes=fire Default=NO DefaultTime=10:00 MaxTime=INFINITE State=UP
```

`prolog.sh`を書き換えてMIG対応処理を追加。
具体的には`SLURM_JOB_PARTITION`で処理を分岐し、MIGを使う場合と使わない場合のそれぞれで必要な処理を書いた。これまでの内容を踏まえて丁寧に処理しているつもり。

```bash:prolog.sh
#!/bin/bash

# fire with MIG
if [ $SLURM_JOB_PARTITION = fire_mig ]; then
        echo "fire" 2>&1 | tee -a /tmp/log_prolog
        # if MIG is disabled, enable it
        ## if sub GPUs are existing, delete them
        if [ `nvidia-smi -L | grep MIG | wc -l` = '0' ]; then
                echo "sub GPUs are not existing" 2>&1 | tee -a /tmp/log_prolog
        else
                echo "sub GPUs are existing, delete them" 2>&1 | tee -a /tmp/log_prolog
                # delete MIG
                sudo /usr/bin/nvidia-smi mig -dci
                sudo /usr/bin/nvidia-smi mig -dgi
        fi
        ## enable MIG
        mig=`nvidia-smi|head -n 11|tail -n 1|grep Enabled`
        # echo $mig 2>&1 | tee -a /tmp/log_prolog
        if [[ -n $mig ]]; then
                echo "MIG is enabled" 2>&1 | tee -a /tmp/log_prolog
        else
                echo "MIG is disabled, enable it" 2>&1 | tee -a /tmp/log_prolog
                sudo /usr/bin/systemctl stop nvidia-dcgm
                sudo /usr/bin/nvidia-smi -mig 1
                sudo /usr/bin/nvidia-smi -r
                sudo /usr/bin/systemctl start nvidia-dcgm
        fi
        ## configure MIG
        sudo nvidia-smi mig -cgi 1g.5gb,1g.5gb,1g.5gb,1g.5gb,1g.5gb,1g.5gb,1g.5gb -C
        # enable PM
        sudo /usr/bin/nvidia-smi -pm 1
fi

# fire without MIG
if [ $SLURM_JOB_PARTITION = fire ]; then
        echo "fire" 2>&1 | tee -a /tmp/log_prolog
        # if MIG is enabled, disable it
        ## if sub GPUs are existing, delete them
        if [ `nvidia-smi -L | grep MIG | wc -l` = '0' ]; then
                echo "MIG is disabled" 2>&1 | tee -a /tmp/log_prolog
        else
                echo "MIG is enabled" 2>&1 | tee -a /tmp/log_prolog
                # delete MIG
                sudo /usr/bin/nvidia-smi mig -dci
                sudo /usr/bin/nvidia-smi mig -dgi
        fi
        ## disable MIG
        mig=`nvidia-smi|head -n 11|tail -n 1|grep Enabled`
        # echo $mig 2>&1 | tee -a /tmp/log_prolog
        if [[ -n $mig ]]; then
                echo "MIG is enabled, disable it" 2>&1 | tee -a /tmp/log_prolog
                sudo /usr/bin/systemctl stop nvidia-dcgm
                sudo /usr/bin/nvidia-smi -mig 0
                sudo /usr/bin/nvidia-smi -r
                sudo /usr/bin/systemctl start nvidia-dcgm
        else
                echo "MIG is disabled" 2>&1 | tee -a /tmp/log_prolog
        fi
        # enable PM
        sudo /usr/bin/nvidia-smi -pm 1
fi
```

これで、fireパーティションに投入されたジョブはMIGなしで実行され、fire_migパーティションに投入されたジョブはMIGにより7つのサブGPUが有効となった状態で実行されるようになった。MIGの構成を色々と弄りたい場合はこれでは駄目だけど、構成が限定されるなら便利に使えるだろう。

これでMIGのテストも行いやすくなったはず！


# おわりに
**私以外に誰が嬉しいのかは知らない。**