You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "DaveT1991 (via GitHub)" <gi...@apache.org> on 2026/04/10 04:57:57 UTC
[PR] Persist Celery external executor ids on send [airflow]
DaveT1991 opened a new pull request, #64997:
URL: https://github.com/apache/airflow/pull/64997
## Summary
Persist the Celery task id to `task_instance.external_executor_id` as soon as a workload is successfully submitted.
Today the Celery executor stores the returned `task_id` in `event_buffer`, and the scheduler later writes it to the database while handling the queued/running event. If the scheduler dies in that window, the replacement scheduler cannot adopt the in-flight task because `external_executor_id` is still `NULL`, so the task may be reset and re-queued instead of adopted.
This change persists the Celery task id immediately after a successful send for task-instance workloads, while keeping the existing event-buffer path as the normal reconciliation flow.
Closes #64971
## Testing
- Added regression coverage in `providers/celery/tests/unit/celery/executors/test_celery_executor.py`
- `git diff --check`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Persist Celery external executor ids on send [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #64997:
URL: https://github.com/apache/airflow/pull/64997#issuecomment-4300900248
@DaveT1991 This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **Failing CI**: `CI image checks / Static checks`. Please investigate and fix.
- :x: **Unresolved review comments** (1 thread): please walk through each unresolved review thread. **Even if a suggestion looks incorrect or irrelevant** — and some of them will be, especially any comments left by automated reviewers like GitHub Copilot — it is still the author's responsibility to respond: apply the fix, reply in-thread with a brief explanation of why the suggestion does not apply, or resolve the thread if the feedback is no longer relevant. **Once you believe a thread is resolved — whether by pushing a fix or by explaining why the suggestion doesn't apply — please mark it as resolved yourself by clicking the 'Resolve conversation' button at the bottom of each thread.** Reviewers don't auto-close their own threads, so an addressed-but-unresolved thread reads as 'still waiting on the author' and keeps the PR from moving forward.
**What to do next:**
- Address the items above.
- Make sure static checks pass locally: `prek run --from-ref main --stage pre-commit`.
- Mark the PR as "Ready for review" when you are done.
Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Persist Celery external executor ids on send [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #64997:
URL: https://github.com/apache/airflow/pull/64997#discussion_r3066465827
##########
providers/celery/src/airflow/providers/celery/executors/celery_executor.py:
##########
@@ -39,6 +39,8 @@
from airflow.exceptions import AirflowProviderDeprecationWarning
from airflow.executors.base_executor import BaseExecutor
+from airflow.models.taskinstance import TaskInstance
+from airflow.models.taskinstancekey import TaskInstanceKey
from airflow.providers.celery.executors import (
Review Comment:
Importing ORM models (TaskInstance/TaskInstanceKey) at module import time makes `airflow.providers.celery.executors.celery_executor` heavier to import. This module is imported by Celery CLI commands to access `app` (see `providers/celery/.../celery_command.py`), and the module already uses `__getattr__` to avoid import-time costs. Consider lazily importing these models inside `_persist_task_external_executor_id` (or another narrow code path) so CLI/worker paths that only need `app` don't pay the ORM import cost / risk DB initialization side effects.
##########
providers/celery/src/airflow/providers/celery/executors/celery_executor.py:
##########
@@ -225,12 +225,39 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]):
result.backend = cached_celery_backend
self.running.add(key)
self.workloads[key] = result
+ # Persist the Celery task_id before scheduler event handling so a replacement
+ # scheduler can still adopt the task after a crash/restart.
+ self._persist_task_external_executor_id(key, result.task_id)
# Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task
# has another event, but that is fine, because the only other events are success/failed at
# which point we don't need the ID anymore anyway.
self.event_buffer[key] = (TaskInstanceState.QUEUED, result.task_id)
+ def _persist_task_external_executor_id(self, key: WorkloadKey, task_id: str) -> None:
+ """Persist Celery task ids for task workloads so they survive scheduler restarts."""
+ from sqlalchemy import update
+
+ from airflow.utils.session import create_session
+
+ if not isinstance(key, TaskInstanceKey):
+ return
+
+ with create_session() as session:
+ result = session.execute(
+ update(TaskInstance)
+ .where(
+ TaskInstance.dag_id == key.dag_id,
Review Comment:
If the DB update/commit fails here (e.g., transient DB outage/lock), the exception will propagate even though the Celery task has already been successfully sent. That can break the scheduler/executor loop and leave the task in-flight but not tracked. Consider catching SQLAlchemy/DB exceptions around this block, logging at warning/debug, and continuing so task submission isn’t treated as failed solely due to the best-effort persistence step.
##########
providers/celery/tests/unit/celery/executors/test_celery_executor.py:
##########
@@ -285,6 +285,45 @@ def test_try_adopt_task_instances(self, clean_dags_dagruns_and_dagbundles, testi
assert executor.workloads == {key_1: AsyncResult("231"), key_2: AsyncResult("232")}
assert not_adopted_tis == []
+ @pytest.mark.backend("mysql", "postgres")
+ def test_send_workloads_persists_external_executor_id(
+ self, clean_dags_dagruns_and_dagbundles, testing_dag_bundle
+ ):
+ from sqlalchemy import select
+
+ from airflow.utils.session import create_session
+
Review Comment:
These imports are inside the test function but don’t appear to be conditional/lazy. Moving them to the module import section keeps test code consistent and avoids repeated imports when the test is collected/executed multiple times.
##########
providers/celery/src/airflow/providers/celery/executors/celery_executor.py:
##########
@@ -225,12 +225,39 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]):
result.backend = cached_celery_backend
self.running.add(key)
self.workloads[key] = result
+ # Persist the Celery task_id before scheduler event handling so a replacement
+ # scheduler can still adopt the task after a crash/restart.
+ self._persist_task_external_executor_id(key, result.task_id)
Review Comment:
`_persist_task_external_executor_id` opens and commits a new DB session for every workload in `_send_workloads`. On busy schedulers this adds extra per-task latency and DB load (in addition to the existing event-buffer reconciliation write). Consider batching these updates (e.g., collect task keys/task_ids and run a single session/transaction per `_send_workloads` call, or use an executemany-style update) to reduce overhead.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Persist Celery external executor ids on send [airflow]
Posted by "DaveT1991 (via GitHub)" <gi...@apache.org>.
DaveT1991 commented on code in PR #64997:
URL: https://github.com/apache/airflow/pull/64997#discussion_r3066098434
##########
providers/celery/src/airflow/providers/celery/executors/celery_executor.py:
##########
@@ -225,12 +225,41 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]):
result.backend = cached_celery_backend
self.running.add(key)
self.workloads[key] = result
+ # Persist the Celery task_id before scheduler event handling so a replacement
+ # scheduler can still adopt the task after a crash/restart.
+ self._persist_task_external_executor_id(key, result.task_id)
# Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task
# has another event, but that is fine, because the only other events are success/failed at
# which point we don't need the ID anymore anyway.
self.event_buffer[key] = (TaskInstanceState.QUEUED, result.task_id)
+ def _persist_task_external_executor_id(self, key: WorkloadKey, task_id: str) -> None:
+ """Persist Celery task ids for task workloads so they survive scheduler restarts."""
+ from sqlalchemy import update
+
+ from airflow.models.taskinstance import TaskInstance as TI
+ from airflow.models.taskinstancekey import TaskInstanceKey
Review Comment:
Of course, can you check now and let me know?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Persist Celery external executor ids on send [airflow]
Posted by "DaveT1991 (via GitHub)" <gi...@apache.org>.
DaveT1991 commented on code in PR #64997:
URL: https://github.com/apache/airflow/pull/64997#discussion_r3066407976
##########
providers/celery/src/airflow/providers/celery/executors/celery_executor.py:
##########
@@ -225,12 +225,41 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]):
result.backend = cached_celery_backend
self.running.add(key)
self.workloads[key] = result
+ # Persist the Celery task_id before scheduler event handling so a replacement
+ # scheduler can still adopt the task after a crash/restart.
+ self._persist_task_external_executor_id(key, result.task_id)
# Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task
# has another event, but that is fine, because the only other events are success/failed at
# which point we don't need the ID anymore anyway.
self.event_buffer[key] = (TaskInstanceState.QUEUED, result.task_id)
+ def _persist_task_external_executor_id(self, key: WorkloadKey, task_id: str) -> None:
+ """Persist Celery task ids for task workloads so they survive scheduler restarts."""
+ from sqlalchemy import update
+
+ from airflow.models.taskinstance import TaskInstance as TI
+ from airflow.models.taskinstancekey import TaskInstanceKey
Review Comment:
@Prab-27 I did some minor import refactoring to solve the failed testing issue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Persist Celery external executor ids on send [airflow]
Posted by "ashb (via GitHub)" <gi...@apache.org>.
ashb closed pull request #64997: Persist Celery external executor ids on send
URL: https://github.com/apache/airflow/pull/64997
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Persist Celery external executor ids on send [airflow]
Posted by "seanmuth (via GitHub)" <gi...@apache.org>.
seanmuth commented on PR #64997:
URL: https://github.com/apache/airflow/pull/64997#issuecomment-4246660363
@DaveT1991 got a failing test in this one.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Persist Celery external executor ids on send [airflow]
Posted by "Prab-27 (via GitHub)" <gi...@apache.org>.
Prab-27 commented on code in PR #64997:
URL: https://github.com/apache/airflow/pull/64997#discussion_r3066045570
##########
providers/celery/src/airflow/providers/celery/executors/celery_executor.py:
##########
@@ -225,12 +225,41 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]):
result.backend = cached_celery_backend
self.running.add(key)
self.workloads[key] = result
+ # Persist the Celery task_id before scheduler event handling so a replacement
+ # scheduler can still adopt the task after a crash/restart.
+ self._persist_task_external_executor_id(key, result.task_id)
# Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task
# has another event, but that is fine, because the only other events are success/failed at
# which point we don't need the ID anymore anyway.
self.event_buffer[key] = (TaskInstanceState.QUEUED, result.task_id)
+ def _persist_task_external_executor_id(self, key: WorkloadKey, task_id: str) -> None:
+ """Persist Celery task ids for task workloads so they survive scheduler restarts."""
+ from sqlalchemy import update
+
+ from airflow.models.taskinstance import TaskInstance as TI
+ from airflow.models.taskinstancekey import TaskInstanceKey
Review Comment:
@DaveT1991 These imports statemtns are defined above. Could we use them here ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org