Data Pipelines, Labeling, and Training Workflows

A model is only as good as the data it was trained on, and data quality is a function of the pipeline that produces it. Ad-hoc scripts that pull data, munge it in notebooks, and write CSVs are fine for exploration. They are not fine for production training jobs that run weekly, need to be audited, and must be reproducible six months later when a model regresses.

This module covers the engineering infrastructure for reliable ML data pipelines and annotation workflows.

What ML Data Pipelines Must Guarantee

Production training data pipelines must satisfy constraints that analytics pipelines do not:

1. Point-in-time correctness: features computed as of the training example's timestamp, not today's values
2. No target leakage: features cannot use information that would not have been available at prediction time
3. Reproducibility: re-running the pipeline on the same date range must produce the same outputs
4. Versioning: different training runs must reference specific, immutable snapshots

These constraints push toward a different architecture than a standard ETL pipeline.

Orchestration with Prefect

Prefect provides task-level retry, caching, observability, and scheduling for Python-first data workflows:

pythonCopy
from prefect import flow, task
from prefect.tasks import task_input_hash
from datetime import timedelta
import pandas as pd

@task(cache_key_fn=task_input_hash, cache_expiration=timedelta(days=1))
def extract_raw_events(start_date: str, end_date: str) -> pd.DataFrame:
    """Pull raw user events from the warehouse. Cached to avoid redundant queries."""
    query = f"""
        SELECT user_id, event_type, properties, created_at
        FROM events
        WHERE created_at BETWEEN '{start_date}' AND '{end_date}'
        ORDER BY user_id, created_at
    """
    return pd.read_sql(query, con=get_db_connection())

@task
def compute_features(events: pd.DataFrame, label_date: str) -> pd.DataFrame:
    """Compute features using only data available before label_date."""
    events = events[events["created_at"] < label_date]
    features = (
        events.groupby("user_id")
        .agg(
            n_events=("event_type", "count"),
            n_distinct_events=("event_type", "nunique"),
            days_active=("created_at", lambda x: x.dt.normalize().nunique()),
        )
        .reset_index()
    )
    return features

@task
def attach_labels(features: pd.DataFrame, label_date: str) -> pd.DataFrame:
    """Join labels computed strictly after the feature window."""
    labels = pd.read_sql(
        f"SELECT user_id, converted FROM conversions WHERE conversion_date = '{label_date}'",
        con=get_db_connection()
    )
    return features.merge(labels, on="user_id", how="left").fillna({"converted": 0})

@flow(name="training-data-pipeline")
def build_training_dataset(feature_start: str, feature_end: str, label_date: str) -> str:
    raw_events = extract_raw_events(feature_start, feature_end)
    features = compute_features(raw_events, label_date)
    labeled = attach_labels(features, label_date)
    output_path = f"s3://ml-data/training/{label_date}/features.parquet"
    labeled.to_parquet(output_path)
    return output_path

The cache_key_fn=task_input_hash makes each task output content-addressable - re-running with the same inputs returns the cached result, making pipelines cheap to re-run during debugging.

Feature Stores: Solving the Train-Serve Skew Problem

The most costly class of production ML bug is train-serve skew: features computed differently during training and serving, producing a distribution mismatch invisible in offline evaluation.

A feature store solves this with:

• Online store: low-latency retrieval of latest feature values per entity (Redis, DynamoDB)
• Offline store: historical feature values for training (S3/GCS + Parquet)
• Point-in-time join: materializes feature values as of any historical timestamp for training

pythonCopy
from feast import FeatureStore, FeatureService
from datetime import datetime
import pandas as pd

store = FeatureStore(repo_path="./feature_repo")

# Point-in-time join for training data
entity_df = pd.DataFrame({
    "user_id": [1001, 1002, 1003],
    "event_timestamp": [
        datetime(2024, 3, 1),
        datetime(2024, 3, 5),
        datetime(2024, 3, 10),
    ]
})

training_df = store.get_historical_features(
    entity_df=entity_df,
    features=[
        "user_features:n_events_7d",
        "user_features:n_distinct_events_7d",
        "user_features:days_active_30d",
        "user_features:total_spend_90d",
    ],
).to_df()

# At serving time - same feature definitions, millisecond latency
online_features = store.get_online_features(
    features=["user_features:n_events_7d", "user_features:n_distinct_events_7d"],
    entity_rows=[{"user_id": 1001}],
).to_dict()

Even without a dedicated feature store, the principle is the same: write feature computation logic once, call it from both the training pipeline and the serving layer.

Data Labeling Workflows

Most production ML requires labeled data. Labeling is engineering work:

pythonCopy
# Label schema - version it in code
from pydantic import BaseModel
from enum import Enum

class SentimentLabel(str, Enum):
    POSITIVE = "positive"
    NEGATIVE = "negative"
    NEUTRAL = "neutral"
    AMBIGUOUS = "ambiguous"

class LabelRecord(BaseModel):
    text_id: str
    text: str
    label: SentimentLabel
    annotator_id: str
    annotated_at: str
    confidence: float  # 1 = certain, 0.5 = unsure

# Inter-annotator agreement (Cohen's Kappa)
from sklearn.metrics import cohen_kappa_score

labels_a = ["positive", "negative", "neutral", "positive"]
labels_b = ["positive", "negative", "positive", "positive"]
kappa = cohen_kappa_score(labels_a, labels_b)
print(f"Cohen's Kappa: {kappa:.3f}")  # > 0.6 = acceptable; > 0.8 = excellent

Labeling quality practices:

• Dual annotation: every example labeled by two annotators independently; disagreements resolved by a third
• Annotator calibration: run annotators on a gold set first to detect systematic biases
• Stratified sampling: ensure rare classes and edge cases are proportionally represented
• Reject ambiguous examples: low-confidence labels hurt training quality more than they help

Training Job Infrastructure

For non-trivial model training, wrap your PyTorch training in a reproducible job:

pythonCopy
import mlflow
import hashlib
import json

def train_and_track(config: dict, dataset_path: str) -> str:
    """Returns the MLflow run ID for reproducibility."""
    config_hash = hashlib.sha256(json.dumps(config, sort_keys=True).encode()).hexdigest()[:8]

    with mlflow.start_run(run_name=f"training-{config_hash}") as run:
        mlflow.log_params(config)
        mlflow.log_param("dataset_path", dataset_path)
        mlflow.log_param("dataset_hash", compute_file_hash(dataset_path))

        model = build_model(**config)
        train_losses, val_losses = [], []

        for epoch in range(config["n_epochs"]):
            train_loss = run_epoch(model, "train")
            val_loss = run_epoch(model, "val")
            train_losses.append(train_loss)
            val_losses.append(val_loss)
            mlflow.log_metrics({"train_loss": train_loss, "val_loss": val_loss}, step=epoch)

        mlflow.pytorch.log_model(model, artifact_path="model")
        return run.info.run_id

MLflow tracking gives you: reproducible experiment runs, hyperparameter comparison, metric curves, and model artifacts linked to the exact training data and code that produced them.

Common Mistakes and Bad Instincts

Using df.sort_values and df.iloc[:n] for train/test splits on time-series data. Temporal data must be split chronologically - never randomly. Random splits allow future data to leak into the training set, producing dramatically overoptimistic offline metrics.

Not versioning data alongside model artifacts. An MLflow run that logs the model but not the training dataset SHA is unreproducible. When a model regresses, you cannot tell whether the issue was the model, the data, or the features.

Running labeling at the wrong granularity. Labeling full documents when you need sentence-level labels produces noisy training data. Match the label granularity to the prediction granularity.

Where to Go Next

• serving-models-and-llm-systems-in-production: deploy the models trained by these pipelines into low-latency serving infrastructure
• mlops-and-cicd-for-ml-teams: automate the pipeline with CI/CD so retraining triggers on data drift or a code change
• observability-drift-feedback-loops-and-llm-evals: detect when production data has drifted from your training distribution