Merge pull request #14 from rhel-lightspeed/perf/lifespan-deferred-imports-3047

perf: defer heavy imports and adopt FastMCP lifespan

Author: major

CHANGELOG.md

@@ -24,6 +24,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
+- Deferred heavy `docs2db-api` imports (torch, transformers) in `engine.py` to reduce module load time
+- Migrated startup health check and engine shutdown to FastMCP lifespan hook for proper async lifecycle management
+- Simplified `__main__.py` — removed manual `asyncio.run()` calls for health check and cleanup
 - Upgraded fastmcp from 2.x to 3.x (`>=3.3.1, <4`)
 - Removed standalone `mcp` dependency (pulled in transitively by fastmcp 3.x)
 - Updated smoke tests to match FastMCP 3.x decorator behavior (`@mcp.tool` now returns the original function)

pyproject.toml

@@ -105,6 +105,3 @@ force-single-line = true
 lines-after-imports = 2
 lines-between-types = 1
 order-by-type = false
-
-
-

src/docs2db_mcp/__main__.py

@@ -1,6 +1,5 @@
 """Main entry point for docs2db MCP server."""
 
-import asyncio
 import logging
 import os
 import sys
@@ -28,43 +27,10 @@
 
 # Import configuration (lightweight, doesn't trigger heavy imports)
 from docs2db_mcp.config import CONFIG  # noqa: E402
-
-
-logger = logging.getLogger(__name__)
-
-
-def _configure_structlog_for_stdio() -> None:
-    """Override docs2db-api's structlog configuration to redirect stdout to stderr.
-
-    docs2db-api configures structlog at import time to write to stdout, but stdio transport
-    requires stdout to be reserved exclusively for MCP protocol messages.
-    """
-    import structlog
-
-    # Reconfigure structlog to output to stderr instead of stdout
-    # Keep existing processors, just redirect the output stream
-    structlog.configure(
-        logger_factory=structlog.PrintLoggerFactory(file=sys.stderr),
-        cache_logger_on_first_use=False,
-    )
-
-
-# Import server and engine modules
-# Must happen after reading transport to configure structlog before they import docs2db-api
-from docs2db_mcp.engine import health_check  # noqa: E402
-from docs2db_mcp.engine import shutdown_engine  # noqa: E402
 from docs2db_mcp.server import mcp  # noqa: E402
 
 
-# docs2db-api has already configured structlog, now override for stdio transport
-if transport != "sse":
-    _configure_structlog_for_stdio()
-
-
-async def cleanup() -> None:
-    """Cleanup resources on shutdown."""
-    logger.info("Shutting down docs2db MCP server")
-    await shutdown_engine()
+logger = logging.getLogger(__name__)
 
 
 def main() -> None:
@@ -79,25 +45,13 @@ def main() -> None:
         CONFIG.rag_enable_reranking,
     )
 
-    # Perform startup health check
-    try:
-        asyncio.run(health_check())
-    except Exception as e:
-        logger.error("Startup health check failed: %s", e)
-        logger.error("Server will not start - please check database connection and configuration")
-        sys.exit(1)
-
     try:
-        # Run the MCP server
         mcp.run(transport=CONFIG.transport, **CONFIG.transport_kwargs)
     except KeyboardInterrupt:
         logger.info("Received keyboard interrupt")
     except Exception as e:
         logger.error("Server error: %s", e, exc_info=True)
         sys.exit(1)
-    finally:
-        # Cleanup
-        asyncio.run(cleanup())
 
 
 if __name__ == "__main__":

src/docs2db_mcp/engine.py

@@ -1,13 +1,22 @@
-"""RAG engine singleton wrapper."""
+"""RAG engine singleton wrapper.
+
+Heavy dependencies (torch, transformers, docs2db-api) are imported lazily
+inside ``get_engine()`` so that importing this module is near-instant.
+"""
+
+from __future__ import annotations
 
 import logging
 
-from docs2db_api.rag.engine import RAGConfig
-from docs2db_api.rag.engine import UniversalRAGEngine
+from typing import TYPE_CHECKING
 
 from docs2db_mcp.config import CONFIG
 
 
+if TYPE_CHECKING:
+    from docs2db_api.rag.engine import UniversalRAGEngine
+
+
 logger = logging.getLogger(__name__)
 
 _engine: UniversalRAGEngine | None = None
@@ -16,6 +25,10 @@
 async def get_engine() -> UniversalRAGEngine:
     """Get or create the singleton RAG engine instance.
 
+    On first call this imports ``docs2db-api`` (which pulls in torch,
+    transformers, etc.) and initialises the engine.  Subsequent calls
+    return the cached instance.
+
     Returns:
         Initialized UniversalRAGEngine instance
 
@@ -25,6 +38,9 @@ async def get_engine() -> UniversalRAGEngine:
     global _engine
 
     if _engine is None:
+        from docs2db_api.rag.engine import RAGConfig
+        from docs2db_api.rag.engine import UniversalRAGEngine
+
         logger.info("Initializing UniversalRAGEngine")
 
         # Configure database connection (dict format)
@@ -53,7 +69,7 @@ async def get_engine() -> UniversalRAGEngine:
             await _engine.start()
             logger.info("UniversalRAGEngine initialized successfully")
         except Exception as e:
-            logger.error(f"Failed to initialize RAG engine: {e}")
+            logger.error("Failed to initialize RAG engine: %s", e)
             _engine = None
             raise
 
@@ -95,11 +111,12 @@ async def health_check() -> None:
 
         # Verify we got a valid response
         if result is None:
-            raise Exception("Health check query returned None")
+            msg = "Health check query returned None"
+            raise Exception(msg)
 
-        logger.info(f"Test query successful (returned {len(result.documents)} documents)")
+        logger.info("Test query successful (returned %d documents)", len(result.documents))
         logger.info("Health check passed - system is ready")
 
     except Exception as e:
-        logger.error(f"Health check failed: {e}", exc_info=True)
+        logger.error("Health check failed: %s", e, exc_info=True)
         raise Exception(f"Startup health check failed - cannot connect to database or perform queries: {e}") from e

src/docs2db_mcp/server.py

@@ -1,20 +1,55 @@
 """FastMCP server instance and initialization."""
 
 import logging
+import sys
+
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 
 from fastmcp import FastMCP
 
+from docs2db_mcp.config import CONFIG
+from docs2db_mcp.engine import health_check
+from docs2db_mcp.engine import shutdown_engine
+
 
 logger = logging.getLogger(__name__)
 
-# Initialize MCP server
+
+@asynccontextmanager
+async def engine_lifespan(server: FastMCP) -> AsyncIterator[dict]:
+    """Run startup health check and manage engine lifecycle.
+
+    The server will not accept tool calls until the health check passes.
+    If it fails, the server raises and refuses to start.
+    """
+    await health_check()
+
+    # docs2db-api configures structlog to stdout at import time (triggered
+    # by health_check → get_engine).  For non-SSE transports stdout is
+    # reserved for MCP protocol messages, so redirect structlog to stderr.
+    if CONFIG.transport != "sse":
+        import structlog
+
+        structlog.configure(
+            logger_factory=structlog.PrintLoggerFactory(file=sys.stderr),
+            cache_logger_on_first_use=False,
+        )
+
+    try:
+        yield {}
+    finally:
+        await shutdown_engine()
+
+
 mcp = FastMCP(
     "docs2db-rag",
     instructions=(
         "RAG search using docs2db for RHEL documentation. "
         "Use search_documents to find relevant information from "
         "RHEL documentation, knowledge base articles, and guides."
     ),
+    lifespan=engine_lifespan,
 )
 
 # Import tools to register them with the MCP server