📑

Elixir + Docker + EC2 + Livebook な環境でGPUぶんぶん回して機械学習を楽しむ環境構築

2023/06/09に公開

手が滑って普段遊んでいる環境を潰してしまったため、復旧ついでに備忘録
https://twitter.com/buy_sake/status/1666887317636927490

前提

AWSで以下のEC2インスタンスを用意します

  • インスタンスタイプは g4dn.xlarge
  • Ubuntu Server 22.04 LTS
  • ボリュームはとりあえず100GB付けておく
  • Elastic IPを紐付けておく
  • 8080と8081ポートは空けておく
    • 8080: livebookを建てます
    • 8081: VegaLiteでグラフ描画する際にjsとのブリッジに使ってるようです

Dockerインストール

公式に従ってシンプルに。
https://docs.docker.com/engine/install/ubuntu/

そして恒例のsudo無しでもdocker関連のコマンドを実行出来るように。

$ sudo gpasswd -a ubuntu docker
$ sudo systemctl restart docker
$ exit

CUDAドライバをインストール

$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb
$ sudo dpkg -i cuda-keyring_1.0-1_all.deb
$ sudo apt-get update
$ sudo apt-get -y install cuda-drivers
$ sudo reboot # 再起動

cuda ではなく cuda-drivers にすると、ドライバのみ綺麗に入ってくれるとの噂を聴き..

参考:
https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=22.04&target_type=deb_network

$ nvidia-smi
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 530.30.02              Driver Version: 530.30.02    CUDA Version: 12.1     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                  Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf            Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  Tesla T4                        Off| 00000000:00:1E.0 Off |                    0 |
| N/A   63C    P0               31W /  70W|      2MiB / 15360MiB |      4%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

GPUが読み込めていればOK 🎉

nvidia-docker2をインストール

dockerコンテナの中からホストのGPUを触れるようにします

$ curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
$ curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-$ docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list
$ sudo apt-get update
$ sudo apt-get install nvidia-docker2
$ sudo systemctl restart docker

参考:
https://www.ibm.com/docs/ja/maximo-vi/continuous-delivery?topic=planning-installing-docker-nvidia-docker2

docker-compose.ymlを書く

イメージは ghcr.io/livebook-dev/livebook:0.9.2-cuda11.8 を使います
https://github.com/livebook-dev/livebook/pkgs/container/livebook/85524865?tag=0.9.2-cuda11.8

$ mkdir livebook
$ cd livebook
$ mkdir books # .livemd ファイル等入れる場所
$ vim docker-compose.yml
docker-compose.yml
version: '3.2'
services:
  livebook:
    image: ghcr.io/livebook-dev/livebook:0.9.2-cuda11.8
    # インスタンスの外とscpでlivemdを出し入れする事が多いので、
    # booksディレクトリとLIVEBOOK_HOMEは常に同期
    volumes:
      - ./books:/data
      
    # Livebookのログインパスワード(最低12桁)
    environment:
      - LIVEBOOK_PASSWORD=asdf1234asdf
      - LIVEBOOK_HOME=/data
      
    ports:
      - 8080:8080
      - 8081:8081

    deploy:
      resources:
        reservations:
          devices:
            - capabilities: [gpu]

起動

$ docker compose up -d

http://{elastic ip}:8080 を開くとLivebookが起動できてるはずです 🎉

試しにdockerの中からGPUが触れるかの確認

$ docker compose exec livebook /bin/bash
> nvidia-smi

これでコンテナ外からnvidia-smiした時と同様の出力が得られればOKです

Nx + Exla セットアップ (with GPU)

各notebookのSetup欄に

setup
Mix.install(
  [
    {:nx, "~> 0.3.0"},
    {:exla, "~> 0.3.0"},
  ],
  config: [
    nx: [
      default_backend: EXLA.Backend,
      default_defn_options: [compiler: EXLA]
    ]
  ],
  system_env: [
    XLA_TARGET: "cuda111"
  ],
)

これで無事mix installが成功すればOKです🎉

LightGBMセットアップ (with GPU)

tato さんが公開してくれているラッパーライブラリたちを使います。
具体的な使い方については本家の記事を見て頂くとして、この記事ではこれをwith GPUでセットアップしていきます
https://zenn.dev/ta_to/articles/298f5aae4873c9

まずはLightGBM CLIのインストール
lgbm-runtimeの中にCLIをインストールするためのDockerfileを作っていきます

$ cd livebook
$ vim Dockerfile

先程はdocker-composeから直接読んでたlivebookのイメージをベースに、LightGBM CLIのインストールを書き加えたDockerfileを作ります

Dockerfile
FROM ghcr.io/livebook-dev/livebook:0.9.2-cuda11.8

RUN apt-get update && apt-get install -y nvidia-opencl-dev opencl-headers git cmake build-essential libboost-dev libboost-system-dev libboost-filesystem-dev
RUN cd / && mkdir lgbm-runtime
RUN cd /lgbm-runtime && git clone --recursive https://github.com/microsoft/LightGBM
RUN cd /lgbm-runtime/LightGBM && mkdir build && cd build && cmake -DUSE_GPU=1 .. && make -j$(nproc)
RUN mkdir -p /etc/OpenCL/vendors && echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd

これをdocker-compose.ymlから読み込みます (imageを消してbuildを生やしました)

docker-compose.yml
version: '3.2'
services:
  livebook:
    build:
      context: .
    volumes:
      - ./books:/data
    environment:
      - LIVEBOOK_HOME=/data
      - LIVEBOOK_PASSWORD=asdf1234asdf
    ports:
      - 8080:8080
      - 8081:8081

    deploy:
      resources:
        reservations:
          devices:
            - capabilities: [gpu]

これでコンテナのビルド時に /lgbm-runtime/LightGBM にLightGBM CLIが入ります

これをlivebook上で使うにはconfigsystem_envを以下のように渡しつつ2つのライブラリをインストールすれば良さそうです

setup
Mix.install(
  [
    {:lgbm_ex_cli, "0.1.0", git: "https://github.com/tato-gh/lgbm_ex_cli"},
    {:lgbm_ex_capi, "0.1.0", git: "https://github.com/tato-gh/lgbm_ex_capi"},
  ],
  config: [
    lgbm_ex_cli: [lightgbm_cmd: "/lgbm-runtime/LightGBM/lightgbm"]
  ],
  system_env: %{
    "LIGHTGBM_DIR" => "/lgbm-runtime/LightGBM"
  }
)

あとはfitの時にdevice: 'gpu' を指定してあげればGPUが躍動します 🎉

LGBMExCli.fit(..., [ ..., device: 'gpu'])

Discussion