From 18fbe9d7e81d72d61613c64b37f870425e548019 Mon Sep 17 00:00:00 2001
From: pablodanswer <pablo@danswer.ai>
Date: Tue, 17 Sep 2024 14:59:43 -0700
Subject: [PATCH] Warn users of gpu-sensitive operation (#2488)

* warn users of gpu-sensitive operation

* update copy
---
 backend/danswer/main.py                       |  8 ++++
 backend/danswer/server/settings/models.py     |  1 +
 .../admin/embeddings/RerankingFormPage.tsx    | 47 ++++++++++++++++++-
 web/src/app/admin/settings/interfaces.ts      |  1 +
 web/src/components/settings/lib.ts            |  1 +
 5 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/backend/danswer/main.py b/backend/danswer/main.py
index 9a681c39a13a..a5abb8f28c2c 100644
--- a/backend/danswer/main.py
+++ b/backend/danswer/main.py
@@ -110,6 +110,8 @@ from danswer.server.query_and_chat.query_backend import (
 from danswer.server.query_and_chat.query_backend import basic_router as query_router
 from danswer.server.settings.api import admin_router as settings_admin_router
 from danswer.server.settings.api import basic_router as settings_router
+from danswer.server.settings.store import load_settings
+from danswer.server.settings.store import store_settings
 from danswer.server.token_rate_limits.api import (
     router as token_rate_limit_settings_router,
 )
@@ -241,6 +243,12 @@ def update_default_multipass_indexing(db_session: Session) -> None:
         )
         update_current_search_settings(db_session, updated_settings)
 
+        # Update settings with GPU availability
+        settings = load_settings()
+        settings.gpu_enabled = gpu_available
+        store_settings(settings)
+        logger.notice(f"Updated settings with GPU availability: {gpu_available}")
+
     else:
         logger.debug(
             "Existing docs or connectors found. Skipping multipass indexing update."
diff --git a/backend/danswer/server/settings/models.py b/backend/danswer/server/settings/models.py
index e999e7294e99..ae7e7236c8de 100644
--- a/backend/danswer/server/settings/models.py
+++ b/backend/danswer/server/settings/models.py
@@ -37,6 +37,7 @@ class Settings(BaseModel):
     search_page_enabled: bool = True
     default_page: PageType = PageType.SEARCH
     maximum_chat_retention_days: int | None = None
+    gpu_enabled: bool | None = None
 
     def check_validity(self) -> None:
         chat_page_enabled = self.chat_page_enabled
diff --git a/web/src/app/admin/embeddings/RerankingFormPage.tsx b/web/src/app/admin/embeddings/RerankingFormPage.tsx
index eadcadc9654c..5425fc89329f 100644
--- a/web/src/app/admin/embeddings/RerankingFormPage.tsx
+++ b/web/src/app/admin/embeddings/RerankingFormPage.tsx
@@ -1,9 +1,16 @@
-import React, { Dispatch, forwardRef, SetStateAction, useState } from "react";
+import React, {
+  Dispatch,
+  forwardRef,
+  SetStateAction,
+  useContext,
+  useState,
+} from "react";
 import { Formik, Form, FormikProps } from "formik";
 import * as Yup from "yup";
 import {
   RerankerProvider,
   RerankingDetails,
+  RerankingModel,
   rerankingModels,
 } from "./interfaces";
 import { FiExternalLink } from "react-icons/fi";
@@ -15,6 +22,7 @@ import {
 import { Modal } from "@/components/Modal";
 import { Button } from "@tremor/react";
 import { TextFormField } from "@/components/admin/connectors/Field";
+import { SettingsContext } from "@/components/settings/SettingsProvider";
 
 interface RerankingDetailsFormProps {
   setRerankingDetails: Dispatch<SetStateAction<RerankingDetails>>;
@@ -38,10 +46,15 @@ const RerankingDetailsForm = forwardRef<
     },
     ref
   ) => {
+    const [showGpuWarningModalModel, setShowGpuWarningModalModel] =
+      useState<RerankingModel | null>(null);
     const [isApiKeyModalOpen, setIsApiKeyModalOpen] = useState(false);
     const [showLiteLLMConfigurationModal, setShowLiteLLMConfigurationModal] =
       useState(false);
 
+    const combinedSettings = useContext(SettingsContext);
+    const gpuEnabled = combinedSettings?.settings.gpu_enabled;
+
     return (
       <Formik
         innerRef={ref}
@@ -169,6 +182,11 @@ const RerankingDetailsForm = forwardRef<
                             RerankerProvider.LITELLM
                           ) {
                             setShowLiteLLMConfigurationModal(true);
+                          } else if (
+                            !card.rerank_provider_type &&
+                            !gpuEnabled
+                          ) {
+                            setShowGpuWarningModalModel(card);
                           }
 
                           if (!isSelected) {
@@ -225,6 +243,33 @@ const RerankingDetailsForm = forwardRef<
                   })}
                 </div>
 
+                {showGpuWarningModalModel && (
+                  <Modal
+                    onOutsideClick={() => setShowGpuWarningModalModel(null)}
+                    width="w-[500px] flex flex-col"
+                    title="GPU Not Enabled"
+                  >
+                    <>
+                      <p className="text-error font-semibold">Warning:</p>
+                      <p>
+                        Local reranking models require significant computational
+                        resources and may perform slowly without GPU
+                        acceleration. Consider switching to GPU-enabled
+                        infrastructure or using a cloud-based alternative for
+                        better performance.
+                      </p>
+                      <div className="flex justify-end">
+                        <Button
+                          onClick={() => setShowGpuWarningModalModel(null)}
+                          color="blue"
+                          size="xs"
+                        >
+                          Understood
+                        </Button>
+                      </div>
+                    </>
+                  </Modal>
+                )}
                 {showLiteLLMConfigurationModal && (
                   <Modal
                     onOutsideClick={() => {
diff --git a/web/src/app/admin/settings/interfaces.ts b/web/src/app/admin/settings/interfaces.ts
index 2c315f320eca..576b7479baae 100644
--- a/web/src/app/admin/settings/interfaces.ts
+++ b/web/src/app/admin/settings/interfaces.ts
@@ -5,6 +5,7 @@ export interface Settings {
   maximum_chat_retention_days: number | null;
   notifications: Notification[];
   needs_reindexing: boolean;
+  gpu_enabled: boolean;
 }
 
 export interface Notification {
diff --git a/web/src/components/settings/lib.ts b/web/src/components/settings/lib.ts
index 18036346f515..f4e14699f1a6 100644
--- a/web/src/components/settings/lib.ts
+++ b/web/src/components/settings/lib.ts
@@ -42,6 +42,7 @@ export async function fetchSettingsSS(): Promise<CombinedSettings | null> {
     if (!results[0].ok) {
       if (results[0].status === 403) {
         settings = {
+          gpu_enabled: false,
           chat_page_enabled: true,
           search_page_enabled: true,
           default_page: "search",