MLflow inside an Airflow task

When ML lives outside the data platform, model runs are invisible to oncall and the data team. When ML lives inside the platform, an Airflow task trains the model, logs the run to MLflow, and the next DAG promotes the model to the registry. One place, one audit trail, one rollback button.

MLflow as a first-class Airflow task

Train step runs inside the same DAG as the data pipeline. Model registry promotion is a downstream task with its own approvals.

Every model in production traces back to a DAG run.
airflow/dags/ml_training_dag.py
python
import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier


def train_model(**ctx):
    mlflow.set_tracking_uri(MLFLOW_URI)
    mlflow.set_experiment('fraud-detection')

    X_train, y_train = load_training_frame(ctx['ds'])

    with mlflow.start_run() as run:
        params = {'n_estimators': 200, 'max_depth': 12}
        mlflow.log_params(params)
        clf = RandomForestClassifier(**params, random_state=42)
        clf.fit(X_train, y_train)
        mlflow.log_metric('train_accuracy', clf.score(X_train, y_train))
        mlflow.sklearn.log_model(clf, 'model')
        ctx['ti'].xcom_push(key='mlflow_run_id', value=run.info.run_id)

The train task logs parameters, metrics, and the model artifact to the MLflow server. The run_id is pushed to xcom so the promote task knows which run to promote.

Because training depends on the same data the DAG just produced. Running training outside the platform means replicating ingestion, scheduling, and retry logic in a second place. Inside Airflow, the same data contracts feed the model, and the same audit trail covers data runs and model runs. One platform, one story.

Quiz: Quiz

Loading practice…