add background job

app.py

@@ -365,19 +365,21 @@ with demo:
         "detection directly over encrypted values. Once the computation is finished, the server returns "
         "the encrypted results to the client."
     )
-
     gr.Markdown("### Step 4: Send the encrypted image to the server.")
     send_input_button = gr.Button("Send the encrypted image to the server.")
     send_input_checkbox = gr.Checkbox(label="Encrypted image sent.", interactive=False)
 
     gr.Markdown("### Step 5: Run FHE execution.")
     execute_fhe_button = gr.Button("Run FHE execution.")
+    fhe_status = gr.Textbox(label="FHE execution status:", max_lines=1, interactive=False)
     fhe_execution_time = gr.Textbox(
         label="Total FHE execution time (in seconds):", max_lines=1, interactive=False
     )
+    task_id = gr.Textbox(label="Task ID:", visible=False)
 
-    gr.Markdown("### Step 6: Receive the encrypted output from the server.")
-    get_output_button = gr.Button("Receive the encrypted output from the server.")
+    gr.Markdown("### Step 6: Check FHE execution status and receive the encrypted output from the server.")
+    check_status_button = gr.Button("Check FHE execution status")
+    get_output_button = gr.Button("Receive the encrypted output from the server.", interactive=False)
 
     with gr.Row():
         encrypted_output = gr.Textbox(

requirements.txt

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

server.py

@@ -4,7 +4,7 @@ import time
 from typing import List
 from fastapi import FastAPI, File, Form, UploadFile
 from fastapi.responses import JSONResponse, Response
-
+import asyncio
 from common import SERVER_TMP_PATH, SEIZURE_DETECTION_MODEL_PATH
 from client_server_interface import FHEServer
 
@@ -52,7 +52,7 @@ def send_input(
 
 
 @app.post("/run_fhe")
-def run_fhe(
+async def run_fhe(
     user_id: str = Form(),
 ):
     """Execute seizure detection on the encrypted input image using FHE."""
@@ -67,19 +67,42 @@ def run_fhe(
         encrypted_image = encrypted_image_file.read()
         evaluation_key = evaluation_key_file.read()
 
-    # 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)
-
-    # Write the file using the above path
-    with encrypted_output_path.open("wb") as encrypted_output_file:
-        encrypted_output_file.write(encrypted_output)
-
-    return JSONResponse(content=fhe_execution_time)
+    # Run the FHE execution in a background task
+    async def run_fhe_task():
+        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)
+
+        # Write the file using the above path
+        with encrypted_output_path.open("wb") as encrypted_output_file:
+            encrypted_output_file.write(encrypted_output)
+
+        return fhe_execution_time
+
+    # Start the background task
+    task = asyncio.create_task(run_fhe_task())
+
+    # Return a response immediately
+    return JSONResponse(content={"message": "FHE execution started", "task_id": str(id(task))})
+
+@app.get("/fhe_status/{task_id}")
+async def fhe_status(task_id: str):
+    """Check the status of an FHE execution task and return the execution time if completed."""
+    for task in asyncio.all_tasks():
+        if str(id(task)) == task_id:
+            if task.done():
+                try:
+                    execution_time = task.result()
+                    return JSONResponse(content={"status": "completed", "execution_time": execution_time})
+                except Exception as e:
+                    return JSONResponse(content={"status": "error", "message": str(e)})
+            else:
+                return JSONResponse(content={"status": "running"})
+
+    return JSONResponse(content={"status": "not_found"})
 
 
 @app.post("/get_output")