使用 Docker Compose 啟用 GPU 存取

如果 Docker 主機包含 GPU 裝置,且 Docker Daemon 已相應設定,則 Compose 服務可以定義 GPU 裝置預留。 為此,請確保您已安裝 必要條件

以下章節中的範例專注於使用 Docker Compose 提供服務容器存取 GPU 裝置。 您可以使用 docker-composedocker compose 命令。 如需更多資訊,請參閱 遷移至 Compose V2

啟用服務容器的 GPU 存取

在需要 GPU 的服務中,可以使用 Compose Deploy 規格中的 device 屬性在 compose.yml 檔案中參考 GPU。

這可以更精細地控制 GPU 預留,因為可以為以下裝置屬性設定自訂值

  • capabilities。 此值指定為字串清單(例如 capabilities: [gpu])。 您必須在 Compose 檔案中設定此欄位。 否則,它會在服務部署時傳回錯誤。
  • count。 此值指定為整數或值 all,表示應預留的 GPU 裝置數量(前提是主機擁有該數量的 GPU)。 如果 count 設定為 all 或未指定,則預設使用主機上所有可用的 GPU。
  • device_ids。 此值指定為字串清單,表示主機的 GPU 裝置 ID。 您可以在主機上 nvidia-smi 的輸出中找到裝置 ID。 如果未設定 device_ids,則預設使用主機上所有可用的 GPU。
  • driver。 此值指定為字串,例如 driver: 'nvidia'
  • options。 代表驅動程式特定選項的鍵值對。

重要

您必須設定 capabilities 欄位。 否則,它會在服務部署時傳回錯誤。

countdevice_ids 互斥。 您一次只能定義一個欄位。

有關這些屬性的更多資訊,請參閱 Compose Deploy 規格

使用 1 個 GPU 裝置執行服務的 Compose 檔案範例

services:
  test:
    image: nvidia/cuda:12.3.1-base-ubuntu20.04
    command: nvidia-smi
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

使用 Docker Compose 執行

$ docker compose up
Creating network "gpu_default" with the default driver
Creating gpu_test_1 ... done
Attaching to gpu_test_1    
test_1  | +-----------------------------------------------------------------------------+
test_1  | | NVIDIA-SMI 450.80.02    Driver Version: 450.80.02    CUDA Version: 11.1     |
test_1  | |-------------------------------+----------------------+----------------------+
test_1  | | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
test_1  | | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
test_1  | |                               |                      |               MIG M. |
test_1  | |===============================+======================+======================|
test_1  | |   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
test_1  | | N/A   23C    P8     9W /  70W |      0MiB / 15109MiB |      0%      Default |
test_1  | |                               |                      |                  N/A |
test_1  | +-------------------------------+----------------------+----------------------+
test_1  |                                                                                
test_1  | +-----------------------------------------------------------------------------+
test_1  | | Processes:                                                                  |
test_1  | |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
test_1  | |        ID   ID                                                   Usage      |
test_1  | |=============================================================================|
test_1  | |  No running processes found                                                 |
test_1  | +-----------------------------------------------------------------------------+
gpu_test_1 exited with code 0

在裝載多個 GPU 的機器上,可以設定 device_ids 欄位以鎖定特定 GPU 裝置,並可以使用 count 來限制分配給服務容器的 GPU 裝置數量。

您可以在每個服務定義中使用 countdevice_ids。 如果您嘗試組合兩者、指定無效的裝置 ID 或使用高於系統中 GPU 數量的 count 值,則會傳回錯誤。

$ nvidia-smi   
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 450.80.02    Driver Version: 450.80.02    CUDA Version: 11.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1B.0 Off |                    0 |
| N/A   72C    P8    12W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Tesla T4            On   | 00000000:00:1C.0 Off |                    0 |
| N/A   67C    P8    11W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   2  Tesla T4            On   | 00000000:00:1D.0 Off |                    0 |
| N/A   74C    P8    12W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   3  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   62C    P8    11W /  70W |      0MiB / 15109MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

存取特定裝置

若要僅允許存取 GPU-0 和 GPU-3 裝置

services:
  test:
    image: tensorflow/tensorflow:latest-gpu
    command: python -c "import tensorflow as tf;tf.test.gpu_device_name()"
    deploy:
      resources:
        reservations:
          devices:
          - driver: nvidia
            device_ids: ['0', '3']
            capabilities: [gpu]