Gunosyデータ分析ブログ

Gunosyで働くデータエンジニアが知見を共有するブログです。

M1 Mac に挫けない!TensorFlow に躓かない開発環境をつくる

こんにちは、GunosyTechLab MediaML 所属 の suchida です。 急激に冬が近づいて参りましたね。 寒がりなので、ヒーター付き手袋を買いました。 キーボードも打てます。 おすすめです。

こちらの記事は Gunosy Advent Calendar 2022 の 3 日目の記事です。 前回の記事は nagayama さんの「Android の Kotlin Coroutines 導入の第一歩」でした。

tech.gunosy.io

はじめに

弊社では、社員が使っている PC が古くなってきたタイミングで新しいものへ置き換える PC リプレースを行っています。 今回私もリプレースの対象になり、MacBook Pro が Intel(amd64) から M1(arm64) になりました。 その際に開発環境構築でハマった点とどう解決したのかをまとめたいと思います。 類似の過去記事はこちらになります。

tech.gunosy.io

MediaML ではニュース記事のパーソナライズシステムの運用をしています。 その一部に記事やユーザーをベクトル化する Python 製の API があり、内部では TensorFlow を利用しています。 今回、その API を新モデルの検証のために新規構築する機会があり、そこで遭遇した arm64 由来の TensorFlow に関するエラーについてまとめていきます。
本記事では主に TensorFlow について記述していますが、amd64 と arm64 が混在するような開発環境の方には参考になるかと思います。

問題

チーム内では、Poetry を使ってパッケージ管理をしているので、それを前提に話を進めていきます。 Python バージョンは 3.10.6 とします。

M1 Mac に TensorFlow がインストールできない

まず、シンプルに M1 Mac に tensorflow をインストールしてみます。

poetry add tensorflow==2.10.1

すると、以下のようにインストールエラーになります。

Unable to find installation candidates for tensorflow-io-gcs-filesystem (0.28.0)

ISSUES にも挙がっているようです。

M1 Mac でサクッと TensorFlow を使いたい場合は、tensorflow-macos を使う必要があります。 実際に、インストールしてみましょう。

poetry add tensorflow-macos==2.10.0

今度は、無事インストールできました。 pyproject.toml にもしっかり追加されていますね。 ライブラリのインポートは import tensorflow で普段と同じように利用することが可能です。

# pyproject.toml

[tool.poetry]
name = "demo"
version = "0.1.0"
description = ""
authors = ["suchida"]

[tool.poetry.dependencies]
python = ">=3.10,<3.11"
tensorflow-macos = "^2.10.0"

めでたしめでたし...とはなりません。
API は EKS 上で稼働しており、その実行環境は amd64 であるため、arm64 ベースの pyproject.toml は残念ながら利用できません。 また、依然として M1 以前の Mac を利用しているエンジニアも多く、彼らの PC でも同様に環境構築ができなくなってしまいます。

では Docker 環境ではどうでしょうか? Docker 環境であれば、CPU アーキテクチャに依存せず TensorFlow の実行が可能かも?ということで検証してみましょう。

Docker 環境でも厳しい

例えば、以下のような pyproject.toml を poetry install できるようにした Dockerfile を build してみます。 platform は linux/amd64 を指定しています。

# pyproject.toml

[tool.poetry]
name = "demo"
version = "0.1.0"
description = ""
authors = ["suchida"]

[tool.poetry.dependencies]
python = ">=3.10,<3.11"
tensorflow = "^2.10.1"

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"

build は無事成功しました。

では、import tensorflow を実行するコンテナを適当に起動してみましょう。 すると、以下のようなメッセージが表示されコンテナがフリーズしてしまいました。

The TensorFlow library was compiled to use AVX instructions,
but these aren't available on your machine.
qemu: uncaught target signal 6 (Aborted) - core dumped

どうやら、M1 Mac のエミュレータ(qemu)では普段の TensorFlow を動かすことができないようです。 これに関しては、以下の記事が参考になるので、チェックしてみてください。

fdp-blog.hatenablog.com

M1 Mac の Docker で TensorFlow を利用したい場合は、別途 arm64 向けに TensorFlow を build する必要があるようです*1。 ただ build には時間を要するようなので、 Docker Hub で arm64 向け image を探して利用してみましょう*2。 今回は以下の image*3 を利用させていただきました。

FROM sonoisa/deep-learning-coding:pytorch1.12.0_tensorflow2.9.1

WORKDIR /srv
COPY demo.py demo.py

CMD ["python", "demo.py"]

platform は linux/arm64 を指定しています。 また、demo.py には import tensorflow を記述しています。 こちらを実行するとエラーが発生しないことが確認でき、arm64 専用の image を build すれば Docker 環境でも TensorFlow を動かすことができるようです。

めでたしめでたし...とはなりません。
前述したとおり EKS の実行環境は amd64 であるため、arm64 ベースの Dockerfile は残念ながら利用できません。

というわけで、amd64 用の実行環境と arm64 用の実行環境は分ける他なさそうというのが現状です。

CPU アーキテクチャの違いに躓かない開発環境づくり

本番環境である amd64 をベースに考えると、 arm64 では以下にまとめたように開発環境が構築できません。

amd64 向けプログラム  arm64
Poetry 環境 作成不可
Docker 環境 build 可、起動不可

唯一救いなのは、arm64 でも amd64 向け TensorFlow プログラムを build できるという点です。 image の build は M1 Mac 上で行い、skaffold などのデプロイツールを利用して EKS 上でコンテナを実行することで、ローカル環境にとらわれずに開発することが可能です。 それはそれでよいのですが、実行までに数分待つ必要がありガツガツ開発を進めたいときには不向きです。

開発環境

前述したとおり、amd64 と arm64 の両環境で TensorFlow を動かしたければ、個々の環境構築ファイルが必要です。 そのため、以下のようにベースとなる amd64 は最上位階層に展開し、ローカル開発時のみ必要な arm64 用ファイルは docs/arm64 に押し込みます。

.
├── Makefile
├── pyproject.toml // 本番環境(amd64)で必要なものは上位階層に
├── poetry.lock
├── docker-compose.yml
├── Dockerfile
└── docs
    └── arm64 // M1 Mac のローカル開発でのみ利用
        ├── Dockerfile
        ├── docker-compose.yml
        ├── poetry.lock
        ├── pyproject.toml
        └── requirements-tensorflow.txt

また、これらは Poetry 環境と Docker 環境の両方で動作するようにしていますので、amd64/arm64 と合わせて 4 通りの実行環境を用意している状態です。 なお、tensorflow-macos は Docker 環境では動作しないと思われるので、Poetry 管理はせず requirements-tensorflow.txt に記述し別途インストールするようにしています。
また、Docker 環境では arm64 向けに build された image を引っ張ってきて環境構築します。 arm64 用の開発環境は完全に amd64 のそれと一致することはできませんが、別途 EKS 上で動作チェックできる環境があるので特に問題を感じたことはありません。

また、誰でも・どの CPU アーキテクチャでも環境構築がシュッとできるように Makefile にコマンドを整備しています。 以下のように CPU アーキテクチャを確認し、内部で実行コマンドを切り替えることでシームレスに開発環境を構築することが可能です。

UNAME ?= $(shell uname -m)

docker-run-app: # Docker 環境で起動したい人用
ifeq ($(UNAME),arm64)
    docker-compose -f docs/arm64/docker-compose.yml up -d
else
    docker-compose up -d
endif

poetry-run-app: # Poetry 環境で使いたい人用(すでに仮想環境は別コマンドにて作成済みのため共通コマンドでよい)
    poetry run python hogehoge.py

先程述べた requirements-tensorflow.txt で一部分離してライブラリ管理している部分なども、環境構築コマンドを用意してあるので身構える必要はありません。

私は手元に仮想環境を作って、自動フォーマットや Pylance でコード補完をしながら開発を進めたい人なので、Poetry のコマンドを主に利用しています。 また Docker 環境は CI 上でテストを実行する際に、Redis などのインフラ周りを丸ごと記述しておいて make コマンドでシュッとテストができるので、こちらも重宝しています。

番外編: arm64 環境で amd64 のための pyproject.toml を作る方法

M1 Mac では amd64 用の TensorFlow を含んだ Poetry 環境が構築できません。 つまり、ライブラリの依存関係の解決もできず、TensorFlow に限らず新しくライブラリの追加ができません。 気軽に poetry add hogehoge できない訳です。

そこで、新たにライブラリを追加したい場合は platform:linux/amd64 のコンテナを起動し、その中で poetry add hogehoge するようにしています。

docker-poetry-add: # M1 で Intel 用の pyproject.toml を更新するためのコマンド
    docker-compose run api poetry add $(LIBRARY)

これにより、手元に arm64 環境しか用意がなくてもライブラリの更新が可能です。 ただし、私の場合は依存関係の解決に 10 分以上かかっており、激遅です。 ローカルにある Poetry の cache を コンテナ内にボリュームしたら早くなりそうです。 が、新規ライブラリを追加する機会はそこまで多くはないので、今の所対応予定はありません。

おわりに

今回奇跡的に TensorFlow を使った API の新規構築期と私の PC リプレース時期が被っていたので、手元に新旧合わせて 2 台の Mac が存在し、各 CPU アーキテクチャの差分をチェックしながら動作確認することができました。 この環境がなければ、そもそもエラーが arm64 起因なのか別問題なのか特定するにも手間取ることになっていたのではないかと戦々恐々です。 今後、arm64 対応が進んでいき、何も考えなくていい未来になるといいですね!!

次回は yamaguchi さんの JAWS に登壇した話です、お楽しみに!!