New to Kestra?
Use blueprints to kickstart your first workflows.
Gate MLflow model promotion with statistical drift detection
Gate MLflow model promotion using PSI and KS feature drift tests, check metric regression against KV baselines, and promote or block with Slack alerts.
Categories
Getting a model into the registry is the easy part. Knowing whether it is safe to actually serve is harder — especially when the production population has shifted since training. This flow puts a deterministic gate between the MLflow registry and Production: it checks feature drift and metric regression on every promotion attempt, promotes on pass, and blocks with a FAILED execution on failure so CI pipelines get a non-zero exit code automatically.
How it works
fetch_model_versionandfetch_run_metricspull the candidate version's metadata and its training run metrics from the MLflow REST API. No MLflow Python client is needed — Kestra's HTTP task calls the/api/2.0endpoints directly, which works with any MLflow deployment (OSS, Databricks, AWS SageMaker Experiments).compute_driftloads a training-time reference Parquet and a recent production sample Parquet, then runs two independent drift tests on every shared numeric feature:- PSI (Population Stability Index): the standard industry metric for population shift. PSI < 0.1 = stable, 0.1–0.2 = warning, > 0.2 = drifted. The threshold is an input so you can tighten or loosen it per model.
- KS test: a distribution-free two-sample test that catches shape
changes PSI can miss when bin sizes are poorly distributed. A p-value
below
ks_alpha(default 0.05) flags drift. A feature is called DRIFTED only when both tests agree, which dramatically cuts false positives from binning artefacts. The full per-feature report lands indrift_report.jsonand is attached to the execution for review.
check_metric_regressioncompares the candidate's logged training metrics against the metrics stored in the KV store from the last promoted version. A 1 pp tolerance window prevents noise from triggering false regressions. If no KV baseline exists — first ever deploy — the check passes unconditionally.gate_decisionbranches:- PASS:
promote_modeltransitions the version totarget_stageviaPATCH /model-versions/transition-stagewitharchive_existing_versions: true, so the previous Production version is automatically archived. The candidate's metrics are written back to KV as the new baseline. A Slack message names the feature and metric counts. - BLOCK:
fail_executionends the run as FAILED so downstream CI jobs receive a non-zero exit code without any extra curl or wrapper. Slack gets the drifted feature count and regression count. The execution is left intact for post-mortem — the drift report is attached as an output file.
- PASS:
The gate result (PASS or BLOCK with timestamp and counts) is written to the KV store on every run, so a scheduled audit flow can query the history across all models without parsing execution logs.
What you get
- A single Kestra execution that is the authoritative promotion record: inputs, drift report, metric comparison, and the registry API response are all attached and queryable.
- A FAILED execution that stops CI without any extra scripting, identical to
how
dbt-core-approval-gateblocks on test failures. - An audit trail in KV (
mlflow_gate_<model>_v<version>_result) that survives execution pruning. - An updated metrics baseline in KV after every successful promotion so the next candidate is compared to the freshest bar, not the original training run.
- A
dry_runinput for CI preview runs that compute the verdict but skip the actual registry transition and KV writes.
Who it's for
- ML platform teams who own the MLflow registry and need a programmatic gate that their model training pipelines can call without writing bespoke Python.
- Data scientists who want to understand why a model was blocked before retraining — the drift report names the exact features and their PSI/KS values.
- Platform engineers standardising multi-team model promotion with a single configurable flow rather than per-team scripts.
Why orchestrate this with Kestra
MLflow provides no built-in promotion gate. The common alternative is a Python script in a GitHub Action that calls the MLflow client and exits non-zero on failure. That approach has no retry semantics (transient API errors become hard failures), no execution history, and no integration with a secrets backend. Kestra replaces the script with a declarative flow: the schedule, retries, KV baseline, Slack alerts, and HTTP calls are all first-class concerns with a full audit trail and no separate service to run.
The drift computation and metric regression live in a pinned Docker container
(python:3.12-slim) so the gate result is reproducible across environments.
Adding a second gate — for example a canary traffic split check — is a new
task in the same flow, not a new script in a new repo.
Prerequisites
- A running MLflow Tracking Server reachable from the Kestra worker.
- A model registered in the MLflow Model Registry with at least one version.
- The model training pipeline must write:
- A reference Parquet with training-time feature statistics to Kestra internal
storage (or any URI the worker can read as a file input). Column names
must match the production sample. Target/label columns should be named
y,label,target, or end with_label/_target— they are skipped automatically. - A production sample Parquet from the serving population.
- A reference Parquet with training-time feature statistics to Kestra internal
storage (or any URI the worker can read as a file input). Column names
must match the production sample. Target/label columns should be named
- Workers must be able to pull
python:3.12-slimfrom Docker Hub, or replacecontainerImagewith a prebuilt image that already haspandas pyarrow scipy.
Secrets
MLFLOW_GATE_WEBHOOK_KEY: key for the inbound webhook trigger (only needed when CI calls the gate via webhook rather than scheduled execution).SLACK_WEBHOOK_URL: Slack incoming webhook for pass/block/error alerts. Only needed whennotify_slackistrue.
If your MLflow server requires authentication (Databricks PAT, basic auth),
add the Authorization header to fetch_model_version and fetch_run_metrics
via headers, sourced from {{ secret('MLFLOW_TOKEN') }}.
Quick start
- Start MLflow and register a model:
pip install mlflow scikit-learn mlflow server --host 0.0.0.0 --port 5000 & python - <<'EOF' import mlflow, numpy as np from sklearn.datasets import make_classification from sklearn.linear_model import LogisticRegression from sklearn.metrics import f1_score X, y = make_classification(n_samples=1000, n_features=8, random_state=42) mlflow.set_experiment("fraud-detector") with mlflow.start_run(): model = LogisticRegression().fit(X, y) mlflow.log_metric("f1", f1_score(y, model.predict(X))) mlflow.log_metric("precision", 0.87) mlflow.log_metric("recall", 0.84) mlflow.sklearn.log_model(model, "model", registered_model_name="fraud-detector") EOF - Produce a reference Parquet (training data) and a production sample Parquet
and upload them to Kestra's internal storage via the namespace file UI or
the API. For a demo, generate two DataFrames with the same 8 columns:
one from the same
make_classificationcall (reference) and one with slightly shifted means (production). - Add the flow, set
mlflow_tracking_uritohttp://localhost:5000(or your host), and run it withcandidate_version: 1anddry_run: true. - The gate computes drift and metrics, logs the verdict, but does not call the
registry. Check
outputs.parse_drift_verdict.values.drifted_count. - Set
dry_run: falseand run again. The version is promoted to Production and the metrics baseline is stored in KV. - Re-run with a second version trained on heavily shifted data. The gate blocks and the execution is FAILED.
How to extend
- Per-model thresholds: store
psi_thresholdandmax_drifted_featuresper model name in the KV store and read them with{{ kv('mlflow_psi_' + inputs.model_name) | default(0.2) }}so different model families have different sensitivity. - Shadow score: before promoting, deploy the candidate as a Shadow stage and capture its predictions on live traffic for N minutes; add a task that reads those shadow metrics from a table and gates on calibration before the actual transition.
- Multi-model: wrap
fetch_model_version+compute_driftin aio.kestra.plugin.core.flow.Loopover a list of model names to gate an entire model family in one execution. - PagerDuty: replace the Slack task with
io.kestra.plugin.pagerduty.CreateIncidentwhen the gate blocks on a model that owns an SLA. - Schedule: add
io.kestra.plugin.core.trigger.Schedulewith a nightly cron to re-run the gate against the currently Production version — this catches population shift that happened after promotion without a new candidate.
Links
Orchestrate with Kestra
Orchestrate Slack with KestraShare this Blueprint
Create a Blueprint
Go to the developer platformMore AI Blueprints
More Data Blueprints
