fix(langflow): handle KeyboardInterrupt and terminate process in run function for graceful shutdown (#2551)

* fix(langflow): handle KeyboardInterrupt and terminate process in run function for graceful shutdown

* refactor(langflow): remove unnecessary try-except block and sys.exit calls in run_langflow function

* fix(langflow): handle exceptions properly and exit with appropriate status codes

* refactor: update multiprocess imports in langflow/__main__.py

Adds type ignore
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-07-05 12:04:33 -03:00 • committed by GitHub
commit e536272759
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -10,7 +10,8 @@ import click
import httpx import httpx
import typer import typer
from dotenv import dotenv_values, load_dotenv from dotenv import dotenv_values, load_dotenv
from multiprocess import Process, cpu_count # type: ignore from multiprocess import cpu_count # type: ignore
from multiprocess.context import Process # type: ignore
from packaging import version as pkg_version from packaging import version as pkg_version
from rich import box from rich import box
from rich import print as rprint from rich import print as rprint
@ -181,6 +182,7 @@ def run(
# Define an env variable to know if we are just testing the server # Define an env variable to know if we are just testing the server
if "pytest" in sys.modules: if "pytest" in sys.modules:
return return
process: Process | None = None
try: try:
if platform.system() in ["Windows"]: if platform.system() in ["Windows"]:
# Run using uvicorn on MacOS and Windows # Run using uvicorn on MacOS and Windows
@ -195,7 +197,12 @@ def run(
if process: if process:
process.join() process.join()
except KeyboardInterrupt: except KeyboardInterrupt:
pass if process is not None:
process.terminate()
sys.exit(0)
except Exception as e:
logger.exception(e)
sys.exit(1)
def wait_for_server_ready(host, port): def wait_for_server_ready(host, port):
@ -417,29 +424,24 @@ def run_langflow(host, port, log_level, options, app):
""" """
Run Langflow server on localhost Run Langflow server on localhost
""" """
try:
if platform.system() in ["Windows"]:
# Run using uvicorn on MacOS and Windows
# Windows doesn't support gunicorn
# MacOS requires an env variable to be set to use gunicorn
import uvicorn
uvicorn.run( if platform.system() in ["Windows"]:
app, # Run using uvicorn on MacOS and Windows
host=host, # Windows doesn't support gunicorn
port=port, # MacOS requires an env variable to be set to use gunicorn
log_level=log_level.lower(), import uvicorn
loop="asyncio",
)
else:
from langflow.server import LangflowApplication
LangflowApplication(app, options).run() uvicorn.run(
except KeyboardInterrupt: app,
pass host=host,
except Exception as e: port=port,
logger.exception(e) log_level=log_level.lower(),
sys.exit(1) loop="asyncio",
)
else:
from langflow.server import LangflowApplication
LangflowApplication(app, options).run()
@app.command() @app.command()