forked from open-webui/open-webui
refac: close subprocess gracefully
This commit is contained in:
parent
8651bec915
commit
3c382d4c6c
2 changed files with 52 additions and 10 deletions
|
@ -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():
|
||||
|
||||
|
|
|
@ -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()
|
||||
|
|
Loading…
Reference in a new issue