Utiliser les AWS Deep Learning Containers avec SageMaker AI géré MLflow

TL;DR

• Intègre les AWS Deep Learning Containers (DLCs) avec SageMaker AI géré par MLflow pour équilibrer le contrôle de l’infrastructure et une gouvernance ML robuste.
• Utilise un flux TensorFlow pour la prédiction de l’âge des abalones où l’entraînement se fait dans des DLCs, les artefacts sont stockés dans Amazon S3 et les résultats sont enregistrés dans MLflow pour un suivi centralisé.
• SageMaker MLflow géré offre une journalisation automatique en une ligne, des capacités de comparaison avancées et un suivi complet des lignées pour l’audit.
• L’approche offre un équilibre entre personnalisation et standardisation, permettant aux organisations de standardiser les workflows ML tout en conservant la flexibilité des environnements via DLCs ou DLAMIs.
• Les prérequis incluent le retrait d’un conteneur TensorFlow optimisé depuis le registre public AWS ECR et la consultation du README du dépôt GitHub pour provisionner et exécuter.

Contexte et historique

Les organisations qui construisent des modèles ML personnalisés ont souvent des exigences spécialisées que les plateformes standard ne peuvent satisfaire. Par exemple, les entreprises de santé ont besoin d’environnements spécifiques pour protéger les données des patients et respecter HIPAA, les institutions financières requièrent des configurations matérielles particulières pour optimiser des algorithmes de trading propriétaires, et les équipes de recherche ont besoin de flexibilité pour expérimenter des techniques de pointe via des frameworks personnalisés. Ces besoins poussent les équipes à créer des environnements d’entraînement sur mesure qui donnent le contrôle sur le matériel, les versions logiciel et les configurations de sécurité. Cette flexibilité est précieuse mais crée des défis importants pour la gestion du cycle de vie du ML, augmentant les coûts opérationnels et nécessitant des ressources d’ingénierie. Les DLCs et les DLAMIs d’AWS offrent une solution puissante pour répondre à ces besoins : les DLCs fournissent des conteneurs Docker préconfigurés avec des frameworks populaires comme TensorFlow et PyTorch, des pilotes NVIDIA CUDA pour le support GPU et sont optimisés pour les performances sur AWS. Ils sont régulièrement mis à jour avec les dernières versions des frameworks et correctifs et sont conçus pour s’intégrer sans couture aux services AWS pour l’entraînement et l’inférence. Les DLAMIs sont des images AMI préconfigurées pour EC2, avec CUDA, cuDNN et d’autres outils nécessaires, gérées par AWS. Ensemble, DLCs et DLAMIs offrent l’infrastructure et les outils pour accélérer le deep learning dans le cloud à grande échelle. SageMaker MLflow fournit une gestion complète du cycle de vie avec journalisation automatique en une ligne, comparaisons améliorées et traçabilité de lignage, tout en déléguant l’exploitation opérationnelle. La solution présentée ici décrit une mise en œuvre fonctionnelle qui équilibre le contrôle d’infrastructure et la gouvernance en combinant DLCs (ou DLAMIs) avec MLflow géré par SageMaker AI. L’objectif est de démontrer comment suivre les expériences et gérer le cycle de vie du modèle de manière centralisée tout en préservant la flexibilité de l’environnement de développement. Le flux est illustré via une architecture qui montre les interactions entre les services AWS, DLCs et MLflow géré par SageMaker AI, et les étapes pour réaliser un premier entraînement avec un suivi d’expériences complet.

Pourquoi c’est important (impact pour les développeurs et les entreprises)

Pour les organisations ayant des exigences strictes en matière de gouvernance ou des besoins informatiques et de sécurité spécifiques, la capacité de contrôler l’environnement d’entraînement sans sacrifier la gouvernance est cruciale. En associant DLCs (ou DLAMIs) à SageMaker MLflow géré, les équipes peuvent:

• Préserver la flexibilité nécessaire pour répondre à la conformité HIPAA, des configurations matérielles spécialisées ou l’expérimentation avancée, tout en bénéficiant d’une gouvernance centralisée.
• Obtenir un suivi complet des expériences, une traçabilité des modèles et une histoire auditable du processus d’entraînement à la mise en production.
• Accélérer le cycle de vie du ML en réduisant la charge opérationnelle associée au suivi et à la gestion des environnements d’expérimentation.
• Standardiser les workflows ML entre les équipes et les projets, facilitant la reproductibilité et la collaboration tout en maintenant les garde-fous requis dans les secteurs réglementés. Cette approche s’inscrit dans les objectifs plus larges visant à équilibrer personnalisation et gouvernance dans les initiatives d’IA en entreprise, permettant aux data scientists et aux équipes opérationnelles de travailler plus efficacement sans compromettre la sécurité, la traçabilité et la conformité.

Détails techniques ou Mise en œuvre

Les éléments clés de la solution décrite incluent:

• AWS Deep Learning Containers (DLCs) : conteneurs Docker préconfigurés avec des frameworks comme TensorFlow et PyTorch, incluant les pilotes NVIDIA CUDA pour le support GPU; les DLCs sont optimisés pour AWS et mis à jour régulièrement. Ils sont conçus pour s’intégrer aux services AWS pour l’entraînement et l’inférence.
• AWS Deep Learning AMIs (DLAMIs) : AMI préconfigurées pour EC2 avec des frameworks populaires; disponibles pour les instances CPU et GPU; incluent CUDA, cuDNN et d’autres outils nécessaires; les mises à jour sont gérées par AWS.
• SageMaker MLflow géré : offre un journalisation automatique en une ligne, des capacités de comparaison et un suivi de lignage complet dans un environnement SageMaker AI géré.
• Étapes d’implémentation (haut niveau) :
• Extraire un conteneur TensorFlow optimisé depuis le dépôt public ECR d’AWS.
• Configurer une instance EC2 avec accès au serveur de suivi MLflow.
• Exécuter l’entraînement dans le DLC tout en stockant les artefacts du modèle dans Amazon S3 et en enregistrant les résultats dans MLflow.
• Visualiser et comparer les résultats des expériences via l’interface MLflow pour évaluer les performances du modèle.
• Lors de l’enregistrement d’un modèle avec mlflow.tensorflow.log_model() en utilisant le paramètre registered_model_name, le modèle est automatiquement enregistré dans le SageMaker Model Registry, créant une traçabilité de l’entraînement jusqu’à la production.
• Les artefacts du modèle sont stockés dans S3 selon une structure standardisée garantissant que les poids, les configurations et les métadonnées sont préservés et accessibles.
• Exemple pratique : prédiction de l’âge des abalones avec un modèle TensorFlow, démontrant l’intégration des fonctionnalités de gouvernance et de versioning via l’interface et le registre de modèles.
• Référence de mise en œuvre : tous les exemples de code et détails d’implémentation sont disponibles dans le README du dépôt GitHub associé, couvrant le provisionnement, les permissions et l’exécution du premier entraînement avec un suivi complet d’expériences. Comme le montrent les visuels accompagnant le billet, SageMaker MLflow géré offre un hub central pour le suivi des métriques, des paramètres, des artefacts et des modèles enregistrés, fournissant une gouvernance complète et une traçabilité d’audit pour les workflows ML. L’intégration assure que les résultats d’entraînement puissent être liés au modèle déployé en production, offrant traçabilité et conformité tout au long du cycle ML.

Points clés à retenir

• DLCs et DLAMIs offrent le bon équilibre entre personnalisation et infrastructure gérée pour les flux ML sur AWS.
• SageMaker AI géré MLflow apporte des capacités de gouvernance solides, y compris la journalisation automatique, la traçabilité et le registre de modèles.
• Le flux de prédiction d’âge des abalones avec TensorFlow montre une gouvernance de bout en bout : entraînement, stockage des artefacts, comparaison des expériences et déploiement en production.
• Le suivi MLflow centralisé sur SageMaker AI simplifie l’audit, la traçabilité et la collaboration entre les équipes de science des données et des opérations.
• Cette approche réduit les charges opérationnelles tout en préservant la flexibilité pour répondre à des exigences spécialisées.

Références

• https://aws.amazon.com/blogs/machine-learning/use-aws-deep-learning-containers-with-amazon-sagemaker-ai-managed-mlflow/