try to fix notebook

http_wrapper/routes/notebook.py

@@ -122,23 +122,18 @@ def download_notebook(post_req:PostReqModel = Depends(common_params)):
     )
     return resp
 
+branches_to_delete = []
+
+def on_exit():
+    for branch in branches_to_delete:
+        delete_branch(branch)
+
 # Only add run endpoint if the necesary env vars are defined
 if run_now_enabled:
 
     # Do not use async, or else it will be blocking
     def delete_branch(branch: str):
-        from datetime import datetime
-        from http_wrapper.server import kill_event
-        from time import sleep
-
-        start = datetime.now()
 
-        diff_seconds = 0
-        while not kill_event.is_set() and diff_seconds < 600:
-
-            diff_seconds = (datetime.now() - start).seconds
-            sleep(5)
-
         from gitlab import Gitlab
         gl = Gitlab(url=HBP_GITLAB_HOST, private_token=HBP_GITLAB_TOKEN)
         project = gl.projects.get(HBP_GITLAB_PROJECT_ID)
@@ -149,7 +144,7 @@ def redirect():
         return RedirectResponse("https://lab.ebrains.eu/")
 
     @router.get("/run", tags=TAGS)
-    def run_notebook(background_tasks: BackgroundTasks, site: NotebookExecutionSite, post_req:PostReqModel = Depends(common_params)):
+    def run_notebook(site: NotebookExecutionSite, post_req:PostReqModel = Depends(common_params)):
         notebook = get_notebook(post_req)
         notebook_filename = f"{get_notebook_name(post_req)}.ipynb"
 
@@ -173,6 +168,8 @@ def run_notebook(background_tasks: BackgroundTasks, site: NotebookExecutionSite,
             'commit_message': 'creating file'
         })
         redirect_url = get_jupyterhub_link(site, project.attributes.get("http_url_to_repo"), branch_name, notebook_filename)
+
+        branches_to_delete.append(branch_name)
 
-        background_tasks.add_task(delete_branch, branch_name)
+        print(f"DEBUG: Redirecting: {redirect_url}")
         return RedirectResponse(redirect_url)

http_wrapper/server.py

@@ -1,5 +1,5 @@
 from http_wrapper.routes.analysis import router as analysis_router
-from http_wrapper.routes.notebook import router as notebook_router
+from http_wrapper.routes.notebook import router as notebook_router, on_exit
 from http_wrapper.routes.git import router as git_router
 from http_wrapper.jugex_logger import access_logger
 from http_wrapper.conf.siibra_jugex_conf import SIIBRA_TOOLBOX_VIEWER_PLUGIN_STATIC_DIR
@@ -48,12 +48,9 @@ async def access_log(request: Request, call_next):
 
 from threading import Event
 
-kill_event = Event()
-
 @app.on_event("shutdown")
 def shutdown_event():
     # TODO doesn't work quite right
     # shutdown handler isn't called until ctrl+c is hit twice
-
-    kill_event.set()
+    on_exit()