tao-run-on-slurm

Par nvidia · skills

Exécution distante sur cluster GPU SLURM via SSH avec sbatch/srun, conteneurs Pyxis/Enroot et stockage Lustre

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

SLURM

Plateforme de calcul GPU distant pour les clusters gérés par SLURM. Les jobs sont soumis depuis le service TAO ou l'hôte SDK à un nœud de connexion via SSH, mis en place sur un système de fichiers partagé, soumis avec sbatch et exécutés avec le support des conteneurs srun.

Utilisez SLURM quand l'utilisateur a accès à un cluster GPU géré, un stockage Lustre partagé et une allocation GPU gérée par le scheduler. N'utilisez pas SLURM pour les fichiers locaux qui n'existent que sur la machine agent ; les données et les résultats doivent être accessibles depuis le cluster.

Preflight

# 1. SSH to the login node works without a password prompt
SLURM_HOST="${SLURM_HOSTNAME%%,*}"
[ -n "$SLURM_USER" ] && [ -n "$SLURM_HOST" ] || {
  echo "MISSING: set SLURM_USER and SLURM_HOSTNAME (comma-separated for failover) in your env (~/.config/tao/.env)."
  exit 1
}
ssh -o BatchMode=yes -o ConnectTimeout=10 "${SLURM_USER}@${SLURM_HOST}" "true" 2>/dev/null || {
  echo "MISSING: passwordless SSH to ${SLURM_USER}@${SLURM_HOST} not working. See references/ssh-setup.md."
  exit 1
}

# 2. Optional: TAO SDK wrapper for Job handles + S3 wrapping.
# nvidia-tao-sdk is on public PyPI; pin lives in versions.yaml (wheels.tao_sdk_slurm).
PIN=$("${TAO_SKILL_BANK_PATH:?}/scripts/resolve_versions_key.py" wheels.tao_sdk_slurm)
python -c "import tao_sdk" 2>/dev/null || {
  echo "MISSING: nvidia-tao-sdk not installed. Run:"
  echo "  pip install \"$PIN\""
  exit 1
}

Si une vérification échoue, l'agent invite l'utilisateur à autoriser l'installation ou la correction via Bash.

Une troisième étape preflight s'applique uniquement pour les images privées nvcr.io : Pyxis sur les nœuds de calcul a besoin des identifiants enroot persistants dans ~/.config/enroot/.credentials sur le cluster (il NE lit PAS NGC_KEY depuis l'env du job). Sans eux, les pulls authentifiés échouent avec « Could not process JSON input » au démarrage du job. Cela s'exécute une fois par (cluster, utilisateur). Voir references/ssh-setup.md pour la vérification complète et le motif printf | ssh qui garde NGC_KEY hors de l'historique, des fichiers et du chat. Ignorez-le pour les images publiques.

Prerequisites

Avant tout envoi de job, l'hôte exécutant le service TAO ou le SDK doit se connecter à au moins un hôte de SLURM_HOSTNAME via SSH sans prompt interactif de mot de passe. Le gestionnaire exécute sbatch, squeue, sacct, scancel et les lectures de logs de façon non-interactive, donc les prompts de mot de passe ou 2FA feront échouer le job au moment de l'envoi ou du suivi du statut.

Configurez cela une fois par tuple (hôte, nœud de connexion, utilisateur) : créez une paire de clés SSH, installez la clé publique sur chaque hôte de connexion, faites confiance à la clé d'hôte, verrouillez les permissions de la clé privée à chmod 600, et vérifiez avec ssh -o BatchMode=yes .... Voir references/ssh-setup.md pour les étapes complètes (incluant l'alias ~/.ssh/config, la note de montage de clé de conteneur et le fallback 2FA / SSH_AUTH_SOCK). Le même fichier contient le prompt de correction des défaillances SSH à afficher à l'utilisateur quand le SSH sans mot de passe échoue.

Credentials

  • SLURM_USER (requis) : Nom d'utilisateur SSH pour le nœud de connexion. Dans les métadonnées d'espace de travail microservices, c'est cloud_specific_details.slurm_user.
  • SLURM_HOSTNAME (requis) : Noms d'hôtes de connexion séparés par des virgules pour le failover. Le schéma microservices stocke cela comme le champ liste cloud_specific_details.slurm_hostname.
  • SLURM_PARTITION (requis) : Liste de partitions pour la soumission de jobs GPU. Demandez cela dans la liste d'intake SLURM obligatoire. La valeur par défaut fournie est polar,polar3,polar4,grizzly, qui sont traitées comme des queues de 4 heures.
  • SSH_KEY_PATH (préféré et attendu avant le lancement) : Chemin de la clé privée pour l'auth par clé publique non-interactive au nœud de connexion. Si le SSH sans mot de passe échoue, demandez à l'utilisateur SSH_KEY_PATH=/path/to/private_key et montrez les étapes de configuration dans references/ssh-setup.md ; ne cachez pas cela derrière plusieurs choix alternatifs.
  • SSH_AUTH_SOCK (fallback avancé) : Socket de l'agent SSH avec une clé acceptée déjà chargée. Préférez SSH_KEY_PATH dans les prompts de correction orientés utilisateur.
  • SLURM_BASE_RESULTS_DIR (optionnel) : Chemin de base du système de fichiers partagé. La convention par défaut de tao-core est /lustre/fsw/portfolios/edgeai/<your-dir>, où <your-dir> est votre répertoire par utilisateur sur le cluster.
  • SLURM_ACCOUNT (généralement requis par la politique du site) : Compte facturé par #SBATCH --account.

Ne demandez pas SLURM_ACCOUNT ou SLURM_BASE_RESULTS_DIR lors de l'intake initial sauf si l'utilisateur dit que son site exige un compte, veut une racine de résultats personnalisée, ou le workflow ne peut pas avancer sans contourner les défauts.

Backend Details

Utilisez backend_details.backend_type = "slurm" quand vous routez un job vers cette plateforme. Les détails backend supportés du schéma microservices :

{
  "backend_type": "slurm",
  "partition": "polar,polar3,polar4,grizzly",
  "cluster_name": "optional-name"
}

Les métadonnées d'exécution sont stockées sous backend_details.slurm_metadata, notamment slurm_job_id et job_dir. N'inventez pas ces valeurs. Elles sont écrites après que sbatch retourne un id de job du scheduler.

Storage

Les jobs SLURM s'exécutent sur le cluster, donc les chemins locaux de l'hôte API ne sont pas des chemins de dataset valides. Préférez les URIs du système de fichiers partagé :

  • Utilisez lustre:///absolute/path pour les datasets fournis par l'utilisateur sur Lustre.
  • Les chemins slurm:// peuvent apparaître dans les métadonnées microservices et sont convertis en chemins Lustre réels avant que le conteneur démarre.
  • Évitez les URIs de dataset bare /local/path et file:// pour SLURM. La validation dans tao-core rejette les chemins locaux et fichier pour les backends distants.

Acceptez soit les racines de dataset soit les chemins directs spec-key :

  • Mode racine : /lustre/.../<model>/train, que les skills de modèle mappent aux fichiers requis comme <root>/annotations.json et <root> comme chemin média.
  • Mode spec direct : les champs exacts comme custom.train_dataset.annotation_path=/lustre/.../train.json et custom.train_dataset.media_path=/lustre/.../videos.tar.gz.

Après la réussite du SSH sans mot de passe et avant de générer les scripts, validez chaque fichier/chemin de dataset requis depuis le nœud de connexion :

ssh -o BatchMode=yes <SLURM_USER>@<working-login-host> \
  'test -e /lustre/.../annotations.json && test -e /lustre/.../media_or_archive'

Si le test distant échoue, arrêtez et demandez des chemins corrigés ou que les données soient mises en place sur le stockage cluster partagé. Ne créez pas de scripts runner qui échoueront à l'intérieur du premier job d'entraînement.

Les résultats par défaut sont :

/lustre/fsw/portfolios/edgeai/<your-dir>/results/<job_id>

<your-dir> est votre répertoire par utilisateur sur le cluster.

Le runner définit TAO_API_RESULTS_DIR au répertoire parent des résultats car le code du conteneur ajoute l'id du job lors de l'écriture du statut et des artefacts.

Utilisez Lustre, pas S3, pour les entrées de jobs SLURM. Le scheduler SLURM impose un timeout GPU-inactif — un long téléchargement s3:// au début du script peut brûler l'allocation avant que l'entraînement ne commence, et le scheduler peut tuer le job. Mettez d'abord les données d'entraînement en place sur Lustre ; la pré-récupération S3 / HF / NGC est fine seulement pour les petites entrées auxiliaires (checkpoints, configs). Voir references/sdk-usage.md pour la rationale complète.

Container Execution

tao-core utilise le gestionnaire SLURM pour exécuter les conteneurs TAO via Pyxis/Enroot :

  1. Mettez en place les fichiers JSON compacts pour les specs, l'environnement et les métadonnées cloud sous <job_dir>/specs, <job_dir>/env et <job_dir>/meta.
  2. Convertissez optionnellement l'image Docker en image SQSH mise en cache avec srun -n1 -p <conversion_partition> enroot import.
  3. Écrivez un script sbatch sous <job_dir>/sbatch/job_<job_id>.sbatch.
  4. Soumettez sbatch --export=ALL <script>.
  5. Exécutez le conteneur avec srun --container-image=<image> --container-mounts=/lustre.

Formats d'image acceptés par le gestionnaire :

  • /path/to/image.sqsh
  • registry#image:tag
  • docker://registry#image:tag
  • registry/image:tag ordinaire, qui est converti en forme Pyxis si nécessaire

La conversion SQSH est mise en cache par nom d'image. Pour les images :latest, le SQSH en cache est utilisé sauf si force_reconvert_latest est activé.

Resource Mapping

Défauts de tao-core :

  • num_nodes: 1
  • num_gpus: 4
  • max_num_gpus_per_node: 8
  • cpus_per_task: 16
  • time_hours: 4
  • timeout_hours: 3.8
  • max_time_hours: 4
  • container_mounts: /lustre
  • use_requeue: true
  • use_sqsh: true

Quand vous générez des launchers ou des scripts wrapper pour SLURM, définissez les défauts de wall-time explicitement depuis les défauts de ressources de plateforme fournis :

export SLURM_TIME_HOURS="${SLURM_TIME_HOURS:-4}"
export SLURM_TIMEOUT_HOURS="${SLURM_TIMEOUT_HOURS:-3.8}"

Ne défaut pas à 12 heures sur SLURM. Si l'utilisateur fournit un SLURM_TIME_HOURS plus long, vérifiez que la partition sélectionnée le supporte avant de soumettre. Pour la liste de partition par défaut fournie polar,polar3,polar4,grizzly, rejetez les requêtes au-dessus de 4 heures et demandez une partition différente seulement si l'utilisateur veut vraiment un wall time plus long.

Quand num_gpus est supérieur ou égal à max_num_gpus_per_node, le gestionnaire traite la requête comme exclusive par nœud et calcule des nœuds additionnels à partir du nombre total de GPU si nécessaire.

Pour les jobs multi-nœuds (num_nodes > 1), le script sbatch exporte WORLD_SIZE, MASTER_ADDR, MASTER_PORT, NODE_RANK et NUM_GPU_PER_NODE, et Cosmos-RL a une gestion de rôles multi-nœuds spéciale pour les workers controller, policy et rollout. Voir references/multi-node.md pour les directives sbatch complètes, la table de contrat env-var de rendezvous et les exigences cluster.

Monitoring

  • Le statut du scheduler vient de l'id du job SLURM stocké via squeue ou sacct.
  • Le statut terminal TAO vient de status.json dans le dossier des résultats partagés.
  • Si l'utilisateur a activé la surveillance du chat, continuez le polling à l'intervalle demandé tant que le job est PENDING, RUNNING ou autrement non-terminal. Ne vous arrêtez pas après un temps écoulé fixe comme 30 minutes ; les longues attentes de queue sont normales sur les partitions GPU partagées.
  • N'envoyez pas une réponse finale pour un job SLURM non-terminal quand la surveillance du chat est activée. Une réponse finale est une action de détachement ; utilisez-la seulement si l'utilisateur a demandé de se détacher/arrêter ou si le job a atteint l'état terminal.
  • Les logs sont lus via SSH depuis :
<job_dir>/slurm-logs/<slurm_job_name>-<slurm_job_id>/main.out
<job_dir>/slurm-logs/<slurm_job_name>-<slurm_job_id>/main.err

Mapping de statut :

  • PENDING -> Pending
  • RUNNING ou COMPLETING -> Running
  • COMPLETED -> vérifiez status.json
  • FAILED, BOOT_FAIL, DEADLINE, OUT_OF_MEMORY, NODE_FAIL -> revenez en arrière si les logs correspondent à des motifs d'infrastructure rejouables, sinon Error
  • CANCELLED, PREEMPTED, REVOKED -> Canceled
  • TIMEOUT -> Error
  • SUSPENDED, STOPPED -> Paused

Cancellation

Annulez en recherchant backend_details.slurm_metadata.slurm_job_id et en exécutant scancel <slurm_job_id> via SSH. Traitez les jobs SLURM manquants ou déjà terminés comme une annulation réussie.

Multi-node training (distributed)

SLURM est la plateforme de choix pour les grands runs multi-nœuds — passez num_nodes > 1 et le SDK gère les directives sbatch et les vars env PyTorch-distributed automatiquement. Voir references/multi-node.md pour un exemple create_job opérationnel, les directives sbatch générées, la table de contrat env-var de rendezvous (WORLD_SIZE, NUM_GPU_PER_NODE, NODE_RANK, MASTER_ADDR, MASTER_PORT), la note de rôle Cosmos-RL, les exigences cluster (Pyxis/Enroot, InfiniBand/NVLink, Lustre) et les liens de référence en amont.

Running via the TAO SDK

L'installation du SDK est couverte dans Preflight — pip install 'nvidia-tao-sdk[slurm]'. Utilisez-la quand vous voulez des handles Job, que le plumbing sbatch/squeue/sacct soit géré pour vous, la durabilité du dossier de run via ActionWorkflow, ou la commodité du cloud-storage I/O (s3://, hf_model://, ngc://). Sans le SDK, pilotez sbatch et srun vous-même.

La retry automatique est entièrement automatique : un moniteur de fond poll squeue/sacct et re-sbatch le script mis en place sur les défaillances de type infrastructure jusqu'à MAX_JOB_RETRIES = 10, tandis que les défaillances d'entraînement simples remontent immédiatement. De plus, #SBATCH --requeue est défini par défaut (SLURM_USE_REQUEUE, défaut à true). Voir references/sdk-usage.md pour l'exemple de code SlurmSDK / build_entrypoint, la règle Lustre-pas-S3, la classification des défaillances rejouables et le comportement complet de retry automatique et requeue.

Failure Modes

Défaillances courantes : défaillance d'auth SSH, chemin de dataset local rejeté, timeout de conversion SQSH, Pyxis/Enroot indisponible et défaillances de mauvais nœud / GPU transitoires (que le gestionnaire réessaie jusqu'à la limite configurée). Voir references/troubleshooting.md pour le diagnostic et la correction de chacune.

Skills similaires

llm-obs-eval-bootstrap
datadog-labs / agent-skills

Générer une suite d'évaluateurs prêts à l'emploi à partir de traces LLM en production.

127
tao-run-on-lepton
nvidia / skills

Soumettre et gérer des workloads GPU conteneurisés sur la plateforme cloud Lepton.

1 116
tao-launch-workflow
nvidia / skills

Valider et lancer des workflows TAO après vérification complète des prérequis plateforme.

1 116
geofeed-tuner
github / awesome-copilot

Créer et optimiser des feeds de géolocalisation IP au format CSV selon RFC 8805.

34 639
llm-obs-session-classify
datadog-labs / agent-skills

Classifier la satisfaction des sessions et traces LLM Observability Datadog en verdicts automatisés.

127