Implementing Mlops

"""
BentoML Model Serving Example

Demonstrates model serving patterns including:
- Model packaging and versioning
- REST API endpoint creation
- Batching for inference optimization
- Docker containerization
"""

import bentoml
from bentoml.io import JSON, NumpyNdarray
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
from pydantic import BaseModel
from typing import List
import asyncio


# =============================================================================
# Model Training and Saving
# =============================================================================

def train_and_save_model():
    """Train a model and save it to BentoML model store."""
    # Generate sample data
    X, y = make_classification(
        n_samples=1000,
        n_features=20,
        n_informative=15,
        random_state=42
    )

    # Train model
    model = RandomForestClassifier(
        n_estimators=100,
        max_depth=10,
        random_state=42
    )
    model.fit(X, y)

    # Save to BentoML model store
    saved_model = bentoml.sklearn.save_model(
        "fraud_classifier",
        model,
        signatures={
            "predict": {"batchable": True, "batch_dim": 0},
            "predict_proba": {"batchable": True, "batch_dim": 0},
        },
        labels={
            "framework": "sklearn",
            "model_type": "RandomForestClassifier",
            "version": "1.0.0",
        },
        metadata={
            "accuracy": 0.95,
            "training_date": "2025-01-15",
            "features": 20,
        },
        custom_objects={
            "feature_names": [f"feature_{i}" for i in range(20)],
        },
    )

    print(f"Model saved: {saved_model}")
    return saved_model


# =============================================================================
# Service Definition (service.py)
# =============================================================================

# Load the model
fraud_classifier = bentoml.sklearn.get("fraud_classifier:latest")

# Create a runner with batching enabled
fraud_runner = fraud_classifier.to_runner()

# Create the BentoML service
svc = bentoml.Service("fraud_detection_service", runners=[fraud_runner])


# Request/Response schemas
class PredictionRequest(BaseModel):
    features: List[float]


class PredictionResponse(BaseModel):
    prediction: int
    probability: float
    model_version: str


class BatchPredictionRequest(BaseModel):
    instances: List[List[float]]


class BatchPredictionResponse(BaseModel):
    predictions: List[int]
    probabilities: List[float]


# Single prediction endpoint
@svc.api(input=JSON(pydantic_model=PredictionRequest), output=JSON(pydantic_model=PredictionResponse))
async def predict(request: PredictionRequest) -> PredictionResponse:
    """
    Single prediction endpoint.

    BentoML automatically batches concurrent requests for efficiency.
    """
    input_array = np.array([request.features])
    prediction = await fraud_runner.predict.async_run(input_array)
    probability = await fraud_runner.predict_proba.async_run(input_array)

    return PredictionResponse(
        prediction=int(prediction[0]),
        probability=float(probability[0][1]),
        model_version=fraud_classifier.tag.version,
    )


# Batch prediction endpoint
@svc.api(
    input=JSON(pydantic_model=BatchPredictionRequest),
    output=JSON(pydantic_model=BatchPredictionResponse)
)
async def predict_batch(request: BatchPredictionRequest) -> BatchPredictionResponse:
    """
    Batch prediction endpoint for bulk inference.

    More efficient for processing multiple instances at once.
    """
    input_array = np.array(request.instances)
    predictions = await fraud_runner.predict.async_run(input_array)
    probabilities = await fraud_runner.predict_proba.async_run(input_array)

    return BatchPredictionResponse(
        predictions=[int(p) for p in predictions],
        probabilities=[float(p[1]) for p in probabilities],
    )


# Health check endpoint
@svc.api(input=JSON(), output=JSON())
async def health() -> dict:
    """Health check for load balancer probes."""
    return {
        "status": "healthy",
        "model": "fraud_classifier",
        "version": fraud_classifier.tag.version,
    }


# =============================================================================
# Bento Configuration (bentofile.yaml)
# =============================================================================

BENTOFILE_CONFIG = """
service: "service:svc"
labels:
  owner: ml-team
  project: fraud-detection
include:
  - "*.py"
python:
  packages:
    - scikit-learn>=1.0
    - numpy>=1.20
    - pydantic>=2.0
docker:
  distro: debian
  python_version: "3.11"
  cuda_version: null  # Set for GPU models
  env:
    BENTOML_CONFIG: /home/bentoml/configuration.yaml
  setup_script: |
    apt-get update && apt-get install -y curl
"""


# =============================================================================
# Deployment Configurations
# =============================================================================

# Kubernetes Deployment
K8S_DEPLOYMENT = """
apiVersion: apps/v1
kind: Deployment
metadata:
  name: fraud-detection
  labels:
    app: fraud-detection
spec:
  replicas: 3
  selector:
    matchLabels:
      app: fraud-detection
  template:
    metadata:
      labels:
        app: fraud-detection
    spec:
      containers:
      - name: fraud-detection
        image: fraud-detection:latest
        ports:
        - containerPort: 3000
        resources:
          requests:
            memory: "512Mi"
            cpu: "500m"
          limits:
            memory: "1Gi"
            cpu: "1000m"
        livenessProbe:
          httpGet:
            path: /health
            port: 3000
          initialDelaySeconds: 30
          periodSeconds: 10
        readinessProbe:
          httpGet:
            path: /health
            port: 3000
          initialDelaySeconds: 5
          periodSeconds: 5
        env:
        - name: BENTOML_PORT
          value: "3000"
---
apiVersion: v1
kind: Service
metadata:
  name: fraud-detection
spec:
  selector:
    app: fraud-detection
  ports:
  - port: 80
    targetPort: 3000
  type: LoadBalancer
---
apiVersion: autoscaling/v2
kind: HorizontalPodAutoscaler
metadata:
  name: fraud-detection-hpa
spec:
  scaleTargetRef:
    apiVersion: apps/v1
    kind: Deployment
    name: fraud-detection
  minReplicas: 2
  maxReplicas: 10
  metrics:
  - type: Resource
    resource:
      name: cpu
      target:
        type: Utilization
        averageUtilization: 70
"""


# =============================================================================
# CLI Commands
# =============================================================================

CLI_COMMANDS = """
# Build the Bento
bentoml build

# List saved models
bentoml models list

# List built Bentos
bentoml list

# Serve locally for development
bentoml serve service:svc --reload

# Containerize
bentoml containerize fraud_detection_service:latest

# Push to registry
bentoml push fraud_detection_service:latest

# Deploy to BentoCloud
bentoml deploy fraud_detection_service:latest

# Export Bento to directory
bentoml export fraud_detection_service:latest ./export/
"""


if __name__ == "__main__":
    # Train and save model
    saved_model = train_and_save_model()

    print("\n=== BentoML Model Serving ===")
    print(f"Model saved: {saved_model.tag}")
    print("\nTo serve locally:")
    print("  bentoml serve service:svc --reload")
    print("\nTo build and containerize:")
    print("  bentoml build")
    print("  bentoml containerize fraud_detection_service:latest")

"""
Feast Feature Store Example

Demonstrates feature store patterns including:
- Feature definitions and entity relationships
- Online/offline feature serving
- Point-in-time correct feature retrieval
- Feature freshness and materialization
"""

from datetime import datetime, timedelta
from feast import Entity, Feature, FeatureView, FileSource, ValueType, Field
from feast.types import Float32, Int64, String
from feast import FeatureStore
import pandas as pd
import numpy as np


# =============================================================================
# Feature Definitions (feature_repo/features.py)
# =============================================================================

# Define entities (primary keys for feature lookup)
customer = Entity(
    name="customer_id",
    description="Unique customer identifier",
    value_type=ValueType.INT64,
)

product = Entity(
    name="product_id",
    description="Unique product identifier",
    value_type=ValueType.INT64,
)

# Define data source (offline store)
customer_features_source = FileSource(
    path="data/customer_features.parquet",
    timestamp_field="event_timestamp",
    created_timestamp_column="created_timestamp",
)

transaction_features_source = FileSource(
    path="data/transaction_features.parquet",
    timestamp_field="event_timestamp",
)

# Define feature views (logical groupings of features)
customer_features = FeatureView(
    name="customer_features",
    entities=[customer],
    ttl=timedelta(days=90),  # Feature freshness window
    schema=[
        Field(name="total_purchases", dtype=Int64),
        Field(name="avg_order_value", dtype=Float32),
        Field(name="customer_segment", dtype=String),
        Field(name="lifetime_value", dtype=Float32),
        Field(name="days_since_last_purchase", dtype=Int64),
    ],
    source=customer_features_source,
    online=True,  # Enable online serving
)

transaction_features = FeatureView(
    name="transaction_features",
    entities=[customer, product],
    ttl=timedelta(days=30),
    schema=[
        Field(name="purchase_count_7d", dtype=Int64),
        Field(name="purchase_count_30d", dtype=Int64),
        Field(name="avg_quantity", dtype=Float32),
        Field(name="total_spend_30d", dtype=Float32),
    ],
    source=transaction_features_source,
    online=True,
)


# =============================================================================
# Feature Store Operations
# =============================================================================

def create_sample_data():
    """Generate sample feature data for demonstration."""
    np.random.seed(42)
    n_customers = 1000
    n_products = 100

    # Customer features
    customer_data = pd.DataFrame({
        "customer_id": range(1, n_customers + 1),
        "total_purchases": np.random.randint(1, 100, n_customers),
        "avg_order_value": np.random.uniform(20, 500, n_customers).astype(np.float32),
        "customer_segment": np.random.choice(["bronze", "silver", "gold", "platinum"], n_customers),
        "lifetime_value": np.random.uniform(100, 10000, n_customers).astype(np.float32),
        "days_since_last_purchase": np.random.randint(0, 365, n_customers),
        "event_timestamp": datetime.now() - timedelta(hours=1),
        "created_timestamp": datetime.now() - timedelta(days=30),
    })

    # Transaction features (customer x product combinations)
    n_transactions = 5000
    transaction_data = pd.DataFrame({
        "customer_id": np.random.randint(1, n_customers + 1, n_transactions),
        "product_id": np.random.randint(1, n_products + 1, n_transactions),
        "purchase_count_7d": np.random.randint(0, 10, n_transactions),
        "purchase_count_30d": np.random.randint(0, 50, n_transactions),
        "avg_quantity": np.random.uniform(1, 10, n_transactions).astype(np.float32),
        "total_spend_30d": np.random.uniform(0, 1000, n_transactions).astype(np.float32),
        "event_timestamp": datetime.now() - timedelta(hours=1),
    })

    return customer_data, transaction_data


def initialize_feature_store(repo_path: str = "feature_repo"):
    """
    Initialize Feast feature store.

    Run `feast apply` to register feature definitions:
    $ cd feature_repo && feast apply
    """
    store = FeatureStore(repo_path=repo_path)
    return store


def materialize_features(store: FeatureStore, end_date: datetime = None):
    """
    Materialize features from offline to online store.

    This populates the online store (Redis/DynamoDB) for low-latency serving.
    Should be run on a schedule (e.g., hourly) to keep features fresh.
    """
    if end_date is None:
        end_date = datetime.now()

    start_date = end_date - timedelta(days=7)

    store.materialize(
        start_date=start_date,
        end_date=end_date,
    )
    print(f"Features materialized from {start_date} to {end_date}")


def get_online_features(store: FeatureStore, customer_ids: list):
    """
    Retrieve features from online store for real-time inference.

    Typical latency: <10ms for Redis online store
    Use case: Real-time recommendation, fraud detection
    """
    entity_rows = [{"customer_id": cid} for cid in customer_ids]

    features = store.get_online_features(
        features=[
            "customer_features:total_purchases",
            "customer_features:avg_order_value",
            "customer_features:customer_segment",
            "customer_features:lifetime_value",
        ],
        entity_rows=entity_rows,
    ).to_df()

    return features


def get_historical_features(
    store: FeatureStore,
    entity_df: pd.DataFrame,
):
    """
    Retrieve point-in-time correct features for training.

    Entity dataframe must include:
    - Entity columns (customer_id, product_id)
    - event_timestamp column for point-in-time join

    Feast ensures features are retrieved as they existed at event_timestamp,
    preventing data leakage in training.
    """
    training_data = store.get_historical_features(
        entity_df=entity_df,
        features=[
            "customer_features:total_purchases",
            "customer_features:avg_order_value",
            "customer_features:customer_segment",
            "customer_features:lifetime_value",
            "customer_features:days_since_last_purchase",
        ],
    ).to_df()

    return training_data


def create_training_dataset(store: FeatureStore):
    """
    Create training dataset with point-in-time correct features.

    This pattern ensures training data reflects the state of features
    at the time of each historical event, preventing future data leakage.
    """
    # Historical events with timestamps
    events = pd.DataFrame({
        "customer_id": [1, 2, 3, 4, 5] * 100,
        "event_timestamp": pd.date_range(
            start=datetime.now() - timedelta(days=90),
            periods=500,
            freq="H"
        ),
        "label": np.random.randint(0, 2, 500),  # Target variable
    })

    # Get point-in-time correct features
    training_df = get_historical_features(store, events)

    print(f"Training dataset shape: {training_df.shape}")
    print(f"Features: {training_df.columns.tolist()}")

    return training_df


# =============================================================================
# Feature Store Configuration (feature_repo/feature_store.yaml)
# =============================================================================

FEATURE_STORE_CONFIG = """
project: ml_platform
registry: data/registry.db
provider: local

online_store:
  type: redis
  connection_string: "localhost:6379"

offline_store:
  type: file

entity_key_serialization_version: 2
"""


# =============================================================================
# Production Patterns
# =============================================================================

def production_inference_example():
    """
    Production inference pattern using Feast.

    1. Receive prediction request with entity IDs
    2. Fetch features from online store (<10ms)
    3. Combine with request features
    4. Run model inference
    5. Return prediction
    """
    store = initialize_feature_store()

    # Incoming request
    request = {
        "customer_id": 123,
        "product_id": 456,
        "request_features": {
            "device_type": "mobile",
            "time_of_day": "evening",
        }
    }

    # Fetch stored features
    online_features = get_online_features(store, [request["customer_id"]])

    # Combine all features
    model_input = {
        **request["request_features"],
        **online_features.iloc[0].to_dict(),
    }

    print(f"Model input features: {model_input}")
    # model.predict(model_input)


if __name__ == "__main__":
    # Generate sample data
    customer_df, transaction_df = create_sample_data()
    print(f"Customer features shape: {customer_df.shape}")
    print(f"Transaction features shape: {transaction_df.shape}")

    # In production, you would:
    # 1. Save data to parquet files
    # 2. Run `feast apply` to register features
    # 3. Run `feast materialize` to populate online store
    # 4. Use get_online_features for inference

    print("\nFeature store configuration:")
    print(FEATURE_STORE_CONFIG)

"""
Kubeflow Pipeline Example

Demonstrates ML pipeline orchestration patterns including:
- Component definition and containerization
- Pipeline DAG construction
- Artifact passing between components
- Conditional execution and loops
"""

from kfp import dsl
from kfp import compiler
from kfp.dsl import Input, Output, Dataset, Model, Metrics, Artifact
from typing import NamedTuple


# =============================================================================
# Pipeline Components
# =============================================================================

@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas", "scikit-learn", "pyarrow"]
)
def load_data(
    dataset_path: str,
    output_dataset: Output[Dataset],
) -> NamedTuple("Outputs", [("num_samples", int), ("num_features", int)]):
    """Load and validate training data."""
    import pandas as pd
    from collections import namedtuple

    # Load data (in production, this would read from GCS/S3)
    from sklearn.datasets import make_classification
    X, y = make_classification(
        n_samples=10000,
        n_features=20,
        n_informative=15,
        random_state=42
    )

    df = pd.DataFrame(X, columns=[f"feature_{i}" for i in range(20)])
    df["target"] = y

    # Save to output artifact
    df.to_parquet(output_dataset.path)

    outputs = namedtuple("Outputs", ["num_samples", "num_features"])
    return outputs(len(df), len(df.columns) - 1)


@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas", "scikit-learn", "pyarrow"]
)
def preprocess_data(
    input_dataset: Input[Dataset],
    train_dataset: Output[Dataset],
    test_dataset: Output[Dataset],
    test_size: float = 0.2,
):
    """Split data into train/test sets and apply preprocessing."""
    import pandas as pd
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler

    # Load data
    df = pd.read_parquet(input_dataset.path)
    X = df.drop("target", axis=1)
    y = df["target"]

    # Split
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_size, random_state=42, stratify=y
    )

    # Scale features
    scaler = StandardScaler()
    X_train_scaled = pd.DataFrame(
        scaler.fit_transform(X_train),
        columns=X_train.columns
    )
    X_test_scaled = pd.DataFrame(
        scaler.transform(X_test),
        columns=X_test.columns
    )

    # Save outputs
    train_df = X_train_scaled.copy()
    train_df["target"] = y_train.values
    train_df.to_parquet(train_dataset.path)

    test_df = X_test_scaled.copy()
    test_df["target"] = y_test.values
    test_df.to_parquet(test_dataset.path)


@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas", "scikit-learn", "pyarrow", "joblib"]
)
def train_model(
    train_dataset: Input[Dataset],
    model_artifact: Output[Model],
    n_estimators: int = 100,
    max_depth: int = 10,
):
    """Train a RandomForest classifier."""
    import pandas as pd
    from sklearn.ensemble import RandomForestClassifier
    import joblib

    # Load training data
    df = pd.read_parquet(train_dataset.path)
    X = df.drop("target", axis=1)
    y = df["target"]

    # Train model
    model = RandomForestClassifier(
        n_estimators=n_estimators,
        max_depth=max_depth,
        random_state=42,
        n_jobs=-1
    )
    model.fit(X, y)

    # Save model artifact
    model_artifact.metadata["framework"] = "sklearn"
    model_artifact.metadata["model_type"] = "RandomForestClassifier"
    model_artifact.metadata["n_estimators"] = n_estimators
    model_artifact.metadata["max_depth"] = max_depth

    joblib.dump(model, model_artifact.path)


@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["pandas", "scikit-learn", "pyarrow", "joblib"]
)
def evaluate_model(
    model_artifact: Input[Model],
    test_dataset: Input[Dataset],
    metrics: Output[Metrics],
    evaluation_report: Output[Artifact],
) -> NamedTuple("Outputs", [("accuracy", float), ("f1_score", float)]):
    """Evaluate model performance on test set."""
    import pandas as pd
    from sklearn.metrics import (
        accuracy_score, precision_score, recall_score, f1_score,
        classification_report, confusion_matrix
    )
    import joblib
    import json
    from collections import namedtuple

    # Load model and test data
    model = joblib.load(model_artifact.path)
    df = pd.read_parquet(test_dataset.path)
    X_test = df.drop("target", axis=1)
    y_test = df["target"]

    # Predict
    y_pred = model.predict(X_test)

    # Calculate metrics
    accuracy = accuracy_score(y_test, y_pred)
    precision = precision_score(y_test, y_pred)
    recall = recall_score(y_test, y_pred)
    f1 = f1_score(y_test, y_pred)

    # Log metrics for Kubeflow UI
    metrics.log_metric("accuracy", accuracy)
    metrics.log_metric("precision", precision)
    metrics.log_metric("recall", recall)
    metrics.log_metric("f1_score", f1)

    # Save detailed report
    report = {
        "classification_report": classification_report(y_test, y_pred, output_dict=True),
        "confusion_matrix": confusion_matrix(y_test, y_pred).tolist(),
    }

    with open(evaluation_report.path, "w") as f:
        json.dump(report, f, indent=2)

    outputs = namedtuple("Outputs", ["accuracy", "f1_score"])
    return outputs(accuracy, f1)


@dsl.component(
    base_image="python:3.11-slim",
    packages_to_install=["joblib", "google-cloud-storage"]
)
def deploy_model(
    model_artifact: Input[Model],
    model_name: str,
    deployment_env: str,
) -> str:
    """Deploy model to serving infrastructure."""
    import joblib
    import json

    model = joblib.load(model_artifact.path)

    # In production, this would:
    # 1. Push model to model registry
    # 2. Update serving deployment
    # 3. Run canary deployment

    deployment_info = {
        "model_name": model_name,
        "environment": deployment_env,
        "model_metadata": model_artifact.metadata,
        "status": "deployed",
    }

    print(f"Deployed model: {json.dumps(deployment_info, indent=2)}")

    return f"gs://models/{model_name}/{deployment_env}/model.joblib"


# =============================================================================
# Pipeline Definition
# =============================================================================

@dsl.pipeline(
    name="ml-training-pipeline",
    description="End-to-end ML training pipeline with evaluation and deployment"
)
def ml_training_pipeline(
    dataset_path: str = "gs://bucket/data/training_data.parquet",
    n_estimators: int = 100,
    max_depth: int = 10,
    test_size: float = 0.2,
    min_accuracy: float = 0.85,
    model_name: str = "fraud_classifier",
    deployment_env: str = "staging",
):
    """
    ML Training Pipeline

    1. Load and validate data
    2. Preprocess and split data
    3. Train model with hyperparameters
    4. Evaluate model performance
    5. Deploy if accuracy threshold met
    """
    # Load data
    load_task = load_data(dataset_path=dataset_path)

    # Preprocess
    preprocess_task = preprocess_data(
        input_dataset=load_task.outputs["output_dataset"],
        test_size=test_size,
    )

    # Train model
    train_task = train_model(
        train_dataset=preprocess_task.outputs["train_dataset"],
        n_estimators=n_estimators,
        max_depth=max_depth,
    )

    # Evaluate
    eval_task = evaluate_model(
        model_artifact=train_task.outputs["model_artifact"],
        test_dataset=preprocess_task.outputs["test_dataset"],
    )

    # Conditional deployment based on accuracy threshold
    with dsl.If(eval_task.outputs["accuracy"] >= min_accuracy):
        deploy_model(
            model_artifact=train_task.outputs["model_artifact"],
            model_name=model_name,
            deployment_env=deployment_env,
        )


# =============================================================================
# Pipeline with Hyperparameter Tuning
# =============================================================================

@dsl.pipeline(
    name="ml-hyperparameter-tuning",
    description="Pipeline with parallel hyperparameter experiments"
)
def hyperparameter_tuning_pipeline(
    dataset_path: str = "gs://bucket/data/training_data.parquet",
):
    """Run multiple training experiments with different hyperparameters."""

    # Load data once
    load_task = load_data(dataset_path=dataset_path)
    preprocess_task = preprocess_data(
        input_dataset=load_task.outputs["output_dataset"],
    )

    # Hyperparameter configurations
    configs = [
        {"n_estimators": 50, "max_depth": 5},
        {"n_estimators": 100, "max_depth": 10},
        {"n_estimators": 200, "max_depth": 15},
        {"n_estimators": 100, "max_depth": 20},
    ]

    # Run experiments in parallel using ParallelFor
    with dsl.ParallelFor(configs) as config:
        train_task = train_model(
            train_dataset=preprocess_task.outputs["train_dataset"],
            n_estimators=config.n_estimators,
            max_depth=config.max_depth,
        )

        evaluate_model(
            model_artifact=train_task.outputs["model_artifact"],
            test_dataset=preprocess_task.outputs["test_dataset"],
        )


# =============================================================================
# Compile and Run
# =============================================================================

if __name__ == "__main__":
    # Compile pipeline to YAML
    compiler.Compiler().compile(
        pipeline_func=ml_training_pipeline,
        package_path="ml_training_pipeline.yaml"
    )
    print("Pipeline compiled to ml_training_pipeline.yaml")

    compiler.Compiler().compile(
        pipeline_func=hyperparameter_tuning_pipeline,
        package_path="hyperparameter_tuning_pipeline.yaml"
    )
    print("Pipeline compiled to hyperparameter_tuning_pipeline.yaml")

    # To run on Kubeflow:
    # from kfp.client import Client
    # client = Client(host="https://kubeflow.example.com")
    # client.create_run_from_pipeline_package(
    #     "ml_training_pipeline.yaml",
    #     arguments={
    #         "n_estimators": 150,
    #         "max_depth": 12,
    #         "min_accuracy": 0.9,
    #     }
    # )

"""
MLflow Experiment Tracking Example

Demonstrates experiment tracking patterns for ML projects including:
- Logging parameters, metrics, and artifacts
- Model versioning and registration
- Experiment comparison and reproducibility
"""

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
import numpy as np
from datetime import datetime


def setup_mlflow(experiment_name: str, tracking_uri: str = "sqlite:///mlflow.db"):
    """Initialize MLflow tracking with local SQLite backend."""
    mlflow.set_tracking_uri(tracking_uri)
    mlflow.set_experiment(experiment_name)
    return mlflow.get_experiment_by_name(experiment_name)


def train_with_tracking(
    n_estimators: int = 100,
    max_depth: int = 10,
    min_samples_split: int = 2,
    experiment_name: str = "rf_classifier_experiment"
):
    """
    Train RandomForest classifier with full MLflow tracking.

    Logs:
    - Hyperparameters (n_estimators, max_depth, min_samples_split)
    - Metrics (accuracy, precision, recall, f1)
    - Model artifacts (sklearn model, feature importances)
    - Metadata (git commit, timestamp, dataset info)
    """
    setup_mlflow(experiment_name)

    # Generate synthetic classification data
    X, y = make_classification(
        n_samples=1000,
        n_features=20,
        n_informative=15,
        n_redundant=5,
        random_state=42
    )
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    with mlflow.start_run(run_name=f"rf_{datetime.now().strftime('%Y%m%d_%H%M%S')}"):
        # Log parameters
        mlflow.log_param("n_estimators", n_estimators)
        mlflow.log_param("max_depth", max_depth)
        mlflow.log_param("min_samples_split", min_samples_split)
        mlflow.log_param("dataset_size", len(X))
        mlflow.log_param("n_features", X.shape[1])

        # Log tags for organization
        mlflow.set_tag("model_type", "RandomForestClassifier")
        mlflow.set_tag("framework", "sklearn")
        mlflow.set_tag("stage", "development")

        # Train model
        model = RandomForestClassifier(
            n_estimators=n_estimators,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            random_state=42,
            n_jobs=-1
        )
        model.fit(X_train, y_train)

        # Evaluate and log metrics
        y_pred = model.predict(X_test)

        metrics = {
            "accuracy": accuracy_score(y_test, y_pred),
            "precision": precision_score(y_test, y_pred),
            "recall": recall_score(y_test, y_pred),
            "f1": f1_score(y_test, y_pred)
        }

        for metric_name, metric_value in metrics.items():
            mlflow.log_metric(metric_name, metric_value)

        # Log feature importances as artifact
        importances = model.feature_importances_
        np.save("feature_importances.npy", importances)
        mlflow.log_artifact("feature_importances.npy")

        # Log model with signature
        signature = mlflow.models.infer_signature(X_train, model.predict(X_train))
        mlflow.sklearn.log_model(
            model,
            "model",
            signature=signature,
            registered_model_name="rf_classifier"
        )

        print(f"Run completed: {mlflow.active_run().info.run_id}")
        print(f"Metrics: {metrics}")

        return model, metrics


def compare_experiments(experiment_name: str, metric: str = "f1"):
    """Compare runs within an experiment by a specific metric."""
    experiment = mlflow.get_experiment_by_name(experiment_name)
    if not experiment:
        raise ValueError(f"Experiment '{experiment_name}' not found")

    runs = mlflow.search_runs(
        experiment_ids=[experiment.experiment_id],
        order_by=[f"metrics.{metric} DESC"]
    )

    print(f"\nTop runs by {metric}:")
    print(runs[["run_id", f"metrics.{metric}", "params.n_estimators", "params.max_depth"]].head(5))

    return runs


def promote_model_to_production(model_name: str, version: int):
    """
    Transition a model version to production stage.

    Stage transitions:
    - None -> Staging: Initial testing
    - Staging -> Production: Approved for serving
    - Production -> Archived: Deprecated
    """
    client = mlflow.tracking.MlflowClient()

    # Transition to staging first
    client.transition_model_version_stage(
        name=model_name,
        version=version,
        stage="Staging"
    )
    print(f"Model {model_name} v{version} promoted to Staging")

    # After validation, promote to production
    client.transition_model_version_stage(
        name=model_name,
        version=version,
        stage="Production"
    )
    print(f"Model {model_name} v{version} promoted to Production")


if __name__ == "__main__":
    # Run hyperparameter experiments
    experiments = [
        {"n_estimators": 50, "max_depth": 5},
        {"n_estimators": 100, "max_depth": 10},
        {"n_estimators": 200, "max_depth": 15},
        {"n_estimators": 100, "max_depth": None},  # Unlimited depth
    ]

    for params in experiments:
        print(f"\nTraining with params: {params}")
        train_with_tracking(**params)

    # Compare experiment results
    compare_experiments("rf_classifier_experiment", metric="f1")

"""
Model Monitoring Example

Demonstrates model monitoring patterns including:
- Data drift detection
- Prediction drift detection
- Model performance degradation alerts
- Feature importance monitoring
"""

import numpy as np
import pandas as pd
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
from datetime import datetime, timedelta
from scipy import stats
from sklearn.ensemble import IsolationForest
import json
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


# =============================================================================
# Drift Detection
# =============================================================================

@dataclass
class DriftResult:
    """Result of drift detection analysis."""
    feature_name: str
    drift_detected: bool
    drift_score: float
    p_value: Optional[float]
    method: str
    threshold: float
    details: Dict


class DataDriftDetector:
    """
    Detect data drift between reference and production distributions.

    Supports multiple detection methods:
    - KS Test (Kolmogorov-Smirnov): Continuous features
    - Chi-Square Test: Categorical features
    - PSI (Population Stability Index): Distribution shift
    """

    def __init__(
        self,
        reference_data: pd.DataFrame,
        ks_threshold: float = 0.05,
        psi_threshold: float = 0.2,
    ):
        self.reference_data = reference_data
        self.ks_threshold = ks_threshold
        self.psi_threshold = psi_threshold
        self._compute_reference_stats()

    def _compute_reference_stats(self):
        """Pre-compute reference distribution statistics."""
        self.reference_stats = {}
        for col in self.reference_data.columns:
            self.reference_stats[col] = {
                "mean": self.reference_data[col].mean(),
                "std": self.reference_data[col].std(),
                "min": self.reference_data[col].min(),
                "max": self.reference_data[col].max(),
                "quantiles": self.reference_data[col].quantile([0.25, 0.5, 0.75]).to_dict(),
            }

    def ks_test(self, feature: str, production_data: pd.Series) -> DriftResult:
        """
        Kolmogorov-Smirnov test for continuous feature drift.

        Returns True if distributions are significantly different.
        """
        reference = self.reference_data[feature].dropna()
        production = production_data.dropna()

        statistic, p_value = stats.ks_2samp(reference, production)
        drift_detected = p_value < self.ks_threshold

        return DriftResult(
            feature_name=feature,
            drift_detected=drift_detected,
            drift_score=statistic,
            p_value=p_value,
            method="ks_test",
            threshold=self.ks_threshold,
            details={
                "reference_mean": float(reference.mean()),
                "production_mean": float(production.mean()),
                "reference_std": float(reference.std()),
                "production_std": float(production.std()),
            }
        )

    def calculate_psi(
        self,
        feature: str,
        production_data: pd.Series,
        n_bins: int = 10
    ) -> DriftResult:
        """
        Population Stability Index for distribution shift.

        PSI Interpretation:
        - < 0.1: No significant change
        - 0.1 - 0.2: Moderate change, investigate
        - > 0.2: Significant change, action required
        """
        reference = self.reference_data[feature].dropna()
        production = production_data.dropna()

        # Create bins from reference data
        bins = np.percentile(reference, np.linspace(0, 100, n_bins + 1))
        bins[0] = -np.inf
        bins[-1] = np.inf

        # Calculate proportions
        ref_counts = np.histogram(reference, bins=bins)[0]
        prod_counts = np.histogram(production, bins=bins)[0]

        # Avoid division by zero
        ref_props = (ref_counts + 1) / (len(reference) + n_bins)
        prod_props = (prod_counts + 1) / (len(production) + n_bins)

        # Calculate PSI
        psi = np.sum((prod_props - ref_props) * np.log(prod_props / ref_props))
        drift_detected = psi > self.psi_threshold

        return DriftResult(
            feature_name=feature,
            drift_detected=drift_detected,
            drift_score=psi,
            p_value=None,
            method="psi",
            threshold=self.psi_threshold,
            details={
                "reference_bins": ref_props.tolist(),
                "production_bins": prod_props.tolist(),
                "interpretation": self._interpret_psi(psi),
            }
        )

    def _interpret_psi(self, psi: float) -> str:
        if psi < 0.1:
            return "No significant change"
        elif psi < 0.2:
            return "Moderate change - investigate"
        else:
            return "Significant change - action required"

    def detect_all_features(
        self,
        production_data: pd.DataFrame,
        method: str = "ks_test"
    ) -> List[DriftResult]:
        """Run drift detection on all features."""
        results = []
        for feature in self.reference_data.columns:
            if feature in production_data.columns:
                if method == "ks_test":
                    result = self.ks_test(feature, production_data[feature])
                elif method == "psi":
                    result = self.calculate_psi(feature, production_data[feature])
                else:
                    raise ValueError(f"Unknown method: {method}")
                results.append(result)
        return results


# =============================================================================
# Prediction Monitoring
# =============================================================================

class PredictionMonitor:
    """
    Monitor model predictions for anomalies and drift.

    Tracks:
    - Prediction distribution changes
    - Confidence score degradation
    - Anomalous prediction patterns
    """

    def __init__(
        self,
        reference_predictions: np.ndarray,
        reference_probabilities: np.ndarray,
        contamination: float = 0.05,
    ):
        self.reference_predictions = reference_predictions
        self.reference_probabilities = reference_probabilities
        self.contamination = contamination

        # Fit anomaly detector on reference probabilities
        self.anomaly_detector = IsolationForest(
            contamination=contamination,
            random_state=42
        )
        self.anomaly_detector.fit(reference_probabilities.reshape(-1, 1))

        # Store reference statistics
        self.ref_pred_dist = np.bincount(reference_predictions) / len(reference_predictions)
        self.ref_prob_mean = reference_probabilities.mean()
        self.ref_prob_std = reference_probabilities.std()

    def check_prediction_distribution(
        self,
        production_predictions: np.ndarray,
        threshold: float = 0.1
    ) -> Dict:
        """Check if prediction class distribution has shifted."""
        prod_dist = np.bincount(
            production_predictions,
            minlength=len(self.ref_pred_dist)
        ) / len(production_predictions)

        # Calculate distribution shift
        shift = np.abs(prod_dist - self.ref_pred_dist).sum() / 2

        return {
            "distribution_shift": float(shift),
            "drift_detected": shift > threshold,
            "reference_distribution": self.ref_pred_dist.tolist(),
            "production_distribution": prod_dist.tolist(),
        }

    def check_confidence_degradation(
        self,
        production_probabilities: np.ndarray,
        z_threshold: float = 2.0
    ) -> Dict:
        """Detect if model confidence has degraded."""
        prod_mean = production_probabilities.mean()
        prod_std = production_probabilities.std()

        # Z-score for mean shift
        z_score = abs(prod_mean - self.ref_prob_mean) / self.ref_prob_std

        return {
            "mean_shift_z_score": float(z_score),
            "confidence_degraded": z_score > z_threshold,
            "reference_mean": float(self.ref_prob_mean),
            "production_mean": float(prod_mean),
            "reference_std": float(self.ref_prob_std),
            "production_std": float(prod_std),
        }

    def detect_anomalous_predictions(
        self,
        probabilities: np.ndarray
    ) -> Dict:
        """Flag anomalous prediction confidence scores."""
        anomaly_labels = self.anomaly_detector.predict(probabilities.reshape(-1, 1))
        anomaly_indices = np.where(anomaly_labels == -1)[0]

        return {
            "num_anomalies": len(anomaly_indices),
            "anomaly_rate": len(anomaly_indices) / len(probabilities),
            "anomaly_indices": anomaly_indices.tolist()[:100],  # Limit output
            "anomaly_values": probabilities[anomaly_indices][:100].tolist(),
        }


# =============================================================================
# Performance Monitoring
# =============================================================================

class PerformanceMonitor:
    """
    Track model performance over time with alerting.

    Monitors:
    - Accuracy, precision, recall, F1
    - Performance degradation trends
    - Threshold violations
    """

    def __init__(
        self,
        baseline_metrics: Dict[str, float],
        alert_thresholds: Dict[str, float],
    ):
        self.baseline_metrics = baseline_metrics
        self.alert_thresholds = alert_thresholds
        self.metrics_history: List[Dict] = []

    def log_metrics(
        self,
        metrics: Dict[str, float],
        timestamp: Optional[datetime] = None
    ):
        """Log metrics with timestamp."""
        if timestamp is None:
            timestamp = datetime.now()

        entry = {
            "timestamp": timestamp.isoformat(),
            "metrics": metrics,
            "alerts": self._check_alerts(metrics),
        }
        self.metrics_history.append(entry)

        # Log alerts
        for alert in entry["alerts"]:
            logger.warning(f"ALERT: {alert}")

        return entry

    def _check_alerts(self, metrics: Dict[str, float]) -> List[str]:
        """Check if any metrics violate thresholds."""
        alerts = []
        for metric, value in metrics.items():
            if metric in self.alert_thresholds:
                threshold = self.alert_thresholds[metric]
                baseline = self.baseline_metrics.get(metric, threshold)

                # Alert if performance dropped below threshold
                if value < threshold:
                    drop = ((baseline - value) / baseline) * 100
                    alerts.append(
                        f"{metric} dropped to {value:.4f} "
                        f"(baseline: {baseline:.4f}, threshold: {threshold:.4f}, "
                        f"drop: {drop:.1f}%)"
                    )
        return alerts

    def get_performance_trend(
        self,
        metric: str,
        window_size: int = 7
    ) -> Dict:
        """Analyze performance trend for a metric."""
        if len(self.metrics_history) < window_size:
            return {"error": "Insufficient history"}

        recent_values = [
            entry["metrics"].get(metric, 0)
            for entry in self.metrics_history[-window_size:]
        ]

        # Simple linear regression for trend
        x = np.arange(len(recent_values))
        slope, intercept = np.polyfit(x, recent_values, 1)

        return {
            "metric": metric,
            "current_value": recent_values[-1],
            "mean": float(np.mean(recent_values)),
            "std": float(np.std(recent_values)),
            "trend_slope": float(slope),
            "trend_direction": "improving" if slope > 0 else "degrading",
            "values": recent_values,
        }

    def generate_report(self) -> Dict:
        """Generate comprehensive monitoring report."""
        if not self.metrics_history:
            return {"error": "No metrics logged"}

        latest = self.metrics_history[-1]

        # Aggregate alerts
        all_alerts = []
        for entry in self.metrics_history[-7:]:
            all_alerts.extend(entry["alerts"])

        return {
            "report_timestamp": datetime.now().isoformat(),
            "latest_metrics": latest["metrics"],
            "baseline_metrics": self.baseline_metrics,
            "active_alerts": latest["alerts"],
            "alerts_last_7_entries": all_alerts,
            "total_entries": len(self.metrics_history),
        }


# =============================================================================
# Alerting Integration
# =============================================================================

class AlertManager:
    """Manage and dispatch monitoring alerts."""

    def __init__(self, alert_config: Dict):
        self.config = alert_config
        self.alert_history: List[Dict] = []

    def send_alert(
        self,
        alert_type: str,
        severity: str,
        message: str,
        details: Dict,
    ):
        """Send alert to configured channels."""
        alert = {
            "timestamp": datetime.now().isoformat(),
            "type": alert_type,
            "severity": severity,
            "message": message,
            "details": details,
        }

        self.alert_history.append(alert)

        # Log alert
        logger.warning(f"[{severity.upper()}] {alert_type}: {message}")

        # In production, send to:
        # - Slack/Teams webhook
        # - PagerDuty
        # - Email
        # - Prometheus Alertmanager

        return alert

    def check_and_alert(
        self,
        drift_results: List[DriftResult],
        prediction_metrics: Dict,
        performance_metrics: Dict,
    ):
        """Analyze all monitoring results and send appropriate alerts."""

        # Check for data drift
        drifted_features = [r for r in drift_results if r.drift_detected]
        if len(drifted_features) > 0:
            self.send_alert(
                alert_type="data_drift",
                severity="warning" if len(drifted_features) < 3 else "critical",
                message=f"Data drift detected in {len(drifted_features)} features",
                details={
                    "features": [r.feature_name for r in drifted_features],
                    "scores": {r.feature_name: r.drift_score for r in drifted_features},
                }
            )

        # Check for prediction drift
        if prediction_metrics.get("drift_detected"):
            self.send_alert(
                alert_type="prediction_drift",
                severity="warning",
                message="Prediction distribution has shifted",
                details=prediction_metrics,
            )

        # Check for performance degradation
        if prediction_metrics.get("confidence_degraded"):
            self.send_alert(
                alert_type="confidence_degradation",
                severity="critical",
                message="Model confidence has degraded significantly",
                details=prediction_metrics,
            )


# =============================================================================
# Usage Example
# =============================================================================

def main():
    """Demonstrate model monitoring pipeline."""
    np.random.seed(42)

    # Generate reference data (from training)
    n_samples = 1000
    reference_features = pd.DataFrame({
        "feature_1": np.random.normal(0, 1, n_samples),
        "feature_2": np.random.normal(5, 2, n_samples),
        "feature_3": np.random.exponential(2, n_samples),
    })
    reference_predictions = np.random.randint(0, 2, n_samples)
    reference_probabilities = np.random.beta(5, 2, n_samples)

    # Initialize monitors
    drift_detector = DataDriftDetector(reference_features)
    prediction_monitor = PredictionMonitor(
        reference_predictions,
        reference_probabilities
    )
    performance_monitor = PerformanceMonitor(
        baseline_metrics={"accuracy": 0.95, "f1": 0.93, "precision": 0.94},
        alert_thresholds={"accuracy": 0.90, "f1": 0.88, "precision": 0.88}
    )
    alert_manager = AlertManager({})

    # Simulate production data with drift
    production_features = pd.DataFrame({
        "feature_1": np.random.normal(0.5, 1.2, n_samples),  # Mean shifted
        "feature_2": np.random.normal(5, 2, n_samples),       # No change
        "feature_3": np.random.exponential(3, n_samples),     # Distribution changed
    })
    production_predictions = np.random.randint(0, 2, n_samples)
    production_probabilities = np.random.beta(4, 3, n_samples)  # Degraded

    # Run drift detection
    print("=" * 60)
    print("Data Drift Detection")
    print("=" * 60)
    drift_results = drift_detector.detect_all_features(production_features)
    for result in drift_results:
        status = "DRIFT" if result.drift_detected else "OK"
        print(f"  {result.feature_name}: {status} (score={result.drift_score:.4f})")

    # Run prediction monitoring
    print("\n" + "=" * 60)
    print("Prediction Monitoring")
    print("=" * 60)
    pred_dist = prediction_monitor.check_prediction_distribution(production_predictions)
    print(f"  Distribution shift: {pred_dist['distribution_shift']:.4f}")

    conf_check = prediction_monitor.check_confidence_degradation(production_probabilities)
    print(f"  Confidence z-score: {conf_check['mean_shift_z_score']:.4f}")
    print(f"  Degraded: {conf_check['confidence_degraded']}")

    # Log performance metrics
    print("\n" + "=" * 60)
    print("Performance Monitoring")
    print("=" * 60)
    performance_monitor.log_metrics({
        "accuracy": 0.91,
        "f1": 0.89,
        "precision": 0.90
    })
    report = performance_monitor.generate_report()
    print(f"  Current accuracy: {report['latest_metrics']['accuracy']}")
    print(f"  Alerts: {report['active_alerts']}")

    # Send alerts
    print("\n" + "=" * 60)
    print("Alert Summary")
    print("=" * 60)
    alert_manager.check_and_alert(drift_results, conf_check, report)
    print(f"  Total alerts sent: {len(alert_manager.alert_history)}")


if __name__ == "__main__":
    main()

skill: "implementing-mlops"
version: "1.0"
domain: "ai-ml"

base_outputs:
  # Core MLOps infrastructure files ALWAYS produced
  - path: "mlops/config/mlflow_config.yaml"
    must_contain: ["tracking_uri", "experiment_name", "artifact_location"]
    description: "MLflow configuration for experiment tracking and model registry"

  - path: "mlops/training/train.py"
    must_contain: ["mlflow\\.start_run", "mlflow\\.log_param", "mlflow\\.log_metric"]
    description: "Training script with MLflow experiment tracking integration"

  - path: "mlops/model_registry/model_metadata.yaml"
    must_contain: ["model_name", "version", "stage", "metrics"]
    description: "Model registry metadata including versioning and stage management"

  - path: "mlops/monitoring/drift_detection.py"
    must_contain: ["detect_drift", "KolmogorovSmirnov|PSI|chi_square"]
    description: "Data and model drift detection implementation"

  - path: "requirements.txt"
    must_contain: ["mlflow"]
    description: "Python dependencies for MLOps infrastructure"

conditional_outputs:

  maturity:
    starter:
      - path: "mlops/serving/simple_api.py"
        must_contain: ["FastAPI|Flask", "predict"]
        description: "Basic REST API for model serving"

      - path: "mlops/pipelines/basic_pipeline.py"
        must_contain: ["def train", "def evaluate", "def deploy"]
        description: "Simple training pipeline without orchestration"

      - path: "docker-compose.yml"
        must_contain: ["mlflow", "postgres"]
        description: "Docker Compose setup for local MLflow server"

    intermediate:
      - path: "mlops/features/feature_definitions.py"
        must_contain: ["FeatureView|Feature", "online_store|offline_store"]
        description: "Feature store definitions for training/serving consistency"

      - path: "mlops/serving/deployment_config.yaml"
        must_contain: ["resources", "replicas", "autoscaling"]
        description: "Deployment configuration with resource management"

      - path: "mlops/pipelines/orchestrated_pipeline.py"
        must_contain: ["@task|@op|@component", "pipeline|workflow"]
        description: "Orchestrated ML pipeline with dependencies"

      - path: "mlops/validation/model_validation.py"
        must_contain: ["accuracy|precision|recall", "threshold", "validation"]
        description: "Automated model validation suite"

      - path: ".github/workflows/ml_ci_cd.yml"
        must_contain: ["train", "test", "deploy"]
        description: "CI/CD pipeline for model deployment automation"

    advanced:
      - path: "mlops/features/feast_feature_store.py"
        must_contain: ["FeatureStore", "get_online_features", "get_historical_features"]
        description: "Production-grade feature store with online/offline serving"

      - path: "mlops/serving/canary_deployment.yaml"
        must_contain: ["canary|traffic_split", "stable", "candidate"]
        description: "Canary deployment strategy configuration"

      - path: "mlops/pipelines/kubeflow_pipeline.py"
        must_contain: ["@dsl\\.component|@dsl\\.pipeline", "ContainerOp|create_component_from_func"]
        description: "Kubeflow Pipelines for ML orchestration"

      - path: "mlops/monitoring/observability_dashboard.py"
        must_contain: ["prometheus|grafana", "metrics", "drift"]
        description: "Comprehensive monitoring dashboard with drift detection"

      - path: "mlops/governance/model_card.md"
        must_contain: ["Model Details", "Intended Use", "Limitations", "Fairness"]
        description: "Model card for governance and compliance"

      - path: "mlops/optimization/model_quantization.py"
        must_contain: ["quantize|quantization", "int8|float16"]
        description: "Model optimization for inference performance"

  infrastructure:
    kubernetes:
      - path: "k8s/mlflow-deployment.yaml"
        must_contain: ["kind: Deployment", "mlflow", "containerPort"]
        description: "Kubernetes deployment for MLflow server"

      - path: "k8s/model-serving-deployment.yaml"
        must_contain: ["kind: Deployment", "image:", "resources:"]
        description: "Kubernetes deployment for model serving"

      - path: "k8s/seldon-inference-graph.yaml"
        must_contain: ["SeldonDeployment|InferenceService", "predictor"]
        description: "Seldon Core or KServe inference service configuration"

      - path: "k8s/gpu-node-pool.yaml"
        must_contain: ["gpu|accelerator", "nodeSelector|tolerations"]
        description: "GPU node pool configuration for training workloads"

    docker_compose:
      - path: "docker-compose.yml"
        must_contain: ["mlflow", "postgres", "minio|s3"]
        description: "Docker Compose with MLflow, database, and artifact storage"

      - path: "mlops/serving/Dockerfile"
        must_contain: ["FROM", "COPY", "CMD|ENTRYPOINT"]
        description: "Dockerfile for model serving container"

    managed_platform:
      - path: "mlops/cloud/sagemaker_pipeline.py"
        must_contain: ["sagemaker|SageMaker", "Pipeline|TrainingStep"]
        description: "AWS SageMaker pipeline configuration"

      - path: "mlops/cloud/vertex_pipeline.py"
        must_contain: ["vertex_ai|aiplatform", "pipeline"]
        description: "GCP Vertex AI pipeline configuration"

      - path: "terraform/mlops_infrastructure.tf"
        must_contain: ["resource", "provider", "aws|google|azurerm"]
        description: "Terraform infrastructure as code for managed ML platform"

  model_type:
    deep_learning:
      - path: "mlops/training/pytorch_train.py"
        must_contain: ["torch\\.nn|nn\\.Module", "optimizer", "loss"]
        description: "PyTorch deep learning training script"

      - path: "mlops/serving/torchserve_config.yaml"
        must_contain: ["model_store", "inference_address", "management_address"]
        description: "TorchServe configuration for PyTorch model serving"

      - path: "mlops/optimization/onnx_conversion.py"
        must_contain: ["onnx", "export|convert"]
        description: "ONNX model conversion for optimized inference"

    classical_ml:
      - path: "mlops/training/sklearn_train.py"
        must_contain: ["sklearn|scikit-learn", "fit", "predict"]
        description: "Scikit-learn classical ML training script"

      - path: "mlops/features/feature_engineering.py"
        must_contain: ["transform|fit_transform", "StandardScaler|OneHotEncoder"]
        description: "Feature engineering pipeline for classical ML"

  ml_framework:
    pytorch:
      - path: "mlops/models/pytorch_model.py"
        must_contain: ["torch\\.nn\\.Module", "forward"]
        description: "PyTorch model architecture definition"

      - path: "mlops/training/pytorch_lightning_train.py"
        must_contain: ["LightningModule|pl\\.LightningModule", "training_step"]
        description: "PyTorch Lightning training with best practices"

    tensorflow:
      - path: "mlops/models/tensorflow_model.py"
        must_contain: ["tf\\.keras|tensorflow\\.keras", "Model|Sequential"]
        description: "TensorFlow/Keras model architecture"

      - path: "mlops/serving/tfserving_config.yaml"
        must_contain: ["model_config_list", "base_path"]
        description: "TensorFlow Serving configuration"

    sklearn:
      - path: "mlops/models/sklearn_pipeline.py"
        must_contain: ["Pipeline", "fit", "predict"]
        description: "Scikit-learn pipeline with preprocessing and model"

    xgboost:
      - path: "mlops/training/xgboost_train.py"
        must_contain: ["xgboost|xgb", "DMatrix", "train"]
        description: "XGBoost training with hyperparameter tuning"

  experiment_tracking:
    mlflow:
      - path: "mlops/tracking/mlflow_setup.py"
        must_contain: ["mlflow\\.set_tracking_uri", "mlflow\\.create_experiment"]
        description: "MLflow tracking server setup and initialization"

      - path: "mlops/training/mlflow_autolog.py"
        must_contain: ["mlflow\\.autolog|mlflow\\.sklearn\\.autolog|mlflow\\.pytorch\\.autolog"]
        description: "MLflow autologging for automatic experiment tracking"

      - path: "scripts/setup_mlflow_server.sh"
        must_contain: ["mlflow server", "backend-store-uri", "default-artifact-root"]
        description: "Script to start MLflow server with PostgreSQL and S3"

    wandb:
      - path: "mlops/tracking/wandb_setup.py"
        must_contain: ["wandb\\.init", "wandb\\.log", "wandb\\.finish"]
        description: "Weights & Biases experiment tracking integration"

      - path: "mlops/training/wandb_sweep.yaml"
        must_contain: ["program:", "method:", "parameters:"]
        description: "W&B Sweeps configuration for hyperparameter optimization"

    neptune:
      - path: "mlops/tracking/neptune_setup.py"
        must_contain: ["neptune\\.init", "run\\[.*\\]\\s*="]
        description: "Neptune.ai tracking for enterprise ML workflows"

scaffolding:
  - path: "data/raw/"
    reason: "Placeholder for raw training data - populated by data pipelines"

  - path: "data/processed/"
    reason: "Placeholder for processed features - generated during training"

  - path: "models/staging/"
    reason: "Staging area for models under validation before production"

  - path: "models/production/"
    reason: "Production models directory - populated by deployment pipeline"

  - path: "artifacts/plots/"
    reason: "Directory for training plots and visualizations"

  - path: "artifacts/reports/"
    reason: "Directory for model evaluation reports and metrics"

  - path: "logs/"
    reason: "Application and training logs directory"

metadata:
  primary_blueprints: ["ml-pipeline"]
  contributes_to:
    - "ML Model Training Pipeline"
    - "Model Serving Infrastructure"
    - "Experiment Tracking System"
    - "Feature Store Implementation"
    - "Model Monitoring and Drift Detection"
    - "ML Workflow Orchestration"
    - "Model Registry and Versioning"
    - "Production Model Deployment"