SlurmでN枚GPUの刺さったノードにNジョブ割り当てるには

[追記:2021.04.06] この記事のきっかけとなったプロジェクト

Slurmとは&この記事の内容について

Slurmの説明いらないですよね，みんな大好きジョブスケジューラです．
無料で使えますし，Ubuntuを使っている場合はaptで簡単にインストールできますし，ユーザがそれなりいるのでググれば情報はたくさん出てきます．
そういった点でとても使いやすいものとなっているのかなと思います．

それはそうと，世の中には1ノードに複数(N)GPUを搭載した計算機がたくさん存在します．
ノード数が少ないというのは管理の面から考えると楽なのです．
しかし，ジョブスケジューリングを行う場合，1ジョブに1ノードを割り当ててしまうと有効にGPUを利用できない可能性があります．
例えば，もし1ジョブが1GPUしか利用しない場合，残りの(N-1) GPUは稼働しないこととなるのです．
これは明らかに効率的なジョブスケジューリングとは言えません．
今回の話は，SlurmでN GPUをN ジョブに割り当てるにはどうすればいいかという話です．

どうやるの？

自分はこのリポジトリを参考に設定を行いました．
dholt/slurm-gpu - GitHub

SlurmにはGRES (General RESource)と呼ばれる機能があり，これを用いることで今回行いたい複数GPUを複数ジョブに割り当てることができます．
今回はこれを用いて設定していきます．
GRESは他にもNVIDIAのMPS (Multi-Process Service)やIntelのMIC(Many Integrated Core)にも対応しています．

環境

• OS : Ubuntu 20.04
• Slurm : 19.05.5

今回扱うクラスタは「ログインノード」と「計算ノード」からなります．
ログインノードでslurmdbdとslurmctldを立ち上げ，ジョブスケジューリングをさせます．

データベースの準備（on ログインノード）

sudo /usr/bin/mysql_secure_installation

CREATE USER 'slurm'@'localhost' IDENTIFIED BY 'password';
CREATE DATABASE slurm_acct_db;
GRANT ALL ON slurm_acct_db.* TO slurm@localhost;

InnoDB周りの推奨設定があるらしので，念の為/etc/mysql/my.cnfに以下を追記．

[mysqld]
innodb_buffer_pool_size=1024M
innodb_log_file_size=64M
innodb_lock_wait_timeout=900

sudo systemctl restart mysql.service

SlurmDBDのインスト〜ル（on ログインノード）

SlurmDBDというのはジョブの実行状況などをデータベースで管理するためのデーモンです．
別にインストールしなくともテキストベースの管理でSlurmは動きはするのですが，今回はこれをインストールしろと書かれているのでインストールし設定していきます．

aptで入れます．

sudo apt install slurmdbd

slurmdbdの設定ファイル/etc/slurm-llnl/slurmdbd.confを作成します．
passwordはデータベースの権限付与で設定したものに書き換えてください．
この設定ファイルではログファイルの出力先を/var/log/slurm/slurmdbd.logとしているため，このディレクトリを作成しておきます．

sudo mkdir /var/log/slurm

sudo chown slurm: /etc/slurm-llnl/slurmdbd.conf
sudo chmod 600 /etc/slurm-llnl/slurmdbd.conf

最後にslurmdbdの起動&有効化を行ってインストールは完了です．

sudo systemctl enable --now slurmdbd

データベースへのクラスタの登録（on ログインノード）

Slurmのsacctmgrコマンドでデータベースへのクラスタの登録を行います．
cluster_nameは自身で決めたクラスタ名を指定します．

sudo sacctmgr add cluster cluster_name

mysql -p -u slurm slurm_acct_db
mysql> show tables;

cgroupsの設定（on 計算ノード）

GPUの数だけ1ノードにジョブを割り当てるのが今回の目的ですが，同時にCPUやメモリと行った資源も分割して提供できるようにします．
/etc/default/grubにGRUB_CMDLINE_LINUXという変数があるため，この右辺に以下を追記します．

cgroup_enable=memory swapaccount=1

次にcgroupで管理するものリストを作成します．

• /etc/slurm-llnl/cgroup.conf

  CgroupAutomount=yes
  CgroupMountpoint="/sys/fs/cgroup"
  CgroupReleaseAgentDir="/etc/slurm-llnl/cgroup"
  
  ConstrainCores=yes
  ConstrainDevices=yes
  ConstrainRAMSpace=yes
  
  AllowedDevicesFile="/etc/slurm-llnl/allowed_devices.conf"
  

• /etc/slurm-llnl/allowed_devices.conf

  /dev/null
  /dev/zero
  /dev/urandom
  /dev/cpu/*/*
  /dev/nvidia*
  

MUNGEの設定（on ログインノード&計算ノード）

MUNGEはSlurmで使われている認証サービスで，ホスト間の接続の暗号化などを行ってくれるものです．
aptで必要なパッケージをインストールします．

sudo apt install munge libmunge-dev libmunge2

sudo systemctl enable --now munge

sudo apt install slurm-wlm

Name=gpu File=/dev/nvidia0 CPUs=0-7
Name=gpu File=/dev/nvidia1 CPUs=8-15

sudo systemctl enable --now slurmd

sudo systemctl enable --now slurmctld

srun --gres=gpu:1 nvidia-smi

srun --gres=gpu:1 --ntasks-per-node=8 --mem=1000 nvidia-smi