Files
OCR-SPRIN-SERVICE/src/ocr_sprint/main.py
Nama Kamu 9d969e61fd update
2026-04-26 22:08:41 +08:00

108 lines
3.5 KiB
Python

"""FastAPI entrypoint."""
from __future__ import annotations
import threading
from contextlib import asynccontextmanager
from typing import AsyncIterator
from fastapi import FastAPI
from ocr_sprint import __version__
from ocr_sprint.api.errors import register_error_handlers
from ocr_sprint.api.metrics import MetricsMiddleware, metrics_endpoint
from ocr_sprint.api.routes import documents, ground_truth, health
from ocr_sprint.config import get_settings
from ocr_sprint.db import models as _models # noqa: F401 (register ORM tables)
from ocr_sprint.db.base import Base, get_engine
from ocr_sprint.utils.logging import configure_logging, get_logger
_startup_logger = get_logger(__name__)
def _ensure_schema() -> None:
"""Create tables if they don't exist.
Production deploys should run Alembic migrations explicitly; this is a
convenience for local dev / tests so the API works without a manual
`alembic upgrade head` step.
"""
Base.metadata.create_all(bind=get_engine())
def _warmup_models_background() -> None:
"""Load PaddleOCR and PP-Structure models in a background thread.
Running in a thread keeps the lifespan non-blocking so uvicorn can
start accepting health-check requests immediately while the heavy models
load (~5-15s on CPU). Requests that arrive before warmup completes will
wait on the existing _lock in each module rather than racing to load.
"""
from ocr_sprint.config import get_settings as _gs
from ocr_sprint.pipeline import ocr as _ocr
from ocr_sprint.pipeline import table as _table
s = _gs()
try:
_ocr.warmup()
except Exception as exc:
_startup_logger.warning("paddleocr.warmup.failed", error=str(exc))
if s.tables_enabled:
try:
_table.warmup()
except Exception as exc:
_startup_logger.warning("pp_structure.warmup.failed", error=str(exc))
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncIterator[None]:
"""FastAPI lifespan: warm OCR models on startup in a background thread."""
_startup_logger.info("startup.warmup.begin")
t = threading.Thread(target=_warmup_models_background, name="ocr-warmup", daemon=True)
t.start()
yield
# Shutdown: nothing to clean up (models are process-global singletons).
_startup_logger.info("shutdown.complete")
def create_app() -> FastAPI:
"""Application factory — keeps top-level state easy to test."""
settings = get_settings()
configure_logging(settings.app_log_level)
_ensure_schema()
# Support for reverse proxy with path prefix (e.g., /ocr)
root_path = getattr(settings, "root_path", "")
app = FastAPI(
lifespan=lifespan,
title="OCR Sprint Service",
version=__version__,
description="OCR + structured extraction for Indonesian police 'surat sprint' documents.",
docs_url="/docs",
redoc_url="/redoc",
openapi_url="/openapi.json",
root_path=root_path,
)
register_error_handlers(app)
app.add_middleware(MetricsMiddleware)
app.include_router(health.router, prefix="/api/v1")
app.include_router(documents.router, prefix="/api/v1")
app.include_router(ground_truth.router, prefix="/api/v1")
app.add_api_route("/metrics", metrics_endpoint, methods=["GET"], include_in_schema=False)
return app
app = create_app()
def run() -> None:
"""Console-script entrypoint (`ocr-sprint-api`)."""
import uvicorn
s = get_settings()
uvicorn.run("ocr_sprint.main:app", host=s.app_host, port=s.app_port, reload=False)