This is an automated email from the ASF dual-hosted git repository.
kaxil 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 3905dc509bb Cache Celery apps when publishing workloads (#67127)
3905dc509bb is described below
commit 3905dc509bbb85b596ae3a25b8066cf59d0895cb
Author: Anmol Mishra <[email protected]>
AuthorDate: Fri Jun 26 06:38:34 2026 +0530
Cache Celery apps when publishing workloads (#67127)
* Cache Celery apps when publishing workloads
* Address Celery app cache review comments
* fix: clear workload celery app cache in integration test _prepare_app
The @cache decorator on _get_celery_app_for_workload meant that once
an app was cached for a team_name, subsequent calls bypassed the
patched create_celery_app in _prepare_app(). With _sync_parallelism=1
(inline execution), send_workload_to_executor got a stale app from a
previous test instead of the test_app, causing the broker connection
to hang and the CI job to time out.
Clear the cache before entering and after exiting the patched context
so each test gets a fresh app via the patched factory.
* fix: address kaxil's docstring and redundant cache_clear review comments
---------
Co-authored-by: Anmol Mishra <[email protected]>
---
.../celery/executors/celery_executor_utils.py | 43 ++++++----
.../integration/celery/test_celery_executor.py | 6 ++
.../unit/celery/executors/test_celery_executor.py | 99 ++++++++++++++++++++++
3 files changed, 132 insertions(+), 16 deletions(-)
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 9b5a9ae759a..39d52deb1dd 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
@@ -160,6 +160,25 @@ def create_celery_app(team_conf: ExecutorConf |
AirflowConfigParser) -> Celery:
return celery_app
+@cache
+def _get_celery_app_for_workload(team_name: str | None) -> Celery:
+ """
+ Return a Celery app cached by team name for task publishing.
+
+ Publishing workloads may run either inline in the scheduler process or in
a publisher
+ subprocess. Cache the app in whichever process executes the publish path
so result
+ backend resolution is amortized while retaining per-team broker isolation.
+ """
+ if AIRFLOW_V_3_2_PLUS:
+ from airflow.executors.base_executor import ExecutorConf
+
+ _conf = ExecutorConf(team_name)
+ else:
+ # Airflow <3.2 ExecutorConf doesn't exist (at least not with the
required attributes), fall back to global conf.
+ _conf = conf
+ return create_celery_app(_conf)
+
+
# Keep module-level app for backward compatibility.
app = _get_celery_app()
@@ -388,25 +407,17 @@ def send_workload_to_executor(
"""
Send workload to executor (serialized and executed as a Celery task).
- This function is called in ProcessPoolExecutor subprocesses. To avoid
pickling issues with
- team-specific Celery apps, we pass the team_name and reconstruct the
Celery app here.
+ This function runs either inline in the long-lived scheduler process
(single-workload or
+ sync_parallelism=1 path) or in short-lived ProcessPoolExecutor
subprocesses (multi-workload
+ path). To avoid pickling issues with team-specific Celery apps, we pass
the team_name and
+ create the app at call time. The cached app lives for the duration of the
caller process, so
+ the main benefit is the scheduler-inline path where the cache persists
across publish cycles.
+ In the ProcessPoolExecutor path, each subprocess is recreated per publish
batch and the cache
+ only lasts for that single batch.
"""
key, args, queue, team_name = workload_tuple
- # Reconstruct the Celery app from configuration, which may or may not be
team-specific.
- # ExecutorConf wraps config access to automatically use team-specific
config where present.
- if TYPE_CHECKING:
- _conf: ExecutorConf | AirflowConfigParser
- # Check if Airflow version is greater than or equal to 3.2 to import
ExecutorConf.
- if AIRFLOW_V_3_2_PLUS:
- from airflow.executors.base_executor import ExecutorConf
-
- _conf = ExecutorConf(team_name)
- else:
- # Airflow <3.2 ExecutorConf doesn't exist (at least not with the
required attributes), fall back to global conf.
- _conf = conf
- # Create the Celery app with the correct configuration.
- celery_app = create_celery_app(_conf)
+ celery_app = _get_celery_app_for_workload(team_name)
celery_task_id = None
if AIRFLOW_V_3_0_PLUS:
diff --git a/providers/celery/tests/integration/celery/test_celery_executor.py
b/providers/celery/tests/integration/celery/test_celery_executor.py
index 4b7bb08d97f..f068f93375c 100644
--- a/providers/celery/tests/integration/celery/test_celery_executor.py
+++ b/providers/celery/tests/integration/celery/test_celery_executor.py
@@ -116,12 +116,18 @@ def _prepare_app(broker_url=None, execute=None):
session = backend.ResultSession()
session.close()
+ # Clear the per-team workload app cache so the patched create_celery_app
+ # is actually called. Without this, _get_celery_app_for_workload returns
+ # a stale app from a previous test and the patched factory is bypassed.
+ celery_executor_utils._get_celery_app_for_workload.cache_clear()
+
with patch_app, patch_execute, patch_factory:
try:
yield test_app
finally:
# Clear event loop to tear down each celery instance
set_event_loop(None)
+ celery_executor_utils._get_celery_app_for_workload.cache_clear()
def setup_dagrun_with_success_and_fail_workloads(dag_maker):
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 c11ea80a5ba..c349efa32cb 100644
--- a/providers/celery/tests/unit/celery/executors/test_celery_executor.py
+++ b/providers/celery/tests/unit/celery/executors/test_celery_executor.py
@@ -77,6 +77,13 @@ else:
pytestmark = pytest.mark.db_test
[email protected](autouse=True)
+def clear_cached_workload_celery_apps():
+ celery_executor_utils._get_celery_app_for_workload.cache_clear()
+ yield
+ celery_executor_utils._get_celery_app_for_workload.cache_clear()
+
+
FAKE_EXCEPTION_MSG = "Fake Exception"
@@ -551,6 +558,98 @@ def
test_send_workload_uses_external_executor_id_as_celery_task_id():
assert result.task_id == pre_assigned_id
[email protected]("team_name", [None, "team-a"])
+def test_get_celery_app_for_workload_reuses_cache_for_same_team(team_name):
+ first_app = mock.Mock()
+ second_app = mock.Mock()
+
+ with mock.patch(
+
"airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
+ side_effect=[first_app, second_app],
+ ) as mock_create_celery_app:
+ assert celery_executor_utils._get_celery_app_for_workload(team_name)
is first_app
+ assert celery_executor_utils._get_celery_app_for_workload(team_name)
is first_app
+
+ mock_create_celery_app.assert_called_once()
+
+
+def test_get_celery_app_for_workload_keeps_cache_team_scoped():
+ team_a_app = mock.Mock()
+ team_b_app = mock.Mock()
+
+ with mock.patch(
+
"airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
+ side_effect=[team_a_app, team_b_app],
+ ) as mock_create_celery_app:
+ assert celery_executor_utils._get_celery_app_for_workload("team-a") is
team_a_app
+ assert celery_executor_utils._get_celery_app_for_workload("team-b") is
team_b_app
+
+ assert mock_create_celery_app.call_count == 2
+
+
+def test_send_workload_reuses_celery_app_for_same_team():
+ """Publishing multiple workloads for the same team reuses the cached
Celery app."""
+ key = TaskInstanceKey(
+ dag_id="test_dag", task_id="test_task", run_id="test_run",
map_index=-1, try_number=1
+ )
+ mock_result = mock.Mock(task_id="mock-task-id")
+ mock_celery_task = mock.Mock()
+ mock_celery_task.apply_async.return_value = mock_result
+ mock_app = mock.Mock()
+ task_name = "execute_workload" if AIRFLOW_V_3_0_PLUS else "execute_command"
+ mock_app.tasks = {task_name: mock_celery_task}
+
+ if AIRFLOW_V_3_0_PLUS:
+ workload = mock.Mock()
+ workload.ti.external_executor_id = None
+ workload.model_dump_json.return_value = "{}"
+ else:
+ workload = ["airflow", "tasks", "run", "test_dag", "test_task",
"test_run"]
+
+ with mock.patch(
+
"airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
+ return_value=mock_app,
+ ) as mock_create_celery_app:
+ celery_executor_utils.send_workload_to_executor((key, workload,
"default", "team-a"))
+ celery_executor_utils.send_workload_to_executor((key, workload,
"default", "team-a"))
+
+ mock_create_celery_app.assert_called_once()
+ assert mock_celery_task.apply_async.call_count == 2
+
+
+def test_send_workload_keeps_celery_app_cache_team_scoped():
+ """Different teams get distinct cached Celery app instances in the
publisher process."""
+ key = TaskInstanceKey(
+ dag_id="test_dag", task_id="test_task", run_id="test_run",
map_index=-1, try_number=1
+ )
+ mock_result = mock.Mock(task_id="mock-task-id")
+ team_a_task = mock.Mock()
+ team_a_task.apply_async.return_value = mock_result
+ team_b_task = mock.Mock()
+ team_b_task.apply_async.return_value = mock_result
+ task_name = "execute_workload" if AIRFLOW_V_3_0_PLUS else "execute_command"
+ team_a_app = mock.Mock(tasks={task_name: team_a_task})
+ team_b_app = mock.Mock(tasks={task_name: team_b_task})
+
+ if AIRFLOW_V_3_0_PLUS:
+ workload = mock.Mock()
+ workload.ti.external_executor_id = None
+ workload.model_dump_json.return_value = "{}"
+ else:
+ workload = ["airflow", "tasks", "run", "test_dag", "test_task",
"test_run"]
+
+ with mock.patch(
+
"airflow.providers.celery.executors.celery_executor_utils.create_celery_app",
+ side_effect=[team_a_app, team_b_app],
+ ) as mock_create_celery_app:
+ celery_executor_utils.send_workload_to_executor((key, workload,
"default", "team-a"))
+ celery_executor_utils.send_workload_to_executor((key, workload,
"default", "team-b"))
+
+ assert mock_create_celery_app.call_count == 2
+ team_a_task.apply_async.assert_called_once()
+ team_b_task.apply_async.assert_called_once()
+
+
@conf_vars({("celery", "result_backend"):
"rediss://test_user:test_password@localhost:6379/0"})
def test_celery_executor_with_no_recommended_result_backend(caplog):
import importlib
