Kubernetes上のTrueFoundryでDockerイメージのビルドを3～15倍高速化

機械学習（ML）の世界では、効率的なDockerイメージのビルドはもはや贅沢品ではなく、必要不可欠なものです。ほとんどの企業は、ローカルのノートPCまたはCI/CDパイプラインを介してDockerイメージをビルドおよびプッシュするための既存のDevOpsパイプラインを持っています。しかし、MLプロジェクトが複雑化し、依存関係が大きくなり、イテレーションが頻繁になるにつれて、従来のDockerビルドプロセスが大きなボトルネックとなる可能性があります。

‍本記事では、標準的なCIパイプラインと比較して、Truefoundryでビルド時間を5～15倍短縮した方法を解説します。

MLにとって高速なDockerビルドが重要な理由

大量の依存関係

MLプロジェクトでは通常、ディープラーニングフレームワーク（PyTorch、TensorFlow）、科学計算ライブラリ（NumPy、SciPy）、GPUドライバー、CUDAツールキットなど、多数の重い依存関係が含まれます。これらの依存関係により、Dockerイメージのサイズが数ギガバイトになり、ビルド時間が長くなる可能性があります。

迅速なイテレーションサイクル

ML開発では、テストのためにデプロイする必要がある頻繁なコード変更が伴います。データサイエンティストは、ローカルのノートPCでコードを実行するために必要なハードウェアを持っていないことがよくあります。これは、コードがリモートクラスターで実行される必要があり、多くの場合、イメージのビルドが伴うことを意味します。

• モデル改善のための頻繁なコード変更
• 依存関係の定期的な更新
• さまざまなモデルアーキテクチャでの継続的な実験、さまざまな構成でのA/Bテスト
  

Truefoundryでは、開発者が迅速なイテレーションペースで作業できるようにすることを目指しており、そのためDockerビルドを非常に高速化したいと考えていました。ビルド時間を最適化するために何をしたかを理解するために、まずTruefoundryで以前どのようにイメージをビルドしていたかを理解しましょう。

開発者がイメージをビルドしたいと思うたびに、コードはコントロールプレーンにアップロードされ、そこで新しいPodがサイドカーでBuildkitを実行してイメージのビルドを開始します。宛先のDockerレジストリはキャッシングレイヤーとして機能し、最終的なイメージはDockerレジストリにプッシュされます。

このセットアップは、ほとんどのCIビルダーと同じであり、現在のCIセットアップと同じ長所と短所があります。

これには以下の利点がありました。

1. システムは無限にスケーリングでき、すべてのビルドが個別のPodとして実行されるため、任意の数の並行ビルドを同時にサポートできます。
2. すべてのビルドのBuildkitは分離されるため、他のビルドに影響を与えることはありません。
3. Buildkitは、最先端の並列処理と高度なキャッシュ機能を提供します。
   

しかし、このアプローチにはいくつかの欠点がありました。

1. BuildkitのPodは大量のリソースを必要とするため、ビルドランナーの起動に時間がかかります。

2. Dockerレジストリからキャッシュをダウンロードするのに多くの時間がかかり、ビルド時間の遅延につながります。

3. 異なるワークロードのビルド間でキャッシュの再利用ができません。

私たちは、ローカルビルドと比較して、リモートでのイメージ構築において同等（あるいはそれ以上）のエクスペリエンスを提供したいと考えました。

Dockerイメージ構築プロセスの高速化

そこで、まずBuildkitのPodをKubernetes上でサービスとしてホストし、複数のビルダー間で共有できるようにし、ローカルディスクキャッシュを提供することで、Dockerビルドを非常に高速にすることにしました。

しかし、このアプローチにはいくつかの制約があります。
1. Buildkitには根本的な制約があり、 キャッシュファイルシステムはBuildkitの1つのインスタンスによってのみ使用できます。これは、複数のビルドを並行して処理するためにBuildkitの複数のインスタンスを稼働させた場合、それぞれが独自のキャッシュを持つことになり、それらは共有できないことを意味します。

2. 複数のBuildkitインスタンスを稼働させ、それぞれが独自のキャッシュを持つ場合、 キャッシュを効果的に使用できるように、同じワークロードは同じマシンにルーティングされるべきです。これにはカスタムルーティングロジックが必要です。

3. 実行中のビルド数に基づいたBuildkit Podのオートスケーリングは容易ではありません。Buildkit PodのCPU使用率をオートスケーリングのメトリックとして使用することはできません。なぜなら、Kubernetesがそのマシン上で何も実行されていないと判断し、実行中の小さなビルドを終了させる可能性があるためです。

動的な数のBuildkit Podを持ち、ワークロードが同じキャッシュインスタンスにルーティングされるようにすることは、容易ではない問題です。KubernetesではPod間でのボリュームのアタッチとデタッチは非常に遅く、そのためビルドの開始に非常に長い時間がかかります。

上記の制約を克服するため、私たちはハイブリッドなアプローチを考案しました。これにより、ほとんどのビルドは非常に高速に完了しますが、並行ビルドの同時実行性が高いまれなケースでは、以前の、Buildkitをサイドカーで実行するフローに戻るようにしています。

下図に示すアーキテクチャでは、特定のビルド同時実行数を設定し、それ以下のビルドはすべてBuildKitサービスに送られるようにしています。これを説明するために、BuildKitサービス用に4CPUと16GB RAMを搭載したマシンを1台割り当てると仮定しましょう。過去のビルドデータから、このマシンは2つの同時ビルドを処理できることがわかります。したがって、すでに1つのビルドが実行中で、新しいビルドが来た場合、それはBuildKitサービスにルーティングされます。しかし、さらに別のビルドが来た場合、Dockerレジストリに保存されたレイヤーキャッシュを使用し、BuildKitをサイドカーとして実行する以前のモデルにルーティングします。

これにより、99%のワークロードで超高速ビルドを提供できますが、ごく一部のケースでは、標準的なCIパイプラインで通常かかる時間と同じくらいビルドに時間がかかります。

ビルド速度のさらなる改善

ビルドプロセスを高速化するために、他にもいくつかの改善を行いました。その一部を以下に示します。

1. pipをuvに置き換える: uvは、pipと比較してPythonの依存関係解決とパッケージインストールをはるかに高速化します。pipをuvに置き換えることで、ビルド時間を約40%短縮できました。
2. EBSの代わりにNVMeディスクを使用する: NVMeディスクは高速な読み書きアクセスを提供し、キャッシュの書き込みと読み込みを高速化することで、ビルド速度の向上に貢献します。

速度改善のベンチマーク

私たちの実験をベンチマークするために、MLワークロードで最も一般的なシナリオを表す1つのサンプルDockerfileを使用しました。

FROM tfy.jfrog.io/tfy-mirror/python:3.10.2-slim

WORKDIR /app

RUN echo "Starting the build"

COPY ./requirements.txt /app/requirements.txt

RUN pip install -r requirements.txt

COPY . /app/

EXPOSE 8000

CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

requirements.txtは以下の通りです。

fastapi[standard]==0.109.1
huggingface-hub==0.24.6
vllm==0.5.4
transformers==4.43.3

ビルドのベンチマークを以下の3つのシナリオで測定しました。

1. 初回ビルド（キャッシュなし）
2. Pythonコードを変更したが依存関係は変更していない場合の2回目のビルド。
3. requirements.txtの依存関係を変更した後の2回目のビルド。

測定時間には、イメージのビルドとレジストリへのプッシュにかかる時間が含まれており、単位は秒です。

コードのみを変更する2番目のシナリオは、開発者が遭遇する最も一般的なシナリオであり、ご覧のとおり、ビルド時間はほぼ15倍改善されています。

また、TritonをベースイメージとするDockerfileを含むシナリオのベンチマークも測定しました。Tritonははるかに大きなベースイメージです。

FROM nvcr.io/nvidia/tritonserver:24.09-py3

WORKDIR /app

RUN echo "Starting the build"

COPY ./requirements2.txt /app/requirements.txt

RUN pip install -r requirements.txt

RUN echo "Finished the build"

結果は以下のとおりです。

このケースでは、初回ビルドのビルド時間が3倍に、その後のビルドでは9倍に高速化されることが確認できました。

上記の変更により、開発体験が劇的に向上し、最終的に本番環境にデプロイされる方法との整合性を保ちながら、アイデアを高速で試行錯誤できるようになりました。