diff --git a/providers/celery/docs/index.rst b/providers/celery/docs/index.rst index e1b09271f2820..78e500b369f86 100644 --- a/providers/celery/docs/index.rst +++ b/providers/celery/docs/index.rst @@ -113,6 +113,7 @@ You can install such cross-provider dependencies when installing from PyPI. For Dependent package Extra ====================================================================================================================== =================== `apache-airflow-providers-cncf-kubernetes `_ ``cncf.kubernetes`` +`apache-airflow-providers-common-compat `_ ``common.compat`` ====================================================================================================================== =================== Optional dependencies diff --git a/providers/celery/provider.yaml b/providers/celery/provider.yaml index f5506e41122f3..5e862a9c1180e 100644 --- a/providers/celery/provider.yaml +++ b/providers/celery/provider.yaml @@ -227,6 +227,17 @@ config: type: boolean example: ~ default: ~ + task_logs_to_stdout: + description: | + Also forward task subprocess stdout/stderr to the Celery worker's own stdout, + so task logs reach a container-level log collector (e.g. Kubernetes/Loki) in + addition to the task-log handler and the UI. This gives the Celery worker parity + with the LocalExecutor and the KubernetesExecutor per-task pod, which always + forward task logs to stdout. Disabled by default to preserve existing behaviour. + version_added: ~ + type: boolean + example: ~ + default: "False" broker_url: description: | The Celery broker URL. Celery supports multiple broker types. See: diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py index 39d52deb1dd13..afe7ff82f1ea5 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py @@ -240,7 +240,10 @@ def execute_workload(input: str) -> None: log.info("[%s] Executing workload in Celery: %s", celery_task_id, workload) try: - BaseExecutor.run_workload(workload) + BaseExecutor.run_workload( + workload, + subprocess_logs_to_stdout=conf.getboolean("celery", "task_logs_to_stdout", fallback=False), + ) except Exception as e: from airflow.sdk.exceptions import TaskAlreadyRunningError @@ -285,6 +288,7 @@ def _execute_workload_pre_3_3(input: str) -> None: token=workload.token, server=conf.get("core", "execution_api_server_url", fallback=default_execution_api_server), log_path=workload.log_path, + subprocess_logs_to_stdout=conf.getboolean("celery", "task_logs_to_stdout", fallback=False), ) else: raise ValueError(f"CeleryExecutor does not know how to handle {type(workload)}") diff --git a/providers/celery/src/airflow/providers/celery/get_provider_info.py b/providers/celery/src/airflow/providers/celery/get_provider_info.py index 4a79cc58821d5..f79c4a68004b2 100644 --- a/providers/celery/src/airflow/providers/celery/get_provider_info.py +++ b/providers/celery/src/airflow/providers/celery/get_provider_info.py @@ -117,6 +117,13 @@ def get_provider_info(): "example": None, "default": None, }, + "task_logs_to_stdout": { + "description": "Also forward task subprocess stdout/stderr to the Celery worker's own stdout,\nso task logs reach a container-level log collector (e.g. Kubernetes/Loki) in\naddition to the task-log handler and the UI. This gives the Celery worker parity\nwith the LocalExecutor and the KubernetesExecutor per-task pod, which always\nforward task logs to stdout. Disabled by default to preserve existing behaviour.\n", + "version_added": None, + "type": "boolean", + "example": None, + "default": "False", + }, "broker_url": { "description": "The Celery broker URL. Celery supports multiple broker types. See:\nhttps://docs.celeryq.dev/en/stable/getting-started/backends-and-brokers/index.html#broker-overview\n", "version_added": None, diff --git a/providers/celery/tests/unit/celery/executors/test_celery_executor.py b/providers/celery/tests/unit/celery/executors/test_celery_executor.py index 28b0840590235..7543cf50f93e0 100644 --- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py +++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py @@ -1147,6 +1147,91 @@ def test_execute_workload_runs_execute_task_before_airflow_3_3(): assert mock_supervise.call_args.kwargs["log_path"] == "test.log" +@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="execute_workload is only used for Airflow 3+") +@pytest.mark.skipif(AIRFLOW_V_3_3_PLUS, reason="pre-3.3 compatibility path only applies before Airflow 3.3") +@pytest.mark.parametrize( + ("config_value", "expected"), + [("True", True), ("False", False)], +) +def test_execute_workload_forwards_task_logs_to_stdout_before_airflow_3_3(config_value, expected): + """The [celery] task_logs_to_stdout config is passed to supervise on Airflow 3.0-3.2.""" + from airflow.executors import workloads + + workload = workloads.ExecuteTask( + ti=workloads.TaskInstance( + id="00000000-0000-0000-0000-000000000001", + dag_version_id="00000000-0000-0000-0000-000000000002", + task_id="test_task", + dag_id="test_dag", + run_id="test_run", + try_number=1, + map_index=-1, + pool_slots=1, + queue="default", + priority_weight=1, + ), + dag_rel_path="test_dag.py", + bundle_info=workloads.BundleInfo(name="test-bundle", version=None), + token="test-token", + log_path="test.log", + ) + mock_current_task = mock.MagicMock() + mock_current_task.request.id = "test-celery-task-id" + mock_app = mock.MagicMock() + mock_app.current_task = mock_current_task + + with ( + conf_vars({("celery", "task_logs_to_stdout"): config_value}), + mock.patch.object(celery_executor_utils, "app", mock_app), + mock.patch("airflow.sdk.execution_time.supervisor.supervise") as mock_supervise, + ): + celery_executor_utils.execute_workload.__wrapped__(workload.model_dump_json()) + + assert mock_supervise.call_args.kwargs["subprocess_logs_to_stdout"] is expected + + +@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="3.3+ path uses BaseExecutor.run_workload") +@pytest.mark.parametrize( + ("config_value", "expected"), + [("True", True), ("False", False)], +) +def test_execute_workload_forwards_task_logs_to_stdout_on_airflow_3_3_plus(config_value, expected): + """The [celery] task_logs_to_stdout config is passed to run_workload on Airflow 3.3+.""" + from airflow.executors import workloads + + workload = workloads.ExecuteTask( + ti=workloads.TaskInstance( + id="00000000-0000-0000-0000-000000000001", + dag_version_id="00000000-0000-0000-0000-000000000002", + task_id="test_task", + dag_id="test_dag", + run_id="test_run", + try_number=1, + map_index=-1, + pool_slots=1, + queue="default", + priority_weight=1, + ), + dag_rel_path="test_dag.py", + bundle_info=workloads.BundleInfo(name="test-bundle", version=None), + token="test-token", + log_path="test.log", + ) + mock_current_task = mock.MagicMock() + mock_current_task.request.id = "test-celery-task-id" + mock_app = mock.MagicMock() + mock_app.current_task = mock_current_task + + with ( + conf_vars({("celery", "task_logs_to_stdout"): config_value}), + mock.patch.object(celery_executor_utils, "app", mock_app), + mock.patch("airflow.executors.base_executor.BaseExecutor.run_workload") as mock_run_workload, + ): + celery_executor_utils.execute_workload.__wrapped__(workload.model_dump_json()) + + assert mock_run_workload.call_args.kwargs["subprocess_logs_to_stdout"] is expected + + @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="3.3+ path uses BaseExecutor.run_workload") def test_execute_workload_runs_base_executor_workload_on_airflow_3_3_plus(): """execute_workload routes serialized ExecuteTask payloads to BaseExecutor on Airflow 3.3+."""