Update app.py

refactor to try and stop request hangs

app.py

@@ -70,21 +70,22 @@ class ModelManager:
                 model_path = MODELS[model_name]
                 logger.debug(f"Loading tokenizer and model from {model_path}")
                 
-                # Simplified tokenizer loading
                 tokenizer = T5Tokenizer.from_pretrained(
                     model_path,
                     token=HF_TOKEN,
-                    use_fast=True  # Added this
+                    use_fast=True
                 )
                 
-                # Simplified model loading
                 model = T5ForConditionalGeneration.from_pretrained(
                     model_path,
                     token=HF_TOKEN,
-                    torch_dtype=torch.float32
+                    torch_dtype=torch.float32,
+                    low_cpu_mem_usage=True,
+                    device_map='auto'
                 )
                 
                 model.eval()
+                torch.set_num_threads(6)  # Number of CPUs used
                 cls._instances[model_name] = (model, tokenizer)
                 
             except Exception as e:
@@ -154,6 +155,12 @@ async def predict(request: PredictionRequest, background_tasks: BackgroundTasks)
             )
 
         model, tokenizer = await ModelManager.get_model_and_tokenizer(request.model)
+        
+        # Add immediate cleanup of memory before generation
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
         generation_params = DEFAULT_GENERATION_CONFIGS[request.model].copy()
 
         try:
@@ -181,17 +188,23 @@ async def predict(request: PredictionRequest, background_tasks: BackgroundTasks)
         )
 
         async def generate():
-            return model.generate(
-                **inputs,
-                **{k: v for k, v in generation_params.items() if k in [
-                    'max_length', 'min_length', 'do_sample', 'temperature',
-                    'top_p', 'top_k', 'num_beams', 'no_repeat_ngram_size',
-                    'repetition_penalty', 'early_stopping'
-                ]}
-            )
+            try:
+                return model.generate(
+                    **inputs,
+                    **{k: v for k, v in generation_params.items() if k in [
+                        'max_length', 'min_length', 'do_sample', 'temperature',
+                        'top_p', 'top_k', 'num_beams', 'no_repeat_ngram_size',
+                        'repetition_penalty', 'early_stopping'
+                    ]}
+                )
+            finally:
+                # Ensure cleanup happens even if generation fails
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
             
         with torch.inference_mode():
-            outputs = await asyncio.wait_for(generate(), timeout=70.0)
+            outputs = await asyncio.wait_for(generate(), timeout=45.0)  # Reduced timeout
             
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         background_tasks.add_task(cleanup_memory)
@@ -201,14 +214,31 @@ async def predict(request: PredictionRequest, background_tasks: BackgroundTasks)
             selected_model=request.model
         )
 
+    except asyncio.TimeoutError:
+        logger.error("Generation timed out")
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        raise HTTPException(status_code=504, detail="Generation timed out")
     except Exception as e:
         error_msg = f"Error during prediction: {str(e)}\n{traceback.format_exc()}"
         logger.error(error_msg)
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         raise HTTPException(status_code=500, detail=error_msg)
 
 def cleanup_memory():
-    gc.collect()
-    torch.cuda.empty_cache() if torch.cuda.is_available() else None
+    try:
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        
+        # Force Python garbage collection
+        gc.collect(generation=2)
+        
+    except Exception as e:
+        logger.error(f"Error in cleanup: {str(e)}")
 
 if __name__ == "__main__":
     import uvicorn