Pythonの機械学習用Docker imageのサイズ削減方法の紹介

エムスリーエンジニアリンググループ AIチームの笹川です。

バスケと、ロードバイクが趣味なのですが、現在、NBAのplayoffと、Tour de Franceが同時に開催されていて大変嬉しい毎日を過ごしています。特にNBAのplayoffは、連日overtimeとなるような激戦や、giant killingがあったりのアツい戦いが繰り広げられていて最高です。

そういう状況なので（？）、今回は先日取り組んだ、Pythonの機械学習バッチを実行するdocker imageのサイズ削減についてのアツい戦いについて紹介したいと思います。

f:id:hsasakawa:20200918070342j:plain — 膝の上に登って寝る為に、筆者がデスクに戻るのを机の下で待ち構える犬氏（かわいい）

今回の取り組みでは、もともと3GB程度だったPythonのML用のimageを、約2.0GBに削減することができました（それでもなかなかのサイズ。MLのimageは特に大きい印象です）。 Docker imageのサイズ削減については、docker push、pull、buildのコストが下がることによりCIの高速化や、実行時のspin upが速くなるなどのメリットがあります。また、今回はサイズの削減と合わせて、開発用imageと、本番用imageの管理についても合わせて考えてみたいと思います。

Python の docker image のサイズ削減

Docker imageのサイズ削減として代表的な手法としては、Goなどのコンパイルでバイナリファイル生成するタイプのアプリケーションで行われる、multi-stage buildがあります。これはbuild stageでbuildされたbinaryファイルを、ターゲットのstageにコピーすることで、buildにしか使わないpackage managerや、ライブラリなどのファイルを削減する方法です。今回対象とする機械学習バッチのためのPythonのdocker imageに対するサイズの削減では、実行に必要なbuildの成果物が単一バイナリファイルになるわけではないので、その辺りをケアして、構成してやる必要があります。

以下は、実際に弊社のproductionで使われているimageから説明に必要な部分を抽出した例です。主として、以下のpackageが入っている環境をpoetryをpackage managerとして利用し構築することを考えます。

mecab
pandas
xgboost

元のDockerfileは以下です。

FROM python:3.6.8-stretch

# mecabとpoetryのinstallする
RUN apt-get update &&\
    apt-get install -y mecab libmecab-dev mecab-ipadic mecab-ipadic-utf8 &&\
    rm -rf ~/.cache &&\
    apt-get clean all
RUN python -m pip install poetry --use-feature=2020-resolver 

# Pythonの依存パッケージ周りの設定ファイルをホストからコピーする
COPY ./poetry.lock /app/poetry.lock
COPY ./pyproject.toml /app/pyproject.toml

# poetryで依存packageをinstallする
WORKDIR /app
RUN poetry config virtualenvs.create false &&\
    poetry install --no-dev &&\
    rm -rf ~/.cache

# ホストからbatchのためのコードをコピーする
WORKDIR /
COPY ./conf /app/conf
COPY ./batch /app/batch
COPY ./main.py /app/main.py

WORKDIR /app
ENTRYPOINT [ "/bin/bash" ]
VOLUME "/app"

このファイルを見ると、Docker image diet厨の皆さんもそうだと思うのですが、手始めにbaseとなるimageに余計な物が入ってそうな匂いがしてきます。サイズの小さいイメージとして候補になるのがalpine imageですが、以下のブログなどで説明されているように、マイナス面が大きいようです。

medium.com

今回は、いらぬトラブルを避けるため、slim imageに変更する方向で進めます。

ちなみに、上記のDockerfileにおいて、単純にbaseを同じPythonバージョンのslimに変更してみると、packageのインストールでライブラリやコマンドなどが足りずエラーになります。ですので、本節冒頭で紹介したmulti-stage buildを用いて、次のような戦略をとることにします。以下では、slimでないimageを通常imageと呼ぶことにします。

build stageでは、baseに通常 imageを利用してbuild
production stageでは、1の成果物をslim imageに持ってきて実行

上記の戦略を実装すると以下のようになります。

# build stageの定義
FROM python:3.6.8-stretch as build

# Pythonの依存パッケージ周りの設定ファイルをホストからコピーする
COPY ./poetry.lock /app/poetry.lock
COPY ./pyproject.toml /app/pyproject.toml

# mecabをinstallする
RUN apt-get update &&\
    apt-get install -y mecab libmecab-dev mecab-ipadic mecab-ipadic-utf8

# poetryで依存packageをinstallする
WORKDIR /app
RUN python -m pip install --upgrade pip &&\
    python -m pip install poetry --use-feature=2020-resolver &&\
    poetry config virtualenvs.create false &&\
    poetry install --no-dev &&\
    rm -rf ~/.cache


# production stageの定義
FROM python:3.6.8-slim-stretch as production

# build stageでpoetryをつかってinstallされたpackage群を丸ごと持ってくる
COPY --from=build /usr/local/lib/python3.6/site-packages /usr/local/lib/python3.6/site-packages

WORKDIR /
COPY ./conf /app/conf
COPY ./batch /app/batch
COPY ./main.py /app/main.py

WORKDIR /app
ENTRYPOINT [ "/bin/bash" ]
VOLUME "/app"

ほぼDockerfileが全てですが、少し解説してみます。

build stageでは、元のDockerfileとほぼ同じ処理を記述し、packageをinstallしています。次にproduction stageでは、baseをslim imageに変更した上で、 build stageでinstallしたpackageが格納されれるsite-package配下のファイルををまるっとコピーしています。このDockerfileで作成したimageをbuildしてバッチを実行すると以下のエラーが出ました。

xgboost.core.XGBoostError: XGBoost Library (libxgboost.so) could not be loaded.
Likely causes:
  * OpenMP runtime is not installed (vcomp140.dll or libgomp-1.dll for Windows, libgomp.so for UNIX-like OSes)
  * You are running 32-bit Python on a 64-bit OS
Error message(s): ['libgomp.so.1: cannot open shared object file: No such file or directory']

shared objectのファイルが足りていないようです。元のイメージに存在していたファイルを探して、コピーすることにします。 Dockerfileは以下のようになります。

FROM python:3.6.8-stretch as build

COPY ./poetry.lock /app/poetry.lock
COPY ./pyproject.toml /app/pyproject.toml
COPY ./setup.py /app/setup.py

RUN apt-get update &&\
    apt-get install -y mecab libmecab-dev mecab-ipadic mecab-ipadic-utf8

WORKDIR /app
RUN python -m pip install --upgrade pip &&\
    python -m pip install poetry --use-feature=2020-resolver &&\
    poetry config virtualenvs.create false &&\
    poetry install --no-dev &&\
    rm -rf ~/.cache


FROM python:3.6.8-slim-stretch as production

COPY --from=build /usr/local/lib/python3.6/site-packages /usr/local/lib/python3.6/site-packages

# これを追加
COPY --from=build /usr/lib/x86_64-linux-gnu/libgomp.so.1 /usr/lib/x86_64-linux-gnu/  

WORKDIR /
COPY ./conf /app/conf
COPY ./batch /app/batch
COPY ./main.py /app/main.py

WORKDIR /app
ENTRYPOINT [ "/bin/bash" ]
VOLUME "/app"

無事動きました！ここでimageサイズを計測してみると、2.03GBでした。やったね！

ここからさらにサイズを削減するとなると、コマンドなどを精査して不要なものを消していくとか、 Pythonではまだexperimentalですが、distrolessなどshellなども含まれない攻め攻めのimageにする方法が考えられます。しかしながら、かける労力と、得られる成果がどう考えても見合わないので、今回はここまでにしておくことにします。

github.com

本番用と開発用imageの管理について

無事に本番用 imageについてはサイズを削減することができました。ここで問題になるのが、CIなどで利用するために、linterや、testing package、unit test、integration testのコードを含めた開発用のimageの管理についてです。

開発用に別のDockerfileを用意する手もあるのですが、ほぼ同じimageを2つ用意することになり、両者に差分が生まれると、「テストは通るのに本番が落ちる」など本末転倒なことになりかねません。そこで、同一のDockerfile内部で、差分を記述した、development stageを用意することで、上記の問題を解決したいい感じの方法を実現することを考えてみます。ここでも活躍するのがmulti-stage buildです。multi-stage buildには、build時にtargetオプションでbuildの対象となるstageを指定し、そこまでを生成して止める機能があるので、それを利用します。

開発用のstageを含めたDockerfileは以下のようになりました。

# build stageの定義
FROM python:3.6.8-stretch as build

COPY ./poetry.lock /app/poetry.lock
COPY ./pyproject.toml /app/pyproject.toml
COPY ./setup.py /app/setup.py

RUN apt-get update &&\
    apt-get install -y mecab libmecab-dev mecab-ipadic mecab-ipadic-utf8

WORKDIR /app
RUN python -m pip install --upgrade pip &&\
    python -m pip install poetry --use-feature=2020-resolver &&\
    poetry config virtualenvs.create false &&\
    poetry install --no-dev &&\
    rm -rf ~/.cache

# development stageの定義
FROM python:3.6.8-slim-stretch as development

COPY --from=build /usr/local/lib/python3.6/site-packages /usr/local/lib/python3.6/site-packages
COPY --from=build /usr/lib/x86_64-linux-gnu/libgomp.so.1 /usr/lib/x86_64-linux-gnu/

COPY ./conf /app/conf
COPY ./batch /app/batch
COPY ./test /app/test  # テストファイルをホストからコピー

WORKDIR /app
RUN apt-get update &&\
    apt-get install -y git # needed to run poetry install

RUN python -m pip install --upgrade pip &&\
    python -m pip install poetry --use-feature=2020-resolver &&\
    poetry config virtualenvs.create false &&\
    poetry install  # dev dependencyも含めてinstall


# production stageの定義
FROM python:3.6.8-slim-stretch as production

COPY --from=build /usr/local/lib/python3.6/site-packages /usr/local/lib/python3.6/site-packages
COPY --from=build /usr/lib/x86_64-linux-gnu/libgomp.so.1 /usr/lib/x86_64-linux-gnu/  # これを追加

WORKDIR /
COPY ./conf /app/conf
COPY ./batch /app/batch
COPY ./main.py /app/main.py

WORKDIR /app
ENTRYPOINT [ "/bin/bash" ]
VOLUME "/app"

ここでは、productionで定義したstageに追加で、development stageを用意し以下の作業を行っています。

base imageは通常image
build stageでinstallしたpackageとライブラリをコピー
testコードをホストからコピー
poetry installをdev dependency含めてinstall

ここで特に4を実行する際にmulti-stage buildの効果が発揮されます。build stageでほとんどのpackageがinstallされており、4ではdev dependencyのみが追加されるので、このstageのbuildは短時間で完了します。このstageまでのimageを生成するには、以下のコマンドを実行します。

docker build --target development -t [tagname] .

このようにすることで、差分のみの比較的簡潔な記述で、本番imageと開発imageを管理することができるようになりました。

上記で紹介したDocker imageのサイズ削減方法をまとめると以下のようになるかと思います。

build stageを作り、通常imageを用いてpackage installを実行する
production stageではbase imageをslimにし、build stageのsite-package配下をコピーする
test結果などを確認し、必要であれば追加でファイルのコピーなどを逐次実施する
development stageを作り、production stageの処理に加えて、dev dependencyを追加する処理を書く

取り組んだ感想

今回Python imageのサイズ削減に取り組んでみて、元サイズが結構大きいこともあり、削減の効果が大きいなと感じました (docker pullのコストなども体感でわかるほど早くなります)。一方で、利用ライブラリが大きく変わるような、開発の初期には向かない方法だなと感じました。これは、シンプルなimageで、site-packageのコピーのみで済むようなケースであれば最初から導入する手もありますが、個別に必要なライブラリ調査しコピーする必要がある場合に、開発体験が著しく毀損されるな、と思ったことが理由です。なので、今回のようにまずは冒頭で紹介したsingle stageのDockerfileで開発を始め、「結構imageも肥大化してきたし、開発は大体固まってきた」という段階になってから、サイズの削減に取り組み始めるという形がいいのかなと思います。