【Docker】Docker ComposeでマルチGPU AI環境を管理する方法

導入

AIモデルの学習や推論において、複数のGPUを効率的に活用することは、処理時間の短縮や大規模モデルの扱いにおいてほぼ必須となっています。しかし、Docker Composeを使用してマルチGPU環境を構築・管理しようとすると、コンテナが特定のGPUしか認識しない、リソースが適切に割り当てられない、あるいはnvidia-dockerランタイムの設定が複雑でエラーが発生するといった問題に直面することが少なくありません。これらの問題は、開発や本番環境の構築を大きく遅らせる要因となります。

本記事では、Docker Composeを用いて、複数のNVIDIA GPUを備えたサーバー上でAI環境を確実に構築し、管理するための具体的な方法を解説します。シンプルな設定から、複数コンテナへのGPUの分割割り当て、リソース制限に至るまで、実践的なソリューションを提供します。

原因の説明

Docker ComposeでマルチGPU環境が正しく機能しない主な原因は、大きく分けて二つあります。第一に、ランタイムの指定不足または誤りです。デフォルトのDockerランタイムはNVIDIA GPUを認識しません。そのため、明示的にnvidiaランタイムを指定する必要があります。第二に、リソース指定の不備です。単にランタイムを指定するだけでは、すべてのGPUがコンテナに使われてしまい、特定のGPUのみを割り当てたり、メモリ使用量を制限したりすることができません。これらの設定は、Docker Composeのデプロイ設定（deploy.resources）または環境変数を用いて詳細に制御する必要があります。

解決方法

以下、シンプルな方法から高度な設定まで、段階的に解決方法を説明します。

方法1: すべてのGPUをコンテナにアクセス可能にする（最も一般的な解決法）

最も簡単な方法は、ホストマシンのすべてのNVIDIA GPUをコンテナから利用できるようにすることです。これには、docker-compose.ymlファイルでnvidiaランタイムを指定し、環境変数NVIDIA_VISIBLE_DEVICES=allを設定します。

事前に、ホストマシンにNVIDIA Container Toolkitが正しくインストールされていることを確認してください。

# NVIDIA Container Toolkitのインストール確認
docker run --rm --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi

次に、以下のようなdocker-compose.ymlを作成します。

version: '3.8'

services:
  ai-training:
    image: nvidia/cuda:12.1.1-runtime-ubuntu22.04
    container_name: multi_gpu_app
    runtime: nvidia # nvidiaランタイムを明示的に指定
    environment:
      - NVIDIA_VISIBLE_DEVICES=all # すべてのGPUを可視化
    volumes:
      - ./app:/workspace
    working_dir: /workspace
    command: tail -f /dev/null # コンテナを起動したままにする例
    deploy: # Docker Swarmモード以外ではこのセクションは無視される場合があるので注意
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

この設定でコンテナを起動すると、コンテナ内からnvidia-smiコマンドを実行することで、すべてのGPUを認識していることを確認できます。

docker-compose up -d
docker-compose exec ai-training nvidia-smi

方法2: 特定のGPUのみを割り当てる

サーバーに複数のGPUがあり、特定のGPU（例えばGPU 0とGPU 2）のみを特定のコンテナで使用したい場合があります。この場合、環境変数NVIDIA_VISIBLE_DEVICESにGPUのインデックスまたはUUIDを指定します。

version: '3.8'

services:
  ai-service-1:
    image: pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
    container_name: gpu_0_2_app
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=0,2 # GPUインデックス0番と2番を割り当て
    # ... その他の設定

  ai-service-2:
    image: tensorflow/tensorflow:2.13.0-gpu
    container_name: gpu_1_app
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=1 # GPUインデックス1番を割り当て
    # ... その他の設定

これにより、異なるコンテナに異なるGPUリソースを柔軟に割り振ることが可能になります。GPUのUUIDを確認するには、ホストでnvidia-smi -Lコマンドを実行します。

方法3: 高度なリソース制限を適用する (Docker Compose v2.3+ および NVIDIA Container Toolkit v2.0+)

より細かい制御として、GPUメモリの制限や、コンピュート能力（MIG）の指定などが可能です。これにはdeploy.reservations.devicesセクションを詳細に記述します。この記法は、Docker Compose v2.3以上およびNVIDIA Container Toolkit v2.0以上でサポートされています。

version: '3.8'

services:
  precision-training:
    image: nvcr.io/nvidia/pytorch:23.10-py3
    container_name: limited_gpu_app
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=0 # 制限を適用するGPUを指定
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: ['0'] # 割り当てるGPU ID
              capabilities: [gpu]
              # 以下のオプションで詳細なリソース要求を指定
              options:
                # GPUメモリの制限（単位はメガバイト）
                memory: 4096
                # 特定のコンピュート能力を要求（例: '7.0', '8.0'）
                compute: '8.0'
    # ... その他の設定

この設定では、GPU 0に対して最大4GBのメモリ使用制限を設けています。コンテナ内のプロセスがこの制限を超えようとすると、CUDAエラーが発生します。これは、複数の軽量なAIサービスを同一GPU上で安全に共存させたい場合に有効です。

注意点: deployセクションは、従来のdocker-compose upでは完全にはサポートされていない場合があります。この機能を完全に活用するには、docker compose（v2コマンド）を使用し、かつDocker Engineのクラシックスwarmモードではなく、通常のスタンドアローンコンテナとして実行するコンテキストで動作します。挙動が不安定な場合は、環境変数による方法（方法2）を基本とし、メモリ制限はアプリケーション側（例: PyTorchのtorch.cuda.set_per_process_memory_fraction()）で行うことも現実的な選択肢です。

まとめ

Docker ComposeでマルチGPU環境を管理する手順は、以下のように要約できます。

1. 前提条件の確認: ホストにNVIDIA Container Toolkitがインストールされ、--gpus allオプションで動作することを確認する。
2. ランタイムの指定: docker-compose.ymlのサービス定義でruntime: nvidiaを明記する。
3. GPUの選択的割り当て: 環境変数NVIDIA_VISIBLE_DEVICESを用いて、使用するGPUをインデックスまたはUUIDで指定する（例: 0,2）。
4. （必要に応じて）詳細なリソース制限: Docker Compose v2.3+環境では、deploy.resources.reservations.devices.optionsを用いてメモリやコンピュート能力の制限を設定できる。

追加のTips:

• プロダクション環境では、GPU UUIDを使用した指定の方が、物理的なGPUの順序が変わっても確実です。
• 開発時は、docker-compose.ymlと並列して.envファイルを作成し、NVIDIA_VISIBLE_DEVICES=${GPU_IDS}のように環境変数で管理すると、環境ごとの切り替えが容易になります。
• PyTorchやTensorFlowなどのフレームワークは、コンテナ内でCUDA_VISIBLE_DEVICES環境変数も参照します。NVIDIA_VISIBLE_DEVICESでフィルタリングされた後のGPUインデックスがCUDA_VISIBLE_DEVICESに反映されるため、アプリケーション側の設定と混同しないように注意しましょう。