app.py

"""
FastAPI application for Job Failure Prediction and Anomaly Detection.
"""
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from typing import Optional, Dict, List, Any
import pandas as pd
from datetime import datetime
import os

from preprocessing import engineer_features, align_schema_df
from model_utils import JobFailurePredictor, AnomalyDetector

app = FastAPI(
    title="Job Failure Prediction & Anomaly Detection API",
    description="ML service for predicting job failures and detecting anomalies",
    version="1.0.0"
)

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize models (lazy loading)
_predictor = None
_anomaly_detector = None


def get_predictor():
    """Lazy load predictor."""
    global _predictor
    if _predictor is None:
        _predictor = JobFailurePredictor()
    return _predictor


def get_anomaly_detector():
    """Lazy load anomaly detector."""
    global _anomaly_detector
    if _anomaly_detector is None:
        _anomaly_detector = AnomalyDetector()
    return _anomaly_detector


# Request/Response models
class JobFailureRequest(BaseModel):
    """Request model for job failure prediction."""
    zone: Optional[str] = None
    job_nm: Optional[str] = None
    tasksgroup_nm: Optional[str] = None
    round_time: Optional[str] = None
    job_start_time: Optional[str] = None
    job_end_time: Optional[str] = None
    duration: Optional[str] = None
    duration_sec: Optional[float] = None
    status: Optional[str] = None
    err_msg: Optional[str] = None
    zeppelin: Optional[str] = None
    ictrl_dt: Optional[str] = None
    start_ictrl_dt: Optional[str] = None
    end_ictrl_dt: Optional[str] = None
    # Optional pre-computed features
    run_hour: Optional[int] = None
    run_dayofweek: Optional[int] = None
    explain: bool = Field(default=False, description="Include SHAP explanations")


class JobFailureResponse(BaseModel):
    """Response model for job failure prediction."""
    fail_probability: float
    risk_level: str
    top_drivers: Optional[List[Dict[str, Any]]] = None
    recommended_actions: Optional[List[str]] = None
    error: Optional[str] = None


class AnomalyRequest(BaseModel):
    """Request model for anomaly detection."""
    features: Dict[str, float] = Field(
        description="Dictionary of feature values (duration_sec, err_msg_len, etc.)"
    )
    threshold: Optional[float] = Field(
        default=None,
        description="Custom threshold (overrides model default)"
    )


class AnomalyResponse(BaseModel):
    """Response model for anomaly detection."""
    reconstruction_error: float
    is_anomaly: bool
    threshold: float
    top_drivers: List[Dict[str, Any]]
    error: Optional[str] = None


# Endpoints
@app.get("/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "service": "job-failure-prediction",
        "models_loaded": {
            "predictor": _predictor is not None,
            "anomaly_detector": _anomaly_detector is not None
        }
    }


@app.post("/predict/job-fail", response_model=JobFailureResponse)
async def predict_job_failure(request: JobFailureRequest):
    """
    Predict job failure probability.
    
    Accepts partial job data and returns failure probability, risk level,
    and optional SHAP explanations.
    """
    try:
        # Convert request to DataFrame for preprocessing
        data_dict = request.dict(exclude={'explain'})
        
        # Create DataFrame with single row
        df = pd.DataFrame([data_dict])
        
        # If duration_sec not provided but duration is, parse it
        if df['duration_sec'].isna().any() and df['duration'].notna().any():
            from preprocessing import parse_duration
            df['duration_sec'] = df['duration'].apply(parse_duration)
        
        # If time features not provided, compute from job_start_time
        if df['run_hour'].isna().any() and df['job_start_time'].notna().any():
            df['job_start_time'] = pd.to_datetime(df['job_start_time'], errors='coerce')
            df['run_hour'] = df['job_start_time'].dt.hour.fillna(0).astype(int)
            df['run_dayofweek'] = df['job_start_time'].dt.dayofweek.fillna(0).astype(int)
        
        # Engineer features
        df = engineer_features(df)
        
        # Predict
        predictor = get_predictor()
        result = predictor.predict(df, explain=request.explain)
        
        return JobFailureResponse(**result)
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {str(e)}")


@app.post("/detect/anomaly", response_model=AnomalyResponse)
async def detect_anomaly(request: AnomalyRequest):
    """
    Detect anomaly in feature vector.
    
    Accepts a dictionary of feature values and returns reconstruction error,
    anomaly flag, and top contributing features.
    """
    try:
        detector = get_anomaly_detector()
        result = detector.detect(request.features, threshold=request.threshold)
        
        return AnomalyResponse(**result)
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Anomaly detection failed: {str(e)}")


@app.get("/")
async def root():
    """Root endpoint with API information."""
    return {
        "service": "Job Failure Prediction & Anomaly Detection API",
        "version": "1.0.0",
        "endpoints": {
            "health": "/health",
            "predict": "/predict/job-fail",
            "anomaly": "/detect/anomaly"
        },
        "docs": "/docs"
    }


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)