Use Arena to submit distributed PyTorch training jobs - Container Service for Kubernetes: Central de Documentação do Alibaba Cloud

Pré-requisitos

Antes de começar, verifique se você tem:

Um cluster ACK com nós acelerados por GPU. Para mais detalhes, consulte Criar um cluster ACK que contém nós acelerados por GPU.
Acesso à internet habilitado para os nós do cluster. Para mais detalhes, consulte Habilitar o acesso à internet em um cluster ACK existente.
O add-on Arena instalado. Para mais detalhes, consulte Configurar o cliente Arena.
Uma persistent volume claim (PVC) chamada training-data com o conjunto de dados MNIST armazenado em /pytorch_data. Para mais detalhes, consulte Configurar um volume NAS compartilhado.

Contexto

Este tutorial treina um modelo PyTorch no conjunto de dados MNIST usando dois nós, cada um com duas GPUs, totalizando quatro GPUs. O código de treinamento é executado com o torchrun, o launcher nativo do PyTorch para jobs distribuídos.

O fluxo de trabalho obtém o código de treinamento de um repositório git e lê os dados de um volume compartilhado baseado em NAS (PV/PVC). Como referência, consulte main.py.

Etapa 1: Verificar os recursos de GPU disponíveis

Execute o comando a seguir para visualizar as GPUs disponíveis no cluster:

arena top node

Saída esperada:

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%)

O cluster possui dois nós acelerados por GPU, cada um com duas GPUs ociosas, totalizando quatro GPUs disponíveis para treinamento.

Etapa 2: Enviar um job de treinamento distribuído do PyTorch

Execute o comando abaixo para enviar o job. Ele cria dois Pods worker, cada um utilizando duas GPUs.

Três parâmetros controlam a topologia distribuída:

--workers=2 — número total de Pods, incluindo o Pod master
--gpus=2 — GPUs alocadas para cada Pod
--nproc-per-node=2 — processos do torchrun iniciados por Pod; cada processo utiliza uma GPU

arena submit pytorch \
    --name=pytorch-mnist \
    --namespace=default \
    --workers=2 \
    --gpus=2 \
    --nproc-per-node=2 \
    --clean-task-policy=None \
    --working-dir=/root \
    --image=kube-ai-registry.cn-shanghai.cr.aliyuncs.com/kube-ai/pytorch-mnist-example: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"

Saída esperada:

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

Como funciona: nós, funções e variáveis de ambiente

Jobs distribuídos do PyTorch no Arena utilizam dois parâmetros desnecessários em jobs standalone:

--workers — o número total de Pods. Um Pod assume a função master; os demais são Pods worker.
--nproc-per-node — a quantidade de processos do torchrun iniciados em cada Pod. Cada processo corresponde a uma GPU.

Os nomes dos Pods seguem o padrão <job_name>-<role_name>-<index>. Por exemplo, usar --workers=3 e --nproc-per-node=2 em um job chamado pytorch-mnist cria três Pods e inicia dois processos em cada um:

Variável de ambiente	pytorch-mnist-master-0	pytorch-mnist-worker-0	pytorch-mnist-worker-1
`MASTER_ADDR`	pytorch-mnist-master-0	—	—
`MASTER_PORT`	23456	—	—
`WORLD_SIZE`	6	—	—
`RANK`	0	1	2
`PET_MASTER_ADDR`	pytorch-mnist-master-0	—	—
`PET_MASTER_PORT`	23456	—	—
`PET_NNODES`	3	—	—
`PET_NODE_RANK`	0	1	2

O Arena injeta essas variáveis automaticamente em cada Pod. No código de treinamento, utilize RANK para identificar qual processo deve salvar checkpoints ou registrar métricas. Por exemplo, apenas o processo onde RANK=0 (o master) grava os resultados, evitando saídas duplicadas. Utilize WORLD_SIZE para determinar o número total de processos paralelos ao calcular gradientes distribuídos.

Referência de parâmetros

Parâmetro	Obrigatório	Descrição	Padrão
`--name`	Sim	Nome do job, globalmente único	—
`--namespace`	Não	Namespace do Kubernetes	`default`
`--workers`	Não	Número total de Pods worker, incluindo o master. Por exemplo, `--workers=3` cria um Pod master e dois Pods worker	`0`
`--gpus`	Não	GPUs alocadas para cada Pod worker	`0`
`--working-dir`	Não	Diretório onde o comando de treinamento é executado	`/root`
`--image`	Sim	Imagem de contêiner usada para o runtime	—
`--sync-mode`	Não	Modo de sincronização do código-fonte: `git` ou `rsync`	—
`--sync-source`	Não	URL do repositório para sincronização do código-fonte. O código é baixado para `code/` dentro de `--working-dir` (por exemplo, `/root/code/arena`)	—
`--data`	Não	Monta uma PVC no Pod como `<pvc-name>:<mount-path>`. Execute `arena data list` para visualizar as PVCs disponíveis	—
`--tensorboard`	Não	Habilita a visualização no TensorBoard. Requer `--logdir`	—
`--logdir`	Não	Caminho onde o TensorBoard lê os arquivos de eventos. Use em conjunto com `--tensorboard`	`/training_logs`

Usar um repositório git privado

Se o código estiver em um repositório privado, passe as credenciais como variáveis de ambiente:

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"

O Arena utiliza o git-sync para obter o código-fonte. Portanto, todas as variáveis de ambiente definidas no projeto git-sync estão disponíveis.

Não consegue obter o código devido a problemas de rede?

Caso o repositório do GitHub esteja inacessível, baixe o código manualmente para o volume NAS. Coloque-o em code/github.com/kubeflow/arena dentro do NAS. Isso corresponde a /mnt/code/github.com/kubeflow/arena após a montagem do volume. Em seguida, envie o job sem --sync-mode:

arena submit pytorch \
    --name=pytorch-mnist \
    --namespace=default \
    --workers=2 \
    --gpus=2 \
    --nproc-per-node=2 \
    --clean-task-policy=None \
    --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"

Etapa 3: Monitorar o job de treinamento

Após o envio, o job passa por vários estágios antes que os logs apareçam:

Pending — O Kubernetes está agendando os Pods nos nós com GPU.
Preparing — Os Pods baixam a imagem do contêiner e sincronizam o repositório git. Os logs ainda não estão disponíveis.
Running — O torchrun inicia e o treinamento começa. Os logs ficam disponíveis nesta fase.

Utilize os comandos a seguir para acompanhar o progresso.

Listar todos os jobs:

arena list -n default

Saída esperada:

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

Verificar o uso de GPU por job:

arena top job -n default

Saída esperada:

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

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

Verificar a alocação de GPU no nível do cluster:

arena top node

Saída esperada:

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           2
cn-beijing.192.168.xxx.xxx  192.168.xxx.xxx  <none>  Ready   2           2
---------------------------------------------------------------------------------------------------
Allocated/Total GPUs In Cluster:
4/4 (100.0%)

Todas as quatro GPUs agora estão alocadas.

Obter informações detalhadas do job:

arena get pytorch-mnist -n default

Saída esperada:

Name:        pytorch-mnist
Status:      RUNNING
Namespace:   default
Priority:    N/A
Trainer:     PYTORCHJOB
Duration:    1m
CreateTime:  2025-02-12 13:54:51
EndTime:

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

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

O job possui um Pod master e um Pod worker, cada um utilizando duas GPUs. A URL do TensorBoard é exibida apenas quando --tensorboard está habilitado.

Etapa 4: Visualizar os resultados do treinamento no TensorBoard

Importante

O comando kubectl port-forward destina-se apenas a desenvolvimento e depuração. Ele não é confiável, seguro ou escalável para produção. Para redes de produção em clusters ACK, consulte Gerenciamento de Ingress.

Encaminhe a porta do TensorBoard para sua máquina local:

kubectl port-forward -n default svc/pytorch-mnist-tensorboard 9090:6006

Abra http://127.0.0.1:9090 em um navegador.

O código de treinamento de exemplo grava eventos a cada 10 épocas. Se você alterar --epochs , defina um múltiplo de 10. Caso contrário, nenhum evento será gravado e o TensorBoard não exibirá dados.

Etapa 5: Visualizar os logs de treinamento

Visualizar os logs do Pod master:

arena logs -n default pytorch-mnist

Saída esperada:

{'PID': 40, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 0, 'RANK': 0, 'GROUP_RANK': 0, 'ROLE_RANK': 0, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
{'PID': 41, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 1, 'RANK': 1, 'GROUP_RANK': 0, 'ROLE_RANK': 1, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
Using cuda:0.
Using cuda:1.
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
...
Train Epoch: 10 [59520/60000 (99%)]     Loss: 0.007343
Train Epoch: 10 [59520/60000 (99%)]     Loss: 0.007343

Accuracy: 9919/10000 (99.19%)

Accuracy: 9919/10000 (99.19%)

Visualizar os logs de um Pod worker específico:

arena logs -n default -i pytorch-mnist-worker-0 pytorch-mnist

Saída esperada:

{'PID': 39, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 0, 'RANK': 2, 'GROUP_RANK': 1, 'ROLE_RANK': 2, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
{'PID': 40, 'MASTER_ADDR': 'pytorch-mnist-master-0', 'MASTER_PORT': '23456', 'LOCAL_RANK': 1, 'RANK': 3, 'GROUP_RANK': 1, 'ROLE_RANK': 3, 'LOCAL_WORLD_SIZE': 2, 'WORLD_SIZE': 4, 'ROLE_WORLD_SIZE': 4}
Using cuda:0.
Using cuda:1.
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
Train Epoch: 1 [0/60000 (0%)]   Loss: 2.283599
...
Train Epoch: 10 [58880/60000 (98%)]     Loss: 0.051877Train Epoch: 10 [58880/60000 (98%)]       Loss: 0.051877

Train Epoch: 10 [59520/60000 (99%)]     Loss: 0.007343Train Epoch: 10 [59520/60000 (99%)]       Loss: 0.007343

Accuracy: 9919/10000 (99.19%)

Accuracy: 9919/10000 (99.19%)

Flags úteis para logs:

-f — transmite os logs em tempo real
-t N / --tail N — exibe as últimas N linhas
arena logs --help — mostra todas as opções

(Opcional) Etapa 6: Limpar recursos

Exclua o job de treinamento quando ele não for mais necessário. Essa ação remove os Pods e o deployment associado do TensorBoard, liberando recursos de GPU no cluster.

arena delete pytorch-mnist -n default

Saída esperada:

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

Próximos passos

Para saber mais sobre os comandos do Arena, execute arena --help ou visite a documentação do Arena.
Para executar um job de treinamento PyTorch standalone como comparação, consulte o tutorial de PyTorch standalone.
Para configurar redes de nível de produção e expor o TensorBoard, consulte Gerenciamento de Ingress.