Return session deletion before storage cleanup

app.py

@@ -26,11 +26,11 @@ import cv2
 from email.message import Message
 from email.utils import collapse_rfc2231_value
 from dotenv import load_dotenv
-from fastapi import FastAPI, File, Form, HTTPException, Request, UploadFile
+from fastapi import BackgroundTasks, FastAPI, File, Form, HTTPException, Request, UploadFile
 from fastapi.responses import FileResponse, JSONResponse
 from huggingface_hub import HfApi, hf_hub_download, snapshot_download
 from huggingface_hub.errors import RemoteEntryNotFoundError
-from storage_listing import session_files, session_folder_exists
+from storage_listing import session_files
 
 load_dotenv(Path(__file__).with_name(".env"))
 jobs = {}
@@ -731,15 +731,38 @@ def cancel_session_job(session_id: UUID):
                 job["done"].wait(timeout=300)
                 if not job["done"].is_set():
                     raise HTTPException(504, "The analysis upload did not stop in time; deletion was not started")
-        work_path = job.get("workPath")
+    return [(client_id, job.get("workPath")) for client_id, job in matching]
+
+
+def cleanup_deleted_session_files(session_id, storage_path, job_paths):
+    for _, work_path in job_paths:
         if work_path:
             shutil.rmtree(work_path, ignore_errors=True)
-        try:
-            for path in (f"colab-staging/{client_id}", f"colab-results/{client_id}"):
-                hub().delete_folder(repo_id=os.environ["STORAGE_REPO_ID"], repo_type="dataset", path_in_repo=path,
-                    commit_message=f"Clean up cancelled analysis {client_id}")
-        except Exception:
-            logger.exception("Unable to clean cancelled analysis staging files: %s", client_id)
+    folders = ([storage_path] if storage_path else []) + [
+        path for client_id, _ in job_paths
+        for path in (f"colab-staging/{client_id}", f"colab-results/{client_id}")
+    ]
+    api = hub()
+    repo_id = os.environ["STORAGE_REPO_ID"]
+    failed_folders = []
+    for folder in folders:
+        for attempt in range(3):
+            try:
+                api.delete_folder(repo_id=repo_id, repo_type="dataset", path_in_repo=folder,
+                    commit_message=f"Delete session {session_id} files")
+                break
+            except RemoteEntryNotFoundError:
+                break
+            except Exception:
+                if attempt == 2:
+                    logger.exception("Unable to delete session storage folder %s", folder)
+                    failed_folders.append(folder)
+                else:
+                    time.sleep(2 ** attempt)
+    if failed_folders:
+        logger.error("Background storage cleanup incomplete for session %s: %s", session_id, failed_folders)
+    else:
+        logger.info("Finished background storage cleanup for session %s", session_id)
 
 
 def session_row(session_id, user_id, is_admin=False):
@@ -858,22 +881,19 @@ def update_action(session_id: UUID, action_id: UUID, payload: dict, request: Req
 
 
 @app.delete("/api/sessions/{session_id}")
-def delete_session(session_id: UUID, request: Request):
-    cancel_session_job(session_id)
+def delete_session(session_id: UUID, request: Request, background_tasks: BackgroundTasks):
     row = owned_session_row(session_id, request.state.user_id, request.state.is_admin)
     expected_path = f"sessions/{session_id}"
     storage_path = row.get("path")
     if storage_path is not None and storage_path != expected_path:
         raise HTTPException(409, "Session uses a different storage layout")
 
-    if storage_path == expected_path:
-        if session_folder_exists(hub(), os.environ["STORAGE_REPO_ID"], storage_path):
-            hub().delete_folder(
-                repo_id=os.environ["STORAGE_REPO_ID"], repo_type="dataset",
-                path_in_repo=storage_path, commit_message=f"Delete session {session_id}",
-            )
-
+    job_paths = cancel_session_job(session_id)
     backend(f"/{session_id}", request.state.user_id, "DELETE", is_admin=request.state.is_admin)
+    for client_id, _ in job_paths:
+        jobs.pop(client_id, None)
+    # The database deletion makes the session disappear before Hugging Face finishes deleting its files.
+    background_tasks.add_task(cleanup_deleted_session_files, session_id, storage_path, job_paths)
     return {"ok": True}