Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 7 commits into from
Aug 31, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion backend/danswer/file_processing/extract_file_text.py
Original file line number Diff line number Diff line change
Expand Up @@ -283,7 +283,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": read_pdf_file,
".docx": docx_to_text,
Expand Down
32 changes: 32 additions & 0 deletions backend/danswer/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@
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
Expand Down Expand Up @@ -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)
pablonyx marked this conversation as resolved.
Show resolved Hide resolved

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()

Expand Down Expand Up @@ -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

Expand Down
30 changes: 30 additions & 0 deletions backend/danswer/utils/gpu_utils.py
Original file line number Diff line number Diff line change
@@ -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:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No need to change but if this is now repeated functionality, we can pull it out into a function

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
1 change: 1 addition & 0 deletions web/src/app/chat/ChatPage.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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} />
Expand Down
Loading