リモートGPUサーバーに立てたDocker上で、jupyter labを動かしてローカルからアクセスする方法（JAX用）

前提

パスワードを入れなくてもsshのみで入れるようにする。 → 参考

Dockerfileの用意

リモートサーバーの中に作業フォルダを作成してDockerfileを入れる

FROM nvcr.io/nvidia/cuda:11.3.1-cudnn8-devel-ubuntu20.04

ENV TZ=Asia/Tokyo
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone

RUN apt-get update
RUN apt-get -y install vim
RUN apt-get -y install software-properties-common
RUN add-apt-repository ppa:deadsnakes/ppa
RUN apt-get install -y python3.9 python3.9-venv
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 100
RUN update-alternatives --config python3
RUN apt-get install -y python3-pip

RUN pip install --upgrade numpy scipy
RUN pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
RUN pip install jupyterlab matplotlib
RUN pip install tqdm optax

WORKDIR /work
CMD jupyter-lab --no-browser --port={好きなport番号(8888がデフォルト)} --ip=0.0.0.0 --allow-root --NotebookApp.token=''

Docker起動

作業フォルダ配下に/work フォルダを作成

作業フォルダ/work はdocker 内にマウントされる。使いたいnotebook等はここに置く。

# docker のイメージを作成（一番最初だけで良い）
docker build -t {任意のイメージ名} .  

# イメージからdockerを起動 --gpus 2 で使用するGPUの数を指定
nvidia-docker run --name {任意のコンテナ名} -p {指定したport番号}:{指定したport番号} -v {作業フォルダの絶対パス}/work:/work  --gpus 2 {任意のイメージ名}

# 使用するGPUを指定する場合
nvidia-docker run --name {任意のコンテナ名} -p {指定したport番号}:{指定したport番号} -v {作業フォルダの絶対パス}/work:/work  --gpus '"device=0,1"' {任意のイメージ名}

ローカルからアクセスしてGPUの接続テスト

ローカルからsshでjupyterにつなぐ。

ssh -L {指定したport番号}:localhost:{指定したport番号}  {ユーザー名}@{IPアドレス}

ブラウザからjupyter labに接続

localhost:{指定したport番号}

これでjupyter labの画面が出てきたら成功

GPUが使えているかのテスト

適当なipynbファイルを開いて以下を実行。 gpu と出てくればGPUに接続できている。

from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)

GPUでJAXを使う場合の注意点

複数のノートブックで同時にGPU上でjitをしようとするとエラーになる。

複数のノートブックを使う場合は、以下のコマンドで別々のGPUを割り当てる。

import os
# dockerから２つのGPUが見えているのであれば、0 or 1を指定することで、切り替わる。
# 0と1の両方をこのnotebookで使いたいならば、 '0,1'に設定する。
os.environ['CUDA_VISIBLE_DEVICES'] = '0'