Orchestrate ML Pipelines with Prefect
Go from fragile training scripts to robust, observable ML workflows using Prefect. Schedule retraining, handle failures gracefully, and monitor every run.
Most ML workflows start as a single script: train.py. It works fine locally. Then someone asks for scheduled retraining. Then another person wants to know why last Tuesday's run failed. Then you need to retrain only when new data arrives. Suddenly your script has grown into a pile of cron jobs, shell scripts, and Slack alerts held together with duct tape.
Prefect solves exactly this problem. It turns your existing Python functions into observable, retryable, schedulable workflows — without requiring you to learn a new DSL or restructure your code. In this guide I take a monolithic training script and turn it into a Prefect flow with retries, caching, scheduled retraining, and conditional model promotion.
What Prefect Gives You
- Tasks: individual units of work, each with automatic retry logic
- Flows: compositions of tasks that Prefect orchestrates
- Deployments: scheduled or event-triggered runs of your flows
- UI: a built-in dashboard for run history, logs, and failure inspection
pip install prefect scikit-learn pandas joblib
From Training Script to Prefect Flow
Here is a typical ML training script before Prefect:
# Before: monolithic script
import pandas as pd
import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score
df = pd.read_csv('data/churn.csv')
X = df.drop('churn', axis=1)
y = df['churn']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
model = RandomForestClassifier(n_estimators=200, random_state=42)
model.fit(X_train, y_train)
f1 = f1_score(y_test, model.predict(X_test))
print(f"F1: {f1:.4f}")
joblib.dump(model, 'model.joblib')
If this fails at step 3, you have no idea where it failed, no retry logic, and no history. Now with Prefect:
# After: Prefect flow
import pandas as pd
import joblib
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import f1_score, classification_report
from prefect import flow, task, get_run_logger
from prefect.tasks import task_input_hash
from datetime import timedelta
@task(
retries=3,
retry_delay_seconds=10,
cache_key_fn=task_input_hash,
cache_expiration=timedelta(hours=1),
)
def load_data(path: str) -> pd.DataFrame:
logger = get_run_logger()
logger.info(f"Loading data from {path}")
df = pd.read_csv(path)
logger.info(f"Loaded {len(df)} rows, {df.shape[1]} columns")
return df
@task
def preprocess(df: pd.DataFrame) -> tuple[pd.DataFrame, pd.Series]:
logger = get_run_logger()
# Drop rows with missing target
df = df.dropna(subset=['churn'])
logger.info(f"After dropping nulls: {len(df)} rows")
X = df.drop('churn', axis=1)
y = df['churn'].astype(int)
return X, y
@task
def split_data(X: pd.DataFrame, y: pd.Series, test_size: float = 0.2):
return train_test_split(X, y, test_size=test_size, random_state=42, stratify=y)
@task
def train_model(X_train: pd.DataFrame, y_train: pd.Series, n_estimators: int = 200):
logger = get_run_logger()
logger.info(f"Training RandomForest with {n_estimators} trees")
model = RandomForestClassifier(n_estimators=n_estimators, random_state=42, n_jobs=-1)
model.fit(X_train, y_train)
return model
@task
def evaluate_model(model, X_test: pd.DataFrame, y_test: pd.Series) -> dict:
logger = get_run_logger()
y_pred = model.predict(X_test)
f1 = f1_score(y_test, y_pred, average='weighted')
acc = model.score(X_test, y_test)
logger.info(f"F1: {f1:.4f} | Accuracy: {acc:.4f}")
logger.info(f"\n{classification_report(y_test, y_pred)}")
return {'f1': f1, 'accuracy': acc}
@task
def save_model(model, metrics: dict, output_path: str = 'model.joblib'):
logger = get_run_logger()
joblib.dump({'model': model, 'metrics': metrics}, output_path)
logger.info(f"Model saved to {output_path} (F1={metrics['f1']:.4f})")
@flow(name="churn-training-pipeline")
def training_pipeline(
data_path: str = 'data/churn.csv',
n_estimators: int = 200,
model_output: str = 'model.joblib',
):
df = load_data(data_path)
X, y = preprocess(df)
X_train, X_test, y_train, y_test = split_data(X, y)
model = train_model(X_train, y_train, n_estimators)
metrics = evaluate_model(model, X_test, y_test)
save_model(model, metrics, model_output)
return metrics
if __name__ == '__main__':
result = training_pipeline()
print(f"Final metrics: {result}")
Run it directly:
python pipeline.py
Prefect logs each task individually, shows their status, and retries load_data up to 3 times if the file is temporarily unavailable.
Starting the Prefect UI
prefect server start
Open http://localhost:4200. You will see a dashboard with your recent flow runs, their duration, task-level status, and logs — all without extra configuration.
Caching Expensive Tasks
The cache_key_fn=task_input_hash argument on load_data's decorator means that if you run the flow twice with the same data_path, the second run will skip loading entirely and use the cached result. This is invaluable during development when you iterate on training code but the data has not changed.
# First run: loads from disk (slow)
training_pipeline(data_path='data/churn.csv')
# Second run within 1 hour: uses cache (instant)
training_pipeline(data_path='data/churn.csv')
Scheduled Retraining
To run your flow on a schedule, create a deployment with serve:
# deploy.py
from pipeline import training_pipeline
if __name__ == '__main__':
training_pipeline.serve(
name="weekly-churn-retraining",
cron="0 6 * * 1", # Every Monday at 6 AM UTC
parameters={
"data_path": "data/churn_latest.csv",
"n_estimators": 300,
},
)
python deploy.py
The serve command keeps a lightweight process alive that waits for the next scheduled trigger. You can also trigger it manually from the UI or CLI:
prefect deployment run 'churn-training-pipeline/weekly-churn-retraining'
Conditional Retraining Based on Metrics
A common pattern is to only promote a new model if it outperforms the current one:
@task
def load_current_metrics(model_path: str = 'model.joblib') -> dict:
import os
if not os.path.exists(model_path):
return {'f1': 0.0}
saved = joblib.load(model_path)
return saved.get('metrics', {'f1': 0.0})
@flow(name="conditional-retraining")
def conditional_training_pipeline(data_path: str = 'data/churn.csv'):
logger = get_run_logger()
df = load_data(data_path)
X, y = preprocess(df)
X_train, X_test, y_train, y_test = split_data(X, y)
model = train_model(X_train, y_train)
new_metrics = evaluate_model(model, X_test, y_test)
current_metrics = load_current_metrics()
if new_metrics['f1'] > current_metrics['f1'] + 0.005:
save_model(model, new_metrics)
logger.info(f"Model promoted: {new_metrics['f1']:.4f} > {current_metrics['f1']:.4f}")
else:
logger.info(f"Model not promoted: {new_metrics['f1']:.4f} <= {current_metrics['f1']:.4f}")
return new_metrics
If you version your models with MLflow, this promotion step maps naturally onto model registry stages — see Manage Your Machine Learning Lifecycle with MLflow in Python.
Handling Failures with State Hooks
Prefect provides first-class failure handling through retries, notifications, and state hooks:
from prefect import flow
def on_failure(flow, flow_run, state):
# Send alert: Slack, email, PagerDuty...
print(f"Flow {flow.name} failed: {state.message}")
@flow(
name="monitored-training",
on_failure=[on_failure],
)
def monitored_pipeline():
# ...
pass
Conclusion
Prefect transforms an ad-hoc Python script into a production-ready ML workflow in a few dozen lines. The same code that runs locally runs in production, with retries, caching, scheduling, and a UI to inspect every run. As your ML systems mature, the investment in orchestration pays dividends every time you need to debug why a model behaved unexpectedly — or why it silently did not retrain when it should have.
With retraining running reliably, you need to know when a model has started to degrade: ML Model Monitoring in Production with Evidently.
About the author
Rodrigo Arenas is a software architect and machine learning engineer in Medellín, Colombia. He builds products, platforms, and open-source software for AI, Data Engineering, and Machine Learning — creator of Ciaren, sklearn-genetic-opt, and PyWorkforce.
Recommended reading
Rodrigo Arenas — ML Engineer & Open Source Maintainer
Working on something similar? I can help you take it to production.