Compiler TensorRT-LLM sur un cluster SLURM

Soumettez, surveillez et vérifiez une tâche de compilation TensorRT-LLM sur un cluster SLURM à l'aide de conteneurs enroot.

Quand l'utiliser

Trouver l'image Docker

La balise officielle de l'image Docker pour une version donnée de TensorRT-LLM est enregistrée dans le dépôt lui-même :

<repo_dir>/jenkins/current_image_tags.properties

Lisez ce fichier pour trouver l'URL de l'image actuelle (par exemple, urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:pytorch-25.12-py3-aarch64-ubuntu24.04-trt10.14.1.48-skip-tritondevel-202602011118-10901).

Pré-dumper l'image du conteneur (enroot import)

Les clusters SLURM utilisant enroot/pyxis nécessitent une image de conteneur .sqsh. Pour éviter la surcharge de téléchargement au moment de la compilation, pré-dumpez l'image à l'avance en utilisant le script compagnon enroot-import :

# Utilisation basique — soumet une tâche SLURM sur une partition CPU pour importer l'image
enroot-import --partition cpu_datamover --debug <docker_image_url>

Le script soumet une tâche sbatch qui exécute enroot import docker://<image_url> et produit un fichier .sqsh dans le répertoire courant. La sortie sur stdout est l'ID de la tâche SLURM.

Flags enroot-import

Flux de travail enroot-import

1. Lisez la balise de l'image depuis jenkins/current_image_tags.properties dans le dépôt TRT-LLM.
2. Exécutez enroot-import pour soumettre la tâche d'import :

   cd <directory_where_sqsh_should_be_stored>
   <path_to>/enroot-import --partition cpu_datamover --debug <image_url>

   IMPORTANT : Convertissez urm.nvidia.com/sw-tensorrt-docker/tensorrt-llm:xxx en urm.nvidia.com#sw-tensorrt-docker/tensorrt-llm:xxx pour éviter les problèmes d'authentification.

3. Attendez que la tâche d'import se termine (squeue -j <job_id>).
4. Le fichier .sqsh obtenu est l'container_image utilisé à l'étape de compilation.

Prérequis

L'utilisateur doit fournir (ou vous devez demander) ces valeurs :

Paramètres optionnels :

Scripts compagnons

Cette skill inclut trois scripts compagnons dans scripts/ :

Répertoire des scripts : skills/exec-slurm-compile/scripts/

Instructions

Suivez ces étapes dans l'ordre :

Étape 0 : Résoudre l'image du conteneur (si nécessaire)

Si l'utilisateur n'a pas déjà une image de conteneur .sqsh :

1. Lisez la balise de l'image Docker depuis <repo_dir>/jenkins/current_image_tags.properties.
2. Utilisez enroot-import pour la pré-dumper :

   cd <directory_for_sqsh_files>
   <scripts_dir>/enroot-import --partition cpu_datamover --debug <image_url>

3. Surveillez la tâche d'import avec squeue -j <job_id>.
4. Une fois terminée, le chemin du fichier .sqsh devient le paramètre container_image.

Si l'utilisateur a déjà un fichier .sqsh, ignorez cette étape.

Étape 1 : Recueillir les informations

Demandez à l'utilisateur les valeurs de prérequis manquantes listées ci-dessus. Au minimum, vous devez avoir :

• container_image (ou l'URL de l'image Docker — exécutez alors l'étape 0 en premier)
• repo_dir
• mount_dir
• partition et account

Si l'utilisateur a utilisé ce flux de travail avant, vérifiez si les valeurs précédentes sont stockées dans des fichiers de mémoire.

Étape 2 : Préparer le répertoire des scripts

Les scripts de compilation doivent être accessibles de l'intérieur du conteneur (c'est-à-dire sous mount_dir). Soit :

Option A — Copiez les scripts compagnons vers un emplacement sous mount_dir :

scripts_dir=<mount_dir>/<username>/workspace/tensorrt_llm_scripts
mkdir -p ${scripts_dir}/log
cp skills/exec-slurm-compile/scripts/compile.sh ${scripts_dir}/
cp skills/exec-slurm-compile/scripts/compile.slurm ${scripts_dir}/
chmod +x ${scripts_dir}/compile.sh ${scripts_dir}/compile.slurm

Option B — Si l'utilisateur a déjà des scripts à un emplacement connu, utilisez-les directement.

Étape 3 : Soumettre la tâche

Exécutez sbatch depuis le nœud de connexion (ou un nœud avec accès au client SLURM) :

sbatch \
    --nodes=1 --ntasks=1 --ntasks-per-node=1 \
    --gres=gpu:<gpu_count> \
    --partition=<partition> \
    --account=<account> \
    --job-name=<jobname> \
    --time=<time_limit> \
    <scripts_dir>/compile.slurm \
    <container_image> <mount_dir> <scripts_dir> <repo_dir>

Capturez et signalez l'ID de la tâche à partir de la sortie sbatch.

Étape 4 : Surveiller la tâche (Proactif — NE PAS attendre l'utilisateur)

Vous DEVEZ interroger activement la tâche jusqu'à ce qu'elle se termine. Ne soumettez pas et ne vous en allez pas.

# Vérifiez l'état de la tâche (répétez toutes les 30-60 secondes)
squeue -j <job_id> -o "%.18i %.9P %.30j %.8u %.2t %.10M %.6D %R"

# Une fois en cours d'exécution, affichez périodiquement la fin du journal (n'utilisez PAS tail -f, utilisez tail -30 à la place)
tail -30 <scripts_dir>/log/compile_<job_id>.srun.log

Boucle de surveillance :

1. Interrogez squeue -j <job_id> pour vérifier l'état
2. Si PD (en attente) — signalez la raison, continuez à interroger toutes les 30-60s
3. Si R (en cours d'exécution) — affichez la fin du journal de compilation toutes les 30-60s ; cherchez [XX%] Building, les erreurs ou la fin
4. Si la tâche disparaît de squeue, elle s'est terminée — passez à l'étape 5
5. Si F (échouée) — lisez immédiatement le journal complet et signalez l'erreur

Indicateurs de progression à chercher dans le journal :

• [XX%] Building CXX object... — progression de la compilation
• Linking CXX... — phase de liaison
• FAILED:, error:, fatal error: — échec de la compilation
• Successfully built — succès

Étape 5 : Vérifier la compilation

Une fois la tâche terminée, vérifiez le succès :

# Vérifiez le code de sortie SLURM
sacct -j <job_id> --format=JobID,State,ExitCode,Elapsed

# Vérifiez le journal de compilation pour les erreurs
tail -50 <scripts_dir>/log/compile_<job_id>.srun.log

Une compilation réussie se termine par un message comme Successfully built tensorrt_llm ou se termine sans erreur.

Référence des flags de compilation courants

Valeurs d'architecture courantes :

• "100-real" — Blackwell (B200, GB200)
• "90-real" — Hopper (H100, H200)
• "89-real" — Ada Lovelace (L40S)
• "80-real" — Ampere (A100)
• "90;100-real" — Plusieurs architectures

Dépannage

Exemple d'interaction

Utilisateur : « Compiler TRT-LLM sur le cluster OCI »

Actions de l'agent :

1. Demander le chemin de l'image de conteneur, le chemin du dépôt, le répertoire de montage (si non connus)
2. Confirmer la partition/le compte pour le cluster OCI
3. Copier les scripts vers un emplacement accessible sous mount_dir
4. Soumettre avec sbatch
5. Signaler l'ID de la tâche
6. Surveiller avec squeue jusqu'à la fin
7. Vérifier les journaux et signaler le succès/l'échec