Porting signal processing to the task monitor condition for 2.10.1 (#151

) The task monitor needs to process the following 4 MWAA signals for the graceful update project: 1. Termination signal: Graceful termination of the workers when the environment is going through a graceful update 2. Resume signal: Reverting the state of graceful termination and resume work when the environment is going through a rollback after attempting a graceful update 3. Kill signal: Shutting down the worker without waiting for the current Airflow tasks to finish when the environment is going through a forced update 4. Activation signal: Starting consumption of work from the queue after termination protection has enabled on the corresponding Fargate task The processing is gated behind certain environment variables which are either absent or marked as false for an environment which does not have graceful updates enabled. The CR brings the changes that have been merged for 2.9.2 (#137 and #150) to the newer version. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
aws · Oct 3, 2024 · fabe3c1 · fabe3c1
1 parent bd9c236
commit fabe3c1
Show file tree

Hide file tree

Showing 13 changed files with 303 additions and 74 deletions.
diff --git a/images/airflow/2.10.1/BillOfMaterials/2.10.1-dev/Python-Packages-BOM.txt b/images/airflow/2.10.1/BillOfMaterials/2.10.1-dev/Python-Packages-BOM.txt
@@ -6162,10 +6162,10 @@ Copyright 2013-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 /usr/local/airflow/.local/lib/python3.11/site-packages/boto3-1.35.7.dist-info/NOTICE
 
 boto3-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/mypy_boto3_builder
-/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2024 Vlad Emelianov
@@ -6437,10 +6437,10 @@ one at http://mozilla.org/MPL/2.0/.
 /usr/local/airflow/.local/lib/python3.11/site-packages/botocore-1.35.7.dist-info/NOTICE
 
 botocore-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/botocore-stubs
-/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2022 Vlad Emelianov

diff --git a/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer-dev/Python-Packages-BOM.txt b/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer-dev/Python-Packages-BOM.txt
@@ -6162,10 +6162,10 @@ Copyright 2013-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 /usr/local/airflow/.local/lib/python3.11/site-packages/boto3-1.35.7.dist-info/NOTICE
 
 boto3-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/mypy_boto3_builder
-/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2024 Vlad Emelianov
@@ -6437,10 +6437,10 @@ one at http://mozilla.org/MPL/2.0/.
 /usr/local/airflow/.local/lib/python3.11/site-packages/botocore-1.35.7.dist-info/NOTICE
 
 botocore-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/botocore-stubs
-/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2022 Vlad Emelianov

diff --git a/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer-privileged-dev/Python-Packages-BOM.txt b/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer-privileged-dev/Python-Packages-BOM.txt
@@ -6162,10 +6162,10 @@ Copyright 2013-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 /usr/local/airflow/.local/lib/python3.11/site-packages/boto3-1.35.7.dist-info/NOTICE
 
 boto3-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/mypy_boto3_builder
-/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2024 Vlad Emelianov
@@ -6437,10 +6437,10 @@ one at http://mozilla.org/MPL/2.0/.
 /usr/local/airflow/.local/lib/python3.11/site-packages/botocore-1.35.7.dist-info/NOTICE
 
 botocore-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/botocore-stubs
-/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2022 Vlad Emelianov

diff --git a/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer-privileged/Python-Packages-BOM.txt b/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer-privileged/Python-Packages-BOM.txt
@@ -6162,10 +6162,10 @@ Copyright 2013-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 /usr/local/airflow/.local/lib/python3.11/site-packages/boto3-1.35.7.dist-info/NOTICE
 
 boto3-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/mypy_boto3_builder
-/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2024 Vlad Emelianov
@@ -6437,10 +6437,10 @@ one at http://mozilla.org/MPL/2.0/.
 /usr/local/airflow/.local/lib/python3.11/site-packages/botocore-1.35.7.dist-info/NOTICE
 
 botocore-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/botocore-stubs
-/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2022 Vlad Emelianov

diff --git a/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer/Python-Packages-BOM.txt b/images/airflow/2.10.1/BillOfMaterials/2.10.1-explorer/Python-Packages-BOM.txt
@@ -6162,10 +6162,10 @@ Copyright 2013-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 /usr/local/airflow/.local/lib/python3.11/site-packages/boto3-1.35.7.dist-info/NOTICE
 
 boto3-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/mypy_boto3_builder
-/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2024 Vlad Emelianov
@@ -6437,10 +6437,10 @@ one at http://mozilla.org/MPL/2.0/.
 /usr/local/airflow/.local/lib/python3.11/site-packages/botocore-1.35.7.dist-info/NOTICE
 
 botocore-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/botocore-stubs
-/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2022 Vlad Emelianov

diff --git a/images/airflow/2.10.1/BillOfMaterials/2.10.1/Python-Packages-BOM.txt b/images/airflow/2.10.1/BillOfMaterials/2.10.1/Python-Packages-BOM.txt
@@ -6162,10 +6162,10 @@ Copyright 2013-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 /usr/local/airflow/.local/lib/python3.11/site-packages/boto3-1.35.7.dist-info/NOTICE
 
 boto3-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/mypy_boto3_builder
-/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/boto3_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2024 Vlad Emelianov
@@ -6437,10 +6437,10 @@ one at http://mozilla.org/MPL/2.0/.
 /usr/local/airflow/.local/lib/python3.11/site-packages/botocore-1.35.7.dist-info/NOTICE
 
 botocore-stubs
-1.35.21
+1.35.29
 MIT License
 https://github.com/youtype/botocore-stubs
-/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.21.dist-info/LICENSE
+/usr/local/airflow/.local/lib/python3.11/site-packages/botocore_stubs-1.35.29.dist-info/LICENSE
 MIT License
 
 Copyright (c) 2022 Vlad Emelianov

diff --git a/images/airflow/2.10.1/docker-compose.yaml b/images/airflow/2.10.1/docker-compose.yaml
@@ -54,6 +54,9 @@ x-airflow-common: &airflow-common
     MWAA__LOGGING__AIRFLOW_WORKER_LOGS_ENABLED: ${MWAA__LOGGING__AIRFLOW_WORKER_LOGS_ENABLED}
     MWAA__LOGGING__AIRFLOW_WORKER_LOG_GROUP_ARN: ${MWAA__LOGGING__AIRFLOW_WORKER_LOG_GROUP_ARN}
     MWAA__LOGGING__AIRFLOW_WORKER_LOG_LEVEL: ${MWAA__LOGGING__AIRFLOW_WORKER_LOG_LEVEL}
+    MWAA__CORE__TASK_MONITORING_ENABLED: ${MWAA__CORE__TASK_MONITORING_ENABLED}
+    MWAA__CORE__TERMINATE_IF_IDLE: ${MWAA__CORE__TERMINATE_IF_IDLE}
+    MWAA__CORE__MWAA_SIGNAL_HANDLING_ENABLED: ${MWAA__CORE__MWAA_SIGNAL_HANDLING_ENABLED}
 
   volumes:
     - ./dags:/usr/local/airflow/dags

diff --git a/images/airflow/2.10.1/python/mwaa/celery/sqs_broker.py b/images/airflow/2.10.1/python/mwaa/celery/sqs_broker.py
@@ -400,8 +400,8 @@ def __init__(self, *args, **kwargs):
 
         self.hub = kwargs.get("hub") or get_event_loop()
 
-        # Dynamic workers have the MWAA__CORE__TASK_MONITORING_ENABLED set to 'true'.
-        # This will be used to determine if idle worker checks are to be enabled.
+        # MWAA__CORE__TASK_MONITORING_ENABLED is set to 'true' for workers where we want to monitor count of tasks currently getting
+        # executed on the worker. This will be used to determine if idle worker checks are to be enabled.
         self.idle_worker_monitoring_enabled = (
             os.environ.get("MWAA__CORE__TASK_MONITORING_ENABLED", "false") == "true"
         )