Browse Source

refac: close subprocess gracefully

Timothy J. Baek 1 year ago
parent
commit
3c382d4c6c
2 changed files with 51 additions and 9 deletions
  1. 41 8
      backend/apps/litellm/main.py
  2. 10 1
      backend/main.py

+ 41 - 8
backend/apps/litellm/main.py

@@ -42,16 +42,40 @@ app.add_middleware(
 )
 )
 
 
 
 
+# Global variable to store the subprocess reference
+background_process = None
+
+
 async def run_background_process(command):
 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 background_process
+    print("run_background_process")
 
 
-    await process.wait()  # Wait for the subprocess to finish
+    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():
 async def start_litellm_background():
@@ -62,6 +86,15 @@ async def start_litellm_background():
     await run_background_process(command)
     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")
 @app.on_event("startup")
 async def startup_event():
 async def startup_event():
 
 

+ 10 - 1
backend/main.py

@@ -20,7 +20,11 @@ from starlette.middleware.base import BaseHTTPMiddleware
 from apps.ollama.main import app as ollama_app
 from apps.ollama.main import app as ollama_app
 from apps.openai.main import app as openai_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.audio.main import app as audio_app
 from apps.images.main import app as images_app
 from apps.images.main import app as images_app
 from apps.rag.main import app as rag_app
 from apps.rag.main import app as rag_app
@@ -316,3 +320,8 @@ app.mount(
     SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
     SPAStaticFiles(directory=FRONTEND_BUILD_DIR, html=True),
     name="spa-static-files",
     name="spa-static-files",
 )
 )
+
+
+@app.on_event("shutdown")
+async def shutdown_event():
+    await shutdown_litellm_background()