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):
|
# 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():
|
||||||
|
|
||||||
|
|
|
@ -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()
|
||||||
|
|
Loading…
Reference in a new issue