From 76649467473ed69a3b896178b3425124f5e4f78f Mon Sep 17 00:00:00 2001
From: Dicklesworthstone <jeff141421@gmail.com>
Date: Sun, 2 Jun 2024 21:22:33 -0400
Subject: [PATCH] Improve GPU handling

---
 shared_resources.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/shared_resources.py b/shared_resources.py
index f7a555f..97d543d 100644
--- a/shared_resources.py
+++ b/shared_resources.py
@@ -7,7 +7,6 @@
 import asyncio
 import urllib.request
 import os
-import nvgpu
 import glob
 import json
 from filelock import FileLock, Timeout
@@ -17,7 +16,11 @@
 from decouple import config
 from fastapi import HTTPException
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
-
+try:
+    import nvgpu
+    GPU_AVAILABLE = True
+except ImportError:
+    GPU_AVAILABLE = False
 logger = setup_logger()
 
 embedding_model_cache = {} # Model cache to store loaded models
@@ -45,6 +48,14 @@
     scheduler.start()
 
 def is_gpu_available():
+    if not GPU_AVAILABLE:
+        return {
+            "gpu_found": False,
+            "num_gpus": 0,
+            "first_gpu_vram": 0,
+            "total_vram": 0,
+            "error": "nvgpu module not found"
+        }
     try:
         gpu_info = nvgpu.gpu_info()
         num_gpus = len(gpu_info)