FIX: cpu fallback

app.py

@@ -6,11 +6,6 @@ from snowflake import SnowflakeGenerator
 
 CompletionIdGenerator = SnowflakeGenerator(42, timestamp=1741101491595)
 
-from pynvml import *
-
-nvmlInit()
-gpu_h = nvmlDeviceGetHandleByIndex(0)
-
 from typing import List, Optional, Union
 from pydantic import BaseModel, Field
 from pydantic_settings import BaseSettings
@@ -40,6 +35,17 @@ import numpy as np
 import torch
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if device == "cpu" and CONFIG.STRATEGY != "cpu":
+    logger.info(f"Cuda not found, fall back to cpu")
+    CONFIG.STRATEGY = "cpu"
+
+if "cuda" in CONFIG.STRATEGY:
+    from pynvml import *
+
+    nvmlInit()
+    gpu_h = nvmlDeviceGetHandleByIndex(0)
+
 torch.backends.cudnn.benchmark = True
 torch.backends.cudnn.allow_tf32 = True
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -520,19 +526,17 @@ async def chatResponseStream(
     yield "data: [DONE]\n\n"
 
 
-
-
-
 @app.post("/api/v1/chat/completions")
 async def chat_completions(request: ChatCompletionRequest):
     completionId = str(next(CompletionIdGenerator))
     logger.info(f"[REQ] {completionId} - {request.model_dump()}")
 
-    def chatResponseStreamDisconnect():  
-        gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
-        logger.info(
-        f"[STATUS] vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}"
-    )
+    def chatResponseStreamDisconnect():
+        if "cuda" in CONFIG.STRATEGY:
+            gpu_info = nvmlDeviceGetMemoryInfo(gpu_h)
+            logger.info(
+                f"[STATUS] vram {gpu_info.total} used {gpu_info.used} free {gpu_info.free}"
+            )
 
     model_state = None
 
@@ -545,7 +549,6 @@ async def chat_completions(request: ChatCompletionRequest):
     else:
         r = await chatResponse(request, model_state, completionId)
 
-
     return r