Files
ENI-projet-piscine/etl/airflow/dags/ml_train.py
Johan LEROY 19ff152445 refactor(etl): migre les DAGs et leurs tests vers le SDK Airflow 3
DAG et BaseOperator viennent d'airflow.sdk, BashOperator du provider
standard (airflow.operators.bash n'est plus qu'un alias déprécié). DagBag
s'importe depuis airflow.dag_processing et ne prend plus include_examples,
les exemples étant déjà coupés par la configuration de conftest.py.

Les timetables n'exposent plus summary : la planification se lit sur
dag.schedule (None) et timetable.expression (cron normalisé).
2026-09-22 09:01:07 +02:00

44 lines
1.9 KiB
Python

"""DAG d'entrainement du modele LightGBM (issue #115).
Pas de planification : reentrainer est couteux et sa cadence n'est pas une decision prise, en
particulier tant que `train.py` ecrase le modele sans comparer ses metriques a l'ancien (cf.
`docs/architecture/10-infra.md`, section Airflow). Declenchement manuel depuis l'UI ou la CLI
Airflow en attendant. `ml_score` (DAG separe, planifie toutes les heures) reutilise le modele que
ce DAG ecrit, il ne reentraine jamais rien lui-meme.
"""
from __future__ import annotations
from datetime import datetime, timedelta
from airflow.providers.standard.operators.bash import BashOperator
from airflow.sdk import DAG
MODEL_PATH = "/opt/ml/state/models/lightgbm-consumption.txt"
MLFLOW_TRACKING_URI = "sqlite:////opt/ml/state/mlflow.db"
with DAG(
dag_id="ml_train",
description="Entraine le modele LightGBM de prevision de consommation (enervision_ml.train).",
schedule=None,
start_date=datetime(2026, 1, 1),
catchup=False,
# Deux entrainements simultanes ecriraient le meme fichier modele.
max_active_runs=1,
tags=["ml"],
) as dag:
# `--no-sync` : l'environnement `/opt/ml/.venv` est fige a la construction de l'image, `uv run`
# ne le resynchronise pas (sinon `enervision-ml` est reconstruit a chaque tache).
# `env -u VIRTUAL_ENV` : l'image de base positionne celui d'Airflow, que `uv` signale a chaque
# execution sans qu'il change quoi que ce soit.
BashOperator(
task_id="train",
bash_command=(
"cd /opt/ml && env -u VIRTUAL_ENV uv run --no-sync python -m enervision_ml.train "
f"--model-output {MODEL_PATH} --mlflow-tracking-uri {MLFLOW_TRACKING_URI}"
),
# Un entrainement complet dure quelques minutes ; une connexion pendue ne doit pas
# immobiliser un slot du scheduler indefiniment.
execution_timeout=timedelta(hours=1),
)