(airflow) branch main updated: Ensure internal api command get the "ready" prefix on start (#38550)

[ds...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dstandish pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new b255012cf0 Ensure internal api command get the "ready" prefix on start (#38550)
b255012cf0 is described below

commit b255012cf01ff79357a22a9b5795447c66f2feab
Author: Daniel Standish <15...@users.noreply.github.com>
AuthorDate: Tue Apr 2 22:37:17 2024 -0700

    Ensure internal api command get the "ready" prefix on start (#38550)
    
    This prefix is used by _get_num_ready_workers_running to determine num workers ready.
---
 airflow/api_internal/gunicorn_config.py         | 33 +++++++++++++++++++++++++
 airflow/cli/commands/internal_api_command.py    |  2 ++
 tests/cli/commands/test_internal_api_command.py |  2 ++
 3 files changed, 37 insertions(+)

diff --git a/airflow/api_internal/gunicorn_config.py b/airflow/api_internal/gunicorn_config.py
new file mode 100644
index 0000000000..1ed8e1f941
--- /dev/null
+++ b/airflow/api_internal/gunicorn_config.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import setproctitle
+
+from airflow import settings
+
+
+def post_worker_init(_):
+    """
+    Set process title.
+
+    This is used by airflow.cli.commands.internal_api_command to track the status of the worker.
+    """
+    old_title = setproctitle.getproctitle()
+    setproctitle.setproctitle(settings.GUNICORN_WORKER_READY_PREFIX + old_title)
diff --git a/airflow/cli/commands/internal_api_command.py b/airflow/cli/commands/internal_api_command.py
index 00a5c4b840..8c25d1fa5a 100644
--- a/airflow/cli/commands/internal_api_command.py
+++ b/airflow/cli/commands/internal_api_command.py
@@ -115,6 +115,8 @@ def internal_api(args):
             str(access_logfile),
             "--error-logfile",
             str(error_logfile),
+            "--config",
+            "python:airflow.api_internal.gunicorn_config",
         ]
 
         if args.access_logformat and args.access_logformat.strip():
diff --git a/tests/cli/commands/test_internal_api_command.py b/tests/cli/commands/test_internal_api_command.py
index f3940278a4..2c41d16f75 100644
--- a/tests/cli/commands/test_internal_api_command.py
+++ b/tests/cli/commands/test_internal_api_command.py
@@ -205,6 +205,8 @@ class TestCliInternalAPI(_ComonCLIGunicornTestClass):
                     "-",
                     "--error-logfile",
                     "-",
+                    "--config",
+                    "python:airflow.api_internal.gunicorn_config",
                     "--access-logformat",
                     "custom_log_format",
                     "airflow.cli.commands.internal_api_command:cached_app()",