使用Arena提交PyTorch單機訓練作業 - Container Service for Kubernetes

PyTorch是一個開源的深度學習架構，廣泛應用於各種深度學習模型的訓練任務中，本文示範如何使用Arena提交PyTorch單機單卡或單機多卡訓練作業，並通過TensorBoard可視化查看訓練作業。

前提條件

已建立包含GPU的Kubernetes叢集。具體操作，請參見建立包含GPU的Kubernetes叢集。
叢集節點可以訪問公網。具體操作，請參見為已有叢集開啟公網訪問能力。
已安裝Arena工具。具體操作，請參見配置Arena用戶端。
建立一個名為training-data的PVC執行個體，並在路徑/pytorch_data下存放MNIST資料集。具體操作，請參見配置NAS共用儲存。

背景資訊

本文將從遠程Git倉庫中下載訓練代碼，並從共用儲存系統（基於NAS的PV和PVC）中讀取訓練資料。torchrun 是 PyTorch 提供的一個命令列工具，用於簡化和管理分布式訓練任務，本文將使用 torchrun 來啟動 PyTorch 單機單卡或單機多卡訓練任務，訓練代碼詳見 main.py。

操作步驟

步驟一：查看GPU資源

執行以下命令查看叢集中可用的GPU資源。

arena top node

預期輸出：

NAME                        IPADDRESS        ROLE    STATUS  GPU(Total)  GPU(Allocated)
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           0
---------------------------------------------------------------------------------------------------
Allocated/Total GPUs In Cluster:
0/4 (0.0%)

可以看到叢集中有2個GPU節點，每個GPU節點都包含2張閒置GPU卡可用於運行訓練作業。

步驟二：提交PyTorch訓練作業

執行如下命令提交一個單機單卡的PyTorch訓練作業。

arena submit pytorch \
    --name=pytorch-mnist \
    --namespace=default \
    --workers=1 \
    --gpus=1 \
    --working-dir=/root \
    --image=kube-ai-registry.cn-shanghai.cr.aliyuncs.com/kube-ai/pytorch-with-tensorboard:2.5.1-cuda12.4-cudnn9-runtime \
    --sync-mode=git \
    --sync-source=https://github.com/kubeflow/arena.git \
    --env=GIT_SYNC_BRANCH=v0.13.1 \
    --data=training-data:/mnt \
    --tensorboard \
    --logdir=/mnt/pytorch_data/logs \
    "torchrun /root/code/arena/examples/pytorch/mnist/main.py --epochs 10 --backend nccl --data /mnt/pytorch_data  --dir /mnt/pytorch_data/logs"

預期輸出：

service/pytorch-mnist-tensorboard created
deployment.apps/pytorch-mnist-tensorboard created
pytorchjob.kubeflow.org/pytorch-mnist created
INFO[0002] The Job pytorch-mnist has been submitted successfully
INFO[0002] You can run `arena get pytorch-mnist --type pytorchjob -n default` to check the job status

說明

如果您需要執行單機多卡訓練作業，例如執行單機兩卡訓練作業，則可以通過指定 --gpus=2 和 --nproc-per-node=2 參數以分配2張GPU卡並啟動2個訓練進程：

arena submit pytorch \
    --name=pytorch-mnist \
    --namespace=default \
    --workers=1 \
    --gpus=2 \
    --nproc-per-node=2 \
    --working-dir=/root \
    --image=kube-ai-registry.cn-shanghai.cr.aliyuncs.com/kube-ai/pytorch-with-tensorboard:2.5.1-cuda12.4-cudnn9-runtime \
    --sync-mode=git \
    --sync-source=https://github.com/kubeflow/arena.git \
    --env=GIT_SYNC_BRANCH=v0.13.1 \
    --data=training-data:/mnt \
    --tensorboard \
    --logdir=/mnt/pytorch_data/logs \
    "torchrun /root/code/arena/examples/pytorch/mnist/main.py --epochs 10 --backend nccl --data /mnt/pytorch_data  --dir /mnt/pytorch_data/logs"

說明

如果您使用的是非公開Git代碼倉庫，則可以通過配置環境變數GIT_SYNC_USERNAME和GIT_SYNC_PASSWORD的方式來設定Git使用者名稱和密碼。

  arena submit pytorch \
        ...
        --sync-mode=git \
        --sync-source=https://github.com/kubeflow/arena.git \
        --env=GIT_SYNC_BRANCH=v0.13.1 \
        --env=GIT_SYNC_USERNAME=<username> \
        --env=GIT_SYNC_PASSWORD=<password> \
        "torchrun /root/code/arena/examples/pytorch/mnist/main.py --epochs 10 --backend nccl --data /mnt/pytorch_data  --dir /mnt/pytorch_data/logs"

arena命令使用git-sync同步原始碼，因此您可以使用任何在git-sync專案中定義的環境變數。

重要

本文樣本從GitHub倉庫中拉取原始碼，如遇到網路原因等導致代碼無法成功拉取時，可以手動將代碼下載到共用儲存系統中，假設您已經將代碼下載至NAS中的 /code/github.com/kubeflow/arena 路徑下，則可以如下提交訓練作業：

arena submit pytorch \
    --name=pytorch-mnist \
    --namespace=default \
    --workers=1 \
    --gpus=1 \
    --working-dir=/root \
    --image=kube-ai-registry.cn-shanghai.cr.aliyuncs.com/kube-ai/pytorch-with-tensorboard:2.5.1-cuda12.4-cudnn9-runtime \
    --data=training-data:/mnt \
    --tensorboard \
    --logdir=/mnt/pytorch_data/logs \
    "torchrun /mnt/code/github.com/kubeflow/arena/examples/pytorch/mnist/main.py --epochs 10 --backend nccl --data /mnt/pytorch_data  --dir /mnt/pytorch_data/logs"

參數解釋如下表。

參數	是否必選	解釋	預設值
`--name`	必選	作業名稱，全域唯一，不能重複。	無
`--namespace`	可選	作業所屬的命名空間。	`default`
`--workers`	可選	指定作業需要的節點數，它設定的值包含Master節點。例如設定3，表明該作業包含1個Master節點和2個Worker節點。	`0`
`--gpus`	可選	指定作業Worker節點需要使用的GPU卡數。	`0`
`--working-dir`	可選	指定當前執行命令所在的目錄。	`/root`
`--image`	必選	指定訓練環境的鏡像地址。	無
`--sync-mode`	可選	同步代碼的模式，您可以指定git、rsync。本文使用Git模式。	無
`--sync-source`	可選	同步代碼的倉庫地址，需要和--sync-mode一起使用，本文樣本使用Git模式，該參數可以為任何GitHub專案、阿里雲Code專案等支援Git的代碼託管地址。專案代碼將會被下載到--working-dir下的`code/`目錄中，在本樣本中則為`/root/code/arena`。	無
`--data`	可選	掛載共用儲存卷PVC到運行環境中。它由兩部分組成，通過分號`:`分割。分號左側是您已經準備好的PVC名稱。您可以通過`arena data list`查看當前叢集可用的PVC列表；分號右側是您想將PVC的掛載到運行環境中的路徑，也是您訓練代碼要讀取資料的本地路徑。這樣通過掛載的方式，您的代碼就可以訪問PVC的資料。說明執行`arena data list`查看本文樣本當前叢集可用的PVC列表。 `NAME ACCESSMODE DESCRIPTION OWNER AGE training-data ReadWriteMany 35m` 如果沒有可用的PVC，您可建立PVC。詳情請參見配置NAS共用儲存。	無
`--tensorboard`	可選	為訓練任務開啟一個TensorBoard服務，用作資料視覺效果，您可以結合--logdir指定TensorBoard要讀取的event路徑。不指定該參數，則不開啟TensorBoard服務。	無
`--logdir`	可選	需要結合--tensorboard一起使用，該參數表示TensorBoard需要讀取event資料的路徑。	`/training_logs`

步驟三：查看PyTorch訓練作業

執行以下命令查看通過arena提交的訓練作業。

arena list -n default

預期輸出：

NAME           STATUS   TRAINER     DURATION  GPU(Requested)  GPU(Allocated)  NODE
pytorch-mnist  RUNNING  PYTORCHJOB  11s       1               1               192.168.xxx.xxx

執行以下命令查看訓練作業所使用的GPU資源。

arena top job -n default

預期輸出：

NAME           STATUS   TRAINER     AGE  GPU(Requested)  GPU(Allocated)  NODE
pytorch-mnist  RUNNING  PYTORCHJOB  18s  1               1               192.168.xxx.xxx

Total Allocated/Requested GPUs of Training Jobs: 1/1

執行以下命令查看叢集中的GPU資源使用方式。

arena top node

預期輸出：

NAME                        IPADDRESS        ROLE    STATUS  GPU(Total)  GPU(Allocated)
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   0           0
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           1
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           0
---------------------------------------------------------------------------------------------------
Allocated/Total GPUs In Cluster:
1/4 (25.0%)

可以看到叢集中總共有4張GPU，當前已經分配1張GPU。

執行以下命令查看訓練作業詳情。

arena get pytorch-mnist -n default

預期輸出：

Name:        pytorch-mnist
Status:      RUNNING
Namespace:   default
Priority:    N/A
Trainer:     PYTORCHJOB
Duration:    45s
CreateTime:  2025-02-12 11:20:10
EndTime:

Instances:
  NAME                    STATUS   AGE  IS_CHIEF  GPU(Requested)  NODE
  ----                    ------   ---  --------  --------------  ----
  pytorch-mnist-master-0  Running  45s  true      1               cn-beijing.192.168.xxx.xxx

Tensorboard:
  Your tensorboard will be available on:
  http://192.168.xxx.xxx:31949

說明

當在訓練作業中啟用了TensorBoard時，作業詳情中將會展示TensorBoard訪問地址；如果沒有啟用TensorBoard，則不會展示。

步驟四：查看TensorBoard

在本地執行如下命令，將叢集中的TensorBoard的訪問連接埠6006映射到本地的9090連接埠。
重要
請注意kubectl port-forward建立的連接埠轉寄不具備生產層級的可靠性、安全性和擴充性，因此僅適用於開發和調試目的，不適合在生產環境使用。更多關於Kubernetes叢集內生產可用的網路方案的資訊，請參見Ingress管理。
```
kubectl port-forward -n default svc/pytorch-mnist-tensorboard 9090:6006
```
在瀏覽器中訪問 http://127.0.0.1:9090，即可查看TensorBoard。
說明
本文中PyTorch範例程式碼，預設每10個epoch寫入event資訊。如果您修改了--epochs參數，請修改為10的整數倍，否則無法在TensorBoard上看到資料。

步驟五：查看訓練作業日誌

執行如下命令查看訓練作業日誌。

arena logs pytorch-mnist -n default

預期輸出：

Train Epoch: 10 [55680/60000 (93%)]     Loss: 0.025778
Train Epoch: 10 [56320/60000 (94%)]     Loss: 0.086488
Train Epoch: 10 [56960/60000 (95%)]     Loss: 0.003240
Train Epoch: 10 [57600/60000 (96%)]     Loss: 0.046731
Train Epoch: 10 [58240/60000 (97%)]     Loss: 0.010752
Train Epoch: 10 [58880/60000 (98%)]     Loss: 0.010934
Train Epoch: 10 [59520/60000 (99%)]     Loss: 0.065813

Accuracy: 9921/10000 (99.21%)

說明

如果需要即時查看作業日誌，可以添加 -f 參數；
如果僅需要查看最後 N 行日誌，可以添加 -t N 或 --tail N 參數；
更多用法請參見 arena logs --help。

（可選）步驟六：環境清理

訓練作業執行結束後如不再需要，執行如下命令進行刪除：

arena delete pytorch-mnist -n default

預期輸出：

INFO[0001] The training job pytorch-mnist has been deleted successfully