refac: close subprocess gracefully

This commit is contained in:
Timothy J. Baek 2024-04-21 01:46:09 -05:00
parent 8651bec915
commit 3c382d4c6c
2 changed files with 52 additions and 10 deletions

View file

@ -42,16 +42,40 @@ app.add_middleware(
) )
async def run_background_process(command): # Global variable to store the subprocess reference
# Start the process background_process = None
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
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(): 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():

View file

@ -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()