Passer à la navigation
Documents de conception

uv dans NeMo RL

Afficher en Markdown

Nous utilisons le gestionnaire de packages Python uv pour gérer les dépendances dans NeMo RL.

Vue d’ensemble

uv est un outil incroyable qui simplifie notre flux de travail et est extrêmement rapide car il est écrit en Rust. Ce document explique pourquoi nous avons adopté uv pour la gestion de packages dans notre dépôt, en particulier pour NeMo RL, et comment il nous aide à gérer les dépendances dans les clusters Ray.

Pourquoi uv ?

uv apporte les avantages clés suivants à notre workflow de développement Python :

Vitesse et Efficacité

  • Écrit en Rust, le rendant significativement plus rapide que les gestionnaires de packages Python traditionnels.
  • Mécanismes de mise en cache optimisés qui réduisent les téléchargements et installations redondants.
  • Création et commutation d’environnements rapides, permettant des cycles de développement rapides.

Environnements Isolés

  • Crée des environnements Python complètement isolés, empêchant les conflits de dépendances entre les packages système et les packages spécifiques au projet.
  • Évite les situations de dépendances nuancées où un script Python pourrait accidentellement utiliser à la fois les dépendances virtualenv et les dépendances système.
  • Assure un comportement cohérent sur différentes machines et environnements de déploiement.

Gestion des Dépendances dans les Clusters Ray

  • Permet la gestion d’environnements Python hétérogènes dans un cluster Ray.
  • Fournit de la flexibilité pour que chaque acteur (worker) utilise les dépendances Python spécifiques dont il a besoin.
  • Simplifie la propagation des environnements aux nœuds workers sans configuration manuelle sur chaque nœud.

Flexibilité sans Conteneur

  • Nous libère de la nécessité de publier de nombreux conteneurs pour différentes combinaisons de dépendances.
  • Nous permet de définir différents groupes de dépendances et extras et de sélectionner dynamiquement ceux dont nous avons besoin.
  • Réduit la complexité de l’infrastructure et la surcharge de maintenance.

Mise en Œuvre dans NeMo RL

Cette section décrit comment les workers définissent leurs exécutables requis, détaille les configurations prédéfinies disponibles (comme BASE ou VLLM), et explique comment personnaliser ces configurations pour des besoins spécifiques, en assurant la cohérence entre les acteurs.

Configuration du Worker

Dans notre base de code, les workers (classes décorées avec @ray.remote, par exemple PolicyWorker) sont associés à un PY_EXECUTABLE qui spécifie les dépendances dont le worker a besoin. Ceux-ci sont définis dans un registre global dans ACTOR_ENVIRONMENT_REGISTRY. Cela permet à différentes parties de notre application d’avoir leurs propres environnements personnalisés.

Exécutables Python Pris en Charge

Nous fournissons plusieurs configurations d’exécutables Python prédéfinies dans PY_EXECUTABLES :

class PY_EXECUTABLES:
    SYSTEM = sys.executable
    # Use NeMo RL direct dependencies.
    BASE = "uv run --locked"
    # Use NeMo RL direct dependencies and vllm.
    VLLM = "uv run --locked --extra vllm"

Pour assurer des dépendances cohérentes entre les acteurs, nous exécutons avec --locked pour garantir que les dépendances sont cohérentes avec le contenu de uv.lock.

Personnalisation

Si vous avez besoin d’une configuration d’exécutable Python différente, vous pouvez remplacer celle par défaut en passant la vôtre dans RayWorkerBuilder.__call__. Cela offre de la flexibilité pour des cas d’utilisation spéciaux sans modifier les configurations de base.

Comment Cela Fonctionne

Lorsqu’un job NeMo RL est démarré :

  1. Le script pilote crée plusieurs RayWorkerGroup.
  2. Chaque groupe de workers créera leurs workers qui sont encapsulés dans un RayWorkerBuilder où le nom complètement qualifié (FQN) de la classe worker est passé en tant que chaîne.
  3. RayWorkerBuilder lance le worker sous RayWorkerBuilder qui nous permet d’initialiser la classe sans importer des packages non disponibles dans l’environnement de base.
  4. Avant que la classe worker ne soit instanciée par le RayWorkerBuilder, le FQN est utilisé pour rechercher — dans un registre global) — pour déterminer quel membre de PY_EXECUTABLES doit être utilisé pour lancer ce jeu de workers. Si le PY_EXECUTABLES.* choisi commence par uv ; un venv est créé avec toutes les dépendances dont il a besoin et le runtime_env["py_executable"] est remplacé par l’interpréteur python du venv.

Cette approche permet un démarrage rapide et maintient l’isolation des dépendances. Elle a également l’avantage supplémentaire d’avoir tous les environnements virtuels localement sous ./venvs.

Conclusion

Utiliser uv pour la gestion des dépendances dans NeMo RL nous offre un moyen rapide, flexible et fiable de gérer les dépendances Python dans les clusters Ray distribués. Il élimine de nombreux points douloureux traditionnels de la gestion des dépendances dans les systèmes distribués, tout en permettant des environnements hétérogènes qui peuvent être adaptés à des charges de travail spécifiques.