Data Pipelines, Feature Pipelines, and Training Workflows

A notebook is a prototype. A production ML system is a pipeline. The difference is not just code quality - it is repeatability, auditability, and the ability to retrain next week with new data and trust the result. This module covers how data flows from raw sources to trained model artifacts inside a team, and the engineering practices that make that flow reliable.

The Data Pipeline Architecture

Every ML pipeline has three stages:

1. Ingestion: pull data from sources (databases, APIs, event streams, files).
2. Transformation: clean, aggregate, and compute features.
3. Loading: materialize the output for downstream consumption (training, serving, analytics).

These stages need orchestration - a scheduler that runs them in the right order, at the right time, with retries on failure.

pythonCopy
# Prefect: modern Python-native orchestration
from prefect import flow, task
from prefect.tasks import task_input_hash
from datetime import timedelta

@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(hours=1))
def extract_raw_events(date: str) -> pd.DataFrame:
    """Pull events for a given date from the warehouse."""
    return query_warehouse(f"SELECT * FROM events WHERE date = '{date}'")

@task
def compute_features(raw: pd.DataFrame) -> pd.DataFrame:
    """Apply feature transformations."""
    raw['days_since_signup'] = (pd.Timestamp.today() - raw['signup_date']).dt.days
    raw['log_spend'] = np.log1p(raw['total_spend'])
    return raw.dropna(subset=['label'])

@task
def validate_and_save(features: pd.DataFrame, output_path: str) -> str:
    """Validate schema and write to storage."""
    assert features.shape[0] > 1000, "Too few rows - pipeline may have failed"
    assert features['label'].isnull().sum() == 0, "Null labels found"
    features.to_parquet(output_path, index=False)
    return output_path

@flow(name="daily-feature-pipeline")
def feature_pipeline(date: str):
    raw = extract_raw_events(date)
    features = compute_features(raw)
    path = validate_and_save(features, f"gs://my-bucket/features/{date}/features.parquet")
    return path

The cache_key_fn=task_input_hash ensures that if the pipeline is re-run with the same inputs, the expensive extraction step is skipped. This is critical for large datasets.

Feature Stores: Sharing Features Across Teams

A feature store solves the train-serve skew problem and the duplication problem. Instead of each team recomputing the same features, a central store defines them once and serves them both to training jobs and production inference.

pythonCopy
# Feast: open-source feature store
from feast import FeatureStore

store = FeatureStore(repo_path="feature_repo/")

# Training: retrieve point-in-time correct features (no leakage)
training_df = store.get_historical_features(
    entity_df=entity_timestamps,  # user_id + event_timestamp pairs
    features=[
        "user_stats:purchase_count_7d",
        "user_stats:avg_session_duration",
        "user_stats:churn_risk_score",
    ]
).to_df()

# Serving: retrieve latest feature values for online inference
feature_vector = store.get_online_features(
    features=["user_stats:purchase_count_7d", "user_stats:avg_session_duration"],
    entity_rows=[{"user_id": "u_12345"}]
).to_dict()

The get_historical_features method performs a point-in-time join - it returns the feature values as they existed at each event_timestamp, preventing data leakage from future values. This is the most important property of a feature store.

Experiment Tracking and Artifact Management

Every training run should be logged so you can reproduce it and compare runs.

pythonCopy
import mlflow
import mlflow.sklearn

mlflow.set_experiment("churn_prediction_v2")

with mlflow.start_run(run_name="lgbm_baseline"):
    # Log hyperparameters
    mlflow.log_params({
        "n_estimators": 300,
        "learning_rate": 0.05,
        "max_depth": 6,
        "min_child_samples": 20,
    })

    # Train model
    model.fit(X_train, y_train, eval_set=[(X_val, y_val)])

    # Log metrics
    val_auc = roc_auc_score(y_val, model.predict_proba(X_val)[:, 1])
    mlflow.log_metrics({"val_auc": val_auc, "n_train": len(X_train)})

    # Log the model artifact
    mlflow.sklearn.log_model(model, "model",
        registered_model_name="churn_classifier")

    # Log data artifact (feature importance)
    importance_df.to_csv("/tmp/feature_importance.csv", index=False)
    mlflow.log_artifact("/tmp/feature_importance.csv")

print(f"Run logged: val_auc={val_auc:.4f}")

The model registry (registered_model_name) lets you version, stage, and promote models: Staging → Production → Archived. Every promotion is tracked and auditable.

Training Workflow Design Principles

Idempotency: running the pipeline twice with the same input should produce the same output. Avoid DATETIME.NOW() inside pipeline steps - pass the date as a parameter.

Parametric configuration: no hardcoded paths or hyperparameters in code. Use config files loaded at runtime.

pythonCopy
from dataclasses import dataclass
import yaml

@dataclass
class TrainingConfig:
    train_date: str
    val_date: str
    feature_version: str
    model_params: dict
    output_path: str

def load_config(path: str) -> TrainingConfig:
    with open(path) as f:
        raw = yaml.safe_load(f)
    return TrainingConfig(**raw)

Separation of concerns: extract, transform, validate, train, and evaluate should be separate steps with clean interfaces. This makes each step independently testable and re-runnable.

Data versioning: store training datasets alongside model artifacts. When a model behaves unexpectedly in production, you need to be able to reproduce the exact training data. Use DVC or MLflow's artifact logging for dataset snapshots.

Common Mistakes and Bad Instincts

Running the full pipeline on every retrain. Cache intermediate results. Re-extract raw data only when the source changed.

Not testing the pipeline end-to-end. Run the pipeline on a 1% data sample in CI to catch schema changes and transform failures before they happen in production.

Hardcoding the training cutoff date. The training cutoff must be a parameter, not a constant. Every time you retrain, you'll need a different cutoff. Hardcoding it means manually editing code on every retrain.

Not validating feature distributions against baseline. When a pipeline runs, compare key feature statistics (mean, null rate, cardinality) against a baseline snapshot. Unexpected shifts indicate upstream data issues before they corrupt a model.

Where to Go Next

• Module 25 (Model Serving) covers how trained models get deployed as APIs and inference services.
• Module 26 (MLOps CI/CD) covers how to automate testing and validation of these pipelines.
• Module 27 (Observability) covers monitoring the data and models that come out of these pipelines.