Detect GPU on startup for default multi-pass indexing value (#2242)

backend/danswer/file_processing/extract_file_text.py

@@ -294,7 +294,7 @@ def extract_file_text(
     file_name: str | None,
     file: IO[Any],
     break_on_unprocessable: bool = True,
-) -> str:
+) -> tuple[str, dict]:
     extension_to_function: dict[str, Callable[[IO[Any]], str]] = {
         ".pdf": pdf_to_text,
         ".docx": docx_to_text,

backend/danswer/main.py

@@ -109,6 +109,7 @@ from danswer.server.token_rate_limits.api import (
 from danswer.tools.built_in_tools import auto_add_search_tool_to_personas
 from danswer.tools.built_in_tools import load_builtin_tools
 from danswer.tools.built_in_tools import refresh_built_in_tools_cache
+from danswer.utils.gpu_utils import gpu_status_request
 from danswer.utils.logger import setup_logger
 from danswer.utils.telemetry import optional_telemetry
 from danswer.utils.telemetry import RecordType
@@ -192,6 +193,34 @@ def setup_postgres(db_session: Session) -> None:
     auto_add_search_tool_to_personas(db_session)
 
 
+def update_default_multipass_indexing(db_session: Session) -> None:
+    docs_exist = check_docs_exist(db_session)
+    connectors_exist = check_connectors_exist(db_session)
+    logger.debug(f"Docs exist: {docs_exist}, Connectors exist: {connectors_exist}")
+
+    if not docs_exist and not connectors_exist:
+        logger.info(
+            "No existing docs or connectors found. Checking GPU availability for multipass indexing."
+        )
+        gpu_available = gpu_status_request()
+        logger.info(f"GPU availability: {gpu_available}")
+
+        current_settings = get_current_search_settings(db_session)
+
+        logger.notice(f"Updating multipass indexing setting to: {gpu_available}")
+        updated_settings = SavedSearchSettings.from_db_model(current_settings)
+        # Enable multipass indexing if GPU is available or if using a cloud provider
+        updated_settings.multipass_indexing = (
+            gpu_available or current_settings.cloud_provider is not None
+        )
+        update_current_search_settings(db_session, updated_settings)
+
+    else:
+        logger.debug(
+            "Existing docs or connectors found. Skipping multipass indexing update."
+        )
+
+
 def translate_saved_search_settings(db_session: Session) -> None:
     kv_store = get_dynamic_config_store()
 
@@ -371,6 +400,9 @@ async def lifespan(app: FastAPI) -> AsyncGenerator:
                 ),
             )
 
+        # update multipass indexing setting based on GPU availability
+        update_default_multipass_indexing(db_session)
+
     optional_telemetry(record_type=RecordType.VERSION, data={"version": __version__})
     yield
 

backend/danswer/utils/gpu_utils.py

@@ -0,0 +1,30 @@
+import requests
+from retry import retry
+
+from danswer.utils.logger import setup_logger
+from shared_configs.configs import INDEXING_MODEL_SERVER_HOST
+from shared_configs.configs import INDEXING_MODEL_SERVER_PORT
+from shared_configs.configs import MODEL_SERVER_HOST
+from shared_configs.configs import MODEL_SERVER_PORT
+
+logger = setup_logger()
+
+
+@retry(tries=5, delay=5)
+def gpu_status_request(indexing: bool = True) -> bool:
+    if indexing:
+        model_server_url = f"{INDEXING_MODEL_SERVER_HOST}:{INDEXING_MODEL_SERVER_PORT}"
+    else:
+        model_server_url = f"{MODEL_SERVER_HOST}:{MODEL_SERVER_PORT}"
+
+    if "http" not in model_server_url:
+        model_server_url = f"http://{model_server_url}"
+
+    try:
+        response = requests.get(f"{model_server_url}/api/gpu-status", timeout=10)
+        response.raise_for_status()
+        gpu_status = response.json()
+        return gpu_status["gpu_available"]
+    except requests.RequestException as e:
+        logger.error(f"Error: Unable to fetch GPU status. Error: {str(e)}")
+        raise  # Re-raise exception to trigger a retry

web/src/app/chat/ChatPage.tsx

@@ -2095,6 +2095,7 @@ export function ChatPage({
                                   )}
                                 </div>
                               )}
+
                             {/* Some padding at the bottom so the search bar has space at the bottom to not cover the last message*/}
                             <div ref={endPaddingRef} className="h-[95px]" />
                             <div ref={endDivRef} />