Use Arena to submit distributed PyTorch training jobs - Container Service for Kubernetes

Prérequis

Avant de commencer, assurez-vous de disposer des éléments suivants :

Un cluster ACK doté de nœuds accélérés par GPU. Pour plus de détails, consultez Créer un cluster ACK contenant des nœuds accélérés par GPU.
L'accès à Internet activé pour les nœuds du cluster. Pour plus de détails, consultez Activer l'accès à Internet pour un cluster ACK existant.
Vérifiez que le module complémentaire Arena est installé. Pour plus de détails, consultez Configurer le client Arena.
Une revendication de volume persistant (PVC) nommée training-data contenant le jeu de données MNIST stocké dans /pytorch_data. Pour plus de détails, consultez Configurer un volume NAS partagé.

Contexte

Ce tutoriel entraîne un modèle PyTorch sur le jeu de données MNIST en utilisant deux nœuds, chacun exploitant deux GPU, soit quatre GPU au total. Le code d'entraînement s'exécute avec torchrun, le lanceur intégré de PyTorch pour les tâches distribuées.

Le flux de travail récupère le code d'entraînement depuis un dépôt Git et lit les données depuis un volume partagé sauvegardé par NAS (PV/PVC). Pour référence, consultez main.py.

Étape 1 : Vérifier les ressources GPU disponibles

Exécutez la commande suivante pour afficher les GPU disponibles dans votre cluster :

arena top node

Sortie attendue :

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

Le cluster dispose de deux nœuds accélérés par GPU avec deux GPU inactifs chacun, soit quatre GPU disponibles pour l'entraînement.

Étape 2 : Soumettre une tâche d'entraînement PyTorch distribuée

Exécutez la commande suivante pour soumettre la tâche. Elle crée deux Pods de travail, chacun utilisant deux GPU.

Trois paramètres contrôlent la topologie distribuée :

--workers=2 — nombre total de Pods, y compris le Pod maître
--gpus=2 — GPU alloués à chaque Pod
--nproc-per-node=2 — processus torchrun lancés par Pod ; chaque processus utilise un 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"

Sortie attendue :

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

Fonctionnement : nœuds, rôles et variables d'environnement

Les tâches PyTorch distribuées dans Arena utilisent deux paramètres dont les tâches autonomes n'ont pas besoin :

--workers — nombre total de Pods. Un Pod assume le rôle master ; les autres sont des Pods worker.
--nproc-per-node — nombre de processus torchrun lancés sur chaque Pod. Chaque processus correspond à un GPU.

Les noms des Pods suivent le schéma <nom_tâche>-<nom_rôle>-<index>. Par exemple, --workers=3 et --nproc-per-node=2 sur une tâche nommée pytorch-mnist créent trois Pods et lancent deux processus pour chacun :

Variable d'environnement	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

Arena injecte automatiquement ces variables dans chaque Pod. Dans votre code d'entraînement, utilisez RANK pour identifier quel processus doit enregistrer les points de contrôle ou consigner les métriques. Par exemple, seul le processus où RANK=0 (le maître) écrit les résultats, ce qui évite les doublons. Utilisez WORLD_SIZE pour déterminer le nombre total de processus parallèles lors du calcul des gradients distribués.

Référence des paramètres

Paramètre	Requis	Description	Par défaut
`--name`	Oui	Nom de la tâche, globalement unique	—
`--namespace`	Non	Namespace Kubernetes	`default`
`--workers`	Non	Nombre total de Pods de travail, y compris le maître. Par exemple, `--workers=3` crée un Pod maître et deux Pods de travail	`0`
`--gpus`	Non	GPU alloués à chaque Pod de travail	`0`
`--working-dir`	Non	Répertoire dans lequel la commande d'entraînement s'exécute	`/root`
`--image`	Oui	Image de conteneur utilisée pour l'exécution	—
`--sync-mode`	Non	Mode de synchronisation du code source : `git` ou `rsync`	—
`--sync-source`	Non	URL du dépôt pour la synchronisation du code source. Le code est téléchargé dans `code/` sous `--working-dir` (par exemple, `/root/code/arena`)	—
`--data`	Non	Monte un PVC dans le Pod sous la forme `<nom-pvc>:<chemin-montage>`. Exécutez `arena data list` pour afficher les PVC disponibles	—
`--tensorboard`	Non	Active la visualisation TensorBoard. Requiert `--logdir`	—
`--logdir`	Non	Chemin où TensorBoard lit les fichiers d'événements. À utiliser conjointement avec `--tensorboard`	`/training_logs`

Utiliser un dépôt Git privé

Si votre code se trouve dans un dépôt privé, transmettez les identifiants sous forme de variables d'environnement :

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 utilise git-sync pour extraire le code source, donc toutes les variables d'environnement définies dans le projet git-sync sont disponibles.

Impossible d'extraire le code en raison de problèmes réseau ?

Si le dépôt GitHub est inaccessible, téléchargez manuellement le code dans le volume NAS. Placez-le dans code/github.com/kubeflow/arena à l'intérieur du NAS — cela correspond à /mnt/code/github.com/kubeflow/arena après le montage du volume. Ensuite, soumettez la tâche sans --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"

Étape 3 : Surveiller la tâche d'entraînement

Après la soumission, la tâche passe par plusieurs étapes avant que les journaux n'apparaissent :

Pending — Kubernetes planifie les Pods sur les nœuds GPU.
Preparing — Les Pods extraient l'image du conteneur et synchronisent le dépôt Git. Les journaux ne sont pas encore disponibles.
Running — torchrun démarre et l'entraînement commence. Les journaux sont disponibles à cette étape.

Utilisez les commandes suivantes pour suivre la progression.

Répertorier toutes les tâches :

arena list -n default

Sortie attendue :

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

Vérifier l'utilisation des GPU par tâche :

arena top job -n default

Sortie attendue :

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

Vérifier l'allocation des GPU au niveau du cluster :

arena top node

Sortie attendue :

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

Les quatre GPU sont désormais alloués.

Obtenir des informations détaillées sur la tâche :

arena get pytorch-mnist -n default

Sortie attendue :

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

La tâche dispose d'un Pod maître et d'un Pod de travail, chacun utilisant deux GPU. L'URL TensorBoard n'est affichée que si --tensorboard est activé.

Étape 4 : Afficher les résultats de l'entraînement dans TensorBoard

Important

kubectl port-forward est destiné uniquement au développement et au débogage. Cette méthode n'est ni fiable, ni sécurisée, ni évolutive pour la production. Pour la mise en réseau en production dans les clusters ACK, consultez Gestion Ingress.

Redirigez le port de TensorBoard vers votre machine locale :

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

Ouvrez http://127.0.0.1:9090 dans un navigateur.

L'exemple de code d'entraînement écrit les événements toutes les 10 époques. Si vous modifiez --epochs , définissez-le sur un multiple de 10 ; sinon, aucun événement n'est écrit et TensorBoard n'affiche aucune donnée.

Étape 5 : Afficher les journaux d'entraînement

Afficher les journaux du Pod maître :

arena logs -n default pytorch-mnist

Sortie attendue :

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

Afficher les journaux d'un Pod de travail spécifique :

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

Sortie attendue :

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

Options utiles pour les journaux :

-f — diffuser les journaux en temps réel
-t N / --tail N — afficher les N dernières lignes
arena logs --help — voir toutes les options

(Facultatif) Étape 6 : Nettoyer

Supprimez la tâche d'entraînement lorsque vous n'en avez plus besoin. Cela supprime les Pods et le déploiement TensorBoard associé, libérant ainsi les ressources GPU de votre cluster.

arena delete pytorch-mnist -n default

Sortie attendue :

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

Étapes suivantes

Pour en savoir plus sur les commandes Arena, exécutez arena --help ou consultez la documentation Arena.
Pour exécuter une tâche d'entraînement PyTorch autonome à titre de comparaison, consultez le tutoriel sur PyTorch autonome.
Pour une mise en réseau de niveau production afin d'exposer TensorBoard, consultez Gestion Ingress.