from prefect import task, flow
import mlflow
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score

@task(retries=2, retry_delay_seconds=30)
def extract_data(data_path: str):
    """Task to load data from a source."""
    import pandas as pd
    df = pd.read_parquet(data_path)
    return df

@task
def preprocess_data(df):
    """Task for validation, splitting, and feature engineering."""
    # Example validation: ensure target column exists
    assert 'target' in df.columns, "Target column missing"
    X = df.drop('target', axis=1)
    y = df['target']
    # Perform train-test split
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    return X_train, X_test, y_train, y_test

@task(retries=2)
def train_model(X_train, y_train, parameters: dict):
    """Isolated training task with MLflow logging."""
    with mlflow.start_run():
        model = RandomForestClassifier(**parameters)
        model.fit(X_train, y_train)
        # Log parameters, metrics, and model
        mlflow.log_params(parameters)
        train_accuracy = model.score(X_train, y_train)
        mlflow.log_metric("train_accuracy", train_accuracy)
        mlflow.sklearn.log_model(model, "model")
    return model

@task
def evaluate_model(model, X_test, y_test):
    """Task to evaluate model performance."""
    predictions = model.predict(X_test)
    accuracy = accuracy_score(y_test, predictions)
    f1 = f1_score(y_test, predictions, average='weighted')
    print(f"Test Accuracy: {accuracy:.4f}, Test F1-Score: {f1:.4f}")
    return {"accuracy": accuracy, "f1_score": f1}

@flow(name="training_pipeline")
def model_training_flow(data_path: str):
    """Main Prefect flow orchestrating the pipeline."""
    df = extract_data(data_path)
    X_train, X_test, y_train, y_test = preprocess_data(df)
    model = train_model(X_train, y_train, parameters={"n_estimators": 100, "max_depth": 10})
    metrics = evaluate_model(model, X_test, y_test)

# Execute the flow
if __name__ == "__main__":
    model_training_flow("s3://bucket/data/training.parquet")

# feature_engineering.py
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer

def create_features(raw_data_path: str) -> pd.DataFrame:
    """
    Transforms raw data into a feature set.
    Args:
        raw_data_path: Path to the raw data file.
    Returns:
        A pandas DataFrame with engineered features.
    """
    df = pd.read_parquet(raw_data_path)

    # Create new features
    df['avg_transaction_value'] = df['total_spent'] / (df['transaction_count'] + 1e-5)  # Avoid division by zero
    df['customer_tenure_days'] = (pd.to_datetime('today') - pd.to_datetime(df['first_purchase_date'])).dt.days

    # Define preprocessing for numerical and categorical columns
    numeric_features = ['total_spent', 'transaction_count', 'avg_transaction_value']
    categorical_features = ['category', 'region']

    numeric_transformer = StandardScaler()
    categorical_transformer = OneHotEncoder(handle_unknown='ignore', sparse_output=False)

    preprocessor = ColumnTransformer(
        transformers=[
            ('num', numeric_transformer, numeric_features),
            ('cat', categorical_transformer, categorical_features)
        ])

    # Fit and transform the data
    feature_array = preprocessor.fit_transform(df)
    # Get feature names after one-hot encoding
    ohe_categories = preprocessor.named_transformers_['cat'].categories_
    ohe_feature_names = []
    for col, cats in zip(categorical_features, ohe_categories):
        ohe_feature_names.extend([f"{col}_{cat}" for cat in cats])

    all_feature_names = numeric_features + ohe_feature_names
    features_df = pd.DataFrame(feature_array, columns=all_feature_names)
    features_df['customer_id'] = df['customer_id'].values  # Keep identifier

    # Save the preprocessor for consistency in inference
    import joblib
    joblib.dump(preprocessor, 'models/feature_preprocessor.joblib')

    return features_df

# Example execution
if __name__ == "__main__":
    engineered_df = create_features("data/raw/customers.parquet")
    engineered_df.to_parquet("data/processed/features.parquet", index=False)

# Agile: Iterative component for a new Kafka stream using PySpark Structured Streaming
from pyspark.sql import SparkSession
from pyspark.sql.functions import from_json, col, current_timestamp
from pyspark.sql.types import StructType, StructField, StringType, DoubleType, TimestampType

def ingest_kafka_stream(topic: str, checkpoint_path: str, output_table: str):
    """
    Ingests and transforms a Kafka stream, writing to a Delta Lake table.
    This is a modular component that can be developed and tested in one sprint.
    """
    spark = SparkSession.builder \
        .appName(f"KafkaIngest_{topic}") \
        .config("spark.sql.extensions", "io.delta.sql.DeltaSparkSessionExtension") \
        .config("spark.sql.catalog.spark_catalog", "org.apache.spark.sql.delta.catalog.DeltaCatalog") \
        .getOrCreate()

    # Define the expected schema for the JSON payload
    event_schema = StructType([
        StructField("user_id", StringType(), False),
        StructField("event_type", StringType(), False),
        StructField("value", DoubleType(), True),
        StructField("timestamp", TimestampType(), True)
    ])

    # Read stream from Kafka
    raw_df = (spark.readStream
              .format("kafka")
              .option("kafka.bootstrap.servers", "kafka-broker:9092")
              .option("subscribe", topic)
              .option("startingOffsets", "latest")
              .load()
              .selectExpr("CAST(value AS STRING) as json_string"))

    # Parse JSON and add ingestion timestamp
    parsed_df = (raw_df
                 .select(from_json(col("json_string"), event_schema).alias("data"))
                 .select("data.*")
                 .withColumn("ingestion_ts", current_timestamp()))

    # Write stream to Delta Lake with merge schema for flexibility
    query = (parsed_df.writeStream
             .format("delta")
             .outputMode("append")
             .option("checkpointLocation", f"{checkpoint_path}/{topic}")
             .option("mergeSchema", "true")  # Allows schema evolution across sprints
             .trigger(processingTime='30 seconds')
             .start(f"/delta/{output_table}"))

    return query  # Query object can be managed and stopped

# Example of invoking the component
if __name__ == "__main__":
    query = ingest_kafka_stream("user_events", "/checkpoints/kafka", "user_events_delta")
    query.awaitTermination()  # In production, this would be managed by an orchestrator

import pandas as pd
import great_expectations as ge
from airflow.decorators import task
from datetime import datetime

@task
def validate_raw_data(execution_date: str):
    """
    Airflow task to validate incoming raw data.
    """
    # Construct data path based on execution date (common pattern)
    data_path = f"s3://raw-data-bucket/date={execution_date}/data.parquet"
    df = pd.read_parquet(data_path)

    # Create a Great Expectations dataset
    dataset = ge.from_pandas(df)

    # Define a suite of expectations
    expectations = [
        dataset.expect_column_to_exist("user_id"),
        dataset.expect_column_values_to_not_be_null("user_id"),
        dataset.expect_column_values_to_be_between("purchase_amount", min_value=0, max_value=10000),
        dataset.expect_table_row_count_to_be_between(min_value=1000, max_value=1000000)
    ]

    # Validate
    validation_result = dataset.validate(expectations=expectations)

    if not validation_result.success:
        # Log failure details and raise an exception to fail the task
        failed_expectations = [e for e in validation_result.results if not e.success]
        error_msg = f"Data validation failed for {data_path}. Failures: {failed_expectations}"
        raise ValueError(error_msg)

    # If validation passes, push the path to XCom for the next task
    return data_path

@task
def compute_features(validated_data_path: str, execution_date: str):
    """
    Task to generate features from validated data.
    Uses a saved preprocessor for consistency.
    """
    import joblib
    import hashlib

    df = pd.read_parquet(validated_data_path)

    # Load the preprocessor saved during development (versioned)
    preprocessor = joblib.load('dags/utils/feature_preprocessor_v1.joblib')

    # Apply transformation
    feature_array = preprocessor.transform(df)

    # Get feature names
    feature_names = preprocessor.get_feature_names_out()

    features_df = pd.DataFrame(feature_array, columns=feature_names)
    features_df['user_id'] = df['user_id'].values
    features_df['effective_date'] = execution_date

    # Create a deterministic path for the features
    # Hashing the data path ensures unique but reproducible output location
    path_hash = hashlib.md5(validated_data_path.encode()).hexdigest()[:8]
    feature_path = f"s3://feature-bucket/run={execution_date}/features_{path_hash}.parquet"

    features_df.to_parquet(feature_path, index=False)
    return feature_path

@task
def train_model(feature_path: str, experiment_name: str = "churn_prediction"):
    """
    Training task that logs everything to MLflow.
    """
    import mlflow
    from sklearn.ensemble import RandomForestClassifier
    from sklearn.model_selection import train_test_split
    from sklearn.metrics import classification_report, roc_auc_score

    mlflow.set_tracking_uri("http://mlflow-server:5000")
    mlflow.set_experiment(experiment_name)

    df = pd.read_parquet(feature_path)
    X = df.drop(['user_id', 'effective_date', 'target'], axis=1, errors='ignore')
    y = df['target']

    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42)

    with mlflow.start_run(run_name=f"training_run_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}"):
        # Define and train model
        model = RandomForestClassifier(n_estimators=150, max_depth=12, random_state=42)
        model.fit(X_train, y_train)

        # Predict and evaluate
        y_pred = model.predict(X_val)
        y_pred_proba = model.predict_proba(X_val)[:, 1]

        auc = roc_auc_score(y_val, y_pred_proba)
        report_dict = classification_report(y_val, y_pred, output_dict=True)

        # Log parameters, metrics, and model
        mlflow.log_params({"n_estimators": 150, "max_depth": 12})
        mlflow.log_metrics({"auc": auc, "accuracy": report_dict['accuracy']})
        mlflow.sklearn.log_model(model, "model", registered_model_name="ChurnClassifier")

        # Log a small artifact, like the classification report as JSON
        import json
        with open("classification_report.json", "w") as f:
            json.dump(report_dict, f)
        mlflow.log_artifact("classification_report.json")

    run_id = mlflow.active_run().info.run_id
    mlflow.end_run()
    return run_id

# Initialize DVC in your project
$ dvc init

# Add and track a large dataset
$ dvc add data/train.csv
$ git add data/train.csv.dvc .gitignore

# Define and run a pipeline stage
$ dvc run -n train_model \
          -d src/train.py -d data/train.csv \
          -o models/model.pkl -M metrics/accuracy.json \
          python src/train.py --data data/train.csv --output models/model.pkl

# Version the results
$ git add dvc.yaml dvc.lock metrics/accuracy.json
$ git commit -m "Experiment #42: RandomForest with new features, accuracy: 0.92"
$ git tag -a "v1.2-model" -m "Model version 1.2"
$ dvc push  # Push data and models to remote storage

import pandera as pa
from pandera import DataFrameSchema, Column, Check, CheckResult
import pandas as pd
import numpy as np

# Define a rigorous schema with custom checks
schema = DataFrameSchema({
    "user_id": Column(int, checks=[
        Check.greater_than(0),
        Check(lambda s: s.is_unique, element_wise=False, error="user_id must be unique")
    ], nullable=False),
    "purchase_amount": Column(float, checks=[
        Check.in_range(0, 10000, error="Amount out of expected range"),
        Check(lambda x: np.isfinite(x), element_wise=True, error="Amount must be finite")
    ]),
    "timestamp": Column(pa.DateTime, checks=[
        Check(lambda s: s.dt.year >= 2020, element_wise=False, error="Date too old")
    ]),
    "category": Column(str, checks=[
        Check.isin(["electronics", "clothing", "home", "books"], error="Invalid category")
    ])
})

# Validate a DataFrame
try:
    validated_df = schema.validate(raw_df, lazy=True)  # Lazy collects all errors
    print("✅ Data validation passed.")
except pa.errors.SchemaErrors as err:
    print("❌ Data validation failed with the following errors:")
    print(err.failure_cases)  # Detailed dataframe of failures
    raise  # Fail the pipeline

# Example: Defining and materializing features with Feast
from feast import FeatureStore, Entity, ValueType, FeatureView, Field
from feast.types import Float32, Int64
from datetime import timedelta
from feast.infra.offline_stores.file_source import FileSource

# Define an entity
user = Entity(name="user", join_keys=["user_id"])

# Define a data source
user_stats_source = FileSource(
    path="data/user_stats.parquet",
    timestamp_field="event_timestamp",
)

# Define a Feature View
user_stats_fv = FeatureView(
    name="user_statistics",
    entities=[user],
    ttl=timedelta(days=365),
    schema=[
        Field(name="avg_transaction_7d", dtype=Float32),
        Field(name="transaction_count_30d", dtype=Int64),
    ],
    source=user_stats_source,
    online=True  # Make available for real-time serving
)

# Apply definitions to the store
store = FeatureStore(repo_path=".")
store.apply([user, user_stats_source, user_stats_fv])

# Materialize features to the online store for the last 30 days
store.materialize(start_date=datetime.now() - timedelta(days=30), end_date=datetime.now())

import mlflow.pyfunc
import json

class CustomModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, sklearn_model):
        self.model = sklearn_model

    def predict(self, context, model_input):
        return self.model.predict_proba(model_input)[:, 1]

# Log the model with a signature defining input/output schema
from mlflow.models.signature import infer_signature

signature = infer_signature(X_train.head(), model.predict_proba(X_train.head())[:, 1])

with mlflow.start_run():
    mlflow.log_params(model.get_params())
    mlflow.log_metric("auc", auc_score)

    # Log the model with its custom class and environment
    mlflow.pyfunc.log_model(
        artifact_path="propensity_model",
        python_model=CustomModelWrapper(model),
        signature=signature,
        registered_model_name="Propensity_Score",
        conda_env="conda.yaml"
    )

# In the registry, you can now transition this version to "Staging" or "Production"
client = mlflow.tracking.MlflowClient()
latest_versions = client.get_latest_versions("Propensity_Score")
client.transition_model_version_stage(
    name="Propensity_Score",
    version=latest_versions[0].version,
    stage="Staging"
)

This creates a versioned model in a central registry, accessible for deployment or comparison.

# Initialize DVC with remote storage (e.g., Amazon S3)
$ dvc init
$ dvc remote add -d myremote s3://my-dvc-bucket/path

# Add a dataset and commit its metadata
$ dvc add data/raw/training.csv
$ git add data/raw/training.csv.dvc .gitignore
$ git commit -m "Track v1.2 of training dataset"

# Push the actual data files to remote storage
$ dvc push

# Later, to retrieve a specific version:
$ git checkout <commit_hash>  # This gets the correct .dvc file
$ dvc pull data/raw/training.csv  # This retrieves the correct data

# Use a specific, lightweight base image for consistency
FROM python:3.9-slim-buster

# Set environment variables to avoid Python buffering and for reproducibility
ENV PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    PIP_NO_CACHE_DIR=1

WORKDIR /app

# Copy dependency file first for better layer caching
COPY requirements.txt .

# Install dependencies precisely
RUN pip install --no-cache-dir --upgrade pip && \
    pip install --no-cache-dir -r requirements.txt --no-deps

# Copy the rest of the application code
COPY src/ ./src/
COPY configs/ ./configs/

# Define a non-root user for security
RUN useradd -m -u 1000 appuser && chown -R appuser:appuser /app
USER appuser

# Health check for long-running training jobs (optional)
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
  CMD python -c "import sys; sys.exit(0)" || exit 1

# Entry point for the training script, allowing parameter override
ENTRYPOINT ["python", "src/train.py"]
CMD ["--config", "configs/default.yaml"]

# Define a reproducible pipeline in dvc.yaml
# dvc.yaml
stages:
  prepare:
    cmd: python src/prepare.py --input data/raw --output data/prepared
    deps:
      - src/prepare.py
      - data/raw
    outs:
      - data/prepared
  train:
    cmd: python src/train.py --data data/prepared --model models/model.pkl
    deps:
      - src/train.py
      - data/prepared
    outs:
      - models/model.pkl
    metrics:
      - metrics/accuracy.json:
          cache: false  # Metrics are not cached, always reproduced

# Run the pipeline
$ dvc repro
# Push data and models
$ dvc push
# Commit the pipeline state (dvc.lock) to Git
$ git add dvc.lock metrics/accuracy.json
$ git commit -m "Run experiment #45: Updated feature set"

# Use an official Python runtime as a base image with a pinned version
FROM python:3.9.16-slim

# Set the working directory
WORKDIR /app

# Install system dependencies if needed (e.g., for scikit-learn or lightgbm)
RUN apt-get update && apt-get install -y --no-install-recommends \
    gcc g++ \
    && rm -rf /var/lib/apt/lists/*

# Copy the requirements file and install Python packages with exact versions
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy the training script and data
COPY train_model.py .
COPY data.csv .

# Define environment variables for configuration
ENV MODEL_NAME="random_forest" \
    RANDOM_SEED=42

# Run the training script
CMD ["python", "train_model.py", "--model", "${MODEL_NAME}", "--seed", "${RANDOM_SEED}"]

apiVersion: batch/v1
kind: Job
metadata:
  name: model-training-job-20231026
  labels:
    app: model-training
    experiment-id: "exp-42"
spec:
  ttlSecondsAfterFinished: 86400  # Auto-delete job and pods after 24 hours
  backoffLimit: 2  # Number of retries before marking as failed
  template:
    spec:
      containers:
      - name: trainer
        image: model-trainer:v1  # Our containerized training environment
        imagePullPolicy: IfNotPresent
        env:
        - name: MODEL_NAME
          value: "random_forest"
        - name: RANDOM_SEED
          value: "42"
        resources:
          requests:
            memory: "4Gi"
            cpu: "2"
          limits:
            memory: "8Gi"
            cpu: "4"
        volumeMounts:
        - name: data-volume
          mountPath: /app/data
        - name: output-volume
          mountPath: /app/output
      restartPolicy: Never
      volumes:
      - name: data-volume
        persistentVolumeClaim:
          claimName: training-data-pvc
      - name: output-volume
        persistentVolumeClaim:
          claimName: model-output-pvc

# Example Argo Workflow snippet for a multi-step pipeline
apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  generateName: ml-pipeline-
spec:
  entrypoint: ml-pipeline
  templates:
  - name: ml-pipeline
    dag:
      tasks:
      - name: extract-data
        template: extract-data-template
      - name: validate-data
        template: validate-data-template
        dependencies: [extract-data]
        arguments:
          artifacts:
          - name: raw-data
            from: "{{tasks.extract-data.outputs.artifacts.raw-data}}"
      - name: train-model
        template: train-model-template
        dependencies: [validate-data]
        arguments:
          artifacts:
          - name: cleaned-data
            from: "{{tasks.validate-data.outputs.artifacts.cleaned-data}}"

name: Model Training Pipeline
on:
  push:
    branches: [ main ]
  pull_request:
    branches: [ main ]
  schedule:
    # Optional: Schedule a nightly retraining job
    - cron: '0 2 * * *'

jobs:
  train-and-evaluate:
    runs-on: ubuntu-latest
    container:
      # Use a consistent, pre-built Docker image with all dependencies
      image: ghcr.io/org/ml-training:3.9-base
      credentials:
        username: ${{ github.actor }}
        password: ${{ secrets.GITHUB_TOKEN }}
    env:
      MLFLOW_TRACKING_URI: ${{ secrets.MLFLOW_TRACKING_URI }}
      DVC_REMOTE_URL: ${{ secrets.DVC_REMOTE_URL }}

    steps:
      - name: Checkout code and data
        uses: actions/checkout@v3
        with:
          fetch-depth: 0 # Fetch all history for DVC

      - name: Setup DVC
        uses: iterative/setup-dvc@v1
        with:
          dvc_version: '2.40.0'

      - name: Pull versioned data
        run: |
          dvc remote modify --local myremote access_key_id ${{ secrets.AWS_ACCESS_KEY_ID }}
          dvc remote modify --local myremote secret_access_key ${{ secrets.AWS_SECRET_ACCESS_KEY }}
          dvc pull -r myremote

      - name: Run unit tests
        run: pytest tests/unit/ -v

      - name: Train Model
        id: train
        run: |
          python src/train.py \
            --data-path data/processed/ \
            --model-output models/model_${{ github.sha }}.pkl \
            --experiment-name "CI_${{ github.ref_name }}"

      - name: Evaluate Model
        id: evaluate
        run: |
          python src/evaluate.py \
            --model-path models/model_${{ github.sha }}.pkl \
            --test-data data/processed/test.parquet \
            --output metrics/metrics_${{ github.sha }}.json

      - name: Check Performance Gate
        id: gate
        run: |
          python scripts/check_metrics.py \
            --current metrics/metrics_${{ github.sha }}.json \
            --baseline metrics/metrics_baseline.json \
            --threshold accuracy 0.01  # Require at least 1% improvement
        continue-on-error: true # Continue to store metrics even if gate fails

      - name: Log Metrics to MLflow
        if: always() # Run this step even if previous steps fail
        run: |
          python scripts/log_run.py \
            --commit-sha ${{ github.sha }} \
            --metrics-file metrics/metrics_${{ github.sha }}.json \
            --gate-status ${{ steps.gate.outcome }}

      - name: Register and Deploy (on success)
        if: steps.gate.outcome == 'success'
        run: |
          python scripts/promote_model.py \
            --model-path models/model_${{ github.sha }}.pkl \
            --stage "Staging"
          echo "Model promoted to Staging. Manual approval needed for Production."

# scripts/check_metrics.py
import json
import argparse
import sys

def check_improvement(current_metrics, baseline_metrics, threshold_dict):
    """
    Compares current metrics against baseline.
    threshold_dict: {'metric_name': min_absolute_improvement}
    """
    print("🔍 Performance Gate Check")
    all_passed = True

    for metric_name, min_improvement in threshold_dict.items():
        if metric_name not in current_metrics or metric_name not in baseline_metrics:
            print(f"  ⚠️  Metric '{metric_name}' not found in one of the reports.")
            continue

        current_val = current_metrics[metric_name]
        baseline_val = baseline_metrics[metric_name]
        improvement = current_val - baseline_val

        if improvement >= min_improvement:
            print(f"  ✅ {metric_name}: {current_val:.4f} (Baseline: {baseline_val:.4f}, Δ={improvement:+.4f})")
        else:
            print(f"  ❌ {metric_name}: {current_val:.4f} (Baseline: {baseline_val:.4f}, Δ={improvement:+.4f}) - Failed.")
            all_passed = False

    # Additional check for fairness/bias metrics
    if 'disparate_impact' in current_metrics:
        di = current_metrics['disparate_impact']
        if 0.8 <= di <= 1.25:  # Common fairness threshold
            print(f"  ✅ Disparate Impact: {di:.3f} (within acceptable range)")
        else:
            print(f"  ❌ Disparate Impact: {di:.3f} (outside acceptable range 0.8-1.25)")
            all_passed = False

    return all_passed

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--current', required=True, help='Path to current metrics JSON')
    parser.add_argument('--baseline', required=True, help='Path to baseline metrics JSON')
    parser.add_argument('--threshold', nargs=2, action='append', metavar=('METRIC', 'VALUE'),
                        help='Thresholds, e.g., --threshold accuracy 0.01')
    args = parser.parse_args()

    with open(args.current) as f:
        current = json.load(f)
    with open(args.baseline) as f:
        baseline = json.load(f)

    thresholds = {k: float(v) for k, v in (args.threshold or [])}

    passed = check_improvement(current, baseline, thresholds)
    sys.exit(0 if passed else 1)

import great_expectations as ge
from great_expectations.core.batch import RuntimeBatchRequest
from great_expectations.checkpoint import SimpleCheckpoint
import pandas as pd

# Initialize Data Context
context = ge.get_context()

# Define a checkpoint for new incoming data
checkpoint_name = "churn_data_validation"
expectation_suite_name = "churn_data_suite"

# Create a batch request for new data (could be triggered by a new file arrival)
batch_request = RuntimeBatchRequest(
    datasource_name="my_datasource",
    data_connector_name="default_runtime_data_connector",
    data_asset_name="new_churn_data",
    runtime_parameters={"batch_data": new_data_df},  # new_data_df is a pandas DataFrame
    batch_identifiers={"run_id": f"run_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}"},
)

# Run the checkpoint
results = context.run_checkpoint(
    checkpoint_name=checkpoint_name,
    validations=[
        {
            "batch_request": batch_request,
            "expectation_suite_name": expectation_suite_name,
        }
    ],
    run_name=f"validation_run_{pd.Timestamp.now().isoformat()}",
)

# Check results and act
if not results["success"]:
    # Send alert (e.g., to Slack, PagerDuty) and fail the pipeline
    send_alert(
        f"Data validation failed for run {results['run_id']}. "
        f"Please check Data Docs at {results['data_docs_sites']}"
    )
    raise DataValidationError("New data failed quality checks. Pipeline halted.")
else:
    print("✅ Data validation passed. Proceeding to model retraining.")
    # Proceed to next pipeline step, e.g., log the validated data path
    validated_data_path = f"s3://validated-data/{results['run_id']}.parquet"
    new_data_df.to_parquet(validated_data_path)
    return validated_data_path

*Measurable Benefit:* Catches data drift and schema issues early, preventing garbage-in-garbage-out scenarios and saving countless hours of debugging.

import pandas as pd
import mlflow
from datetime import datetime
import matplotlib.pyplot as plt

def analyze_experiment_velocity(experiment_id: str):
    """
    Analyzes cycle time and other metrics for a given MLflow experiment.
    """
    client = mlflow.tracking.MlflowClient()
    runs = client.search_runs(experiment_ids=[experiment_id])

    metrics_data = []
    for run in runs:
        # Extract timestamps and relevant tags/metadata
        start_time = pd.to_datetime(run.info.start_time, unit='ms')
        end_time = pd.to_datetime(run.info.end_time, unit='ms') if run.info.end_time else pd.NaT

        cycle_time = (end_time - start_time).total_seconds() / 3600 if pd.notna(end_time) else None  # in hours

        run_data = {
            'run_id': run.info.run_id,
            'start_time': start_time,
            'end_time': end_time,
            'cycle_time_hours': cycle_time,
            'status': run.info.status,
            'user': run.data.tags.get('mlflow.user', 'unknown'),
            'accuracy': run.data.metrics.get('accuracy', None),
            'experiment_name': run.data.tags.get('mlflow.experiment.name', '')
        }
        metrics_data.append(run_data)

    df = pd.DataFrame(metrics_data)

    # Calculate aggregate statistics
    if not df.empty:
        successful_runs = df[df['status'] == 'FINISHED']
        avg_cycle_time = successful_runs['cycle_time_hours'].mean()
        median_cycle_time = successful_runs['cycle_time_hours'].median()
        deployment_freq = df['start_time'].dt.date.nunique()  # Unique days with deployments

        print(f"📊 Experiment Velocity Report for {experiment_id}")
        print(f"   Average Cycle Time: {avg_cycle_time:.2f} hours")
        print(f"   Median Cycle Time: {median_cycle_time:.2f} hours")
        print(f"   Deployment Frequency: ~{deployment_freq} days with activity in period")

        # Visualize cycle time trend
        plt.figure(figsize=(10, 5))
        plt.scatter(successful_runs['start_time'], successful_runs['cycle_time_hours'], alpha=0.7)
        plt.xlabel('Start Time')
        plt.ylabel('Cycle Time (hours)')
        plt.title(f'Cycle Time Trend for Experiment: {experiment_id}')
        plt.xticks(rotation=45)
        plt.tight_layout()
        plt.savefig(f'cycle_time_trend_{experiment_id}.png')
        plt.show()

    return df

# Example usage
df_runs = analyze_experiment_velocity("123")

# Example: Analyzing A/B test results for a model update
import numpy as np
import scipy.stats as stats

def analyze_ab_test(control_conversions, control_total,
                    treatment_conversions, treatment_total,
                    confidence_level=0.95):
    """
    Performs a proportion test to check if a new model (treatment)
    performs significantly better than the old one (control).
    """
    # Conversion rates
    p_control = control_conversions / control_total
    p_treatment = treatment_conversions / treatment_total

    # Standard error
    se = np.sqrt(p_control * (1 - p_control) / control_total +
                 p_treatment * (1 - p_treatment) / treatment_total)

    # Z-score and p-value for one-tailed test (treatment > control)
    z = (p_treatment - p_control) / se
    p_value = 1 - stats.norm.cdf(z)

    # Confidence interval
    z_critical = stats.norm.ppf((1 + confidence_level) / 2)
    margin = z_critical * se
    ci_lower = (p_treatment - p_control) - margin
    ci_upper = (p_treatment - p_control) + margin

    print(f"A/B Test Results:")
    print(f"  Control Conversion Rate: {p_control:.4%} ({control_conversions}/{control_total})")
    print(f"  Treatment Conversion Rate: {p_treatment:.4%} ({treatment_conversions}/{treatment_total})")
    print(f"  Absolute Lift: {p_treatment - p_control:+.4%}")
    print(f"  Relative Lift: {(p_treatment/p_control - 1):+.2%}")
    print(f"  p-value: {p_value:.6f}")
    print(f"  {confidence_level:.0%} CI for difference: [{ci_lower:.4%}, {ci_upper:.4%}]")

    is_significant = p_value < (1 - confidence_level)
    print(f"  Statistically Significant at {confidence_level:.0%}? {is_significant}")

    return {
        'lift_abs': p_treatment - p_control,
        'lift_rel': p_treatment / p_control - 1,
        'p_value': p_value,
        'significant': is_significant
    }

# Simulated data: new model (treatment) vs old model (control)
results = analyze_ab_test(
    control_conversions=1200, control_total=50000,
    treatment_conversions=1350, treatment_total=50000
)

import pytest
from pydantic import BaseModel, ValidationError, Field
from typing import Optional

# Define the expected schema contract as a Pydantic model
# This acts as a living document and an executable contract.
class UserEventContract(BaseModel):
    """Contract for the user event data from the streaming service."""
    user_id: int = Field(..., gt=0, description="Positive integer user ID")
    event_type: str = Field(..., regex="^(page_view|purchase|signup)$", description="Allowed event types")
    session_duration: Optional[int] = Field(None, ge=0, description="Duration in seconds, nullable")
    timestamp: str = Field(..., regex=r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}Z", description="ISO 8601 UTC timestamp")

    class Config:
        extra = 'forbid'  # Reject any extra fields not in the contract

def test_contract_with_production_data_sample():
    """
    This test should be run as part of CI whenever the upstream service is updated.
    It fetches a recent sample from production (or a test endpoint) and validates it.
    """
    # In reality, this would fetch from a known test endpoint or a recent log sample
    sample_event = {
        "user_id": 101,
        "event_type": "page_view",
        "session_duration": None,  # This is the new nullable field
        "timestamp": "2023-10-26T15:30:00Z"
    }

    # Validation: This should pass with the new nullable field
    validated_event = UserEventContract(**sample_event)
    assert validated_event.session_duration is None

    # Test a breaking change: if 'event_type' gets a new value not in our regex
    breaking_event = sample_event.copy()
    breaking_event["event_type"] = "new_event_type"
    with pytest.raises(ValidationError) as exc_info:
        UserEventContract(**breaking_event)
    assert "event_type" in str(exc_info.value)

    # Log the contract version for traceability
    print(f"Contract validation successful. Schema version: {UserEventContract.schema_json(indent=2)}")

# This test can be automated to run daily, alerting the team of schema drift.

# Example: Automated drift detection trigger for model retraining
import numpy as np
from scipy import stats
from datetime import datetime, timedelta

def check_feature_drift(current_features, reference_features, feature_name, threshold=0.05):
    """
    Uses the Kolmogorov-Smirnov test to detect distribution drift for a single feature.
    Returns True if drift is detected (p-value < threshold).
    """
    stat, p_value = stats.ks_2samp(reference_features, current_features)
    drift_detected = p_value < threshold
    if drift_detected:
        print(f"⚠️  Drift detected for '{feature_name}': KS stat={stat:.3f}, p={p_value:.4f}")
    return drift_detected

def evaluate_model_decay(current_accuracy, baseline_accuracy, decay_threshold=0.02):
    """Checks if model accuracy has decayed beyond an acceptable threshold."""
    decay = baseline_accuracy - current_accuracy
    if decay > decay_threshold:
        print(f"⚠️  Model accuracy decay detected: {current_accuracy:.4f} vs baseline {baseline_accuracy:.4f} (Δ={decay:.4f})")
        return True
    return False

# In a scheduled monitoring job:
def monitoring_job():
    # 1. Load recent production inference data (e.g., from last week)
    recent_data = load_inference_data(datetime.now() - timedelta(days=7), datetime.now())
    # 2. Load reference data (e.g., from when the model was last trained)
    reference_data = load_reference_data()

    drift_flags = []
    for feature in ['amount', 'frequency', 'recency']:
        drift = check_feature_drift(
            recent_data[feature].values,
            reference_data[feature].values,
            feature_name=feature,
            threshold=0.01  # 1% significance level
        )
        drift_flags.append(drift)

    # 3. Check accuracy on a recent holdout set
    current_acc = calculate_current_accuracy()
    baseline_acc = 0.875  # Retrieved from model registry

    accuracy_decay = evaluate_model_decay(current_acc, baseline_acc)

    # 4. Trigger retraining pipeline if any condition is met
    if any(drift_flags) or accuracy_decay:
        print("🚨 Drift or decay detected. Triggering retraining pipeline.")
        trigger_retraining_pipeline(
            reason="drift" if any(drift_flags) else "decay",
            details={
                'features_with_drift': [f for f, d in zip(['amount', 'frequency', 'recency'], drift_flags) if d],
                'accuracy_decay': baseline_acc - current_acc if accuracy_decay else 0.0
            }
        )
    else:
        print("✅ System healthy. No retraining required.")