tao-run-on-lepton

Par nvidia · skills

Plateforme de calcul GPU managée DGX Cloud Lepton avec interface run/status/cancel. À utiliser pour soumettre des jobs TAO.

npx skills add https://github.com/nvidia/skills --skill tao-run-on-lepton

Lepton

Plateforme GPU managée sur DGX Cloud. Les jobs sont soumis sous forme de workloads conteneurisés qui s'exécutent sur des groupes de nœuds GPU dédiés ou partagés. Lepton gère la planification, la récupération d'images, la collecte de logs et le cycle de vie des jobs.

Utilisez Lepton quand vous avez besoin de calcul GPU basé en cloud sans gérer directement l'infrastructure Kubernetes ou SLURM.

Vérification préalable

Lepton est API-first — pas d'équivalent docker-run. Cette skill nécessite le SDK TAO avec l'extra Lepton. nvidia-tao-sdk est sur PyPI public ; la version épinglée se trouve dans versions.yaml (wheels.tao_sdk_lepton), résolue via scripts/resolve_versions_key.py :

PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_lepton)
python -c "import tao_sdk" 2>/dev/null || {
  echo "MISSING: nvidia-tao-sdk not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}
python -c "import leptonai" 2>/dev/null || {
  echo "MISSING: lepton extra not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

Si absent, l'agent invite l'utilisateur à autoriser l'installation via Bash, puis réexécute la vérification préalable avant de continuer.

Identifiants

  • LEPTON_WORKSPACE_ID (requis) : Détermine le cluster et le compte de facturation sur lequel le job s'exécute.
  • LEPTON_AUTH_TOKEN (requis) : Token API pour l'authentification auprès du plan de contrôle Lepton.
  • NGC_KEY (optionnel) : Utilisé pour créer des secrets de récupération d'image pour puller les images de conteneur TAO depuis nvcr.io.
  • ACCESS_KEY / SECRET_KEY (optionnel) : Clés de stockage compatibles S3 pour les URIs de datasets et de checkpoints.
  • S3_ENDPOINT_URL (optionnel) : Point d'accès S3 personnalisé (ex. MinIO ou S3 non-AWS).
  • S3_BUCKET_NAME (optionnel) : Bucket pour les artefacts de sortie du job.
  • CLOUD_REGION (optionnel) : Région de stockage (ex. us-east-1).

Vérification préalable au lancement

Avant de générer des scripts ou de soumettre des jobs :

  1. Vérifiez que LEPTON_WORKSPACE_ID et LEPTON_AUTH_TOKEN sont définis.
  2. Vérifiez que l'API du workspace est accessible avec l'utilitaire packagé : scripts/check_tao_launch_preflight.py --platform lepton ....
  3. Pour les datasets/résultats s3://, vérifiez que ACCESS_KEY et SECRET_KEY sont définis et que les chemins exacts sont lisibles avec aws s3 ls.
  4. Pour les chemins montés NFS/Lustre, exigez une preuve de Lepton concernant les permissions de volume/stockage que le chemin sera monté dans le job. Ne traitez pas un test de système de fichiers local test -e sur l'hôte agent comme preuve pour les jobs Lepton.
  5. Vérifiez les identifiants spécifiques au modèle tels que HF_TOKEN avant le lancement.

Détails du backend

LeptonSDK.create_job accepte ces kwargs spécifiques à Lepton (en plus des arguments agnostiques de plateforme — image, command, gpu_count, env_vars, inputs, outputs, hooks) :

  • resource_shape : ID explicite de la forme de ressource GPU (ex. "gpu.8xh100-sxm"). Quand défini, ignore la résolution automatique depuis gpu_count. Le format est opaque (ce que l'API Lepton retourne comme instance metadata.id) — découvrez les IDs valides via sdk.list_resource_shapes().
  • dedicated_node_group : ID du groupe de nœuds pour une allocation GPU garantie (sans préemption). Omettez pour les ressources partagées.
  • num_nodes : nombre de nœuds pour l'entraînement distribué. Par défaut 1. Quand > 1, active la communication intra-job et l'initialisation distribuée PyTorch (voir Entraînement multi-nœuds).
  • mounts : objets Mount pré-construits pour NFS / Lustre. Auto-détectés du groupe de nœuds quand non définis.

Découverte des formes / volumes du workspace

shapes = sdk.list_resource_shapes()
# {<platform_id>: {"cluster": ..., "gpu_type": "gpu.8xh100-sxm",
#                   "gpu_count": 8, "instance_type": ..., ...}, ...}

volumes = sdk.get_volumes(node_group_id="my-h100-pool")
# [{"name": "lustre", "from_path": "/lustre", "type": "Lustre"}, ...]

prefixes = sdk.get_storage_permissions("lustre", "my-h100-pool")
# ["/lustre/fsw/portfolios/edgeai/...", ...]

Entraînement multi-nœuds (distribué)

Passez num_nodes > 1 à create_job pour l'entraînement distribué multi-nœuds. Le gestionnaire Lepton (tao_sdk/platforms/lepton/handler.py) configure le LeptonJob sous-jacent en définissant intra_job_communication=True (ouvre la communication pod-à-pod), parallelism=num_nodes et completions=num_nodes (Lepton planifie N réplicas), et exporte WORLD_SIZE=num_nodes comme variable d'environnement du conteneur.

Les variables d'environnement natives par-replica de Lepton utilisent des noms spécifiques à Lepton (LEPTON_JOB_WORKER_INDEX, LEPTON_JOB_TOTAL_WORKERS, LEPTON_JOB_WORKER_PREFIX, LEPTON_SUBDOMAIN), donc le gestionnaire ajoute un bootstrap qui source le script de traduction officiel de Lepton :

wget -O init.sh https://raw.githubusercontent.com/leptonai/scripts/main/lepton_env_to_pytorch.sh
chmod +x init.sh
source init.sh
# user command runs here

Après source, les variables d'environnement suivantes sont définies :

Variable d'env Source Valeur
MASTER_ADDR script ${LEPTON_JOB_WORKER_PREFIX}-0.${LEPTON_SUBDOMAIN}
MASTER_PORT script 29400
NNODES script ${LEPTON_JOB_TOTAL_WORKERS}
NODE_RANK script ${LEPTON_JOB_WORKER_INDEX}
WORKER_ADDRS script liste séparée par des virgules des noms d'hôte des workers non-master
WORLD_SIZE gestionnaire TAO SDK num_nodes (convention du conteneur TAO — même valeur que NNODES)
NUM_GPU_PER_NODE gestionnaire TAO SDK gpu_count (lu par le point d'entrée du conteneur TAO)
job = sdk.create_job(
    image='nvcr.io/nvidia/tao/tao-toolkit:6.26.3-pyt',
    command='dino train -e /tmp/spec.yaml',  # TAO entrypoint reads WORLD_SIZE + NUM_GPU_PER_NODE
    gpu_count=8,                          # GPUs per node
    num_nodes=4,                          # 4 × 8 = 32 GPUs total
    dedicated_node_group='my-h100-pool',
    inputs={'/data/train.json': 's3://bucket/coco/train.json'},
    outputs=['/results/'],
)

Pour les commandes torchrun brutes (conteneurs non-TAO) :

command='torchrun --nnodes=$NNODES --nproc-per-node=8 --node-rank=$NODE_RANK '
        '--master-addr=$MASTER_ADDR --master-port=$MASTER_PORT train.py'

Deux façons d'exécuter des jobs distribués sur Lepton

Chemin Quand l'utiliser
TAO SDK create_job(num_nodes=N) (cette skill) Soumission programmatique depuis le code de l'agent ; vous voulez le wrapping S3 du SDK, la surveillance, l'analyse des défaillances et le JobStore.
Type de job Lepton "Torchrun" (UI Lepton / lep CLI) Soumission artisanale via la console Lepton. L'UI de Lepton a un mode "Torchrun" de première classe qui configure le rendezvous pour vous — pas besoin de script de bootstrap. Voir l'exemple officiel.

Lectures de référence

Notes

  • Préférez dedicated_node_group pour multi-nœuds pour garder les réplicas sur la même interconnexion à faible latence (NVLink / InfiniBand).
  • Si un replica est préempté sur un groupe de nœuds partagés, le job entier échoue — Lepton ne redémarre pas élastiquement en v1. Utilisez un groupe de nœuds dédié pour les exécutions longues.
  • Pour les datasets sauvegardés par Lustre, le même mount est exposé à chaque replica — pas de wrapping d'I/O par-replica nécessaire.

Stockage cloud

Même si la plateforme est Lepton, la couche de stockage est compatible S3. Utilisez toujours aws comme clé cloud_metadata et s3:// comme protocole URI pour les datasets et results_dir.

  • Correct : s3://bucket-name/path
  • Incorrect : lepton://bucket-name/path

Le get_cloud_storage_class_object() du conteneur analyse le protocole URI pour rechercher les identifiants dans CLOUD_METADATA[protocol][bucket].

Stockage partagé (NFS/Lustre)

Les groupes de nœuds peuvent avoir des volumes NFS ou Lustre attachés. Le SDK les détecte automatiquement et les monte dans les conteneurs pour le partage de données persistant entre jobs.

Fonctions SDK

  • sdk.get_volumes(node_group_id=None) — retourne les volumes disponibles (name, from_path, type) depuis la spec du groupe de nœuds
  • sdk.get_storage_permissions(volume_name, node_group_id) — retourne les préfixes de chemin autorisés pour un volume

LeptonSDK.create_job() appelle ces fonctions automatiquement pour détecter les mounts et construire les objets Mount appropriés pour les specs de job.

Comment le script runner utilise les mounts

Quand un mount Lustre est disponible :

  • Inputs : les chemins S3 sont mappés à Lustre (s3://bucket/path/mnt/lustre/bucket/path). Si le fichier existe sur Lustre, il est utilisé directement (zéro téléchargement). S'il manque, il est téléchargé de S3 à Lustre et persiste pour les jobs futurs.
  • Outputs : les résultats s'écrivent d'abord à Lustre (rapide, persistent), puis se téléchargent vers S3 (durable). Les jobs en aval (ex. analyse des écarts) peuvent lire les résultats directement depuis Lustre sans aller-retour S3.

Ordre de préférence des volumes

lustre > filestore > premier disponible

Invalidation du cache Lustre

Lustre cache les fichiers de manière persistante entre les jobs. Il n'y a pas d'invalidation intégrée. Si les données en amont changent mais le chemin S3 reste le même, Lustre serve la version cachée obsolète. Pour forcer une absence de cache :

  • Renommez le fichier sur S3 (ex. prompt_v2.txt au lieu de remplacer prompt.txt)
  • Utilisez une nouvelle storage_root entre les itérations pour éviter la obsolescence entre itérations
  • Utilisez un nouveau chemin pour tout artefact régénéré

Surveillance

État du job

Utilisez sdk.get_job_status(job_id) pour l'état de haut niveau (Pending, Running, Complete, Error).

État du replica

Utilisez sdk.get_job_replicas(job_id) lors du démarrage pour les infos détaillées au niveau replica. Chaque replica est un dict :

replicas = sdk.get_job_replicas(job_id)
for r in replicas:
    node = r["status"]["node"]["name"]           # ex. "node-ip-10-50-111-24"
    node_group = r["status"]["node"]["node_group_id"]
    cpu = r["status"]["cpu"]                      # ex. 2
    memory_mb = r["status"]["memory_in_mb"]       # ex. 8192
    readiness = r["status"].get("readiness_issue")
    if readiness:
        reason = readiness["reason"]   # "InProgress", "Failed", "ConfigError"
        message = readiness["message"] # "Pulling image", "Mount point not found", etc.

Patterns clés de readiness_issue :

  • reason="InProgress", message="Pulling image" — récupération d'image en cours (normal pour les grandes images)
  • reason="Failed" — échec de la récupération d'image (vérifiez NGC_KEY)
  • reason="ConfigError" — problème de nœud (échec du mount, erreur GPU)
  • Pas de readiness_issue — le replica s'exécute

L'état du replica est particulièrement utile quand un job est bloqué en Pending — il révèle si le problème vient de la récupération d'image, de la planification des ressources ou de la santé du nœud.

Logs du job

Utilisez sdk.get_job_logs(job_id, tail=N) pour les N lignes de log les plus récentes. Les logs sont récupérés du service de collecte de logs de Lepton.

Jobs parallèles

Pour les étapes de workflow qui s'exécutent en parallèle (ex. génération vidéo x8) :

  1. Lancement : Appelez execute_step(plan, step_id, extra_args={"split_id": i}) pour chaque split. Chaque appel retourne immédiatement avec un job_id.
  2. Surveillance : Pollisez tous les jobs : sdk.get_job_status(job_id) pour chacun. Utilisez get_job_replicas(job_id) pour la diagnostique au démarrage.
  3. Achèvement : Tous les jobs sont terminés quand chaque statut est Complete ou Error.
  4. Défaillance partielle : Réessayez seulement les splits échoués — les splits réussis n'ont pas besoin de réexécution. Passez le même split_id à execute_step.

Analyse des défaillances

Quand un job échoue, utilisez sdk.get_failure_analysis(job_id) pour la détection automatique de la cause racine :

analysis = sdk.get_failure_analysis(job_id)
if analysis:
    print(analysis["err_class"])    # ex. "ERR_PROGRAM"
    print(analysis["suggestion"])   # Fix lisible
    for event in analysis.get("job_failure_by_node_event", []):
        print(event["node_event_name"], event["message"])
        # ex. "OOM", "OOM encountered, victim process: cosmos-rl-evalu, pid: 3368483"

Retourne :

  • err_class : Classification d'erreur (ERR_PROGRAM, ERR_INFRA, etc.)
  • suggestion : Ce qui s'est probablement mal passé et comment corriger
  • job_failure_by_node_event : Événements au niveau nœud (kills OOM, erreurs GPU, défaillances de mount)
  • log_streams : Snippets de log pertinents avec contexte d'erreur

Appelez toujours cela sur les jobs échoués avant de réessayer — cela distingue les erreurs utilisateur (mauvaise config, OOM) des problèmes infrastructure (défaillance de nœud, éviction).

Modes de défaillance

OOM killed : Le conteneur a dépassé la mémoire GPU ou système. Détection : get_failure_analysis() retourne node_event_name: "OOM". Causes communes : evaluation.batch_size trop élevé, max_length trop grand pour le KV cache disponible. Récupération : réduisez batch_size, ajoutez des GPUs avec parallélisme tensoriel, ou réduisez max_length.

Échec de la récupération d'image : L'image du conteneur TAO ne peut pas être pullée depuis nvcr.io. Généralement causé par un secret de récupération d'image manquant ou expiré. Le SDK provisionne automatiquement le secret depuis NGC_KEY, mais si NGC_KEY est invalide, le job échouera. Détection : vérifiez get_job_replicas()readiness_issue.reason affichera InProgress avec message = "Pulling image" pendant de longues périodes, ou Failed si le pull échoue. Récupération : vérifiez que NGC_KEY est valide.

Ressource indisponible : La forme de GPU demandée n'est pas disponible. Le job entre dans l'état Queueing indéfiniment. Détection : Pending > 15 minutes, les replicas n'affichent pas d'assignation de nœud. Récupération : essayez une autre resource_shape ou dedicated_node_group, ou attendez les ressources.

Défaillance d'authentification : LEPTON_AUTH_TOKEN invalide ou expiré. Tous les appels API échouent avec 401/403. Détection : la création du job lève une exception immédiatement. Récupération : actualisez le token et réinitialisez le SDK.

Nœud malsain : Le nœud assigné a des problèmes infrastructure (défaillances de mount, erreurs GPU, problèmes réseau). Détection : vérifiez get_job_replicas()readiness_issue.reason = "ConfigError" avec des messages comme "Mount point not found". Le job reste Pending indéfiniment sur le mauvais nœud. Récupération : annulez le job et resoumettez — Lepton planifiera sur un nœud différent. Si le problème se reproduit, essayez un autre dedicated_node_group ou resource_shape.

Éviction du job : Sur les groupes de nœuds partagés, Lepton peut évincer les jobs sous pression de ressources. Détection : le job effectue une transition inattendue de Running à Error. Récupération : réessayez, ou utilisez un dedicated_node_group.

Skills similaires