Revert "memory fix"

This reverts commit 4d2771fa3a148324cd57eb98e371add1bceb96e7.

app.py

@@ -238,33 +238,32 @@ def send_input(user_id):
         return response.ok
 
 def run_fhe(user_id):
-    """Start the FHE execution and poll for results."""
+    """Apply the seizure detection model on the encrypted image previously sent using FHE."""
     data = {"user_id": user_id}
     url = SERVER_URL + "run_fhe"
     
     try:
-        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")
+        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.")
     except requests.exceptions.RequestException as e:
-        logger.error(f"Error during FHE execution: {str(e)}")
-        raise gr.Error(f"An error occurred during FHE execution: {str(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)}")
 
 def get_output(user_id):
     """Retrieve the encrypted output (boolean).

requirements.txt

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

server.py

@@ -4,10 +4,8 @@ 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, BackgroundTasks
+from fastapi import FastAPI, File, Form, UploadFile, HTTPException
 from fastapi.responses import JSONResponse, Response
 from fastapi.middleware.cors import CORSMiddleware
 
@@ -40,12 +38,6 @@ 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."""
@@ -70,62 +62,45 @@ 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")
-async def run_fhe(background_tasks: BackgroundTasks, user_id: str = Form()):
-    """Start the FHE execution in the background with memory management."""
+def run_fhe(user_id: str = Form()):
+    """Execute seizure detection on the encrypted input image using FHE."""
     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)
 
-        if not encrypted_image_path.exists() or not evaluation_key_path.exists():
-            raise FileNotFoundError("Required files not found")
-
-        # Start the FHE execution in the background
-        background_tasks.add_task(execute_fhe_with_memory_management, user_id, encrypted_image_path, evaluation_key_path)
-
-        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)}")
+        logger.info(f"Looking for encrypted_image at: {encrypted_image_path}")
+        logger.info(f"Looking for evaluation_key at: {evaluation_key_path}")
 
-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()
+        # 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}")
 
-        available_memory = get_available_memory()
-        estimated_memory_needed = len(encrypted_image) * 2  # Rough estimate, adjust as needed
+        # 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()
 
-        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)
+        # Run the FHE execution
+        start = time.time()
+        encrypted_output = FHE_SERVER.run(encrypted_image, evaluation_key)
+        fhe_execution_time = round(time.time() - start, 2)
 
+        # Retrieve the encrypted output path
         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}")
+        # Write the file using the above path
+        with encrypted_output_path.open("wb") as encrypted_output_file:
+            encrypted_output_file.write(encrypted_output)
+
+        logger.info(f"FHE execution completed for user {user_id} in {fhe_execution_time} seconds")
+        return JSONResponse(content=fhe_execution_time)
     except Exception as 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"})
+        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)}")
 
 @app.post("/get_output")
 def get_output(user_id: str = Form()):