Commit 3c382d4c authored by Timothy J. Baek's avatar Timothy J. Baek
Browse files

refac: close subprocess gracefully

parent 8651bec9
......@@ -42,16 +42,40 @@ app.add_middleware(
)
# Global variable to store the subprocess reference
background_process = None
async def run_background_process(command):
# Start the process
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
)
# Read output asynchronously
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()) # Print stdout line by line
print(line.decode().strip())
await process.wait() # Wait for the subprocess to finish
# 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():
......
......@@ -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()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment