TeamGenKI/Inference-API

1"""2LLM Inference Server main application using LitServe framework.3"""4from sys import platform5 6import litserve as ls7import logging8import os9from fastapi.middleware.cors import CORSMiddleware10from huggingface_hub import login11from .routes import router, init_router12from .api import InferenceApi13from .utils import load_config14 15# Store process list globally so it doesn't get garbage collected16_WORKER_PROCESSES = []17_MANAGER = None18 19# Load configuration20config = load_config()21 22 23def setup_logging():24    """Set up basic logging configuration"""25    logging.basicConfig(26        level=logging.DEBUG,27        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'28    )29    return logging.getLogger(__name__)30 31 32def create_app():33    """Create and configure the application instance."""34    global _WORKER_PROCESSES, _MANAGER, config35 36    logger = setup_logging()37 38    # Log into Hugging Face Hub39    access_token = os.environ.get("InfAPITokenWrite")40    if access_token:41        try:42            login(token=access_token)43            logger.info("Successfully logged into Hugging Face Hub")44        except Exception as e:45            logger.error(f"Failed to login to Hugging Face Hub: {str(e)}")46    else:47        logger.warning("No Hugging Face access token found")48 49    server_config = config.get('server', {})50 51    # Initialize API with config52    api = InferenceApi(config)53 54    # Initialize router with API instance55    init_router(api, config)56 57    if platform == "darwin":  # Darwin is macOS58        server = ls.LitServer(59            api,60            timeout=server_config.get('timeout', 60),61            max_batch_size=server_config.get('max_batch_size', 1),62            track_requests=True,63            accelerator="cpu"  # Force CPU on Mac64        )65    else:66        server = ls.LitServer(67            api,68            timeout=server_config.get('timeout', 60),69            max_batch_size=server_config.get('max_batch_size', 1),70            track_requests=True71        )72 73    # Launch inference workers (assuming single uvicorn worker for now)74    _MANAGER, _WORKER_PROCESSES = server.launch_inference_worker(num_uvicorn_servers=1)75 76    # Get the FastAPI appls77 78    app = server.app79 80    # Add CORS middleware81    app.add_middleware(82        CORSMiddleware,83        allow_origins=["*"],84        allow_credentials=True,85        allow_methods=["*"],86        allow_headers=["*"],87    )88 89    # Add routes with configured prefix90    api_prefix = config.get('llm_server', {}).get('api_prefix', '/api/v1')91    app.include_router(router, prefix=api_prefix)92 93    # Set the response queue ID for the app94    app.response_queue_id = 0  # Since we're using a single worker95 96    return app97 98# Create the app instance for uvicorn99app = create_app()100 101if __name__ == "__main__":102    # Run the app with uvicorn103    import uvicorn104    host = config["server"]["host"]105    port = config["server"]["port"]106    uvicorn.run(107        app,108        host=host,109        port=port,110        log_level=config["logging"]["level"].lower()111    )