使用するもの

NixOS

Linuxディストリビューションの１つ
- 関数型パッケージ管理ツールNixを使用している
- パッケージ自体はnixpkgsのリポジトリで管理・開発されている(個々のパッケージはhttps://nixos.org/nixos/packages.html# で検索できる)
環境の構築が楽にできる
- １つのファイルでシステム全体の設定ができる -> 管理が楽
- システムのアップデート -> システムのロールバックということが楽にできる
- sudoとかしなくてもパッケージを用意できる(たとえばpythonでnumpyとmatplotlibを使いたい時は以下のようにする)

$ nix-shell -p python3 -p pythonPackages.numpy -p pythonPackages.matplotlib

NixOps

NixOSの環境をリモートマシンにdeployするためのツール
基本的な使い方は以下のような感じ

登録
$ nixops create [deployする環境を書いた.nixファイル(複数可)] -d [環境の名前]

インスタンスの作成＆deploy
$ nixops deploy -d [環境の名前]

インスタンスの破棄
$ nixops destroy -d [環境の名前]

Chainer

ニューラルネットワーク記述用のpythonフレームワーク

ChainerMN

Chainerで複数のGPU上での分散処理を行うための追加パッケージ

はじめに

NixOpsでは、Nixのコードを書くことで、Nixパッケージマネージャを利用した環境をリモートにdeployすることができます。
これを使ってGPUクラスタをデプロイできないかということでやってみました。

利用環境

AWSのEC2のp2.xlargeのインスタンスを使用しました。
GPUはTesla K80になります。

CUDAはCUDA 9.0、cudnnはv7.0、NCCLはv2.1.2を使用します。
MPIはOpenMPIのv1.10です。
これらはすべて利用したnixpkgsに含まれているものを使用します。
(これを書いている時点ではnixpkgsは自前で用意しています)

コード

NixOpsで使用するコード
https://github.com/hyphon81/ec2-gpu-cluster-with-NixOps

./ec2-gpu-cluster-with-NixOps というディレクトリがある想定で書きます。

nixpkgs
https://github.com/hyphon81/nixpkgs/tree/add/chainermn

これについても、 ./nixpkgs が存在する前提で書きます。

事前準備

AWSの準備

※これらの手順等についてはここでは説明しません。

ファイルの準備

NixOpsのマニュアルの説明に従い、以下のどちらかの用意をします。

~/.ec2-keys

<アクセスキーID> <アクセスキー> <IAMユーザ名>

~/.aws/credentials

[<IAMユーザ名>]
aws_access_key_id = <アクセスキーID>
aws_secret_access_key = <アクセスキー>

./ec2-gpu-cluster-with-NixOps/ec2-info.nix を次のように用意します。

./ec2-gpu-cluster-with-NixOps/ec2-info.nix

{ config, pkgs, ... }:

with pkgs.lib;

{
  deployment.ec2.accessKeyId = "<アクセスキーID>";
  deployment.ec2.keyPair = "<AWS EC2のsshキーペア名>";
  deployment.ec2.privateKey = "<ssh秘密鍵ファイルのパス>";
}

インスタンス間のssh接続に使用するsshキーペア(ed25519)を用意します。

$ ssh-keygen -t ed25519
...
$ cp ~/.ssh/id_ed25519* ./ec2-gpu-cluster-with-NixOps/

NixOpsによるデプロイ

①. nixops create コマンドを実行し、設定を作ります。

$ nixops create ./ec2-gpu-cluster-with-NixOps/ec2gpu-cluster.nix -d ec2-gpu-cluster

②. 早速 nixops deploy コマンドでAWS EC2インスタンスの作成→環境のデプロイを行います。

$ nixops deploy -d ec2-gpu-cluster -I nixpkgs=./nixpkgs

以上です。
コマンドは2行とあっけないですが、デプロイに結構時間がかかります。
私の場合は1〜2時間ぐらい待ちました。
NVIDIAのドライバやCUDAなどを含むため、デプロイするファイルサイズの合計が５GBぐらいになるようです。

デプロイする環境の設定について

今回用意したファイルは以下のようになっています。

./ec2-gpu-cluster-with-NixOps/ec2gpu-cluster-conf.nix

{ region, ami, instanceType, securityGroups, publicKey, secretKey, hosts }:

{ config, pkgs, resources, ...}:
let
  ## この辺は使用するパッケージの設定です。
  python = pkgs.python3;
  cudatoolkit = pkgs.cudatoolkit;
  mpi = pkgs.openmpi.overrideDerivation (attrs: {
    configureFlags = [ "--with-cuda=${pkgs.cudatoolkit}" ];
  });
  gpu-test1 = pkgs.callPackage ./gpu-test1.nix {
    mpi = mpi;
  };
  gpu-test2 = pkgs.callPackage ./gpu-test2.nix {
    python = python;
    cython = python.pkgs.cython;
    mpi = mpi;
    cudnnSupport = true;
    cudatoolkit = cudatoolkit;
    cudnn = pkgs.cudnn;
    nccl = pkgs.nccl;
  };
in
{ imports = [ ./ec2-info.nix ];
  deployment.targetEnv = "ec2";
  deployment.ec2.ami = ami;
  deployment.ec2.region = region;
  deployment.ec2.instanceType = instanceType;
  deployment.ec2.securityGroups = securityGroups;
  deployment.ec2.ebsInitialRootDiskSize = 20;

  nixpkgs.config.allowUnfree = true;

  environment.systemPackages = [
    python
    mpi
    cudatoolkit
    gpu-test1
    gpu-test2
  ];

  ## 以下が環境設定です。
  services.xserver = {
    enable = true;
    videoDrivers = [
      "nvidia"
    ];
  };

  services.openssh = {
    enable = true;
    hostKeys = [
      { bits = 4096; path = "/etc/ssh/ssh_host_rsa_key"; type = "rsa"; }
      { path = "/etc/ssh/ssh_host_ed25519_key"; type = "ed25519"; }
      { path = "/etc/root/ssh/id_ed25519"; type = "ed25519"; }
    ];
    knownHosts = hosts;
    passwordAuthentication = false;
    challengeResponseAuthentication = false;
  };

  environment.etc = {
    "/root/ssh/id_ed25519" = {
      enable = true;
      user = "root";
      group = "root";
      mode = "0600";
      source = secretKey;
    };
  };

  ## LD_LIBRARY_PATHに必要なライブラリを含めるためにこんなことやってしまっています。
  ## もっとかっこいい方法をご存じの方がいたらぜひ教えてください。
  hardware.opengl = {
    extraPackages = [
      mpi
      pkgs.linuxPackages.nvidia_x11
      cudatoolkit
      cudatoolkit.lib
    ];
  };

  networking.firewall.enable = false;

  users.users.root = {
    openssh.authorizedKeys.keyFiles = [ publicKey ];
  };

  programs.bash.shellInit = ''
    if [ ! -e ~/.ssh/id_ed25519 ] ; then
      if [ -f /etc/root/ssh/id_ed25519 ] ; then
        /run/current-system/sw/bin/ln -s /etc/root/ssh/id_ed25519 .ssh/id_ed25519
      fi
    fi
  '';
}

./ec2-gpu-cluster-with-NixOps/ec2gpu-cluster.nix

let
  ## ここでデプロイする先のAWS EC2の設定を記述しています
  region = "ap-northeast-1"; ## 東京リージョン
  ami = "ami-89b921ef"; ## 東京リージョンのNixOSのAMI https://nixos.org/nixos/download.html
  instanceType = "p2.xlarge"; ## インスタンスタイプ
  ec2 = import ./ec2gpu-cluster-conf.nix;

  securityGroups = [ "allow-ssh" ]; ## インスタンスが所属するセキュリティグループ名を入れてください
  secretKey = ./id_ed25519;
  publicKey = ./id_ed25519.pub;

  hosts = [
    { hostNames = [ "node1" ]; publicKeyFile = publicKey; }
    { hostNames = [ "node2" ]; publicKeyFile = publicKey; }
    { hostNames = [ "node3" ]; publicKeyFile = publicKey; }
    { hostNames = [ "node4" ]; publicKeyFile = publicKey; }
  ];
in
{
  network.description = "An ec2 gpu cluster create test";

  node1 = ec2 {
    inherit region ami instanceType securityGroups publicKey secretKey hosts;
  };
  node2 = ec2 {
    inherit region ami instanceType securityGroups publicKey secretKey hosts;
  };
  node3 = ec2 {
    inherit region ami instanceType securityGroups publicKey secretKey hosts;
  };
  node4 = ec2 {
    inherit region ami instanceType securityGroups publicKey secretKey hosts;
  };
  ## 4ノード作成します
}

デプロイ完了後の操作(動作確認)

①. NVIDIAドライバインストールがあったため、ドライバの読み込み等必要になるので、一度全ノードをrebootしたほうが良いです。

$ nixops ssh-for-each -d ec2-gpu-cluster -p reboot

②. その後少し待ってnode1に接続します。

$ nixops ssh -d ec2-gpu-cluster node1
...
[root@node1]#

③. 今回用意したnixコードでは以下のパッケージを含み、テストができます。
(こちらに記載されているテストコードをそのまま使用しています。)

// gpu-mpi-test
// mpiパッケージがCUDA-Awareに対応しているか確認するプログラムです。
// 問題なければOKと出るだけです。
[root@node1]# gpu-mpi-test
OK.

// check_mpi4py
// mpi4pyが動作するかの確認です。
// node1〜node4が表示されれば問題ありません。
[root@node1]# mpiexec --allow-run-as-root -n 4 -host node1,node2,node3,node4 check_mpi4py
node1 0
node2 1
node3 2
node4 3

まぁ、Nixなので一度動作を確認できた環境なら同じコード使えば同じように動くはずなんですが・・・

ChainerMNを使用したMNISTの学習の動作

コード自体はここに従って用意しています。

実際に動作させてみたときの動画がこちらです
※ 6分程度と長いのでBGMつけました。音が出ます。

一応、動作していました。

余談

ChainerMNでMNISTの学習が動作できたのですが、
ノードを1つだけ使ったときのほうが計算が早く終わりました。

ChainerMNの論文を見る限り、
ノード増やすとほぼリニアにスピードアップするグラフが乗っているので、
何か間違ったかと思っていました。
https://arxiv.org/abs/1710.11351

ですが、こちらのスライドではそもそもMNISTの学習は
ChainerMNの速度の測定に使用するには向かない旨の注意が書いてありました。
https://www.slideshare.net/pfi/20171128chainermn

というわけで、ノード増やしてスピードアップすることを確認するには、
ちゃんとしたベンチマークになりうるプログラムを用意する必要がありそうです。

とりあえず、NixOpsでChainerMNが動作する環境の構築はできたということで、
今回はここまでです。