diff --git a/src/backend/base/langflow/__main__.py b/src/backend/base/langflow/__main__.py index 53407b223..95368d5b3 100644 --- a/src/backend/base/langflow/__main__.py +++ b/src/backend/base/langflow/__main__.py @@ -1,6 +1,7 @@ import asyncio import inspect import platform +import signal import socket import sys import time @@ -24,9 +25,7 @@ from sqlmodel import select from langflow.logging.logger import configure, logger from langflow.main import setup_app -from langflow.services.database.models.folder.utils import ( - create_default_folder_if_it_doesnt_exist, -) +from langflow.services.database.models.folder.utils import create_default_folder_if_it_doesnt_exist from langflow.services.database.utils import session_getter from langflow.services.deps import async_session_scope, get_db_service, get_settings_service from langflow.services.settings.constants import DEFAULT_SUPERUSER @@ -77,6 +76,13 @@ def set_var_for_macos_issue() -> None: logger.debug("Set OBJC_DISABLE_INITIALIZE_FORK_SAFETY to YES to avoid error") +def handle_sigterm(signum, frame): # noqa: ARG001 + """Handle SIGTERM signal gracefully.""" + logger.info("Received SIGTERM signal. Performing graceful shutdown...") + # Raise SystemExit to trigger graceful shutdown + sys.exit(0) + + @app.command() def run( *, @@ -150,6 +156,9 @@ def run( ), ) -> None: """Run Langflow.""" + # Register SIGTERM handler + signal.signal(signal.SIGTERM, handle_sigterm) + if env_file: load_dotenv(env_file, override=True) @@ -211,13 +220,20 @@ def run( click.launch(f"http://{host}:{port}") if process: process.join() - except KeyboardInterrupt: + except (KeyboardInterrupt, SystemExit) as e: + logger.info("Shutting down server...") if process is not None: process.terminate() - sys.exit(0) - except Exception as e: # noqa: BLE001 + process.join(timeout=15) # Wait up to 15 seconds for process to terminate + if process.is_alive(): + logger.warning("Process did not terminate gracefully, forcing...") + process.kill() + raise typer.Exit(0) from e + except Exception as e: logger.exception(e) - sys.exit(1) + if process is not None: + process.terminate() + raise typer.Exit(1) from e def wait_for_server_ready(host, port) -> None: @@ -359,9 +375,6 @@ def print_banner(host: str, port: int) -> None: def run_langflow(host, port, log_level, options, app) -> None: """Run Langflow server on localhost.""" if platform.system() == "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( @@ -374,7 +387,28 @@ def run_langflow(host, port, log_level, options, app) -> None: else: from langflow.server import LangflowApplication - LangflowApplication(app, options).run() + server = LangflowApplication(app, options) + + def graceful_shutdown(signum, frame): # noqa: ARG001 + """Gracefully shutdown the server when receiving SIGTERM.""" + # Suppress click exceptions during shutdown + import click + + click.echo = lambda *args, **kwargs: None # noqa: ARG005 + + logger.info("Gracefully shutting down server...") + # For Gunicorn workers, we raise SystemExit to trigger graceful shutdown + raise SystemExit(0) + + # Register signal handlers + signal.signal(signal.SIGTERM, graceful_shutdown) + signal.signal(signal.SIGINT, graceful_shutdown) + + try: + server.run() + except (KeyboardInterrupt, SystemExit): + # Suppress the exception output + sys.exit(0) @app.command() @@ -499,10 +533,7 @@ def api_key( ) return None from langflow.services.database.models.api_key import ApiKey, ApiKeyCreate - from langflow.services.database.models.api_key.crud import ( - create_api_key, - delete_api_key, - ) + from langflow.services.database.models.api_key.crud import create_api_key, delete_api_key api_key = (await session.exec(select(ApiKey).where(ApiKey.user_id == superuser.id))).first() if api_key: @@ -568,4 +599,8 @@ def main() -> None: if __name__ == "__main__": - main() + try: + main() + except Exception as e: + logger.exception(e) + raise typer.Exit(1) from e diff --git a/src/backend/base/langflow/main.py b/src/backend/base/langflow/main.py index 344646d15..6edd353f4 100644 --- a/src/backend/base/langflow/main.py +++ b/src/backend/base/langflow/main.py @@ -97,6 +97,7 @@ def get_lifespan(*, fix_migration=False, version=None): @asynccontextmanager async def lifespan(_app: FastAPI): configure(async_file=True) + # Startup message if version: rprint(f"[bold green]Starting Langflow v{version}...[/bold green]") @@ -108,15 +109,20 @@ def get_lifespan(*, fix_migration=False, version=None): await asyncio.to_thread(create_or_update_starter_projects, all_types_dict) telemetry_service.start() await asyncio.to_thread(load_flows_from_directory) + yield + except Exception as exc: if "langflow migration --fix" not in str(exc): logger.exception(exc) raise - # Shutdown message - rprint("[bold red]Shutting down Langflow...[/bold red]") - await teardown_services() - await logger.complete() + finally: + # Clean shutdown + logger.info("Cleaning up resources...") + await teardown_services() + await logger.complete() + # Final message + rprint("[bold red]Langflow shutdown complete[/bold red]") return lifespan diff --git a/src/backend/base/pyproject.toml b/src/backend/base/pyproject.toml index b7e95c5c0..dd62d2ef6 100644 --- a/src/backend/base/pyproject.toml +++ b/src/backend/base/pyproject.toml @@ -119,7 +119,7 @@ dependencies = [ "langchain-experimental>=0.0.61", "pydantic>=2.7.0", "pydantic-settings>=2.2.0", - "typer>=0.12.0", + "typer>=0.13.0", "cachetools>=5.3.1", "platformdirs>=4.2.0", "python-multipart>=0.0.12", diff --git a/uv.lock b/uv.lock index b98f771d7..4f2145ecd 100644 --- a/uv.lock +++ b/uv.lock @@ -4082,7 +4082,7 @@ requires-dist = [ { name = "setuptools", specifier = ">=70" }, { name = "spider-client", specifier = ">=0.0.27" }, { name = "sqlmodel", specifier = "==0.0.18" }, - { name = "typer", specifier = ">=0.12.0" }, + { name = "typer", specifier = ">=0.13.0" }, { name = "types-google-cloud-ndb", marker = "extra == 'dev'", specifier = ">=2.2.0.0" }, { name = "types-markdown", marker = "extra == 'dev'", specifier = ">=3.7.0.20240822" }, { name = "types-passlib", marker = "extra == 'dev'", specifier = ">=1.7.7.13" }, @@ -8060,7 +8060,7 @@ wheels = [ [[package]] name = "typer" -version = "0.12.5" +version = "0.13.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -8068,9 +8068,9 @@ dependencies = [ { name = "shellingham" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/c5/58/a79003b91ac2c6890fc5d90145c662fd5771c6f11447f116b63300436bc9/typer-0.12.5.tar.gz", hash = "sha256:f592f089bedcc8ec1b974125d64851029c3b1af145f04aca64d69410f0c9b722", size = 98953 } +sdist = { url = "https://files.pythonhosted.org/packages/e7/87/9eb07fdfa14e22ec7658b5b1147836d22df3848a22c85a4e18ed272303a5/typer-0.13.0.tar.gz", hash = "sha256:f1c7198347939361eec90139ffa0fd8b3df3a2259d5852a0f7400e476d95985c", size = 97572 } wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/2b/886d13e742e514f704c33c4caa7df0f3b89e5a25ef8db02aa9ca3d9535d5/typer-0.12.5-py3-none-any.whl", hash = "sha256:62fe4e471711b147e3365034133904df3e235698399bc4de2b36c8579298d52b", size = 47288 }, + { url = "https://files.pythonhosted.org/packages/18/7e/c8bfa8cbcd3ea1d25d2beb359b5c5a3f4339a7e2e5d9e3ef3e29ba3ab3b9/typer-0.13.0-py3-none-any.whl", hash = "sha256:d85fe0b777b2517cc99c8055ed735452f2659cd45e451507c76f48ce5c1d00e2", size = 44194 }, ] [[package]]