Zenn
🕹️

Slurm + Pyxis/enroot環境でJupyter Notebookを起動させる手順

2024/10/25に公開
Docker
Python
NVIDIA
jupyter
GPU
tech

この手順書では、Slurmによるリソース配分の恩恵を受けながら、コンテナを利用してJupyter Notebookを起動する方法を解説します。コンテナを使用することで、環境の依存関係や破壊リスクが少なくなり、簡単にJupyterの実行環境を構築できます。

前提条件と注意事項

  • Pyxis と enroot は、GPUクラスタ上でコンテナを扱うためのツールです。本手順では、これらを用いてコンテナをSlurm上で実行し、Jupyter Notebookを起動します。
  • 必要な権限：ヘッドノードにrootでアクセスできることが前提です。
  • 利用するコンテナ：Jupyterの公式コンテナイメージを使用しますが、利用者の好みに応じて他のコンテナも選択可能です。

ステップ1：DockerHub等から使いたいコンテナを探す

ステップ2：利用者のhomeディレクトリにコンテナをダウンロード

以下のコマンドで、enrootを使ってJupyterコンテナを取得し、利用者のhomeディレクトリに保存します。

# DockerHubからダウンロードする場合
enroot import docker://jupyter/base-notebook

# Quay.ioからダウンロードする場合
enroot import docker://quay.io#jupyter/base-notebook

DockerHub以外からコンテナをダウンロードする場合は、enrootのオプションが少し変わります。
quay.io#のような記法で、落としてくるサイトを指定する必要があります。

実行例では、「jupyter+base-notebook.sqsh」というファイルが生成されます。適切なディレクトリ（例：/home/user01/）にこのファイルを移動してください。

ステップ3：利用者毎にJupyter Notebook用のバッチファイルを作成する

以下は、4GPUを割り当てる場合のバッチファイル例です。必要に応じてGPU数やファイルパスを変更してください。

$ vi jupyter-notebook.sbatch

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --gpus=4
#SBATCH --job-name=jupyter-notebook
#SBATCH --output=%x_%j.out
#SBATCH --error=%x_%j.err
set -euxo pipefail

CONT='/home/user01/jupyter+base-notebook.sqsh'  # コンテナのパスを環境に合わせて変更

srun --container-image="${CONT}" jupyter lab --notebook-dir=/home/user01

wait

各パラメータの説明

  • --gpus=4：利用したいGPU数を指定します。
  • CONT='/home/user01/jupyter+base-notebook.sqsh'：コンテナイメージのパスを指定します。利用者に応じてパスを変更してください。

ステップ4：バッチファイルを実行してJupyterサーバを立ち上げる

作成したバッチファイルをsbatchコマンドで実行し、Jupyter Notebookサーバを起動します。

$ sbatch jupyter-notebook.sbatch

実行結果の確認

このコマンドを実行すると、jupyter-notebook_xx.errというファイルがhomeディレクトリに生成され、Jupyterサーバのトークンコードを含むURLが出力されます。DGXノードに直接アクセスできない場合、ヘッドノードにログインして、SSHポートフォワーディングを設定してアクセスします。

SSHポートフォワーディング例

# ヘッドノードにログインしてポートフォワーディング
ssh -L 8888:localhost:8888 user@headnode

上記を設定後、http://localhost:8888 にアクセスし、出力されたトークンコードでJupyter Notebookにログインします。

ステップ5：終了時はscancelコマンドでジョブを停止

Jupyterサーバの利用が終わったら、以下のコマンドでジョブを終了します。

$ scancel <ジョブID>

参考情報

よくあるエラーと対処方法

  • エラー例1: enroot import でコンテナがダウンロードできない場合
    対処：docker://のURLが正しいか確認し、インターネット接続も確認してください。

  • エラー例2: sbatchコマンドでのエラー
    対処：バッチファイルのパスやGPU設定が正しいか再確認してください。また、Slurmの設定ファイルが影響する場合もあるため、管理者に確認が必要な場合もあります。

公式ドキュメントリンク

Discussion