launch-nemo-rl — exécuter les recettes NeMo-RL sur Kubernetes via nrl-k8s

Ceci est le playbook pour le CLI nrl-k8s situé dans infra/nrl_k8s/. Suivez-le quand l'utilisateur demande de lancer / itérer / déboguer une recette NeMo-RL sur un cluster Kubernetes. Vérifiez l'état actuel (kubectl, git log, les fichiers recette + infra) avant d'agir — le cluster est partagé et le coût d'une mauvaise action est élevé.

1. Une commande, deux modes

Il existe une seule commande de soumission de haut niveau : nrl-k8s run. Elle a deux modes de cycle de vie.

Demandez-vous : Ai-je besoin de ce cluster après l'exécution ? Si oui, utilisez --raycluster. Sinon, utilisez le mode par défaut (éphémère).

Le reste du CLI est observabilité / contrôle étape par étape :

2. Paire recette + infra

Chaque lancement nécessite deux fichiers. Passez l'infra avec --infra, pas fusionné en ligne :

nrl-k8s run infra/nrl_k8s/examples/<recipe>.yaml \
  --infra infra/nrl_k8s/examples/<recipe>.<profile>.infra.yaml

• Recette (ex. qwen3_30b_math_8n_4gpu.yaml) — config NeMo-RL : modèle, boutons GRPO/SFT, cluster.{gpus_per_node,num_nodes}. Utilise defaults: pour hériter de examples/configs/recipes/llm/....
• Infra (ex. *.<profile>.infra.yaml) — shape K8s/Ray : namespace, image, service account, spec RayCluster sous kuberay:, Deployments optionnels sous deployments:, submit.submitter, launch.{mode,codeSource,codePath,entrypoint}. Les noms de paires suivent <recipe>.<profile>[.prod].infra.yaml où <profile> nomme la cible matérielle (ex. gb300).

Des paires d'exemples dans infra/nrl_k8s/examples/ — lisez les fichiers voisins pour voir les conventions actuelles pour le profile cible.

3. Flags du mode long-lived

Trois dimensions indépendantes. --mode est une macro qui choisit les défauts ; les flags individuels la remplacent.

--mode interactive   → --submitter portForward  --code-source upload  (suit les logs)
--mode batch         → --submitter exec         --code-source image   (retourne après nohup)

• Submitter : portForward utilise kubectl port-forward + Ray Job SDK (obtient un submission_id que le dashboard suit). exec utilise kubectl exec + nohup sur le pod head (pas de submission_id ; le driver apparaît comme type=DRIVER dans le dashboard).
• Code source : upload stage un working_dir depuis l'ordinateur portable (limite Ray 100 MiB). image / lustre attendent le code sur le filesystem du pod — couplé avec --code-path (typiquement /opt/nemo-rl), qui est un subPath du mount PVC filesystem partagé dans les exemples d'infra standard.
• Wait : --wait suit les logs jusqu'au terminal ; --no-wait retourne dès que le driver s'exécute.

Autres flags long-lived uniquement :

• --replace — arrêtez tout job d'entraînement / daemon en cours avant de soumettre de nouveaux (ajoute un timestamp aux submissionIds des daemons pour que Ray accepte la resoumission).
• --recreate — supprimez + réappliquez un RayCluster dont la spec live a divergé du manifest rendu (par défaut avertir + réutiliser).
• --skip-daemons — apportez tous les clusters déclarés mais soumettez uniquement l'entraînement. À utiliser sur les recettes disaggregées où gym/generation sont déjà sains.

Piège : sur l'infra où le entrypoint fait cd /opt/nemo-rl (ou un autre chemin en image / Lustre) et charge la recette à partir de là, --code-source upload N'ÉCRASE PAS la recette sur le pod — le working_dir uploadé se situe dans /tmp/ray/... mais le entrypoint se cd ailleurs. Pour vraiment tester un changement local de recette, soit synchronisez vos modifications avec le filesystem partagé monté dans les pods, soit activez les overrides Hydra dans le entrypoint.

4. Flags du mode éphémère (--rayjob)

Quand --rayjob est défini, run branche dans le chemin du code RayJob. Flags pertinents :

• --rayjob-name NAME — nom des métadonnées du RayJob (par défaut le nom du cluster d'entraînement).
• --shutdown / --no-shutdown — par défaut true : KubeRay supprime le RayCluster une fois que le Ray Job atteint un état terminal.
• --ttl SECONDS — par défaut 3600s : gardez l'objet RayJob autour après la fin de l'exécution pour l'accès aux logs post-mortem.
• --wait / --no-wait — par défaut wait : scrutez jobDeploymentStatus jusqu'à Complete/Failed. --no-wait retourne dès que le RayJob est appliqué.
• --timeout SECONDS — par défaut 86400s (24h) : bornez le scrute --wait.
• --dry-run — affichez le manifest RayJob et imprimez-le ; ne l'appliquez pas.

--replace / --recreate / --skip-daemons sont silencieusement ignorés en mode --rayjob (KubeRay possède le cycle de vie).

5. Itérer sur une config sans toucher au filesystem partagé

Quand la recette sur le filesystem du pod a la mauvaise valeur pour votre expérience, utilisez plutôt les overrides Hydra sur le entrypoint au lieu de forker la recette. Motif :

entrypoint: |
  set -eu
  cd /opt/nemo-rl
  RUN_ID="\${RAY_JOB_SUBMISSION_ID:-\${NRL_K8S_RUN_ID:-$(date -u +%Y%m%d-%H%M%S)}}"
  python -u examples/run_grpo.py \
    --config infra/nrl_k8s/examples/<recipe>.yaml \
    logger.wandb_enabled=true \
    logger.wandb.project=<project> \
    "logger.wandb.name=<run-name>-\${RUN_ID}"

Échappez ${…} avec un backslash. OmegaConf interprète sinon comme interpolation et erreur sur le style shell ${VAR:-default}. RUN_ID se résout à RAY_JOB_SUBMISSION_ID (injecté par KubeRay en mode rayjob) → NRL_K8S_RUN_ID (injecté par le CLI en mode long-lived) → timestamp local — donc le nom est unique sur l'un ou l'autre chemin.

6. Préoccupations par profile (matériel + scheduler + DRA)

Chaque YAML infra code un profile matériel/scheduler. Les exemples concrets dans infra/nrl_k8s/examples/ sont faisant autorité pour les profiles qu'ils ciblent — lisez le fichier infra voisin avant d'en écrire un nouveau. Les choses qui varient couramment :

• GPUs par nœud (ex. 4 vs 8) — doit correspondre à cluster.gpus_per_node dans la recette, sinon les workers restent Pending.
• Node selectors — les pods head atterrissent généralement sur un node pool CPU-only ; les workers GPU correspondent sur nvidia.com/gpu.product ou un label de node-group.
• Scheduler — KAI (schedulerName: kai-scheduler + label kai.scheduler/queue) avec annotations de topologie (kai.scheduler/topology, kai.scheduler/topology-required-placement) planifie les workers dans une seule clique. Sans cela, les pods peuvent atterrir sur différents racks et NVLink/RoCE ne les traversera pas.
• DRA claims — ComputeDomain + RoCE sont attachés via resourceClaims référençant ResourceClaimTemplates. Le CLI crée/supprime automatiquement ceux-ci quand la spec du pod worker contient des références de DRA claim — pas de setup manuel nécessaire.
• Secrets — toujours via secretKeyRef (wandb-api-key, image pull secret). Ne les intégrez jamais.
• Shared filesystem mounts — typiquement une Lustre PVC montée deux fois : une fois au chemin du code (ex. /opt/nemo-rl avec un subPath limité à l'utilisateur) et une fois à la racine du workspace (ex. /mnt/rl-workspace) pour les datasets, HF cache, et les checkpoints.

Avant d'appliquer une infra, vérifiez que les prérequis existent dans le namespace cible :

kubectl get pvc <workspace-pvc>
kubectl get secret <wandb-secret> <image-pull-secret>
kubectl get sa <service-account>

7. Workflows de bout en bout

7a. Exécution unique fraîche (rayjob)

# Depuis la racine du repo NeMo-RL :
nrl-k8s check <recipe> --infra <infra>                               # valider d'abord
nrl-k8s run <recipe> --infra <infra> --rayjob --dry-run              # afficher le manifest RayJob
nrl-k8s run <recipe> --infra <infra> --rayjob --no-wait              # appliquer, retour rapide

Regardez le statut + teardown (fonctionne même après déconnexion du laptop car KubeRay possède le cycle de vie) :

kubectl get rayjob -n default <name> -w
kubectl get raycluster -n default                                    # vide = teardown réussi

7b. Boucle de développement (long-lived)

nrl-k8s run <recipe> --infra <infra> --run-id $(date +%Y%m%d-%H%M%S)
# Modifications dans la recette ? Réexécutez juste — réutilise le cluster live.
# Pod spec changée ? Ajoutez --recreate pour supprimer + réappliquer.
# Recette disaggregée avec gym/gen déjà sains ? --skip-daemons.

7c. Premier apport disaggregated

nrl-k8s run <recipe> --infra <disagg-infra> --mode batch --code-source image

7d. Cycle de vie cluster uniquement

nrl-k8s cluster up   <recipe> --infra <infra> --target kuberay.training --wait
nrl-k8s cluster up   <recipe> --infra <infra> --target kuberay.training --dry-run   # afficher le manifest
nrl-k8s cluster down <recipe> --infra <infra> --target kuberay.training --wait
nrl-k8s cluster down <recipe> --infra <infra>                                       # tear down all
nrl-k8s cluster list -n default
nrl-k8s cluster dashboard <cluster-name>                                  # port-forward + browser

7e. Deployments (ex. nemo-skills sandbox)

# Apportez juste le deployment
nrl-k8s cluster up <recipe> --infra <infra> --target deployments.nemo_skills
# Supprimez juste le deployment
nrl-k8s cluster down <recipe> --infra <infra> --target deployments.nemo_skills
# Supprimez tout (RayClusters + Deployments)
nrl-k8s cluster down <recipe> --infra <infra>

La section deployments: dans le YAML infra déclare les Deployments Kubernetes gérés aux côtés des RayClusters. Le CLI patche image, imagePullSecrets, et serviceAccountName à partir des clés infra de haut niveau (pareil que RayClusters). Les Deployments démarrent en parallèle avec le bring-up du cluster — aucune dépendance d'ordre.

8. Surveiller une exécution

# Statut
nrl-k8s status <recipe> --infra <infra>
kubectl get rayjob,raycluster -n default

# Suivez le driver
nrl-k8s job list <recipe> --infra <infra> --role training
nrl-k8s job logs <run-id> <recipe> --infra <infra> --role training -f

Quand le subprocess nrl-k8s job logs -f meurt (kubectl port-forward i/o timeout après ~15 min idle), réexécutez-le. Le job d'entraînement continue.

Pour récupérer les logs du driver d'un job terminal (SUCCEEDED/FAILED) ou un RayJob via l'API du dashboard :

RC=$(kubectl get rayjob -n default <rayjob-name> -o jsonpath='{.status.rayClusterName}')
kubectl port-forward -n default svc/${RC}-head-svc 18266:8265 &
curl -s http://localhost:18266/api/jobs/                              # lister les jobs, trouver submission_id
curl -s "http://localhost:18266/api/jobs/<submission_id>/logs"        # log du driver complet

type=DRIVER avec submission_id=null signifie une exécution exec-submitter (aucun endpoint de log du dashboard — utilisez nrl-k8s job logs à la place). type=SUBMISSION a submission_id défini et /api/jobs/<id>/logs fonctionne.

L'URL Wandb apparaît dans le log du driver sur le premier appel wandb.init ; grep grep -oE 'https://wandb\.ai/[A-Za-z0-9_./-]+'.

9. Arrêter les choses

Confirmez avant de supprimer l'infra partagée. Le coût de cluster down sur le cluster de quelqu'un d'autre est élevé.

10. Vérifier le teardown du RayJob

Après qu'un run --rayjob se complète avec --shutdown (défaut), KubeRay devrait supprimer le RayCluster :

kubectl get rayjob   -n default <rayjob-name>                        # jobDeploymentStatus = Complete
kubectl get raycluster -n default | grep <rayjob-name>               # pas de sortie = détruit

L'objet RayJob lui-même reste pendant --ttl secondes (défaut 3600s) pour que vous puissiez toujours récupérer les logs.

11. Pièges courants

• Interpolation OmegaConf dévore ${VAR} dans recipe/infra YAML. Échappez les variables shell avec \${VAR} pour qu'OmegaConf les passe au shell du pod verbatim.
• Configs optimiseur Megatron ne comportent pas foreach / fused. Les overrides comme ~policy.optimizer.kwargs.foreach ~policy.optimizer.kwargs.fused (valides pour les configs DTensor) cassent sur les recettes Megatron. Omettez-les pour Megatron.
• DTensor vs Megatron — les recettes MoE utilisent généralement megatron_cfg.enabled=true ; assurez-vous dtensor_cfg.enabled=false dans les défauts hérités.
• Shared filesystem vs git divergence — codeSource: image|lustre lit depuis le filesystem du pod. Si vos modifications locales ne sont pas sur le filesystem partagé que les pods montent, l'exécution teste la version sur disque, pas la vôtre. Synchronisez soit via un pod helper (l'exec du pod head est souvent bloqué) soit override via des flags Hydra.
• Ephemeral-storage + readinessProbe sont injectés par les webhooks kuberay/CDI au moment de l'application du pod. N'LES AJOUTEZ PAS à la spec RayCluster en ligne.
• Node taints varient par cluster. tolerations: [{operator: Exists}] sur les workers est défensif et vaut le coup d'être conservé.
• Page vide du dashboard — Ray 2.52 installe les assets du dashboard comme symlinks par défaut ; nrl-k8s cluster dashboard <name> réinstalle automatiquement ray[default] --link-mode=copy pour corriger. Intégrez ENV UV_LINK_MODE=copy dans l'image pour éviter cela entièrement.
• kubectl exec est généralement bloqué en automation — contournez avec kubectl get ... -o yaml, kubectl logs, et kubectl port-forward + Ray dashboard APIs.

12. Checklist avant de dire qu'une exécution est "terminée"

Avant de signaler un lancement comme réussi, vérifiez :

1. kubectl get rayjob/raycluster -n default affiche les objets attendus.
2. nrl-k8s job list (ou curl /api/jobs/) affiche le job en RUNNING / SUCCEEDED.
3. Le log du driver contient wandb.ai/<project>/runs/<id> (si wandb est activé) — partagez l'URL avec l'utilisateur.
4. Au moins une ligne Processed prompts: 100% apparaît (confirme que la génération est câblée).
5. Pour le mode --rayjob uniquement : après jobDeploymentStatus=Complete, confirmez que kubectl get raycluster | grep <name> est vide (teardown a fonctionné).

13. Dev pod

nrl-k8s dev gère un pod CPU léger sur le cluster pour la synchro du code, le débogage, et l'exécution de kubectl/nrl-k8s depuis le cluster.

# Unique : configurer les secrets (HF token, wandb, SSH key, rclone)
nrl-k8s dev setup-secrets --ssh-key ~/.ssh/id_rsa --add-rclone

# Créer le pod et exec (idempotent — réutilise le pod existant)
nrl-k8s dev connect

# Changer l'image (doit d'abord arrêter — le changement d'image est averti mais pas appliqué automatiquement)
nrl-k8s dev stop
nrl-k8s dev connect --image nvcr.io/nvidian/nemo-rl:v0.7.0

# Détruire
nrl-k8s dev stop

Le dev pod :

• Tourne sur un nœud CPU-only (anti-affinity aux nœuds GPU)
• Monte la PVC rl-workspace partagée à /mnt/rl-workspace
• Définit la var env USER au nom d'utilisateur nrl-k8s (donc $USER et getpass.getuser() fonctionnent correctement malgré l'exécution en tant que root)
• Installe kubectl, rclone (si configuré) au premier boot
• Injecte les clés SSH et tokens via envFrom sur un Secret K8s par utilisateur

Le service account default du pod a besoin d'une RoleBinding edit dans le namespace pour que kubectl fonctionne dedans. dev connect vérifie cela et imprime le YAML requis s'il manque.

14. Où les choses vivent dans le repo

• Code CLI : infra/nrl_k8s/src/nrl_k8s/ (cli.py, orchestrate.py, manifest.py, rayjob.py, k8s.py, submitters/, schema.py).
• Tests : infra/nrl_k8s/tests/unit/ — exécutez avec uv run --extra test pytest -x -q depuis infra/nrl_k8s/.
• Exemples recette + infra : infra/nrl_k8s/examples/.
• Recettes de base que cet outil enveloppe : examples/configs/recipes/llm/… et examples/nemo_gym/….