[Python, Docker] CPU 環境と GPU 環境の使い分けが面倒なあなたに

TL;DR

• CPU 環境で実装，軽いテスト -> GPU 環境で実行の流れが（個人的に）よくある．
• しかし，この時パッケージや Dockerfile が CPU 向けと GPU 向けで別れていたり， コマンドの引数が変わったりと面倒．
• そこで GNU Make を用いて CPU/GPU 環境を意識せず，同一コマンドで実行できるようにした．
• サンプルレポジトリ（tensorflow 環境）

とりあえず環境が欲しい場合

Docker (>= 19.03) および GNU Make のインストールを済ました上で以下のコマンドでコンテナ内の bash に入れます．
また，GPU 環境が必要な場合は別途 nvidia-drivers および nvidia-docker を先にセットアップしておいてください．

$ git clone https://github.com/hiro-o918/dev-python-tensorflow-dockerash
$ cd dev-python-tensorflow-dockerash
$ make build
$ make bash

仕組み

環境設定

Dockerfile を GPU 環境版用のものと CPU 環境版用のものを用意します．
また，tensorflow のような GPU/CPU 環境でパッケージが異なるようなものを用いる場合は，ベースイメージを各環境に対応したものを用いるか，Dockerfile 側でインストールするようにします．

また，GPU/CPU 環境に依存しないパッケージは Pipenv や poetry を使って管理します．

GNU Make

上記のように GPU/CPU 環境の両方を用意した場合は Docker 周りのコマンドが各環境で異なるので面倒なことになります．例えば，イメージを build する際は以下の 2 つのコマンドを使い分けなければなりません．

$ docker build ./docker/Dockerfile.cpu  # CPU ver
$ docker build ./docker/Dockerfile.gpu  # GPU ver

そこで，以下のように GNU Make を使ってこの使い分けを隠蔽します．
詳しい流れは後述します．

NVIDIA_SMI_PATH := $(shell which nvidia-smi)
IMAGE_NAME := python_tensorflow
CONTAINER_NAME := python.tensorflow
WORKINGDIR := /var/www
PWD := $(shell pwd)

ifdef NVIDIA_SMI_PATH
    DOCKER_GPU_PARAMS := --gpus all
endif

.PHONY: _build/cpu
_build/cpu:
	@docker build --tag $(IMAGE_NAME) -f $(PWD)/docker/Dockerfile.cpu .

.PHONY: _build/gpu
_build/gpu:
	@docker build --tag $(IMAGE_NAME) -f $(PWD)/docker/Dockerfile.gpu .

.PHONY: build
build:
    ifdef NVIDIA_SMI_PATH
	    @$(MAKE) _build/gpu
    else
	    @$(MAKE) _build/cpu
    endif

.PHONY: run
run:
	@docker run \
		--rm -it \
		$(DOCKER_GPU_PARAMS) \
		--name $(CONTAINER_NAME) \
		--volume $(PWD):$(WORKINGDIR) \
		$(IMAGE_NAME) \
		$(ARGS)

GPU 環境であるかのフラグ設定

まず，ホストマシンで nvidia-smi にパスが通っているかを確認します．
これによって nvidia-driver が適切に入っているかを検証しています．

NVIDIA_SMI_PATH := $(shell which nvidia-smi)

イメージの build

次に，以下の箇所にて先程設定した NVIDIA_SMI_PATH に値があるかどうかを見ることで， GPU が使える環境かどうかを判別し，build する Dockerfile を使い分けています．

.PHONY: build
build:
    ifdef NVIDIA_SMI_PATH
        @$(MAKE) _build/gpu
    else
        @$(MAKE) _build/cpu
    endif

このため開発者は以下のコマンドでホスト環境にマッチしたイメージのビルドが可能になります．

$ make build

GPU 環境特有のパラメータ設定

また，GPU 版のコンテナ実行時に必要なオプションは以下の部分で設定しています．
デフォルトではホストの GPU をすべて使えるようにしました．

ifdef NVIDIA_SMI_PATH
    DOCKER_GPU_PARAMS := --gpus all
endif

コンテナ実行

最後にコンテナ実行箇所は以下になります．
先程 DOCKER_GPU_PARAMS として設定した値を投げることで，GPU 環境特有の設定を渡すことができます．
環境変数なので設定していない場合は空白となり，CPU 環境での実行時に影響を与えることはありません．
コマンド内の $(ARGS) は実行コマンドを override するために使います．

.PHONY: run
run:
	@docker run \
		--rm -it \
		$(DOCKER_GPU_PARAMS) \
		--name $(CONTAINER_NAME) \
		--volume $(PWD):$(WORKINGDIR) \
		$(IMAGE_NAME) \
		$(ARGS)

よって，以下のコマンドで環境気にすることなくコンテナの実行が可能です．

$ make run

カスタムコマンド

この結果，以下のようにカスタムコマンドも簡単に追加でき，1 つにまとめることができます．

.PHONY: bash
bash: ARGS=bash
export ARGS
bash:
	@$(MAKE) run

以前はこのようなカスタムコマンドを bash/gpu, bash/cpu のように設定していたため，すごいうっとおしかった....

まとめ

個人的にはよく CPU 環境と　GPU 環境を行き来するので，実験がしやすくなったと思っています．サンプルレポジトリ では簡単な CI 環境も整えているので，適当に使って貰えると喜びます．