Replacing gunicornmontor with uvicorn.run()

airflow/cli/commands/local_commands/fastapi_api_command.py

@@ -20,25 +20,16 @@
 
 import logging
 import os
-import signal
 import subprocess
-import sys
 import textwrap
-from contextlib import suppress
-from pathlib import Path
-from time import sleep
-from typing import NoReturn
 
-import psutil
-from lockfile.pidlockfile import read_pid_from_pidfile
-from uvicorn.workers import UvicornWorker
+import uvicorn
+from gunicorn.util import daemonize
+from setproctitle import setproctitle
 
 from airflow import settings
-from airflow.cli.commands.local_commands.daemon_utils import run_command_with_daemon_option
-from airflow.cli.commands.local_commands.webserver_command import GunicornMonitor
 from airflow.exceptions import AirflowConfigException
 from airflow.utils import cli as cli_utils
-from airflow.utils.cli import setup_locations
 from airflow.utils.providers_configuration_loader import providers_configuration_loaded
 
 log = logging.getLogger(__name__)
@@ -47,8 +38,6 @@
 # This shouldn't be necessary but there seems to be an issue in uvloop that causes bad file descriptor
 # errors when shutting down workers. Despite the 'closed' status of the issue it is not solved,
 # more info here: https://github.com/benoitc/gunicorn/issues/1877#issuecomment-1911136399
-AirflowUvicornWorker = UvicornWorker
-AirflowUvicornWorker.CONFIG_KWARGS = {"loop": "asyncio", "http": "auto"}
 
 
 @cli_utils.action_cli
@@ -59,18 +48,13 @@ def fastapi_api(args):
 
     apps = args.apps
     access_logfile = args.access_logfile or "-"
-    error_logfile = args.error_logfile or "-"
     access_logformat = args.access_logformat
     num_workers = args.workers
     worker_timeout = args.worker_timeout
 
-    worker_class = "airflow.cli.commands.local_commands.fastapi_api_command.AirflowUvicornWorker"
-
-    from airflow.api_fastapi.app import create_app
-
     if args.debug:
         print(f"Starting the FastAPI API server on port {args.port} and host {args.hostname} debug.")
-        log.warning("Running in dev mode, ignoring gunicorn args")
+        log.warning("Running in dev mode, ignoring uvicorn args")
 
         run_args = [
             "fastapi",
@@ -93,124 +77,35 @@ def fastapi_api(args):
             process.wait()
         os.environ.pop("AIRFLOW_API_APPS")
     else:
+        if args.daemon:
+            daemonize()
+            log.info("Daemonized the FastAPI API server process PID: %s", os.getpid())
+
         log.info(
             textwrap.dedent(
                 f"""\
-                Running the Gunicorn Server with:
+                Running the uvicorn with:
                 Apps: {apps}
-                Workers: {num_workers} {worker_class}
+                Workers: {num_workers}
                 Host: {args.hostname}:{args.port}
                 Timeout: {worker_timeout}
-                Logfiles: {access_logfile} {error_logfile}
+                Logfiles: {access_logfile}
                 Access Logformat: {access_logformat}
                 ================================================================="""
             )
         )
-
-        pid_file, _, _, _ = setup_locations("fastapi-api", pid=args.pid)
-
-        run_args = [
-            sys.executable,
-            "-m",
-            "gunicorn",
-            "--workers",
-            str(num_workers),
-            "--worker-class",
-            str(worker_class),
-            "--timeout",
-            str(worker_timeout),
-            "--bind",
-            args.hostname + ":" + str(args.port),
-            "--name",
-            "airflow-fastapi-api",
-            "--pid",
-            pid_file,
-            "--access-logfile",
-            str(access_logfile),
-            "--error-logfile",
-            str(error_logfile),
-            "--config",
-            "python:airflow.api_fastapi.gunicorn_config",
-        ]
-
         ssl_cert, ssl_key = _get_ssl_cert_and_key_filepaths(args)
-        if ssl_cert and ssl_key:
-            run_args += ["--certfile", ssl_cert, "--keyfile", ssl_key]
-
-        if args.access_logformat and args.access_logformat.strip():
-            run_args += ["--access-logformat", str(args.access_logformat)]
-
-        if args.daemon:
-            run_args += ["--daemon"]
-
-        run_args += [f"airflow.api_fastapi.app:cached_app(apps='{apps}')"]
-
-        # To prevent different workers creating the web app and
-        # all writing to the database at the same time, we use the --preload option.
-        # With the preload option, the app is loaded before the workers are forked, and each worker will
-        # then have a copy of the app
-        run_args += ["--preload"]
-
-        def kill_proc(signum: int, gunicorn_master_proc: psutil.Process | subprocess.Popen) -> NoReturn:
-            log.info("Received signal: %s. Closing gunicorn.", signum)
-            gunicorn_master_proc.terminate()
-            with suppress(TimeoutError):
-                gunicorn_master_proc.wait(timeout=30)
-            if isinstance(gunicorn_master_proc, subprocess.Popen):
-                still_running = gunicorn_master_proc.poll() is not None
-            else:
-                still_running = gunicorn_master_proc.is_running()
-            if still_running:
-                gunicorn_master_proc.kill()
-            sys.exit(0)
-
-        def monitor_gunicorn(gunicorn_master_proc: psutil.Process | subprocess.Popen) -> NoReturn:
-            # Register signal handlers
-            signal.signal(signal.SIGINT, lambda signum, _: kill_proc(signum, gunicorn_master_proc))
-            signal.signal(signal.SIGTERM, lambda signum, _: kill_proc(signum, gunicorn_master_proc))
-
-            # These run forever until SIG{INT, TERM, KILL, ...} signal is sent
-            GunicornMonitor(
-                gunicorn_master_pid=gunicorn_master_proc.pid,
-                num_workers_expected=num_workers,
-                master_timeout=120,
-                worker_refresh_interval=30,
-                worker_refresh_batch_size=1,
-                reload_on_plugin_change=False,
-            ).start()
-
-        def start_and_monitor_gunicorn(args):
-            if args.daemon:
-                subprocess.Popen(run_args, close_fds=True)
-
-                # Reading pid of gunicorn master as it will be different that
-                # the one of process spawned above.
-                gunicorn_master_proc_pid = None
-                while not gunicorn_master_proc_pid:
-                    sleep(0.1)
-                    gunicorn_master_proc_pid = read_pid_from_pidfile(pid_file)
-
-                # Run Gunicorn monitor
-                gunicorn_master_proc = psutil.Process(gunicorn_master_proc_pid)
-                monitor_gunicorn(gunicorn_master_proc)
-            else:
-                with subprocess.Popen(run_args, close_fds=True) as gunicorn_master_proc:
-                    monitor_gunicorn(gunicorn_master_proc)
-
-        if args.daemon:
-            # This makes possible errors get reported before daemonization
-            os.environ["SKIP_DAGS_PARSING"] = "True"
-            create_app(apps)
-            os.environ.pop("SKIP_DAGS_PARSING")
-
-        pid_file_path = Path(pid_file)
-        monitor_pid_file = str(pid_file_path.with_name(f"{pid_file_path.stem}-monitor{pid_file_path.suffix}"))
-        run_command_with_daemon_option(
-            args=args,
-            process_name="fastapi-api",
-            callback=lambda: start_and_monitor_gunicorn(args),
-            should_setup_logging=True,
-            pid_file=monitor_pid_file,
+        setproctitle(f"airflow fastapi_api -- host:{args.hostname} port:{args.port}")
+        uvicorn.run(
+            "airflow.api_fastapi.main:app",
+            host=args.hostname,
+            port=args.port,
+            workers=num_workers,
+            timeout_keep_alive=worker_timeout,
+            timeout_graceful_shutdown=worker_timeout,
+            ssl_keyfile=ssl_key,
+            ssl_certfile=ssl_cert,
+            access_log=access_logfile,
         )
 
 

tests/cli/commands/local_commands/test_fastapi_api_command.py

@@ -16,10 +16,6 @@
 # under the License.
 from __future__ import annotations
 
-import os
-import subprocess
-import sys
-import time
 from unittest import mock
 
 import pytest
@@ -37,73 +33,9 @@
 class TestCliFastAPI(_CommonCLIGunicornTestClass):
     main_process_regexp = r"airflow fastapi-api"
 
-    @pytest.mark.execution_timeout(210)
-    def test_cli_fastapi_api_background(self, tmp_path):
-        parent_path = tmp_path / "gunicorn"
-        parent_path.mkdir()
-        pidfile_fastapi_api = parent_path / "pidflow-fastapi-api.pid"
-        pidfile_monitor = parent_path / "pidflow-fastapi-api-monitor.pid"
-        stdout = parent_path / "airflow-fastapi-api.out"
-        stderr = parent_path / "airflow-fastapi-api.err"
-        logfile = parent_path / "airflow-fastapi-api.log"
-        try:
-            # Run fastapi-api as daemon in background. Note that the wait method is not called.
-            console.print("[magenta]Starting airflow fastapi-api --daemon")
-            env = os.environ.copy()
-            proc = subprocess.Popen(
-                [
-                    "airflow",
-                    "fastapi-api",
-                    "--daemon",
-                    "--pid",
-                    os.fspath(pidfile_fastapi_api),
-                    "--stdout",
-                    os.fspath(stdout),
-                    "--stderr",
-                    os.fspath(stderr),
-                    "--log-file",
-                    os.fspath(logfile),
-                ],
-                env=env,
-            )
-            assert proc.poll() is None
-
-            pid_monitor = self._wait_pidfile(pidfile_monitor)
-            console.print(f"[blue]Monitor started at {pid_monitor}")
-            pid_fastapi_api = self._wait_pidfile(pidfile_fastapi_api)
-            console.print(f"[blue]FastAPI API started at {pid_fastapi_api}")
-            console.print("[blue]Running airflow fastapi-api process:")
-            # Assert that the fastapi-api and gunicorn processes are running (by name rather than pid).
-            assert self._find_process(r"airflow fastapi-api --daemon", print_found_process=True)
-            console.print("[blue]Waiting for gunicorn processes:")
-            # wait for gunicorn to start
-            for _ in range(30):
-                if self._find_process(r"^gunicorn"):
-                    break
-                console.print("[blue]Waiting for gunicorn to start ...")
-                time.sleep(1)
-            console.print("[blue]Running gunicorn processes:")
-            assert self._find_all_processes("^gunicorn", print_found_process=True)
-            console.print("[magenta]fastapi-api process started successfully.")
-            console.print(
-                "[magenta]Terminating monitor process and expect "
-                "fastapi-api and gunicorn processes to terminate as well"
-            )
-            self._terminate_multiple_process([pid_fastapi_api, pid_monitor])
-            self._check_processes(ignore_running=False)
-            console.print("[magenta]All fastapi-api and gunicorn processes are terminated.")
-        except Exception:
-            console.print("[red]Exception occurred. Dumping all logs.")
-            # Dump all logs
-            for file in parent_path.glob("*"):
-                console.print(f"Dumping {file} (size: {file.stat().st_size})")
-                console.print(file.read_text())
-            raise
-
     def test_cli_fastapi_api_debug(self, app):
         with (
             mock.patch("subprocess.Popen") as Popen,
-            mock.patch.object(fastapi_api_command, "GunicornMonitor"),
         ):
             port = "9092"
             hostname = "somehost"
@@ -130,7 +62,6 @@ def test_cli_fastapi_api_env_var_set_unset(self, app):
         """
         with (
             mock.patch("subprocess.Popen") as Popen,
-            mock.patch.object(fastapi_api_command, "GunicornMonitor"),
             mock.patch("os.environ", autospec=True) as mock_environ,
         ):
             apps_value = "core,execution"
@@ -172,8 +103,7 @@ def test_cli_fastapi_api_args(self, ssl_cert_and_key):
         cert_path, key_path = ssl_cert_and_key
 
         with (
-            mock.patch("subprocess.Popen") as Popen,
-            mock.patch.object(fastapi_api_command, "GunicornMonitor"),
+            mock.patch("uvicorn.run") as mock_run,
         ):
             args = self.parser.parse_args(
                 [
@@ -192,39 +122,16 @@ def test_cli_fastapi_api_args(self, ssl_cert_and_key):
             )
             fastapi_api_command.fastapi_api(args)
 
-            Popen.assert_called_with(
-                [
-                    sys.executable,
-                    "-m",
-                    "gunicorn",
-                    "--workers",
-                    "4",
-                    "--worker-class",
-                    "airflow.cli.commands.local_commands.fastapi_api_command.AirflowUvicornWorker",
-                    "--timeout",
-                    "120",
-                    "--bind",
-                    "0.0.0.0:9091",
-                    "--name",
-                    "airflow-fastapi-api",
-                    "--pid",
-                    "/tmp/x.pid",
-                    "--access-logfile",
-                    "-",
-                    "--error-logfile",
-                    "-",
-                    "--config",
-                    "python:airflow.api_fastapi.gunicorn_config",
-                    "--certfile",
-                    str(cert_path),
-                    "--keyfile",
-                    str(key_path),
-                    "--access-logformat",
-                    "custom_log_format",
-                    "airflow.api_fastapi.app:cached_app(apps='core')",
-                    "--preload",
-                ],
-                close_fds=True,
+            mock_run.assert_called_with(
+                "airflow.api_fastapi.main:app",
+                host="0.0.0.0",
+                port=9091,
+                workers=4,
+                timeout_keep_alive=120,
+                timeout_graceful_shutdown=120,
+                ssl_keyfile=str(key_path),
+                ssl_certfile=str(cert_path),
+                access_log="-",
             )
 
     @pytest.mark.parametrize(