diff --git a/airflow/providers/celery/executors/celery_executor.py b/airflow/providers/celery/executors/celery_executor.py index 2a75be91da56e..0b4293cde728f 100644 --- a/airflow/providers/celery/executors/celery_executor.py +++ b/airflow/providers/celery/executors/celery_executor.py @@ -30,10 +30,12 @@ import time from collections import Counter from concurrent.futures import ProcessPoolExecutor +from importlib.metadata import version as importlib_version from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, Optional, Sequence, Tuple from celery import states as celery_states +from packaging.version import Version try: from airflow.cli.cli_config import ( @@ -178,11 +180,19 @@ def __getattr__(name): action="store_true", ) +AIRFLOW_VERSION = Version(importlib_version("apache-airflow")) + +CELERY_CLI_COMMAND_PATH = ( + "airflow.providers.celery.cli.celery_command" + if AIRFLOW_VERSION >= Version("2.8.0") + else "airflow.cli.commands.celery_command" +) + CELERY_COMMANDS = ( ActionCommand( name="worker", help="Start a Celery worker node", - func=lazy_load_command("airflow.providers.celery.cli.celery_command.worker"), + func=lazy_load_command(f"{CELERY_CLI_COMMAND_PATH}.worker"), args=( ARG_QUEUES, ARG_CONCURRENCY, @@ -203,7 +213,7 @@ def __getattr__(name): ActionCommand( name="flower", help="Start a Celery Flower", - func=lazy_load_command("airflow.providers.celery.cli.celery_command.flower"), + func=lazy_load_command(f"{CELERY_CLI_COMMAND_PATH}.flower"), args=( ARG_FLOWER_HOSTNAME, ARG_FLOWER_PORT, @@ -222,7 +232,7 @@ def __getattr__(name): ActionCommand( name="stop", help="Stop the Celery worker gracefully", - func=lazy_load_command("airflow.providers.celery.cli.celery_command.stop_worker"), + func=lazy_load_command(f"{CELERY_CLI_COMMAND_PATH}.stop_worker"), args=(ARG_PID, ARG_VERBOSE), ), )