streamlit_healthcheck.server

  1from fastapi import FastAPI, HTTPException
  2from fastapi.responses import JSONResponse
  3from typing import Dict, Any, Optional
  4import uvicorn
  5from .healthcheck import HealthCheckService
  6import logging
  7import argparse
  8from contextlib import asynccontextmanager
  9"""
 10This module implements a FastAPI server for monitoring the health of a Streamlit application.
 11Features:
 12- Provides REST API endpoints to report health status of system resources, dependencies, and Streamlit pages.
 13- Uses a configurable health check service to gather health data.
 14- Supports startup and shutdown events for proper initialization and cleanup of the health check service.
 15- Includes endpoints:
 16    - `/health`: Returns complete health status.
 17    - `/health/system`: Returns system resource health (CPU, Memory, Disk).
 18    - `/health/dependencies`: Returns dependencies health status.
 19    - `/health/pages`: Returns Streamlit pages health status.
 20- Configurable via command-line arguments for host, port, and health check config file.
 21- Uses logging for operational visibility.
 22Global Variables:
 23- `health_service`: Instance of HealthCheckService, manages health checks.
 24- `config_file`: Path to health check configuration file.
 25Usage:
 26Run as a standalone server or import as a module.
 27"""
 28
 29
 30# Set up logging
 31logging.basicConfig(
 32    level=logging.INFO,
 33    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
 34    handlers=[
 35        logging.StreamHandler()  # Prints to console
 36    ]
 37)
 38logger = logging.getLogger(__name__)
 39
 40# Initialize health check service as global variable
 41health_service: Optional[HealthCheckService] = None
 42config_file: Optional[str] = None
 43
 44@asynccontextmanager
 45async def lifespan(app: FastAPI):
 46    """Lifespan context manager for FastAPI application"""
 47    # Startup
 48    global health_service
 49    logger.info("Initializing health check service")
 50    config_path = config_file if config_file is not None else "health_check_config.json"
 51    health_service = HealthCheckService(config_path=config_path)
 52    health_service.start()
 53    
 54    yield
 55    
 56    # Shutdown
 57    if health_service:
 58        logger.info("Stopping health check service")
 59        health_service.stop()
 60
 61app = FastAPI(
 62    title="Streamlit Health Check API",
 63    description="API endpoints for monitoring Streamlit application health",
 64    version="1.0.0",
 65    lifespan=lifespan
 66)
 67
 68# Initialize health check service and config file path as global variables
 69# health_service and config_file are already defined above
 70
 71
 72@app.get("/health", response_model=Dict[str, Any])
 73async def get_health_status():
 74    """
 75    Get complete health check status including system resources, 
 76    dependencies, custom checks, and Streamlit pages
 77    """
 78    global health_service
 79    if not health_service:
 80        raise HTTPException(status_code=503, detail="Health check service not initialized")
 81    
 82    try:
 83        health_data = health_service.get_health_data()
 84        return JSONResponse(content=health_data)
 85    except Exception as e:
 86        logger.error(f"Error getting health data: {str(e)}")
 87        raise HTTPException(status_code=500, detail=str(e))
 88
 89@app.get("/health/system", response_model=Dict[str, Any])
 90async def get_system_health():
 91    """Get system resource health status (CPU, Memory, Disk)"""
 92    global health_service
 93    if not health_service:
 94        raise HTTPException(status_code=503, detail="Health check service not initialized")
 95    
 96    try:
 97        health_data = health_service.get_health_data()
 98        return JSONResponse(content={"system": health_data.get("system", {})})
 99    except Exception as e:
100        logger.error(f"Error getting system health data: {str(e)}")
101        raise HTTPException(status_code=500, detail=str(e))
102
103@app.get("/health/dependencies", response_model=Dict[str, Any])
104async def get_dependencies_health():
105    """Get dependencies health status"""
106    global health_service
107    if not health_service:
108        raise HTTPException(status_code=503, detail="Health check service not initialized")
109    
110    try:
111        health_data = health_service.get_health_data()
112        return JSONResponse(content={"dependencies": health_data.get("dependencies", {})})
113    except Exception as e:
114        logger.error(f"Error getting dependencies health data: {str(e)}")
115        raise HTTPException(status_code=500, detail=str(e))
116
117@app.get("/health/pages", response_model=Dict[str, Any])
118async def get_pages_health():
119    """Get Streamlit pages health status"""
120    global health_service
121    if not health_service:
122        raise HTTPException(status_code=503, detail="Health check service not initialized")
123    
124    try:
125        health_data = health_service.get_health_data()
126        return JSONResponse(content={"streamlit_pages": health_data.get("streamlit_pages", {})})
127    except Exception as e:
128        logger.error(f"Error getting pages health data: {str(e)}")
129        raise HTTPException(status_code=500, detail=str(e))
130
131def start_api_server(host: str = "0.0.0.0", port: int = 8000, config: str = "health_check_config.json"):
132    """Start the FastAPI server"""
133    global config_file
134    config_file = config
135    uvicorn.run(app, host=host, port=port)
136
137def parse_args():
138    """Parse command line arguments"""
139    parser = argparse.ArgumentParser(description="Streamlit Health Check API Server")
140    parser.add_argument("--host", default="0.0.0.0", help="Host address to bind")
141    parser.add_argument("--port", type=int, default=8000, help="Port to run the server on")
142    parser.add_argument(
143        "--config", 
144        default="health_check_config.json",
145        help="Path to health check configuration file"
146    )
147    return parser.parse_args()
148
149if __name__ == "__main__":
150    args = parse_args()
151    logger.info(f"Starting server with config file: {args.config}")
152    start_api_server(host=args.host, port=args.port, config=args.config)