memory fix

app.py

@@ -238,32 +238,33 @@ def send_input(user_id):
         return response.ok
 
 def run_fhe(user_id):
-    """Apply the seizure detection model on the encrypted image previously sent using FHE."""
+    """Start the FHE execution and poll for results."""
     data = {"user_id": user_id}
     url = SERVER_URL + "run_fhe"
     
     try:
-        logger.info(f"Sending request to {url} with user_id: {user_id}")
-        with requests_retry_session().post(url=url, data=data, timeout=300) as response:
-            logger.info(f"Received response with status code: {response.status_code}")
-            response.raise_for_status()  # Raises an HTTPError for bad responses
-            if response.ok:
-                return response.json()
-            else:
-                logger.error(f"Server responded with status code {response.status_code}")
-                raise gr.Error(f"Server responded with status code {response.status_code}")
-    except requests.exceptions.Timeout:
-        logger.error("The request timed out. The server might be overloaded.")
-        raise gr.Error("The request timed out. The server might be overloaded.")
-    except requests.exceptions.ConnectionError as e:
-        logger.error(f"Failed to connect to the server. Error: {str(e)}")
-        raise gr.Error("Failed to connect to the server. Please check your network connection.")
+        with requests_retry_session().post(url=url, data=data, timeout=30) as response:
+            response.raise_for_status()
+            logger.info("FHE execution started successfully")
+        
+        # Poll for FHE execution status
+        status_url = SERVER_URL + f"fhe_status/{user_id}"
+        max_attempts = 60  # Adjust this value based on expected execution time
+        attempt = 0
+        while attempt < max_attempts:
+            with requests_retry_session().get(status_url, timeout=10) as status_response:
+                status_response.raise_for_status()
+                status = status_response.json()["status"]
+                if status == "completed":
+                    logger.info("FHE execution completed successfully")
+                    return "FHE execution completed successfully"
+            time.sleep(5)  # Wait for 5 seconds before next poll
+            attempt += 1
+        
+        raise gr.Error("FHE execution timed out")
     except requests.exceptions.RequestException as e:
-        logger.error(f"An error occurred: {str(e)}")
-        raise gr.Error(f"An error occurred: {str(e)}")
-    except Exception as e:
-        logger.error(f"An unexpected error occurred: {str(e)}")
-        raise gr.Error(f"An unexpected error occurred: {str(e)}")
+        logger.error(f"Error during FHE execution: {str(e)}")
+        raise gr.Error(f"An error occurred during FHE execution: {str(e)}")
 
 def get_output(user_id):
     """Retrieve the encrypted output (boolean).

requirements.txt

@@ -1,3 +1,4 @@
 concrete-ml
 gradio
-fastapi
+fastapi
+psutil

server.py

@@ -4,8 +4,10 @@ import time
 import logging
 from pathlib import Path
 from typing import List
+import psutil
+import numpy as np
 
-from fastapi import FastAPI, File, Form, UploadFile, HTTPException
+from fastapi import FastAPI, File, Form, UploadFile, HTTPException, BackgroundTasks
 from fastapi.responses import JSONResponse, Response
 from fastapi.middleware.cors import CORSMiddleware
 
@@ -38,6 +40,12 @@ def get_server_file_path(file_type: str, user_id: str) -> Path:
         file_type = "evaluation"
     return Path(__file__).parent / "server_tmp" / f"{file_type}_{user_id}"
 
+def get_available_memory():
+    return psutil.virtual_memory().available
+
+def chunk_data(data, num_chunks):
+    return np.array_split(data, num_chunks)
+
 @app.post("/send_input")
 async def send_input(user_id: str = Form(), files: List[UploadFile] = File(...)):
     """Receive the encrypted input image and the evaluation key from the client."""
@@ -62,45 +70,62 @@ async def send_input(user_id: str = Form(), files: List[UploadFile] = File(...))
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
     
 @app.post("/run_fhe")
-def run_fhe(user_id: str = Form()):
-    """Execute seizure detection on the encrypted input image using FHE."""
+async def run_fhe(background_tasks: BackgroundTasks, user_id: str = Form()):
+    """Start the FHE execution in the background with memory management."""
     logger.info(f"Starting FHE execution for user {user_id}")
     try:
-        # Retrieve the encrypted input image and the evaluation key paths
         encrypted_image_path = get_server_file_path("encrypted", user_id)
         evaluation_key_path = get_server_file_path("evaluation", user_id)
 
-        logger.info(f"Looking for encrypted_image at: {encrypted_image_path}")
-        logger.info(f"Looking for evaluation_key at: {evaluation_key_path}")
+        if not encrypted_image_path.exists() or not evaluation_key_path.exists():
+            raise FileNotFoundError("Required files not found")
 
-        # Check if files exist
-        if not encrypted_image_path.exists():
-            raise FileNotFoundError(f"Encrypted image file not found at {encrypted_image_path}")
-        if not evaluation_key_path.exists():
-            raise FileNotFoundError(f"Evaluation key file not found at {evaluation_key_path}")
+        # Start the FHE execution in the background
+        background_tasks.add_task(execute_fhe_with_memory_management, user_id, encrypted_image_path, evaluation_key_path)
 
-        # Read the files using the above paths
-        with encrypted_image_path.open("rb") as encrypted_image_file, evaluation_key_path.open("rb") as evaluation_key_file:
-            encrypted_image = encrypted_image_file.read()
-            evaluation_key = evaluation_key_file.read()
+        return JSONResponse(content={"message": "FHE execution started"})
+    except Exception as e:
+        logger.error(f"Error in run_fhe for user {user_id}: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-        # Run the FHE execution
-        start = time.time()
-        encrypted_output = FHE_SERVER.run(encrypted_image, evaluation_key)
-        fhe_execution_time = round(time.time() - start, 2)
+async def execute_fhe_with_memory_management(user_id: str, encrypted_image_path: Path, evaluation_key_path: Path):
+    try:
+        with encrypted_image_path.open("rb") as f, evaluation_key_path.open("rb") as g:
+            encrypted_image = f.read()
+            evaluation_key = g.read()
 
-        # Retrieve the encrypted output path
-        encrypted_output_path = get_server_file_path("encrypted_output", user_id)
+        available_memory = get_available_memory()
+        estimated_memory_needed = len(encrypted_image) * 2  # Rough estimate, adjust as needed
 
-        # Write the file using the above path
-        with encrypted_output_path.open("wb") as encrypted_output_file:
-            encrypted_output_file.write(encrypted_output)
+        if estimated_memory_needed > available_memory:
+            num_chunks = (estimated_memory_needed // available_memory) + 1
+            encrypted_image_chunks = chunk_data(np.frombuffer(encrypted_image, dtype=np.uint8), num_chunks)
+            
+            results = []
+            for chunk in encrypted_image_chunks:
+                chunk_result = FHE_SERVER.run(chunk.tobytes(), evaluation_key)
+                results.append(chunk_result)
+            
+            encrypted_output = b''.join(results)
+        else:
+            encrypted_output = FHE_SERVER.run(encrypted_image, evaluation_key)
 
-        logger.info(f"FHE execution completed for user {user_id} in {fhe_execution_time} seconds")
-        return JSONResponse(content=fhe_execution_time)
+        encrypted_output_path = get_server_file_path("encrypted_output", user_id)
+        with encrypted_output_path.open("wb") as f:
+            f.write(encrypted_output)
+
+        logger.info(f"FHE execution completed for user {user_id}")
     except Exception as e:
-        logger.error(f"Error in run_fhe for user {user_id}: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        logger.error(f"Error during FHE execution for user {user_id}: {str(e)}")
+
+@app.get("/fhe_status/{user_id}")
+async def fhe_status(user_id: str):
+    """Check the status of the FHE execution."""
+    encrypted_output_path = get_server_file_path("encrypted_output", user_id)
+    if encrypted_output_path.exists():
+        return JSONResponse(content={"status": "completed"})
+    else:
+        return JSONResponse(content={"status": "in_progress"})
 
 @app.post("/get_output")
 def get_output(user_id: str = Form()):