""" FastAPI inference service for candlestick pattern prediction. Provides REST API endpoints for model serving, health checks, and prediction. """ import logging import threading import uuid as uuid_lib from pathlib import Path from typing import Optional, Dict, Any, List from datetime import datetime import json from fastapi import FastAPI, HTTPException, status from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field import numpy as np import pandas as pd import joblib import mlflow import mlflow.sklearn import mlflow.xgboost from sqlalchemy import update as sa_update, desc from app.config import load_config, PipelineConfig, get_default_config from app.db import get_db, TrainingRun, init_db from app.preprocessing import preprocess_candles, extract_feature_columns from app.patterns import ( TALIB_PATTERNS, get_available_patterns, validate_pattern_names, detect_patterns, ) # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) # FastAPI app app = FastAPI( title="Candle Pattern Inference API", description="ML inference service for candlestick pattern recognition", version="1.0.0" ) # CORS middleware app.add_middleware( CORSMiddleware, allow_origins=["*"], # In production, specify actual origins allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # Global state class AppState: """Application state container.""" model: Optional[Any] = None model_info: Optional[Dict[str, Any]] = None pipeline_config: Optional[PipelineConfig] = None feature_columns: Optional[List[str]] = None label_encoder: Optional[Dict[int, str]] = None # Training thread management active_training_run_id: Optional[str] = None training_lock: threading.Lock = None def __init__(self): self.training_lock = threading.Lock() state = AppState() # --- Pydantic Models --- class CandleData(BaseModel): """Single candle data point.""" time: int = Field(..., description="Unix timestamp in seconds") open: float high: float low: float close: float volume: Optional[float] = None class PredictRequest(BaseModel): """Request model for /predict endpoint.""" pair: Optional[str] = Field(None, description="Trading pair (e.g., EURUSD)") timeframe: Optional[str] = Field(None, description="Timeframe (e.g., 1H, 4H, 1D)") candles: List[CandleData] = Field(..., min_length=1, description="Array of candle data") class PredictionResult(BaseModel): """Single candle prediction.""" time: int label: str confidence: float = Field(..., ge=0.0, le=1.0) class PredictionSpan(BaseModel): """Grouped prediction span.""" start_time: int end_time: int label: str avg_confidence: float = Field(..., ge=0.0, le=1.0) candle_count: int class ModelInfo(BaseModel): """Model metadata.""" model_name: str model_version: Optional[str] = None model_type: str trained_at: Optional[str] = None dataset_version: Optional[str] = None feature_engineering_enabled: bool class PredictResponse(BaseModel): """Response model for /predict endpoint.""" predictions: List[PredictionResult] spans: List[PredictionSpan] model_info: ModelInfo class BatchPredictRequest(BaseModel): """Request model for /predict/batch endpoint.""" pair: str timeframe: str start_date: str = Field(..., description="ISO format date (YYYY-MM-DD)") end_date: str = Field(..., description="ISO format date (YYYY-MM-DD)") class LabelInfo(BaseModel): """Pattern label with display color.""" name: str color: str class PerClassMetrics(BaseModel): """Per-class performance metrics.""" label: str precision: float recall: float f1_score: float support: int class ModelInfoResponse(BaseModel): """Extended model information with metrics.""" model_name: str model_version: Optional[str] = None model_type: str trained_at: Optional[str] = None dataset_version: Optional[str] = None feature_engineering_enabled: bool labels: List[str] per_class_metrics: List[PerClassMetrics] class HealthResponse(BaseModel): """Health check response.""" status: str = Field(..., description="healthy, degraded, or unhealthy") model_loaded: bool mlflow: str = Field(..., description="connected or disconnected") database: str = Field(..., description="connected or disconnected") # --- Model Loading Functions --- def load_model_from_mlflow(model_name: str, stage: str) -> tuple[Any, Dict[str, Any]]: """ Load model from MLflow model registry. Args: model_name: Name of the registered model stage: Model stage (Production, Staging, None) Returns: Tuple of (model, model_info) Raises: Exception: If model not found or MLflow unavailable """ logger.info(f"Loading model from MLflow: {model_name} stage={stage}") try: # Load model from registry model_uri = f"models:/{model_name}/{stage}" model = mlflow.pyfunc.load_model(model_uri) # Get model version details client = mlflow.tracking.MlflowClient() model_versions = client.get_latest_versions(model_name, stages=[stage]) if not model_versions: raise ValueError(f"No model found for {model_name} at stage {stage}") model_version = model_versions[0] run_id = model_version.run_id # Get run metadata run = client.get_run(run_id) # Extract model info model_info = { "model_name": model_name, "model_version": model_version.version, "model_type": run.data.params.get("model_type", "unknown"), "trained_at": datetime.fromtimestamp(run.info.start_time / 1000).isoformat(), "dataset_version": run.data.params.get("dataset_version", None), "feature_engineering_enabled": run.data.params.get("feature_engineering", "true") == "true", "labels": json.loads(run.data.params.get("labels", "[]")), "per_class_metrics": [] } # Extract per-class metrics for key, value in run.data.metrics.items(): if key.startswith("precision_"): label = key.replace("precision_", "") recall = run.data.metrics.get(f"recall_{label}", 0.0) f1 = run.data.metrics.get(f"f1_{label}", 0.0) support = int(run.data.params.get(f"support_{label}", 0)) model_info["per_class_metrics"].append({ "label": label, "precision": value, "recall": recall, "f1_score": f1, "support": support }) logger.info(f"Successfully loaded model: {model_name} v{model_version.version}") return model, model_info except Exception as e: logger.error(f"Failed to load model from MLflow: {e}") raise def load_model_from_local(model_path: str) -> tuple[Any, Dict[str, Any]]: """ Load model from local file using joblib. Args: model_path: Path to .pkl model file Returns: Tuple of (model, model_info) Raises: FileNotFoundError: If model file doesn't exist Exception: If model loading fails """ model_path = Path(model_path) logger.info(f"Loading model from local file: {model_path}") if not model_path.exists(): raise FileNotFoundError(f"Model file not found: {model_path}") try: # Load model with joblib model_data = joblib.load(model_path) # Extract model and metadata if isinstance(model_data, dict): model = model_data.get("model") metadata = model_data.get("metadata", {}) else: model = model_data metadata = {} # Extract labels from model if not in metadata labels = metadata.get("labels", []) if not labels: # Try to get class labels from the model itself if hasattr(model, 'classes_'): labels = [str(c) for c in model.classes_] elif hasattr(model, 'model') and hasattr(model.model, 'classes_'): labels = [str(c) for c in model.model.classes_] # Build model info model_info = { "model_name": model_path.stem, "model_version": metadata.get("version", "local"), "model_type": metadata.get("model_type", "unknown"), "trained_at": metadata.get("trained_at", None), "dataset_version": metadata.get("dataset_version", None), "feature_engineering_enabled": metadata.get("feature_engineering_enabled", True), "labels": labels, "per_class_metrics": metadata.get("per_class_metrics", []) } logger.info(f"Successfully loaded local model: {model_path.name}") return model, model_info except Exception as e: logger.error(f"Failed to load model from local file: {e}") raise # --- Startup Event --- @app.on_event("startup") async def startup_event(): """ Load model and pipeline config on startup. """ logger.info("Starting inference service...") # Ensure training_runs table exists init_db() # Mark any stale "running" records as failed — they belong to a previous # process and will never complete. try: with get_db() as db: stmt = ( sa_update(TrainingRun) .where(TrainingRun.status == "running") .values( status="failed", completed_at=datetime.utcnow(), metrics_summary={"error": "Service restarted while training was in progress"}, ) ) result = db.execute(stmt) db.commit() if result.rowcount: logger.warning( f"Marked {result.rowcount} stale training run(s) as failed on startup" ) except Exception as exc: logger.error(f"Failed to clean up stale training runs: {exc}") # Load pipeline config config_path = Path("config/pipeline.yaml") if not config_path.exists(): logger.warning(f"Config file not found: {config_path}. Using defaults.") return try: state.pipeline_config = load_config(config_path) logger.info(f"Loaded pipeline config from {config_path}") except Exception as e: logger.error(f"Failed to load pipeline config: {e}") return # Load model based on config inference_config = state.pipeline_config.stages.inference if not inference_config.enabled: logger.warning("Inference stage is disabled in config") return # Load model try: if inference_config.model_source == "mlflow": # Configure MLflow tracking URI mlflow.set_tracking_uri(state.pipeline_config.stages.training.mlflow.tracking_uri) state.model, state.model_info = load_model_from_mlflow( model_name=inference_config.mlflow_model_name, stage=inference_config.mlflow_model_stage ) elif inference_config.model_source == "local": state.model, state.model_info = load_model_from_local( model_path=inference_config.local_model_path ) else: logger.error(f"Unknown model_source: {inference_config.model_source}") return logger.info("Model loaded successfully") logger.info(f"Model info: {state.model_info['model_name']} " f"v{state.model_info['model_version']} " f"({state.model_info['model_type']})") except Exception as e: logger.error(f"Failed to load model: {e}") logger.warning("Service will start without a model. Use /health to check status.") # --- Health Check --- @app.get("/health", response_model=HealthResponse) async def health_check(): """ Health check endpoint. Returns service status, model loaded status, and dependency health. """ model_loaded = state.model is not None # Check MLflow connection mlflow_status = "disconnected" try: # TODO: Actually check MLflow connection mlflow_status = "connected" except Exception: pass # Check database connection db_status = "disconnected" try: # TODO: Actually check database connection db_status = "connected" except Exception: pass # Determine overall status if model_loaded and mlflow_status == "connected" and db_status == "connected": overall_status = "healthy" elif model_loaded: overall_status = "degraded" else: overall_status = "unhealthy" return HealthResponse( status=overall_status, model_loaded=model_loaded, mlflow=mlflow_status, database=db_status ) # --- Model Info Endpoints (Stubs) --- @app.get("/model/info", response_model=ModelInfoResponse) async def get_model_info(): """ Get detailed model information and per-class metrics. Returns HTTP 503 if no model is loaded. """ if state.model is None or state.model_info is None: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="No model available" ) # Build response from loaded model info return ModelInfoResponse( model_name=state.model_info["model_name"], model_version=state.model_info.get("model_version"), model_type=state.model_info["model_type"], trained_at=state.model_info.get("trained_at"), dataset_version=state.model_info.get("dataset_version"), feature_engineering_enabled=state.model_info["feature_engineering_enabled"], labels=state.model_info.get("labels", []), per_class_metrics=[ PerClassMetrics(**metric) for metric in state.model_info.get("per_class_metrics", []) ] ) @app.get("/model/labels", response_model=List[LabelInfo]) async def get_model_labels(): """ Get all pattern labels the model can predict with display colors. """ if state.model is None: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="No model available" ) # TODO: Load label colors from config or database # For now, return placeholder labels = state.model_info.get("labels", []) if state.model_info else [] return [ LabelInfo(name=label, color="#888888") for label in labels ] # --- Prediction Helper Functions --- def group_prediction_spans( predictions: List[PredictionResult] ) -> List[PredictionSpan]: """ Group consecutive predictions with the same label into spans. Only groups non-"O" (non-background) predictions. "O" predictions are treated as background and not grouped into spans. Args: predictions: List of per-candle predictions Returns: List of prediction spans """ if not predictions: return [] spans = [] current_span = None for pred in predictions: # Skip "O" (background) predictions if pred.label == "O": if current_span is not None: # End current span spans.append(current_span) current_span = None continue # Start new span or continue current span if current_span is None or current_span.label != pred.label: # End previous span if exists if current_span is not None: spans.append(current_span) # Start new span current_span = PredictionSpan( start_time=pred.time, end_time=pred.time, label=pred.label, avg_confidence=pred.confidence, candle_count=1 ) else: # Continue current span current_span.end_time = pred.time current_span.candle_count += 1 # Update running average confidence total_confidence = current_span.avg_confidence * (current_span.candle_count - 1) current_span.avg_confidence = (total_confidence + pred.confidence) / current_span.candle_count # Don't forget last span if current_span is not None: spans.append(current_span) logger.info(f"Grouped {len(predictions)} predictions into {len(spans)} spans") return spans # --- Prediction Endpoints --- @app.post("/predict", response_model=PredictResponse) async def predict(request: PredictRequest): """ Predict candlestick patterns for provided candles. Accepts OHLCV candles, runs preprocessing, and returns predictions. """ if state.model is None: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="No model available" ) if not request.candles: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Candles array cannot be empty" ) if state.pipeline_config is None: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Pipeline configuration not loaded" ) logger.info(f"Predict request: {request.pair or 'unknown'} {request.timeframe or 'unknown'}, {len(request.candles)} candles") try: # Convert candles to list of dicts candles_data = [candle.model_dump() for candle in request.candles] # Preprocess candles (feature engineering + windowing) X, window_times = preprocess_candles(candles_data, state.pipeline_config) # Get predictions and probabilities if hasattr(state.model, 'predict_proba'): y_pred = state.model.predict(X) y_proba = state.model.predict_proba(X) # Get confidence as max probability confidences = np.max(y_proba, axis=1) else: # Fallback for models without predict_proba y_pred = state.model.predict(X) confidences = np.ones(len(y_pred)) # Default confidence of 1.0 # Get label names (handle both string and int predictions) if state.label_encoder is not None: labels = [state.label_encoder.get(int(pred), f"unknown_{pred}") for pred in y_pred] else: labels = [str(pred) for pred in y_pred] # Build per-window predictions (each window maps to its last candle time) predictions = [ PredictionResult( time=int(time), label=label, confidence=float(conf) ) for time, label, conf in zip(window_times, labels, confidences) ] # Group into spans spans = group_prediction_spans(predictions) # Build model info for response model_info = ModelInfo( model_name=state.model_info["model_name"], model_version=state.model_info.get("model_version"), model_type=state.model_info["model_type"], trained_at=state.model_info.get("trained_at"), dataset_version=state.model_info.get("dataset_version"), feature_engineering_enabled=state.model_info["feature_engineering_enabled"] ) logger.info( f"Prediction complete: {len(predictions)} windows, " f"{len(spans)} spans, {len([p for p in predictions if p.label != 'O'])} patterns" ) return PredictResponse( predictions=predictions, spans=spans, model_info=model_info ) except ValueError as e: logger.error(f"Prediction validation error: {e}") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=str(e) ) except Exception as e: logger.error(f"Prediction failed: {e}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Prediction failed: {str(e)}" ) @app.post("/predict/batch", response_model=PredictResponse) async def predict_batch(request: BatchPredictRequest): """ Batch prediction for a date range. Loads data from storage and returns predictions for the full range. """ if state.model is None: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="No model available" ) if state.pipeline_config is None: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Pipeline configuration not loaded" ) logger.info( f"Batch predict: {request.pair} {request.timeframe} " f"from {request.start_date} to {request.end_date}" ) try: # Load OHLCV data from raw data path raw_path = Path(state.pipeline_config.data.raw_path) if not raw_path.exists(): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"No data found for {request.pair} {request.timeframe}" ) # Load data df_raw = pd.read_csv(raw_path) # Filter by date range if time column exists if 'time' in df_raw.columns: # Parse dates to timestamps start_ts = int(pd.Timestamp(request.start_date).timestamp()) end_ts = int(pd.Timestamp(request.end_date).timestamp()) # Filter data df_filtered = df_raw[ (df_raw['time'] >= start_ts) & (df_raw['time'] <= end_ts) ].copy() if len(df_filtered) == 0: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"No data found in date range {request.start_date} to {request.end_date}" ) logger.info(f"Loaded {len(df_filtered)} candles from {request.start_date} to {request.end_date}") else: df_filtered = df_raw logger.warning("No 'time' column found, using all data") # Get batch size from config batch_size = state.pipeline_config.stages.inference.batch_size # Process in batches all_predictions = [] all_spans = [] for i in range(0, len(df_filtered), batch_size): batch_df = df_filtered.iloc[i:i + batch_size] logger.info(f"Processing batch {i // batch_size + 1}: rows {i} to {i + len(batch_df)}") # Convert batch to candles format batch_candles = batch_df.to_dict('records') # Preprocess (feature engineering + windowing) X, window_times = preprocess_candles(batch_candles, state.pipeline_config) # Predict if hasattr(state.model, 'predict_proba'): y_pred = state.model.predict(X) y_proba = state.model.predict_proba(X) confidences = np.max(y_proba, axis=1) else: y_pred = state.model.predict(X) confidences = np.ones(len(y_pred)) # Get labels if state.label_encoder is not None: labels = [state.label_encoder.get(int(pred), f"unknown_{pred}") for pred in y_pred] else: labels = [str(pred) for pred in y_pred] # Build predictions for this batch batch_predictions = [ PredictionResult( time=int(time), label=label, confidence=float(conf) ) for time, label, conf in zip(window_times, labels, confidences) ] all_predictions.extend(batch_predictions) # Group all predictions into spans all_spans = group_prediction_spans(all_predictions) # Build model info model_info = ModelInfo( model_name=state.model_info["model_name"], model_version=state.model_info.get("model_version"), model_type=state.model_info["model_type"], trained_at=state.model_info.get("trained_at"), dataset_version=state.model_info.get("dataset_version"), feature_engineering_enabled=state.model_info["feature_engineering_enabled"] ) logger.info( f"Batch prediction complete: {len(all_predictions)} candles, " f"{len(all_spans)} spans" ) return PredictResponse( predictions=all_predictions, spans=all_spans, model_info=model_info ) except HTTPException: raise except Exception as e: logger.error(f"Batch prediction failed: {e}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Batch prediction failed: {str(e)}" ) # --------------------------------------------------------------------------- # Pattern Detection Endpoints # --------------------------------------------------------------------------- class PatternInfo(BaseModel): """A single supported CDL pattern.""" function_name: str display_name: str class DetectPatternsRequest(BaseModel): """Request model for POST /patterns/detect.""" candles: List[CandleData] = Field(..., min_length=1, description="Array of candle data") patterns: List[str] = Field( default=[], description="CDL function names to run. Empty list means run all.", ) class SpanAnnotation(BaseModel): """A span annotation returned by pattern detection.""" start_time: int end_time: int label: str confidence: float = Field(..., ge=0.0, le=1.0) source: str notes: str class DetectPatternsResponse(BaseModel): """Response model for POST /patterns/detect.""" annotations: List[SpanAnnotation] metadata: Dict[str, Any] @app.get("/patterns/available", response_model=List[PatternInfo]) async def patterns_available(): """ Return all supported CDL pattern names with display names. """ return get_available_patterns() @app.post("/patterns/detect", response_model=DetectPatternsResponse) async def patterns_detect(request: DetectPatternsRequest): """ Detect TA-Lib CDL patterns on provided candle data. - Empty ``patterns`` list runs all available CDL functions. - Invalid pattern names return HTTP 400. """ # 1.4 – Validate pattern names if request.patterns: invalid = validate_pattern_names(request.patterns) if invalid: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid pattern name(s): {', '.join(invalid)}. " f"Use GET /patterns/available to see supported patterns.", ) candles_data = [c.model_dump() for c in request.candles] try: raw_annotations = detect_patterns(candles_data, request.patterns or None) except RuntimeError as exc: raise HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=str(exc), ) annotations = [SpanAnnotation(**ann) for ann in raw_annotations] return DetectPatternsResponse( annotations=annotations, metadata={"source": "talib", "count": len(annotations)}, ) # --------------------------------------------------------------------------- # Training Endpoints # --------------------------------------------------------------------------- SUPPORTED_MODEL_TYPES = ["random_forest", "xgboost"] class TrainingStartRequest(BaseModel): """Request model for POST /training/start.""" model_type: str = Field( "random_forest", description="Model type: random_forest or xgboost", ) class TrainingStartResponse(BaseModel): """Response model for POST /training/start.""" run_id: str status: str class TrainingRunInfo(BaseModel): """Summary of a single training run.""" run_id: str model_type: str status: str experiment_name: Optional[str] = None created_at: Optional[str] = None completed_at: Optional[str] = None metrics_summary: Optional[Dict[str, Any]] = None class TrainingRunsResponse(BaseModel): """Response model for GET /training/runs.""" runs: List[TrainingRunInfo] class DatasetInfoResponse(BaseModel): """Response model for GET /training/dataset-info.""" path: str exists: bool size_bytes: Optional[int] = None last_modified: Optional[str] = None row_count: Optional[int] = None def build_dataset_from_db(config: PipelineConfig) -> dict: """ Build the labeled training dataset directly from the database. Steps: 1. Export candles from PostgreSQL to raw CSV 2. Run feature engineering (TA-Lib indicators, candle features) 3. Run annotation ingestion from DB (span_annotations -> labeled CSV) Returns: dict with keys: chart_name, n_candles, n_annotations, n_samples, labeled_path """ from app.data_access import DataAccess from app.annotation_ingestion import AnnotationIngestion from features.engineer import run_feature_engineering_stage data_access = DataAccess() # Find all charts, use the first one (single-chart app) charts_df = data_access.get_all_charts() if charts_df.empty: raise ValueError("No charts found in database. Upload candle data first.") chart = charts_df.iloc[0] chart_name = chart["name"] chart_id = int(chart["id"]) logger.info(f"Building dataset for chart: {chart_name} (id={chart_id})") # Step 1: Export candles to raw CSV candles_df = data_access.get_candles(chart_id) if candles_df.empty: raise ValueError(f"No candles found for chart: {chart_name}") raw_path = Path(config.data.raw_path) raw_path.parent.mkdir(parents=True, exist_ok=True) # Ensure 'time' column is suitable for feature engineering export_df = candles_df[["time", "open", "high", "low", "close"]].copy() export_df.to_csv(raw_path, index=False) logger.info(f"Exported {len(export_df)} candles to {raw_path}") # Step 2: Run feature engineering run_feature_engineering_stage(config) enriched_path = Path(config.data.enriched_path) logger.info(f"Feature engineering complete: {enriched_path}") # Step 3: Run annotation ingestion from database enriched_df = pd.read_csv(enriched_path, parse_dates=["time"]) ingestion = AnnotationIngestion(config.stages.annotation_ingestion) labeled_df = ingestion.process_from_db(enriched_df, chart_name, source="human") if labeled_df.empty: raise ValueError( f"No labeled samples produced. " f"Ensure you have span annotations on chart '{chart_name}'." ) # Write labeled dataset labeled_path = Path(config.data.labeled_path) labeled_path.parent.mkdir(parents=True, exist_ok=True) labeled_df.to_csv(labeled_path, index=False) result = { "chart_name": chart_name, "n_candles": len(export_df), "n_samples": len(labeled_df), "n_features": len([c for c in labeled_df.columns if c != "label"]), "labeled_path": str(labeled_path), } logger.info(f"Dataset built: {result}") return result def _run_training_background(run_id: str, model_type: str, config: PipelineConfig) -> None: """ Background thread target: build dataset then train a model. Uses the pre-inserted TrainingRun record identified by ``run_id``. """ logger.info(f"Training thread started: run_id={run_id}, model_type={model_type}") try: # Import training utilities here to avoid circular import issues from training.train import create_model, temporal_split from sklearn.metrics import accuracy_score, f1_score # Build dataset from database (feature engineering + annotation ingestion) logger.info("Building dataset from database...") build_dataset_from_db(config) labeled_path = Path(config.data.labeled_path) if not labeled_path.exists(): raise FileNotFoundError(f"Labeled dataset not found: {labeled_path}") # Load dataset df = pd.read_csv(labeled_path) if "label" not in df.columns: raise ValueError("Labeled dataset must have 'label' column") feature_cols = [ col for col in df.columns if col not in ("label", "time", "timestamp") and not col.startswith("label_programmatic_") ] X = df[feature_cols].values y = df["label"].values logger.info(f"Loaded {len(X)} samples, {len(feature_cols)} features") # Split data training_cfg = config.stages.training X_train, X_val, X_test, y_train, y_val, y_test = temporal_split( X, y, training_cfg.test_split, training_cfg.validation_split ) # Train model model_instance = create_model( model_type, training_cfg.hyperparameters, training_cfg.class_weights ) model_instance.fit(X_train, y_train) logger.info("Model training complete") # Evaluate y_val_pred = model_instance.predict(X_val) y_test_pred = model_instance.predict(X_test) metrics = { "val_accuracy": float(accuracy_score(y_val, y_val_pred)), "val_f1_macro": float( f1_score(y_val, y_val_pred, average="macro", zero_division=0) ), "test_accuracy": float(accuracy_score(y_test, y_test_pred)), "test_f1_macro": float( f1_score(y_test, y_test_pred, average="macro", zero_division=0) ), "n_samples": int(len(X)), "n_features": int(X.shape[1]), } # Save model locally models_dir = Path("models") models_dir.mkdir(exist_ok=True) model_path = models_dir / f"{run_id}.pkl" model_data = { "model": model_instance, "metadata": { "model_type": model_type, "trained_at": datetime.utcnow().isoformat(), "run_id": run_id, "feature_columns": feature_cols, "labels": ( [str(c) for c in model_instance.model.classes_] if hasattr(model_instance, "model") and hasattr(model_instance.model, "classes_") else [] ), }, } joblib.dump(model_data, model_path) logger.info(f"Model saved to {model_path}") # Update DB: completed with get_db() as db: stmt = ( sa_update(TrainingRun) .where(TrainingRun.run_id == run_id) .values( status="completed", completed_at=datetime.utcnow(), metrics_summary=metrics, ) ) db.execute(stmt) db.commit() logger.info(f"Training completed successfully: run_id={run_id}") except Exception as exc: logger.error(f"Training failed for run_id={run_id}: {exc}", exc_info=True) try: with get_db() as db: stmt = ( sa_update(TrainingRun) .where(TrainingRun.run_id == run_id) .values( status="failed", completed_at=datetime.utcnow(), metrics_summary={"error": str(exc)}, ) ) db.execute(stmt) db.commit() except Exception as db_exc: logger.error(f"Failed to update DB for failed run {run_id}: {db_exc}") finally: with state.training_lock: if state.active_training_run_id == run_id: state.active_training_run_id = None logger.info(f"Training thread exiting: run_id={run_id}") @app.post("/training/start", response_model=TrainingStartResponse) async def training_start(request: TrainingStartRequest): """ Start a training run in a background thread. Returns immediately with run_id and status "running". Rejects concurrent runs with HTTP 409. """ # Validate model type if request.model_type not in SUPPORTED_MODEL_TYPES: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Unsupported model type. Available: {', '.join(SUPPORTED_MODEL_TYPES)}", ) # Reject concurrent runs (atomic check-and-set) with state.training_lock: if state.active_training_run_id is not None: raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail={ "error": "Training already in progress", "run_id": state.active_training_run_id, }, ) run_id = str(uuid_lib.uuid4()) state.active_training_run_id = run_id config = state.pipeline_config or get_default_config() # Compute config hash (best-effort) config_hash = "unknown" try: from training.train import compute_config_hash config_hash = compute_config_hash(config) except Exception: pass # Pre-insert the run record so callers can track it immediately try: with get_db() as db: training_run = TrainingRun( run_id=run_id, model_type=request.model_type, experiment_name=config.stages.training.mlflow.experiment_name, pipeline_config_hash=config_hash, status="running", created_at=datetime.utcnow(), metrics_summary={}, ) db.add(training_run) db.commit() except Exception as exc: with state.training_lock: state.active_training_run_id = None logger.error(f"Failed to insert training run record: {exc}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to create training run record: {exc}", ) # Launch background thread (daemon so it doesn't block process exit) thread = threading.Thread( target=_run_training_background, args=(run_id, request.model_type, config), daemon=True, name=f"training-{run_id[:8]}", ) thread.start() logger.info(f"Training started: run_id={run_id}, model_type={request.model_type}") return TrainingStartResponse(run_id=run_id, status="running") @app.get("/training/runs", response_model=TrainingRunsResponse) async def training_runs(): """ Return training run history from the database, sorted by date descending. """ try: from sqlalchemy import select with get_db() as db: stmt = select(TrainingRun).order_by(desc(TrainingRun.created_at)) rows = db.execute(stmt).scalars().all() runs = [ TrainingRunInfo( run_id=row.run_id, model_type=row.model_type, status=row.status, experiment_name=row.experiment_name, created_at=row.created_at.isoformat() if row.created_at else None, completed_at=row.completed_at.isoformat() if row.completed_at else None, metrics_summary=row.metrics_summary, ) for row in rows ] except Exception as exc: logger.error(f"Failed to fetch training runs: {exc}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to fetch training runs: {exc}", ) return TrainingRunsResponse(runs=runs) class ActiveTrainingResponse(BaseModel): """Response model for GET /training/active.""" active: bool run_id: Optional[str] = None @app.get("/training/active", response_model=ActiveTrainingResponse) async def training_active(): """ Return whether a training run is currently active and its run_id. """ with state.training_lock: run_id = state.active_training_run_id return ActiveTrainingResponse(active=run_id is not None, run_id=run_id) class DeleteRunResponse(BaseModel): """Response model for DELETE /training/runs/{run_id}.""" run_id: str deleted: bool @app.delete("/training/runs/{run_id}", response_model=DeleteRunResponse) async def delete_training_run(run_id: str): """ Delete a training run record and its model artifact. Returns HTTP 409 if the run is currently active. Returns HTTP 404 if the run_id doesn't exist. """ from sqlalchemy import select, delete as sa_delete # Reject deletion of the active run with state.training_lock: if state.active_training_run_id == run_id: raise HTTPException( status_code=status.HTTP_409_CONFLICT, detail="Cannot delete an active training run", ) try: with get_db() as db: stmt = select(TrainingRun).where(TrainingRun.run_id == run_id) row = db.execute(stmt).scalar_one_or_none() if row is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Training run not found: {run_id}", ) db.execute(sa_delete(TrainingRun).where(TrainingRun.run_id == run_id)) db.commit() except HTTPException: raise except Exception as exc: logger.error(f"Failed to delete training run {run_id}: {exc}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to delete training run: {exc}", ) # Remove model artifact if it exists model_path = Path("models") / f"{run_id}.pkl" if model_path.exists(): try: model_path.unlink() logger.info(f"Deleted model artifact: {model_path}") except Exception as exc: logger.warning(f"Could not delete model artifact {model_path}: {exc}") logger.info(f"Deleted training run: {run_id}") return DeleteRunResponse(run_id=run_id, deleted=True) @app.get("/training/dataset-info", response_model=DatasetInfoResponse) async def training_dataset_info(): """ Return information about the labeled training dataset. Includes file path, existence, size, last modified date, and row count. """ config = state.pipeline_config or get_default_config() labeled_path = Path(config.data.labeled_path) if not labeled_path.exists(): return DatasetInfoResponse(path=str(labeled_path), exists=False) try: stat = labeled_path.stat() size_bytes = stat.st_size last_modified = datetime.fromtimestamp(stat.st_mtime).isoformat() row_count = None try: # Read only one column for efficiency df_head = pd.read_csv(labeled_path, usecols=[0]) row_count = len(df_head) except Exception: pass return DatasetInfoResponse( path=str(labeled_path), exists=True, size_bytes=size_bytes, last_modified=last_modified, row_count=row_count, ) except Exception as exc: logger.error(f"Failed to get dataset info: {exc}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to get dataset info: {exc}", ) class BuildDatasetResponse(BaseModel): """Response model for POST /training/build-dataset.""" chart_name: str n_candles: int n_samples: int n_features: int labeled_path: str @app.post("/training/build-dataset", response_model=BuildDatasetResponse) async def training_build_dataset(): """ Build the labeled training dataset from database annotations. Exports candles, runs feature engineering, and ingests span annotations into a labeled CSV ready for training. """ config = state.pipeline_config or get_default_config() try: result = build_dataset_from_db(config) return BuildDatasetResponse(**result) except ValueError as exc: raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) except Exception as exc: logger.error(f"Failed to build dataset: {exc}", exc_info=True) raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to build dataset: {exc}", ) # --------------------------------------------------------------------------- # Model Loading Endpoint # --------------------------------------------------------------------------- class ModelLoadRequest(BaseModel): """Request model for POST /model/load.""" run_id: str = Field(..., description="Training run ID to load model from") class ModelLoadResponse(BaseModel): """Response model for POST /model/load.""" run_id: str model_type: str status: str # Lock protecting model hot-swap _model_swap_lock = threading.Lock() @app.post("/model/load", response_model=ModelLoadResponse) async def model_load(request: ModelLoadRequest): """ Load a trained model from a completed training run. Looks up the run_id in the training_runs table, loads the model artifact, and replaces the active model in AppState with a brief lock to prevent conflicts with in-flight prediction requests. """ from sqlalchemy import select # 1. Look up the training run try: with get_db() as db: stmt = select(TrainingRun).where(TrainingRun.run_id == request.run_id) row = db.execute(stmt).scalar_one_or_none() except Exception as exc: logger.error(f"DB lookup failed for run_id={request.run_id}: {exc}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Database error: {exc}", ) if row is None: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Training run not found: {request.run_id}", ) if row.status != "completed": raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"Training run is not completed (status={row.status})", ) # 2. Resolve model artifact path model_path = Path("models") / f"{request.run_id}.pkl" if not model_path.exists(): raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"Model artifact not found at {model_path}", ) # 3. Load model (outside lock – can be slow) try: new_model, new_model_info = load_model_from_local(str(model_path)) except Exception as exc: logger.error(f"Failed to load model for run_id={request.run_id}: {exc}") raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to load model: {exc}", ) # 4. Thread-safe model swap (3.2 – brief lock) with _model_swap_lock: state.model = new_model state.model_info = new_model_info logger.info(f"Model hot-swapped: run_id={request.run_id}, type={row.model_type}") return ModelLoadResponse( run_id=request.run_id, model_type=row.model_type, status="loaded", ) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8001)