# app.py
import os
from traceloop.sdk import Traceloop
from traceloop.sdk.decorators import workflow, task

# Configure SigNoz endpoint
os.environ["TRACELOOP_BASE_URL"] = "http://localhost:4318"

# Initialize tracing
Traceloop.init(app_name="my-ml-app")

@workflow(name="prediction_pipeline")
def run_prediction(data):
    """Main prediction workflow."""
    preprocessed = preprocess_data(data)
    prediction = model_inference(preprocessed)
    result = postprocess_result(prediction)
    return result

@task(name="preprocess")
def preprocess_data(data):
    # Your preprocessing logic
    return processed_data

@task(name="inference")
def model_inference(data):
    # Model prediction
    return prediction

@task(name="postprocess")
def postprocess_result(prediction):
    # Format result
    return result

# deployment.yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: ml-app
spec:
  template:
    spec:
      containers:
      - name: app
        image: my-ml-app:latest
        env:
        - name: TRACELOOP_BASE_URL
          value: "http://my-release-signoz-otel-collector.default:4318"
        - name: OTEL_SERVICE_NAME
          value: "ml-app-production"

# main.py
from fastapi import FastAPI
from traceloop.sdk import Traceloop
import os

os.environ["TRACELOOP_BASE_URL"] = "http://localhost:4318"
Traceloop.init(app_name="fastapi-ml-service")

app = FastAPI()

@app.post("/predict")
async def predict(request: PredictionRequest):
    # Automatically traced
    result = model.predict(request.data)
    return {"prediction": result}

# metrics.py
from prometheus_client import Counter, Histogram, start_http_server

# Define metrics
prediction_count = Counter(
    'ml_predictions_total',
    'Total number of predictions',
    ['model', 'status']
)

prediction_duration = Histogram(
    'ml_prediction_duration_seconds',
    'Time spent on prediction',
    ['model']
)

input_feature_values = Histogram(
    'ml_input_feature_value',
    'Distribution of input feature values',
    ['feature_name']
)

# Start metrics server
start_http_server(8000)

# Use in your code
def predict(data):
    with prediction_duration.labels(model='v1').time():
        try:
            result = model.predict(data)
            prediction_count.labels(model='v1', status='success').inc()
            return result
        except Exception as e:
            prediction_count.labels(model='v1', status='error').inc()
            raise

# servicemonitor.yaml
apiVersion: monitoring.coreos.com/v1
kind: ServiceMonitor
metadata:
  name: ml-app-metrics
  labels:
    release: monitoring  # Must match Prometheus selector
spec:
  selector:
    matchLabels:
      app: ml-app
  endpoints:
  - port: metrics
    interval: 30s
    path: /metrics

# airflow_dag.py
from airflow import DAG
from airflow.operators.python import PythonOperator
from evidently.test_suite import TestSuite
from evidently.tests import TestColumnDrift
import pandas as pd
from datetime import datetime

def check_drift(**context):
    """Check for data drift."""
    # Load reference data
    reference_df = pd.read_parquet("s3://bucket/reference_data.parquet")
    
    # Load recent production data
    current_df = pd.read_parquet("s3://bucket/production_data_latest.parquet")
    
    # Define tests
    tests = TestSuite(tests=[
        TestColumnDrift(column_name="age"),
        TestColumnDrift(column_name="income"),
        TestColumnDrift(column_name="education"),
    ])
    
    # Run tests
    tests.run(reference_data=reference_df, current_data=current_df)
    
    # Check results
    results = tests.as_dict()
    if not results["summary"]["all_passed"]:
        # Send alert
        send_slack_alert("Drift detected in production data!")
        
        # Save report
        tests.save_html(f"drift_report_{datetime.now().isoformat()}.html")
        
        raise AirflowException("Drift detected - manual review required")

with DAG('ml_pipeline_with_drift', schedule_interval='@daily') as dag:
    drift_check = PythonOperator(
        task_id='check_drift',
        python_callable=check_drift
    )
    
    retrain = PythonOperator(
        task_id='retrain_model',
        python_callable=retrain_model
    )
    
    drift_check >> retrain

# pipeline.yaml
apiVersion: mlops.seldon.io/v1alpha1
kind: Pipeline
metadata:
  name: ml-pipeline-with-drift
spec:
  steps:
    - name: preprocessor
    - name: model
      inputs:
      - preprocessor
    - name: drift-detector
      inputs:
      - preprocessor
      batch:
        size: 100  # Check drift every 100 requests
  output:
    steps:
    - model
    - drift-detector.outputs.drift_score
---
apiVersion: mlops.seldon.io/v1alpha1
kind: Model
metadata:
  name: drift-detector
spec:
  storageUri: "gs://my-bucket/drift-detector"
  requirements:
    - mlserver
    - alibi-detect

# drift_component.py
from kfp import dsl
from kfp.dsl import component, Input, Output, Dataset, Metrics

@component(
    base_image="python:3.10",
    packages_to_install=["evidently", "pandas"]
)
def check_drift(
    reference_data: Input[Dataset],
    current_data: Input[Dataset],
    metrics: Output[Metrics]
):
    """Check for drift between datasets."""
    import pandas as pd
    from evidently.report import Report
    from evidently.metric_preset import DataDriftPreset
    
    ref_df = pd.read_parquet(reference_data.path)
    curr_df = pd.read_parquet(current_data.path)
    
    report = Report(metrics=[DataDriftPreset()])
    report.run(reference_data=ref_df, current_data=curr_df)
    
    result = report.as_dict()
    drift_score = result["metrics"][0]["result"]["drift_score"]
    
    metrics.log_metric("drift_score", drift_score)
    
    if drift_score > 0.1:
        raise ValueError(f"Drift detected: {drift_score}")

# Use in pipeline
@dsl.pipeline(name="ML Pipeline with Drift Check")
def ml_pipeline():
    drift_task = check_drift(
        reference_data=load_reference_data(),
        current_data=load_production_data()
    )
    
    with dsl.Condition(drift_task.outputs["drift_score"] > 0.1):
        retrain_model()