Skip to content

Commit ea486c9

Browse files
committed
Preserve Dag processor work across retries and reparses
Delayed claim responses and repeated reparses must not consume newer work or lose the request that triggered a later import. API-mode processors also need consistent callback delivery semantics and bounded cleanup when recovering from interrupted work.
1 parent 6baf202 commit ea486c9

6 files changed

Lines changed: 174 additions & 104 deletions

File tree

‎airflow-core/docs/security/jwt_token_authentication.rst‎

Lines changed: 18 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -620,21 +620,27 @@ observed server revision; conflicting refreshes must
620620
rediscover before retrying. Imports carry the accepted revision, so an older snapshot cannot
621621
publish after a changed inventory. Only accepted publications establish missing Dags within
622622
a file; rejected requests and incomplete discovery do not establish absence. Every
623-
``[scheduler] parsing_cleanup_interval``, the scheduler marks the Dags of inactive or unassigned
624-
bundles stale, removes stale warnings, and ends expired processor Jobs, in bounded batches and in
623+
``[scheduler] parsing_cleanup_interval``, the scheduler marks the Dags of inactive bundles stale,
624+
removes stale warnings, and ends expired processor Jobs, in bounded batches and in
625625
both modes.
626626

627627
Callback and priority requests remain stored while claimed. A lost claim response or
628-
acknowledgment can be retried, and requests from a retired Job become available to another
629-
processor. A callback acknowledgment confirms delivery, not successful execution of user
630-
code. As with other at-least-once delivery, a crash after a callback's side effect and before
631-
acknowledgment can repeat that side effect. A graceful shutdown acknowledges finished work
632-
before the processor's Job ends. When a callback's bundle version cannot be
633-
prepared, the processor retries with backoff for up to ``[dag_processor] job_heartbeat_timeout``
634-
and then acknowledges the callback as failed. Priority requests are assigned to an import
635-
started after their requested refresh. Only that import's publication or terminal failure
636-
acknowledges the request; an older in-flight import cannot consume it. Discovery after the
637-
requested refresh can also acknowledge a request if it confirms the definition is absent.
628+
acknowledgment can be retried. Claims use increasing sequences per Job and work type; the
629+
client waits for a claim response before sending the next sequence for that type.
630+
Retrying an empty claim remains empty, even if work arrives
631+
afterward, and retrying an older claim cannot consume newly requested work. Requests from a
632+
retired Job become available to another processor. A callback acknowledgment confirms
633+
delivery, not successful execution of user code. As with other at-least-once delivery, a crash
634+
after a callback's side effect and before acknowledgment can repeat that side effect. A graceful
635+
shutdown acknowledges finished work before the processor's Job ends. When a callback's bundle
636+
version cannot be prepared, the processor retries with backoff for up to
637+
``[dag_processor] job_heartbeat_timeout``
638+
and then acknowledges the request without running the callback. Priority requests are assigned
639+
to an import started after their requested refresh. Only that import's publication or terminal
640+
failure acknowledges the request; an older in-flight import cannot consume it. Discovery after
641+
the requested refresh can also acknowledge a request if it confirms the definition is absent.
642+
A new request for a file already being imported remains pending for a subsequent import,
643+
even if the earlier import has not yet been acknowledged.
638644

639645
The manager retains a bounded queue of pending publications. Each loop it sends queued
640646
publications and acknowledgments for up to a second, stopping at the first failure that may

‎airflow-core/src/airflow/config_templates/config.yml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3519,12 +3519,12 @@ dag_processor:
35193519
or listing a single bundle, so set this above the longest of those operations; the local
35203520
health check must also tolerate them, through ``health_check_threshold`` and the probe's
35213521
failure budget. It also bounds how long a processor retries preparing the bundle
3522-
version of a claimed callback before acknowledging the callback as failed.
3522+
version of a claimed callback before acknowledging it without running it.
3523+
A callback that ran before its acknowledgment was recorded can run again after the Job
3524+
times out, repeating its side effects. Callbacks should tolerate repeated execution.
35233525
The scheduler reads this value to end Jobs and the processor reads it to stop, so set the
35243526
same value for both; changing it only in the processor's configuration has no effect on
35253527
when the scheduler ends the Job.
3526-
A callback that ran before its acknowledgment was recorded can run again after the Job
3527-
times out, repeating its side effects. Callbacks should tolerate repeated execution.
35283528
version_added: 3.5.0
35293529
type: integer
35303530
example: ~

‎airflow-core/src/airflow/dag_processing/api_client.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ def __init__(
146146
self._completion_state: TerminalJobState | None = None
147147
self._completed = False
148148
self._pending_claims: dict[str, ProcessorWorkClaimBody] = {}
149+
self._claim_sequences: dict[str, int] = {}
149150
super().__init__(base_url=base_url, token="", **kwargs)
150151

151152
@property
@@ -411,6 +412,7 @@ def claim_work(
411412
kind,
412413
ProcessorWorkClaimBody(
413414
claim_id=uuid7(),
415+
claim_sequence=self._claim_sequences.get(kind, 0) + 1,
414416
bundle_names=bundle_names,
415417
limit=min(limit, 100),
416418
),
@@ -424,13 +426,12 @@ def claim_work(
424426
timeout=self._get_bounded_timeout(),
425427
)
426428
work = [ProcessorWorkItem.model_validate(item) for item in response.json()]
429+
self._claim_sequences[kind] = body.claim_sequence
427430
del self._pending_claims[kind]
428431
return work
429432

430-
def acknowledge_work(
431-
self, kind: Literal["callbacks", "priority"], work: ProcessorWorkItem, *, failed: bool = False
432-
) -> None:
433-
body = ProcessorWorkAckBody(claim_id=work.claim_id, state="failed" if failed else "success")
433+
def acknowledge_work(self, kind: Literal["callbacks", "priority"], work: ProcessorWorkItem) -> None:
434+
body = ProcessorWorkAckBody(claim_id=work.claim_id)
434435
self.request(
435436
"POST",
436437
f"jobs/{self._require_job_id()}/requested-work/{kind}/{work.id}/ack",

‎airflow-core/src/airflow/dag_processing/manager.py‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -341,8 +341,8 @@ class DagFileProcessorManager(LoggingMixin):
341341
_deferred_api_callbacks: list[CallbackRequest] = attrs.field(factory=list, init=False)
342342
_deferred_callback_retries: dict[int, DeferredCallback] = attrs.field(factory=dict, init=False)
343343
_callback_attempts: dict[UUID, list[ProcessorWorkItem]] = attrs.field(factory=dict, init=False)
344-
_priority_claims: dict[str, PriorityParseClaim] = attrs.field(factory=dict, init=False)
345-
_pending_work_acks: list[tuple[Literal["callbacks", "priority"], ProcessorWorkItem, bool]] = attrs.field(
344+
_priority_claims: dict[tuple[str, UUID], PriorityParseClaim] = attrs.field(factory=dict, init=False)
345+
_pending_work_acks: list[tuple[Literal["callbacks", "priority"], ProcessorWorkItem]] = attrs.field(
346346
factory=list, init=False
347347
)
348348
_num_run: int = attrs.field(default=0, init=False)
@@ -810,7 +810,7 @@ def claim_priority_files(self) -> list[DagFileInfo]:
810810
available = (
811811
self._parallelism
812812
- len(self._priority_claims)
813-
- sum(kind == "priority" for kind, _, _ in self._pending_work_acks)
813+
- sum(kind == "priority" for kind, _ in self._pending_work_acks)
814814
)
815815
if available <= 0:
816816
return []
@@ -822,7 +822,7 @@ def claim_priority_files(self) -> list[DagFileInfo]:
822822
return []
823823
files = []
824824
for item in work:
825-
self._priority_claims[item.id] = PriorityParseClaim(
825+
self._priority_claims[item.id, item.claim_id] = PriorityParseClaim(
826826
work=item, refresh_generation=self._bundle_refresh_generations[item.bundle_name] + 1
827827
)
828828
bundle = bundles[item.bundle_name]
@@ -897,7 +897,7 @@ def fetch_callbacks(self) -> list[CallbackRequest]:
897897
self.max_callbacks_per_loop
898898
- len(self._callback_claims)
899899
- sum(len(work) for work in self._callback_attempts.values())
900-
- sum(kind == "callbacks" for kind, _, _ in self._pending_work_acks)
900+
- sum(kind == "callbacks" for kind, _ in self._pending_work_acks)
901901
)
902902
if available <= 0:
903903
return list(self._deferred_api_callbacks)
@@ -922,9 +922,9 @@ def _acknowledge_requested_work(self, *, deadline: float | None = None) -> None:
922922
if deadline is None:
923923
deadline = time.monotonic() + _API_SEND_BUDGET_SECONDS
924924
while self._pending_work_acks and time.monotonic() < deadline:
925-
kind, item, failed = self._pending_work_acks[0]
925+
kind, item = self._pending_work_acks[0]
926926
try:
927-
self.api_client.acknowledge_work(kind, item, failed=failed)
927+
self.api_client.acknowledge_work(kind, item)
928928
except httpx2.HTTPStatusError as error:
929929
if error.response.status_code >= 500:
930930
self.log.warning("Unable to acknowledge requested work; retaining the claim")
@@ -944,13 +944,13 @@ def _acknowledge_requested_work(self, *, deadline: float | None = None) -> None:
944944

945945
def _finish_callback_attempt(self, proc: BaseDagFileProcessorProcess) -> None:
946946
for item in self._callback_attempts.pop(proc.id, []):
947-
self._pending_work_acks.append(("callbacks", item, proc._exit_code != 0))
947+
self._pending_work_acks.append(("callbacks", item))
948948

949-
def _finish_priority_claims(self, attempt_id: UUID, *, failed: bool) -> None:
949+
def _finish_priority_claims(self, attempt_id: UUID) -> None:
950950
"""Acknowledge only requests assigned to this import, never newer requests for its file."""
951951
for key, claim in list(self._priority_claims.items()):
952952
if claim.attempt_id == attempt_id:
953-
self._pending_work_acks.append(("priority", claim.work, failed))
953+
self._pending_work_acks.append(("priority", claim.work))
954954
del self._priority_claims[key]
955955

956956
def _get_priority_claims(self, file: DagFileInfo) -> list[PriorityParseClaim]:
@@ -1065,7 +1065,7 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None:
10651065
if get_claiming_importer(request.filepath, request.bundle_name) is not None:
10661066
self._log_dropped_lang_sdk_callback(request)
10671067
if item := self._callback_claims.pop(id(request), None):
1068-
self._pending_work_acks.append(("callbacks", item, True))
1068+
self._pending_work_acks.append(("callbacks", item))
10691069
self._remove_deferred_callback(request)
10701070
return
10711071
retry = self._deferred_callback_retries.get(id(request))
@@ -1098,7 +1098,7 @@ def _remove_deferred_callback(self, request: CallbackRequest) -> None:
10981098
self._deferred_callback_retries.pop(id(request), None)
10991099

11001100
def _defer_claimed_callback(self, request: CallbackRequest) -> None:
1101-
"""Retry preparing a claimed callback's bundle with backoff, then acknowledge it as failed."""
1101+
"""Retry preparing a claimed callback's bundle with backoff, then drop the unpreparable request."""
11021102
now = time.monotonic()
11031103
timeout = conf.getint("dag_processor", "job_heartbeat_timeout")
11041104
retry = self._deferred_callback_retries.setdefault(id(request), DeferredCallback(now + timeout))
@@ -1111,7 +1111,7 @@ def _defer_claimed_callback(self, request: CallbackRequest) -> None:
11111111
request.bundle_version,
11121112
timeout,
11131113
)
1114-
self._pending_work_acks.append(("callbacks", self._callback_claims.pop(id(request)), True))
1114+
self._pending_work_acks.append(("callbacks", self._callback_claims.pop(id(request))))
11151115
self._remove_deferred_callback(request)
11161116
return
11171117
retry.attempts += 1
@@ -1214,7 +1214,7 @@ def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set
12141214
and item.relative_fileloc not in body.files
12151215
and not is_source_discovery_failed(item.relative_fileloc, failed_sources)
12161216
):
1217-
self._pending_work_acks.append(("priority", item, True))
1217+
self._pending_work_acks.append(("priority", item))
12181218
del self._priority_claims[key]
12191219

12201220
def _discard_unusable_bundle_source(self, bundle_name: str) -> None:
@@ -1232,7 +1232,7 @@ def _discard_unusable_bundle_source(self, bundle_name: str) -> None:
12321232
)
12331233
for key, claim in list(self._priority_claims.items()):
12341234
if claim.work.bundle_name == bundle_name and claim.attempt_id not in active_attempts:
1235-
self._pending_work_acks.append(("priority", claim.work, True))
1235+
self._pending_work_acks.append(("priority", claim.work))
12361236
del self._priority_claims[key]
12371237

12381238
def _invalidate_bundle_parse_source(self, bundle_name: str) -> None:
@@ -1882,11 +1882,11 @@ def handle_parsing_result(
18821882
file.bundle_name,
18831883
)
18841884
self._record_failed_publication(file, next_stat)
1885-
self._finish_priority_claims(proc.id, failed=True)
1885+
self._finish_priority_claims(proc.id)
18861886
return
18871887

18881888
if proc.parsing_result is None and not is_callback_only:
1889-
self._finish_priority_claims(proc.id, failed=True)
1889+
self._finish_priority_claims(proc.id)
18901890
self._file_stats[file] = next_stat
18911891

18921892
def _record_failed_publication(self, file: DagFileInfo, stat: DagFileStat) -> None:
@@ -1964,15 +1964,14 @@ def _publish_pending_results(self, *, deadline: float | None = None) -> None:
19641964
continue
19651965
self.log.exception("Failed to publish parsing result for %s", file)
19661966
self._record_failed_publication(file, pending.stat)
1967-
self._finish_priority_claims(pending.body.attempt_id, failed=True)
19681967
else:
19691968
self._file_stats[file] = pending.stat
1970-
self._finish_priority_claims(pending.body.attempt_id, failed=False)
1969+
self._finish_priority_claims(pending.body.attempt_id)
19711970
del self._pending_publications[file]
19721971

19731972
def _build_rejection_report(self, pending: PendingDagPublication, reason: str) -> PendingDagPublication:
19741973
"""Replace a result that cannot be published with an import error users can see."""
1975-
self._finish_priority_claims(pending.body.attempt_id, failed=True)
1974+
self._finish_priority_claims(pending.body.attempt_id)
19761975
self._dispatch_sequence += 1
19771976
if len(reason) > _MAX_REJECTION_REASON_LENGTH:
19781977
reason = reason[:_MAX_REJECTION_REASON_LENGTH] + "..."
@@ -2505,7 +2504,7 @@ def _kill_timed_out_processors(self):
25052504
)
25062505
self._file_stats[file] = stat
25072506
if not processor.had_callbacks:
2508-
self._finish_priority_claims(processor.id, failed=True)
2507+
self._finish_priority_claims(processor.id)
25092508

25102509
# Clean up `self._processors` after iterating over it
25112510
for proc in processors_to_remove:

‎airflow-core/tests/unit/dag_processing/test_api_client.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from airflow.api_fastapi.execution_api.datamodels.dag_parsing import (
3131
DagBundleInventoryBody,
3232
DagParseResultBody,
33+
ProcessorWorkItem,
3334
)
3435
from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody, JobState, TerminalJobState
3536
from airflow.dag_processing import api_client
@@ -1079,3 +1080,59 @@ def test_inventory_waits_for_a_slow_write_but_connects_quickly(make_client):
10791080
timeout = requests[-1].extensions["timeout"]
10801081
assert timeout["connect"] <= 1
10811082
assert timeout["read"] == client.timeout.read > 1
1083+
1084+
1085+
@pytest.mark.parametrize("kind", ["callbacks", "priority"])
1086+
@pytest.mark.parametrize(
1087+
"failure",
1088+
[
1089+
httpx2.ReadError("Lost claim response"),
1090+
httpx2.Response(503),
1091+
httpx2.Response(200, json=[{"id": "incomplete"}]),
1092+
],
1093+
ids=["lost_response", "server_error", "invalid_response"],
1094+
)
1095+
def test_claim_sequence_survives_retries_and_renewal_then_advances_per_kind(make_client, kind, failure):
1096+
client, requests = make_client(
1097+
make_registration_response(),
1098+
failure,
1099+
make_registration_response(),
1100+
httpx2.Response(200, json=[]),
1101+
httpx2.Response(200, json=[]),
1102+
httpx2.Response(200, json=[]),
1103+
)
1104+
client.register_job()
1105+
with pytest.raises((httpx2.HTTPError, ValueError)):
1106+
client.claim_work(kind, ["bundle-a"], 2)
1107+
1108+
client.register_job()
1109+
assert client.claim_work(kind, ["bundle-b"], 5) == []
1110+
other_kind = "callbacks" if kind == "priority" else "priority"
1111+
assert client.claim_work(other_kind, ["bundle-a"], 2) == []
1112+
assert client.claim_work(kind, ["bundle-b"], 5) == []
1113+
1114+
assert requests[1].content == requests[3].content
1115+
original = json.loads(requests[1].content)
1116+
assert original["claim_sequence"] == 1
1117+
assert original["bundle_names"] == ["bundle-a"]
1118+
assert original["limit"] == 2
1119+
assert json.loads(requests[4].content)["claim_sequence"] == 1
1120+
subsequent = json.loads(requests[5].content)
1121+
assert subsequent["claim_sequence"] == 2
1122+
assert subsequent["claim_id"] != original["claim_id"]
1123+
assert subsequent["bundle_names"] == ["bundle-b"]
1124+
assert subsequent["limit"] == 5
1125+
1126+
1127+
@pytest.mark.parametrize("kind", ["callbacks", "priority"])
1128+
def test_work_acknowledgment_sends_only_the_claim_identity(make_client, kind):
1129+
client, requests = make_client(make_registration_response(), httpx2.Response(204))
1130+
client.register_job()
1131+
work = ProcessorWorkItem(
1132+
id=str(uuid7()), claim_id=uuid7(), bundle_name="bundle-a", relative_fileloc="dag.py"
1133+
)
1134+
1135+
client.acknowledge_work(kind, work)
1136+
1137+
assert requests[-1].url.path == f"/execution/jobs/1/requested-work/{kind}/{work.id}/ack"
1138+
assert json.loads(requests[-1].content) == {"claim_id": str(work.claim_id)}

0 commit comments

Comments
 (0)