From 3c382d4c6cbea0352a4ad4bc3a90ed8f339a148b Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Sun, 21 Apr 2024 01:46:09 -0500
Subject: [PATCH] refac: close subprocess gracefully

---
 backend/apps/litellm/main.py | 51 +++++++++++++++++++++++++++++-------
 backend/main.py              | 11 +++++++-
 2 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/backend/apps/litellm/main.py b/backend/apps/litellm/main.py
index 5a8b37f47..68e48858b 100644
--- a/backend/apps/litellm/main.py
+++ b/backend/apps/litellm/main.py
@@ -42,16 +42,40 @@ app.add_middleware(
 )
 
 
-async def run_background_process(command):
-    # Start the process
-    process = await asyncio.create_subprocess_exec(
-        *command.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE
-    )
-    # Read output asynchronously
-    async for line in process.stdout:
-        print(line.decode().strip())  # Print stdout line by line
+# Global variable to store the subprocess reference
+background_process = None
 
-    await process.wait()  # Wait for the subprocess to finish
+
+async def run_background_process(command):
+    global background_process
+    print("run_background_process")
+
+    try:
+        # Log the command to be executed
+        print(f"Executing command: {command}")
+        # Execute the command and create a subprocess
+        process = await asyncio.create_subprocess_exec(
+            *command.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
+        background_process = process
+        print("Subprocess started successfully.")
+
+        # Capture STDERR for debugging purposes
+        stderr_output = await process.stderr.read()
+        stderr_text = stderr_output.decode().strip()
+        if stderr_text:
+            print(f"Subprocess STDERR: {stderr_text}")
+
+        # Print output line by line
+        async for line in process.stdout:
+            print(line.decode().strip())
+
+        # Wait for the process to finish
+        returncode = await process.wait()
+        print(f"Subprocess exited with return code {returncode}")
+    except Exception as e:
+        log.error(f"Failed to start subprocess: {e}")
+        raise  # Optionally re-raise the exception if you want it to propagate
 
 
 async def start_litellm_background():
@@ -62,6 +86,15 @@ async def start_litellm_background():
     await run_background_process(command)
 
 
+async def shutdown_litellm_background():
+    print("shutdown_litellm_background")
+    global background_process
+    if background_process:
+        background_process.terminate()
+        await background_process.wait()  # Ensure the process has terminated
+        print("Subprocess terminated")
+
+
 @app.on_event("startup")
 async def startup_event():
 
diff --git a/backend/main.py b/backend/main.py
index 48e14f1dd..579ff2ee0 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -20,7 +20,11 @@ from starlette.middleware.base import BaseHTTPMiddleware
 from apps.ollama.main import app as ollama_app
 from apps.openai.main import app as openai_app
 
-from apps.litellm.main import app as litellm_app, start_litellm_background
+from apps.litellm.main import (
+    app as litellm_app,
+    start_litellm_background,
+    shutdown_litellm_background,
+)
 from apps.audio.main import app as audio_app
 from apps.images.main import app as images_app
 from apps.rag.main import app as rag_app
@@ -316,3 +320,8 @@ app.mount(
     SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
     name="spa-static-files",
 )
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    await shutdown_litellm_background()