garbage collection

server.py

@@ -11,6 +11,8 @@ from fastapi.middleware.cors import CORSMiddleware
 
 from concrete.ml.deployment import FHEModelServer
 
+import gc
+
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -61,6 +63,12 @@ async def send_input(user_id: str = Form(), files: List[UploadFile] = File(...))
         logger.error(f"Error in send_input: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
     
+def get_memory_usage():
+    """Get current memory usage in GB"""
+    with open('/proc/self/status') as f:
+        memusage = f.read().split('VmRSS:')[1].split('\n')[0][:-3]
+    return int(memusage.strip()) / 1024 / 1024  # Convert KB to GB
+    
 @app.post("/run_fhe")
 def run_fhe(user_id: str = Form()):
     """Execute seizure detection on the encrypted input image using FHE."""
@@ -84,11 +92,34 @@ def run_fhe(user_id: str = Form()):
             encrypted_image = encrypted_image_file.read()
             evaluation_key = evaluation_key_file.read()
 
+        memory_before = get_memory_usage()
+        logger.info(f"Memory usage before FHE execution: {memory_before:.2f} GB")
+
+        # Force garbage collection before FHE execution
+        gc.collect()
+
         # Run the FHE execution
         start = time.time()
-        encrypted_output = FHE_SERVER.run(encrypted_image, evaluation_key)
+        try:
+            encrypted_output = FHE_SERVER.run(encrypted_image, evaluation_key)
+        except MemoryError:
+            logger.error("FHE execution failed due to insufficient memory")
+            raise HTTPException(status_code=503, detail="Insufficient memory during FHE execution")
+        except Exception as e:
+            logger.error(f"FHE execution failed: {str(e)}")
+            raise HTTPException(status_code=500, detail="FHE execution failed")
+        finally:
+            # Force garbage collection after FHE execution
+            gc.collect()
+
         fhe_execution_time = round(time.time() - start, 2)
 
+        # Check memory usage after FHE execution
+        memory_after = get_memory_usage()
+        logger.info(f"Memory usage after FHE execution: {memory_after:.2f} GB")
+        logger.info(f"Memory increase during FHE execution: {memory_after - memory_before:.2f} GB")
+
+
         # Retrieve the encrypted output path
         encrypted_output_path = get_server_file_path("encrypted_output", user_id)