use of debugpy to run airflow components via breeze

dev/breeze/src/airflow_breeze/global_constants.py

@@ -370,6 +370,12 @@ def get_default_platform_machine() -> str:
 SSH_PORT = "12322"
 VITE_DEV_PORT = "5173"
 WEB_HOST_PORT = "28080"
+BREEZE_DEBUG_SCHEDULER_PORT = "50231"
+BREEZE_DEBUG_DAG_PROCESSOR_PORT = "50232"
+BREEZE_DEBUG_API_SERVER_PORT = "50233"
+BREEZE_DEBUG_CELERY_WORKER_PORT = "50234"
+BREEZE_DEBUG_EDGE_PORT = "50235"
+BREEZE_DEBUG_FLOWER_PORT = "50236"
 
 CELERY_BROKER_URLS_MAP = {"rabbitmq": "amqp://guest:guest@rabbitmq:5672", "redis": "redis://redis:6379/0"}
 SQLITE_URL = "sqlite:////root/airflow/sqlite/airflow.db"

dev/breeze/src/airflow_breeze/params/shell_params.py

@@ -37,6 +37,12 @@
     ALLOWED_POSTGRES_VERSIONS,
     ALLOWED_PYTHON_MAJOR_MINOR_VERSIONS,
     APACHE_AIRFLOW_GITHUB_REPOSITORY,
+    BREEZE_DEBUG_API_SERVER_PORT,
+    BREEZE_DEBUG_CELERY_WORKER_PORT,
+    BREEZE_DEBUG_DAG_PROCESSOR_PORT,
+    BREEZE_DEBUG_EDGE_PORT,
+    BREEZE_DEBUG_FLOWER_PORT,
+    BREEZE_DEBUG_SCHEDULER_PORT,
     CELERY_BROKER_URLS_MAP,
     CELERY_EXECUTOR,
     DEFAULT_CELERY_BROKER,
@@ -644,6 +650,13 @@ def env_variables_for_docker_commands(self) -> dict[str, str]:
         _set_var(_env, "WEB_HOST_PORT", None, WEB_HOST_PORT)
         _set_var(_env, "_AIRFLOW_RUN_DB_TESTS_ONLY", self.run_db_tests_only)
         _set_var(_env, "_AIRFLOW_SKIP_DB_TESTS", self.skip_db_tests)
+        _set_var(_env, "BREEZE_DEBUG_SCHEDULER_PORT", None, BREEZE_DEBUG_SCHEDULER_PORT)
+        _set_var(_env, "BREEZE_DEBUG_DAG_PROCESSOR_PORT", None, BREEZE_DEBUG_DAG_PROCESSOR_PORT)
+        _set_var(_env, "BREEZE_DEBUG_API_SERVER_PORT", None, BREEZE_DEBUG_API_SERVER_PORT)
+        _set_var(_env, "BREEZE_DEBUG_CELERY_WORKER_PORT", None, BREEZE_DEBUG_CELERY_WORKER_PORT)
+        _set_var(_env, "BREEZE_DEBUG_EDGE_PORT", None, BREEZE_DEBUG_EDGE_PORT)
+        _set_var(_env, "BREEZE_DEBUG_FLOWER_PORT", None, BREEZE_DEBUG_FLOWER_PORT)
+
         self._generate_env_for_docker_compose_file_if_needed(_env)
 
         _target_env: dict[str, str] = os.environ.copy()

scripts/ci/docker-compose/base-ports.yml

@@ -21,3 +21,9 @@ services:
       - "${SSH_PORT}:22"
       - "${WEB_HOST_PORT}:8080"
       - "${FLOWER_HOST_PORT}:5555"
+      - "${BREEZE_DEBUG_SCHEDULER_PORT}:50231"
+      - "${BREEZE_DEBUG_DAG_PROCESSOR_PORT}:50232"
+      - "${BREEZE_DEBUG_API_SERVER_PORT}:50233"
+      - "${BREEZE_DEBUG_CELERY_WORKER_PORT}:50234"
+      - "${BREEZE_DEBUG_EDGE_PORT}:50235"
+      - "${BREEZE_DEBUG_FLOWER_PORT}:50236"

scripts/in_container/bin/run_tmux

@@ -24,7 +24,7 @@ if [[ ${BACKEND} != "sqlite"  ]]; then
     export AIRFLOW__CORE__EXECUTOR=${AIRFLOW__CORE__EXECUTOR:-LocalExecutor}
 fi
 
-#this is because I run docker in WSL - Hi Nadella!
+
 export TMUX_TMPDIR=~/.tmux/tmp
 if [ -e ~/.tmux/tmp ]; then
     rm -rf ~/.tmux/tmp
@@ -57,70 +57,112 @@ tmux send-keys -t 'Main' 'bash' C-m 'clear' C-m
 
 tmux split-window -v
 tmux select-pane -t 1
-tmux set-option -p @airflow_component Scheduler
-tmux send-keys 'airflow scheduler' C-m
+if [[ ${BREEZE_DEBUG_SCHEDULER} == "true" ]]; then
+    tmux set-option -p @airflow_component "Scheduler(Debug Mode)"
+    tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_SCHEDULER_PORT} --wait-for-client -m airflow scheduler" C-m
+else
+    tmux set-option -p @airflow_component Scheduler
+    tmux send-keys "airflow scheduler" C-m
+fi
 
 if [[ ! ${USE_AIRFLOW_VERSION=} =~ ^2\..*  ]]; then
-  tmux split-window -h
-  tmux select-pane -t 2
-  tmux set-option -p @airflow_component "API Server"
-
-    if [[ ${DEV_MODE=} == "true" ]]; then
-        tmux send-keys 'airflow api-server -d' C-m
+    tmux split-window -h
+    tmux select-pane -t 2
+    if [[ ${BREEZE_DEBUG_APISERVER} == "true" ]]; then
+        tmux set-option -p @airflow_component "API Server(Debug Mode)"
+        tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_APISERVER_PORT} --wait-for-client -m airflow api-server -d" C-m
     else
-        tmux send-keys 'airflow api-server' C-m
+        tmux set-option -p @airflow_component "API Server"
+        if [[ ${DEV_MODE=} == "true" ]]; then
+            tmux send-keys 'airflow api-server -d' C-m
+        else
+            tmux send-keys 'airflow api-server' C-m
+        fi
     fi
 else
-  tmux split-window -h
-  tmux select-pane -t 3
-  tmux set-option -p @airflow_component Webserver
-  if [[ ${DEV_MODE=} == "true" ]]; then
-      tmux send-keys 'airflow webserver -d' C-m
-  else
-      tmux send-keys 'airflow webserver' C-m
-  fi
+    tmux split-window -h
+    tmux select-pane -t 3
+    if [[ ${BREEZE_DEBUG_WEBSERVER} == "true" ]]; then
+        tmux set-option -p @airflow_component "Webserver(Debug Mode)"
+        tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_WEBSERVER_PORT} --wait-for-client -m airflow webserver" C-m
+    else
+        tmux set-option -p @airflow_component Webserver
+        if [[ ${DEV_MODE=} == "true" ]]; then
+            tmux send-keys 'airflow webserver -d' C-m
+        else
+            tmux send-keys 'airflow webserver' C-m
+        fi
+    fi
 fi
 
 tmux select-pane -t 0
 tmux split-window -h
-tmux set-option -p @airflow_component Triggerer
-tmux send-keys 'airflow triggerer' C-m
+if [[ ${BREEZE_DEBUG_TRIGGERER} == "true" ]]; then
+    tmux set-option -p @airflow_component "Triggerer(Debug Mode)"
+    tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_TRIGGERER_PORT} --wait-for-client -m airflow triggerer" C-m
+else
+    tmux set-option -p @airflow_component Triggerer
+    tmux send-keys 'airflow triggerer' C-m
+fi
 
 if [[ ${INTEGRATION_CELERY} == "true" ]]; then
     tmux select-pane -t 0
     tmux split-window -h
-    tmux set-option -p @airflow_component "Celery Worker"
-    tmux send-keys 'airflow celery worker' C-m
+    if [[ ${BREEZE_DEBUG_CELERY_WORKER} == "true" ]]; then
+        tmux set-option -p @airflow_component "Celery Worker(Debug Mode)"
+        tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_CELERY_WORKER_PORT} --wait-for-client -m airflow celery worker" C-m
+    else
+        tmux set-option -p @airflow_component "Celery Worker"
+        tmux send-keys 'airflow celery worker' C-m
+    fi
 fi
+
 if [[ ${INTEGRATION_CELERY} == "true" && ${CELERY_FLOWER} == "true" ]]; then
     tmux select-pane -t 3
     tmux split-window -h
-    tmux set-option -p @airflow_component Flower
-    tmux send-keys 'airflow celery flower' C-m
+    if [[ ${BREEZE_DEBUG_FLOWER} == "true" ]]; then
+        tmux set-option -p @airflow_component "Flower(Debug Mode)"
+        tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_FLOWER_PORT} --wait-for-client -m airflow celery flower" C-m
+    else
+        tmux set-option -p @airflow_component Flower
+        tmux send-keys 'airflow celery flower' C-m
+    fi
 fi
+
 if [[ ${AIRFLOW__CORE__EXECUTOR} == "airflow.providers.edge3.executors.edge_executor.EdgeExecutor" ]]; then
     tmux select-pane -t 0
     tmux split-window -h
-    tmux set-option -p @airflow_component EdgeExec
+    if [[ ${BREEZE_DEBUG_EDGE} == "true" ]]; then
+        tmux set-option -p @airflow_component "EdgeExec(Debug Mode)"
+        tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_EDGE_PORT} --wait-for-client -m airflow edge worker --edge-hostname breeze --queues default" C-m
+    else
+        tmux set-option -p @airflow_component EdgeExec
 
     # Ensure we are not leaking any DB connection information to Edge Worker process
-    tmux send-keys 'unset AIRFLOW__DATABASE__SQL_ALCHEMY_CONN' C-m
-    tmux send-keys 'unset AIRFLOW__CELERY__RESULT_BACKEND' C-m
-    tmux send-keys 'unset POSTGRES_HOST_PORT' C-m
-    tmux send-keys 'unset BACKEND' C-m
-    tmux send-keys 'unset POSTGRES_VERSION' C-m
+        tmux send-keys 'unset AIRFLOW__DATABASE__SQL_ALCHEMY_CONN' C-m
+        tmux send-keys 'unset AIRFLOW__CELERY__RESULT_BACKEND' C-m
+        tmux send-keys 'unset POSTGRES_HOST_PORT' C-m
+        tmux send-keys 'unset BACKEND' C-m
+        tmux send-keys 'unset POSTGRES_VERSION' C-m
 
     # Ensure logs are smelling like Remote and are not visible to other components
-    tmux send-keys 'export AIRFLOW__LOGGING__BASE_LOG_FOLDER=edge_logs' C-m
+        tmux send-keys 'export AIRFLOW__LOGGING__BASE_LOG_FOLDER=edge_logs' C-m
 
     # Start Edge Worker and make a "breeze" hostname, let it pick only default queue
-    tmux send-keys 'airflow edge worker --edge-hostname breeze --queues default' C-m
+        tmux send-keys 'airflow edge worker --edge-hostname breeze --queues default' C-m
+    fi
 fi
+
 if [[ ${STANDALONE_DAG_PROCESSOR} == "true" ]]; then
     tmux select-pane -t 3
     tmux split-window -h
-    tmux set-option -p @airflow_component "DAG Processor"
-    tmux send-keys 'airflow dag-processor' C-m
+    if [[ ${BREEZE_DEBUG_DAG_PROCESSOR} == "true" ]]; then
+        tmux set-option -p @airflow_component "DAG Processor(Debug Mode)"
+        tmux send-keys "debugpy --listen 0.0.0.0:${BREEZE_DEBUG_DAG_PROCESSOR_PORT} --wait-for-client -m airflow dag-processor" C-m
+    else
+        tmux set-option -p @airflow_component "DAG Processor"
+        tmux send-keys 'airflow dag-processor' C-m
+    fi
 fi
 
 # Attach Session, on the Main window