MLflow pas à pas : instrumenter un script d’entraînement

Un script ElasticNet qui n’affiche ses résultats que dans le terminal, puis les six lignes qui le rendent comparable, reproductible et prêt pour le registre de modèles.

10 min de lecturemlflowmlopsmachine-learningpython

La meilleure façon de comprendre ce que MLflow apporte est de partir d’un script qui n’en a pas, de le faire souffrir, puis d’ajouter les lignes une à une en nommant ce que chacune achète.

Voici donc un script d’entraînement complet et honnête : il charge un jeu de données, entraîne un ElasticNet, prédit, évalue, et affiche. Rien de plus.

Le point de départ

python
import argparse
import warnings

import numpy as np
import pandas as pd
from sklearn.linear_model import ElasticNet
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import train_test_split

parser = argparse.ArgumentParser()
parser.add_argument('--alpha', type=float, default=0.7)
parser.add_argument('--l1_ratio', type=float, default=0.7)
args = parser.parse_args()


def eval_metrics(actual, pred):
    rmse = np.sqrt(mean_squared_error(actual, pred))
    mae = mean_absolute_error(actual, pred)
    r2 = r2_score(actual, pred)

    return rmse, mae, r2


if __name__ == '__main__':
    warnings.filterwarnings('ignore')

    data = pd.read_csv('data/red-wine-quality.csv')

    train, test = train_test_split(data, test_size=0.25, random_state=42)

    train_x = train.drop(['quality'], axis=1)
    test_x = test.drop(['quality'], axis=1)
    train_y = train['quality']
    test_y = test['quality']

    alpha = args.alpha
    l1_ratio = args.l1_ratio

    model = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
    model.fit(train_x, train_y)

    predicted_qualities = model.predict(test_x)

    rmse, mae, r2 = eval_metrics(test_y, predicted_qualities)

    print(f'ElasticNet — alpha={alpha}, l1_ratio={l1_ratio}')
    print(f'  RMSE : {rmse:.4f}')
    print(f'  MAE  : {mae:.4f}')
    print(f'  R²   : {r2:.4f}')

Le trajet est linéaire :

Le faire souffrir

bash
python train.py
python train.py --alpha 0.3 --l1_ratio 0.8
python train.py --alpha 0.1 --l1_ratio 0.2

Trois exécutions, trois blocs de texte dans le terminal. Et voilà le problème, qui n’a rien de théorique : au troisième essai, le premier a déjà défilé hors de l’écran. Demain, il ne reste rien.

Le script répond aux bonnes questions et ne conserve aucune réponse. Or ce sont exactement les quatre questions qu’on se posera dans trois jours : quel réglage était le meilleur, sur quelles données, mesuré comment, et avec quelle version du code.

Trois remarques sur ce script avant de l’instrumenter

Elles valent d’être faites, parce que deux d’entre elles se retrouvent dans énormément de scripts d’entraînement.

np.random.seed(40) ne sert à rien ici. On le voit souvent en tête de ces scripts, et il y donne une fausse impression de reproductibilité. Dans ce code, les deux seules sources d’aléa sont train_test_split et ElasticNet, et toutes deux reçoivent un random_state explicite. La graine globale de NumPy n’a donc aucun effet — ce qui est sans danger, mais trompeur : ce qui rend ce script reproductible, ce sont les deux random_state=42, et c’est eux qu’il faudra consigner.

warnings.filterwarnings('ignore') masque un message utile. Avec une régularisation forte, ElasticNet peut ne pas converger dans le nombre d’itérations imparti et le dit. Faire taire tous les avertissements pour obtenir une sortie propre revient à supprimer le seul indice disponible.

Le découpage est figé, et c’est une chance. Comme random_state=42 et test_size=0.25 sont écrits en dur, toutes les exécutions portent sur le même découpage : leurs scores sont donc réellement comparables. Le jour où quelqu’un change cette valeur pour « voir », les comparaisons deviennent silencieusement fausses. C’est pourquoi ces deux valeurs seront enregistrées comme des paramètres à part entière, au même titre que alpha.

Un mot enfin sur les résultats attendus : avec une régularisation aussi appuyée, le R² reste faible sur ce jeu de données — de l’ordre de 0,1, et il varie selon le découpage. Ce n’est pas un bug mais un alpha trop grand, qui écrase les coefficients et rapproche le modèle d’une prédiction constante ; ce que règlent exactement alpha et l1_ratio l’explique en détail. Autrement dit : le besoin de comparer apparaît dès la deuxième exécution.

Les six lignes qui changent tout

python
import mlflow

mlflow.set_experiment('vin-rouge-elasticnet')

with mlflow.start_run(run_name=f'alpha-{alpha}-l1-{l1_ratio}'):
    model = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
    model.fit(train_x, train_y)

    predicted_qualities = model.predict(test_x)
    rmse, mae, r2 = eval_metrics(test_y, predicted_qualities)

    mlflow.log_params(
        {
            'alpha': alpha,
            'l1_ratio': l1_ratio,
            'random_state': 42,
            'test_size': 0.25,
        }
    )
    mlflow.log_metrics({'rmse': rmse, 'mae': mae, 'r2': r2})
    mlflow.sklearn.log_model(model, name='modele')

Ce que chaque appel achète, dans l’ordre :

  • set_experiment regroupe les exécutions sous un nom, plutôt que de les laisser tomber dans un fourre-tout par défaut.
  • start_run délimite une exécution. Le run_name construit à partir des réglages est un petit geste qui rend la liste lisible sans l’ouvrir.
  • log_params répond à « avec quels réglages ». Noter la présence de random_state et test_size : ils ne sont pas des hyperparamètres du modèle, mais ils définissent le protocole, et sans eux deux lignes du tableau ne sont pas comparables.
  • log_metrics répond à « pour quel résultat ».
  • log_model enregistre l’objet entraîné, rechargeable plus tard sans réentraînement.

Le print peut rester : voir les chiffres en direct est agréable, et le journal n’en dépend plus.

Ce que le bloc with délimite exactement

start_run ouvre un dossier, et tout ce qui est enregistré à l’intérieur appartient à cet essai.

text
Début de l'exécution
┌──────────────────────────────┐
│ création du modèle           │
│ entraînement                 │
│ prédictions                  │
│ calcul des métriques         │
│ enregistrement des paramètres│
│ enregistrement des métriques │
│ enregistrement du modèle     │
└──────────────────────────────┘
Fin de l'exécution

Le with sert à fermer ce dossier automatiquement : quand Python sort du bloc indenté, MLflow termine l’exécution. Sans lui, il faut appeler mlflow.end_run() — et l’oublier a une conséquence désagréable, car l’exécution reste active et les enregistrements suivants viennent s’y ajouter. On se retrouve avec un essai qui contient les paramètres de deux entraînements différents, sans aucun message d’erreur.

Le même mécanisme explique un phénomène déroutant : appeler log_param alors qu’aucune exécution n’est ouverte n’échoue pas. MLflow en démarre une implicitement, ce qui produit ces exécutions mystérieuses que personne ne se souvient avoir lancées.

Faut-il créer le modèle à l’intérieur du bloc ? Techniquement non — ceci fonctionne :

python
lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)

with mlflow.start_run():
    lr.fit(train_x, train_y)
    mlflow.log_param('alpha', alpha)

Mais deux raisons de tout mettre dedans. La première est conceptuelle : une exécution MLflow correspond à une tentative complète d’entraînement, et il est plus clair que le bloc couvre exactement le cycle création, entraînement, évaluation, enregistrement.

La seconde est pratique, et elle décide. Si fit lève une exception, le gestionnaire de contexte termine quand même l’exécution et la marque en échec. Vous obtenez donc dans le journal la trace d’un essai qui a échoué, avec ses paramètres — ce qui est précisément l’information utile. Une exception hors du bloc laisse une exécution ouverte ou aucune trace du tout.

set_experiment ou experiment_id ?

Les deux désignent l’expérience dans laquelle ranger l’essai, par deux chemins.

python
mlflow.set_experiment('vin-rouge-elasticnet')  # par le nom, créée si absente

with mlflow.start_run():
    ...
python
exp = mlflow.get_experiment_by_name('vin-rouge-elasticnet')

with mlflow.start_run(experiment_id=exp.experiment_id):
    ...

set_experiment est le plus court et convient à un script : il crée l’expérience si elle n’existe pas et la fixe pour tout le processus. Passer un experiment_id explicite devient utile quand un même programme écrit dans plusieurs expériences, ou quand l’identifiant vient d’une configuration extérieure — auquel cas il vaut mieux vérifier que l’expérience existe plutôt que de supposer un nom.

Le piège de log_model

Écrivez bien name='modele', et non le deuxième argument positionnel qu’on trouve dans la plupart des exemples en circulation :

python
mlflow.sklearn.log_model(model, 'modele')       # avertissement de dépréciation
mlflow.sklearn.log_model(model, name='modele')  # correct

Ce deuxième argument correspond à artifact_path, déprécié au profit de name. Le code fonctionne encore, avec un avertissement — et si vous passez les deux, MLflow lève une erreur. C’est le genre de détail qui fait douter de son installation alors que seul l’exemple recopié était périmé.

Et si je veux tout enregistrer sans y penser ?

Une ligne, placée avant l’entraînement :

python
mlflow.sklearn.autolog()

Elle capte les hyperparamètres du modèle, les métriques d’entraînement et le modèle lui-même à chaque appel de fit. C’est le meilleur point de départ, parce qu’elle ne demande de ne rien oublier.

Elle ne connaît en revanche que ce que scikit-learn lui expose : elle ignore votre test_size, l’origine des données et vos métriques calculées à la main. La combinaison gagnante est donc l’enregistrement automatique plus quelques appels explicites pour ce qui vient de votre code.

Comparer, enfin

bash
mlflow ui

L’interface s’ouvre sur http://localhost:5000, et les exécutions apparaissent en tableau : un réglage par ligne, une métrique par colonne, triables. C’est le tableau que vous auriez tenu à la main, sans le tenir.

Pour lui donner de quoi montrer, un balayage :

bash
for a in 0.1 0.3 0.5 0.7 0.9; do python train.py --alpha "$a" --l1_ratio 0.5; done

Sous PowerShell :

powershell
foreach ($a in 0.1, 0.3, 0.5, 0.7, 0.9) { python train.py --alpha $a --l1_ratio 0.5 }

Cinq exécutions, et la relation entre alpha et le RMSE devient lisible d’un coup d’œil. C’est le moment où le suivi cesse d’être une corvée d’hygiène pour devenir un instrument de mesure.

Ce qui manque encore

Trois compléments transforment un journal correct en journal exploitable dans six mois.

L’origine des données. data/red-wine-quality.csv désigne un chemin, pas un contenu. Une empreinte règle la question :

python
import hashlib
from pathlib import Path

CHEMIN = Path('data/red-wine-quality.csv')

mlflow.log_param('donnees_empreinte', hashlib.sha256(CHEMIN.read_bytes()).hexdigest()[:12])
mlflow.log_param('donnees_lignes', len(data))

MLflow propose aussi un mécanisme dédié aux jeux de données, qui conserve la source et le schéma :

python
jeu = mlflow.data.from_pandas(data, source=str(CHEMIN), name='vin-rouge')
mlflow.log_input(jeu, context='training')

La version du code. Bonne nouvelle : si le script tourne depuis un dépôt Git, MLflow enregistre de lui-même le commit et le nom du fichier source. C’est l’un des bénéfices les moins connus, et il suffit à répondre à « avec quel code ce chiffre a-t-il été obtenu ». Encore faut-il que le dépôt soit propre au moment de l’exécution : un commit enregistré alors que trois fichiers sont modifiés localement désigne un état qui n’a jamais existé.

L’intention. Une étiquette d’une phrase, qui vaut plus que le score six mois plus tard :

python
mlflow.set_tag('intention', 'chercher le meilleur alpha à l1_ratio fixé')

Le raisonnement complet sur ce qu’il faut consigner, et pourquoi la dispersion compte autant que la moyenne, est dans le suivi des expériences.

Où vivent ces données

Par défaut, MLflow écrit dans un dossier mlruns/ à côté du script. Trois conséquences pratiques :

  • À exclure de Git. Une ligne dans .gitignore : ces fichiers sont des résultats, pas des sources, et les artefacts de modèles pèsent vite lourd.
  • Local, donc solitaire. Vos collègues ne voient rien. Le partage passe par un serveur de suivi, désigné par MLFLOW_TRACKING_URI, avec une base pour les métadonnées et un stockage d’objets pour les artefacts.
  • Périssable. Un dossier de travail nettoyé emporte l’historique. Dès que les résultats comptent, ils appartiennent à un serveur.

Dans un cluster, l’adresse du serveur et ses identifiants viennent de la configuration et non du code : un ConfigMap pour l’URI, un Secret pour les accès.

Du script au registre

Une exécution reste une exécution. Le jour où l’une d’elles devient le modèle servi, elle change de statut : il lui faut un nom et une version.

python
mlflow.sklearn.log_model(
    model,
    name='modele',
    registered_model_name='QualiteVinRouge',
)

Puis désigner celle qui sert :

python
from mlflow import MlflowClient

MlflowClient().set_registered_model_alias('QualiteVinRouge', 'champion', 3)

Et la recharger sans savoir quel numéro tourne :

python
modele = mlflow.pyfunc.load_model('models:/QualiteVinRouge@champion')

C’est la frontière entre expérimenter et servir, et le rôle exact du registre — détaillé dans MLOps et MLflow, y compris le piège de l’alias déplacé qui ne redémarre aucun service.

Ce qu’il faut retenir

Le script de départ n’était pas mauvais : il était muet. Six lignes lui rendent la parole, et aucune ne touche à sa logique — l’entraînement, les prédictions et les métriques sont exactement les mêmes.

Trois idées à emporter. Enregistrer le protocole (random_state, test_size) autant que les hyperparamètres, sans quoi les lignes du tableau ne se comparent pas. Écrire name= dans log_model, parce que l’argument positionnel est déprécié. Et laisser MLflow capter le commit Git, à condition de travailler dans un dépôt propre.

Pour la théorie derrière la pratique : ce que vous réglez et ce que le modèle apprend pour les hyperparamètres alpha et l1_ratio de ce script, et quand réentraîner pour la suite du cycle.

Continuer sur le même sujet