Slurmを使った際のhostfile/machinefileの自動生成について

背景/目的

最近、ジョブスケジューラといえばPBSばかりを使っているため、自身、Slurmの振り返りを兼ねて記事を書いてみます。

複数のノードに跨ってMPI計算を実行する際、計算に使用するマシンの「ホスト名 or IPアドレス」を列挙したファイルを作成することが一般的です。このファイルをhostfileやmachinefileと呼び、mpirun実行時に引数として渡すと、ファイル内で記述した計算ノードを使って、ノード跨ぎの計算を行います。

OpenPBSやPBS Professionalを使う場合、hostfile/machinefileはジョブ実行中に環境変数として登録され、$PBS_NODEFILE の名前で自動的に利用するノードの一覧を参照できるようになっています。詳細はこちら(参考[1])の記事が分かりやすいです。

当然Slurmも複数ノードを使って計算実行したい場合、$PBS_NODEFILE のように「計算を実行するホストをまとめて生成してくれる機能」があるかと期待するのが自然だと思います。が、残念ながらそのような機能は無いようです (参考[2])

このような背景から、機能が無いなら既存の他ツールを使ってみよう！そして不足部分は自前で補ってしまおう！というのが趣旨です。

hostfile/machinefileの記述について

MPIと一口に言ってもOpenMPIやIntelMPIやら様々な種類のMPIが存在しております。それぞれのMPIごとに各計算ノードへのプロセス割り付けに関する記述が微妙に異なってきます。

そのため、本記事内では統一してOpenMPIにおける記述に則ってhostfile/machinefileを生成します。具体的には下記のような具合で「slots=」として表現します。

compute01 slots=2 # 計算ノードcompute01に2プロセスの割り当て
compute02 slots=2 # 計算ノードcompute02に2プロセスの割り当て
compute03 slots=2 # 計算ノードcompute03に2プロセスの割り当て

ここでプロセス割り当ての記述について言及した意図として、slots値をhostfile/machinefile内に記述したいと考えたためです。ニーズとしては例えば、各計算ノードに均等にプロセスを割り当てたり、プロセスのメモリ使用量の都合から、使う計算ノードの台数を増やす代わりに、1台あたりに割り当てるプロセス数を減らしたりするケースが考えられます。

本記事のゴールイメージ

ここまでの内容から、最終的に作りたいものとしては下記を想定しています。

Slurmのジョブスクリプトの中で自動的にhostfile / machinefileを生成する
hostfile / machinefileの中に各計算ノードが使用するプロセス数を記載できる
hostfile / machinefileファイル生成の仕組みを簡単に利用できる

HPC利用者が手間暇かけずに使える仕組みが望ましいと思い、最後の項目を追記しました

実行環境

今回使用したOSや各種ソフトウェアのバージョンは下記の通りです。

OS: Ubuntu 22.04 LTS
AWS ParallelCluster: 3.11.1 (HeadNode / Compute共にvCPU 1個のt2.microを使用)
Slurm: 23.11.10
OpenMPI: 4.1.6
Python: 3.9.20

以降では実際のツールの導入やその使い方について説明します。また、これらの操作は全てAWS ParallelClusterのHeadNode上で実施している操作となります。

python-hostlistの導入と挙動

今回、Slurm上でhostfile/machinefileの生成に使ってみたのは「python-hostlist」(参考[3]) というツールです。実際にインストール作業を進めてみます。以下はHPCユーザ目線での作業例ということで、root権限を使用しないユーザレベルでのインストール例となります。

### ソース一式のclone
> git clone git://www.nsc.liu.se/~kent/python-hostlist.git
> cd python-hostlist/
### READMEに従い、make prepareを事前に済ませる
> make prepare; cd versioned
### 下記にてビルドやらファイルのパーミッションの変更やらを実施している模様
> python setup.py build

インストールは3分もあれば完了するくらい簡単に実行できました。ここまで出来たら動作確認兼使い方の確認を行ってみます。python-hostlistのサイトトップに掲載してある利用例を参考に、自身のクラスタの設定値を入力します。

### Interactive shellで実行した例
>>> import hostlist
### テストとしてslurmのsinfoコマンドで取得できるNODELIST情報を引数に与えてみる
>>> hosts = hostlist.expand_hostlist("normal-dy-t2micro-[1-10]")
### 各計算ノードのホスト名をリストで返してくれる
>>> hosts
['normal-dy-t2micro-1', 'normal-dy-t2micro-2', 'normal-dy-t2micro-3', 'normal-dy-t2micro-4', 'normal-dy-t2micro-5', 'normal-dy-t2micro-6', 'normal-dy-t2micro-7', 'normal-dy-t2micro-8', 'normal-dy-t2micro-9', 'normal-dy-t2micro-10']

問題無く計算ノードのホスト名が返ってくることを確認できました。

少々脱線しますが、上記の例で与えた「”normal-dy-t2micro-[1-10]”」のような表記を入力し、各々のホスト名に変換できると、何が嬉しいかというと、実はジョブ自身が割り当てられた計算ノードを格納した環境変数 $SLURM_JOB_NODELIST と、今回例として与えた引数は同じ形式で記述されるためです。文字だけだと状況が分かりにくいかと思いますので、以下で簡単に $SLURM_JOB_NODELIST の参照例を示しておきます。

#!/bin/bash
### 計算ノードを2台 割り当てるように記述したジョブスクリプト

#SBATCH --job-name=env_sample # ジョブ名
#SBATCH --output=env_sample.out # 出力ファイル名の指定
#SBATCH --ntasks=2
#SBATCH --cpus-per-task=1

### 以下の結果がジョブ名「sample」に対して割り当てられた計算ノードになる
echo $SLURM_JOB_NODELIST

### 以下がechoの出力でsinfoコマンドの出力normal-dy-t2micro-[1-10]と同様の形式になる
normal-dy-t2micro-[2-3]

上記の出力を hostlist.expand_hostlistの引数に入れればジョブが使用する計算ノードのリストを扱えるようになり、hostfile/machinefileが作りやすくなるという具合です。

python-hostlistの利用例

さて、インストールが完了したので次はジョブスクリプト内でpython-hostlistの出力を取得できるようにすること、そして取得した出力結果に各計算ノードのプロセス割り当て情報を加える必要があります。

やり方は多々あるかと思いますが、私は以下の実行フローを元に実装してみました。

$SLURM_JOB_NODELISTの取得+プロセス数割り当ての設定をジョブスクリプト内で実施
1.で扱った2つの値を1つの変数にまとめる
2.でまとめた変数をPythonスクリプト(hostfile_gen.py)に与え、hostfile/machinefileを生成
3.で生成したhostfile/machinefileをMPI実行時の引数として与える

簡単に私が作成してみたジョブスクリプト & Pythonスクリプトと、その実行例を示します

#!/bin/bash

#SBATCH --job-name=hostfile-test
#SBATCH --output=result.log
#SBATCH --error=error.log
#SBATCH --ntasks=2
#SBATCH --cpus-per-task=1
#SBATCH --time=01:00:00
### 下記で各ノードに割り当てるプロセス数を設定する
PROC_IN_EACH_NODE=1

### Pythonスクリプトに与える2つの変数を1つにまとめておく
### OpenMPI以外を使うときはここでslotsの表記を修正するとよい
SLURM_JOB_NODELIST_AND_SLOTS="`echo $SLURM_JOB_NODELIST` slots=$PROC_IN_EACH_NODE"

### Pythonスクリプトの呼び出し 以下ではジョブを実行したディレクトリにhostfileを生成
python hostfile_gen.py $SLURM_JOB_NODELIST_AND_SLOTS > $SLURM_SUBMIT_DIR/hostfile

### 引数に生成したhostfileを与えて、MPIを使ったHello Worldを実行してみる
### プログラムは各自が使いたいプログラムを与えてください
mpirun -np $SLURM_NTASKS --hostfile $SLURM_SUBMIT_DIR/hostfile /efs/mpi_hello

上記ジョブスクリプト内において、ハイライトしているPythonスクリプトは以下の通り。

import os
import sys
import hostlist # import pathにご注意ください

### コマンドライン引数を利用する
### 第1引数にslurmから割り当てられたノード群を与える
### 第2引数に各ノードに割り当てるプロセス数を与える
args = sys.argv
hosts = hostlist.expand_hostlist(args[1])
slots = args[2]
for host in hosts:
   print(host, slots)

これを使って、バッチジョブを投入した際の結果を以下に示します。

### 上記のジョブスクリプトを使ってバッチジョブを投入する
> sbatch hostfile-test.sh 

### 生成されたhostfileの記述を確認 想定通りに記述されている
> cat hostfile 
normal-dy-t2micro-1 slots=1
normal-dy-t2micro-2 slots=1

### 実行結果も想定通りで、無事に各ノードで処理が実行されている
> cat result.log
Hello world from normal-dy-t2micro-1! Rank is 0.
Hello world from normal-dy-t2micro-2! Rank is 1.

いずれも想定通りに実行されており、無事に実装完了できました！めでたしめでたし。

感想

以前はノードを跨いだ計算を行う利用者が少ない計算基盤を担当していたため、あまり注目していなかった問題でした。昨年に書いた「Apptainer+MPI」利用をHPC基盤上で実行する際に、ジョブスケジューラとしてSlurmを採用していたら直面する問題なのだろうと思います。

またhostfile/machinefileとは別に、Slurm上で使える環境変数をまとめているサイトを参照していると、新しい発見があったりするなという印象でした。こちらのサイト(参考[4])や、Slurmの公式ページ(参考[5])をボーっと眺めてみるもの悪くないなと思いました。

もっと楽にhostfile/machinefileを生成する術を知っている方はお気軽にコメントください。

参考にさせていただいた情報

[1] OpenPBS/PBS_NODEFILE – Chaperone
[2] slurm_scripts [Arkansas High Performace Computing Center [hpcwiki]]
[3] python-hostlist
[4] Common SLURM environment variables — Sheffield HPC Documentation
[5] Slurm Workload Manager – sbatch

以上

Slurmを使った際のhostfile/machinefileの自動生成について

背景/目的

hostfile/machinefileの記述について

本記事のゴールイメージ

実行環境

python-hostlistの導入と挙動

python-hostlistの利用例

感想

参考にさせていただいた情報

Comments

コメントを残すコメントをキャンセル

Slurmを使った際のhostfile/machinefileの自動生成について

背景/目的

hostfile/machinefileの記述について

本記事のゴールイメージ

実行環境

python-hostlistの導入と挙動

python-hostlistの利用例

感想

参考にさせていただいた情報

Comments

コメントを残す コメントをキャンセル

コメントを残すコメントをキャンセル