kaxil commented on code in PR #74794:
URL: https://github.com/apache/airflow/pull/74794#discussion_r4501826311
##########
airflow-core/src/airflow/cli/commands/dag_processor_command.py:
##########
@@ -72,13 +86,56 @@ def _create_dag_processor_job_runner(args: Any) ->
DagProcessorJobRunner:
@enable_memray_trace(component=MemrayTraceComponents.dag_processor)
def _run_dag_processor_job(job_runner: DagProcessorJobRunner) -> None:
- run_job(job=job_runner.job, execute_callable=job_runner._execute)
+ if not (token_file := conf.get("dag_processor",
"execution_api_token_file", fallback=None)):
+ run_job(job=job_runner.job, execute_callable=job_runner._execute)
+ return
+
+ # Create the connection pool after the daemon or hot-reload fork.
Review Comment:
The import doesn't create the pool, the `DagProcessorAPIClient(...)` call
below does, so this comment gives the wrong reason. Is the reason import weight
(the client pulls in the Execution API versions)? The local imports at
`manager.py:965` and `api_requested_work.py:99` have no reason comment at all.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3452,6 +3478,8 @@ dag_processor:
DAG is loaded. The absolute maximum that this could take is
``[dag_processor] dag_file_processor_timeout``, but when you have a
long timeout configured,
it results in a significant delay in the deactivation of stale dags.
+ When ``[dag_processor] execution_api_token_file`` is set, Dags missing
from an accepted parse
+ are marked stale at once; the threshold applies only to the Dags of a
file whose import failed.
Review Comment:
In API mode the processor skips `_scan_stale_dags`, and the API server reads
this value when it accepts a parse result (`routes/dag_parsing.py:229`). Could
this say it is set on the API server, the way `health_check_threshold` names
who reads it?
##########
airflow-core/src/airflow/dag_processing/processor.py:
##########
@@ -183,6 +193,20 @@ class DagFileParsingResult(BaseModel):
]
[email protected]()
+class _DagParseCommsDecoder(CommsDecoder[ToDagProcessor, ToManager]):
+ api_request_failed: bool = attrs.field(default=False, init=False)
+ """Whether the Dag processor could not answer a request of this process
through the Execution API."""
+
+ def send(self, msg: ToManager) -> ToDagProcessor | None:
Review Comment:
`CommsDecoder.asend` calls `_from_frame` directly, so an `API_SERVER_ERROR`
from an import-time `asyncio.run(Variable.aget(...))` never goes through this
override and `api_request_failed` stays false. The child then exits 0 and the
callback is acked without running, which is the case exit code 75 is meant to
retry. Would overriding `_from_frame` cover both paths? The error-type check
itself isn't pinned either: replacing it with `if True:` keeps every processor
test passing, so a decoder test that feeds a `GENERIC_ERROR` and an
`API_SERVER_ERROR` frame through `send` and `asend` would cover it.
##########
airflow-core/src/airflow/dag_processing/api_client.py:
##########
@@ -438,6 +443,44 @@ def get_bundles(self, *, retry: bool = False) ->
list[DagBundleStateResponse]:
)
return [DagBundleStateResponse.model_validate(item) for item in
response.json()]
+ def claim_work(
+ self, kind: Literal["callbacks", "priority"], bundle_names: list[str],
limit: int
+ ) -> list[ProcessorWorkItem]:
+ if not bundle_names:
+ return []
+ body = self._pending_claims.setdefault(
+ kind,
+ ProcessorWorkClaimBody(
+ claim_id=uuid7(),
+ claim_sequence=self._claim_sequences.get(kind, 0) + 1,
+ bundle_names=bundle_names,
+ limit=min(limit, 100),
+ ),
+ )
+ response = self.request(
+ "POST",
+ f"jobs/{self._require_job_id()}/requested-work/{kind}/claim",
+ content=body.model_dump_json(),
+ headers=_JOB_API_HEADERS,
+ retry=False,
+ timeout=self._get_bounded_timeout(),
Review Comment:
Claims and acks cap every timeout phase at 1s, but the claim route assigns
the rows and commits before it responds. If the API takes 1 to 5 seconds to
answer claims (a loaded API server or database), every claim times out after
the rows are already owned, the replay of the same `claim_id` is just as slow,
and no callback or reparse is ever delivered while heartbeats succeed and
`--check-ready` stays green. Could `claim_work` and `acknowledge_work` use
`_get_response_timeout()` like publication does?
##########
airflow-core/docs/security/jwt_token_authentication.rst:
##########
@@ -557,14 +587,168 @@ parsing subprocess can without this client.
``Variable.set`` upserts on the key
multi-team mode a write can overwrite another team's Variable with the same
key, as it can with task
tokens.
-The client attempts renewal at 80% of token lifetime. Transient early-renewal
failures leave
-the existing token usable until expiry. A retired registration signals that
its owner must
-restart. Completion retains its original identity and outcome across retries.
-
-File and attempt identity are carried in the token but are not currently used
for access
-checks or request attribution. Archives retain one file identity under the
current processor
-model. These credentials do not isolate hostile code from other credentials
accessible on
-the processor host; the process and deployment limitations above still apply.
+For each ``dag_parse`` token, the manager chooses the bundle, bundle-relative
file location,
+and parsing-attempt UUID.
+
+The manager publishes completed imports through ``POST
/jobs/{job_id}/parse-results``
+using its Job credential. Parsing tokens cannot publish. A request contains
one file or
+container's serialized Dags, diagnostics, source text (or an explicit
unavailable marker),
+and the bundle version captured at dispatch. The API never reads the
processor's source
+paths. It checks bundle access, existing Dag ownership and team-scoped plugin
classes
+before deserializing results and rechecks ownership on the rows used for
persistence.
+A Dag can move between files within its bundle. Recovery from a removed bundle
is
+limited to stale legacy rows with neither a relative file location nor a Dag
version.
+
+Each request carries an attempt UUID and a dispatch sequence within its Job. A
publication
+receipt and the metadata writes commit together. Replaying the latest
identical request returns
+the saved receipt; conflicting payloads and superseded sequences are rejected.
Publication
+receipts retain one entry per Job and source and are removed with Job cleanup.
There is no total
+ordering between different Jobs publishing a source without versioning; an
already-accepted
+request's retry does not overwrite another Job's result. Listener side effects
are not
+transactional and may repeat after a failed transaction.
+
+Publication and inventory request bodies are limited to 16 MiB, and the API
refuses a
+larger body before reading the rest of it. A rejected publication does not
fall back
+to direct database writes, and does not prevent other files from being
processed. Transport
+failures and server errors are retried; exhausting those retries preserves the
last accepted
+result. When the API refuses a result for its content, or the result is over
the limit, the
+processor publishes an import error for the file instead, so the problem is
visible.
+Dags whose definition path matches an import error are marked stale until a
later parse
+publishes them again. Dags missing from an accepted parse are marked stale at
once: the
+publication is atomic, so there is no gap between file and Dag parse times to
wait out. For a
+container-level error, missing member Dags are marked stale only when their
last successful
+parse is older than ``[dag_processor] stale_dag_threshold``, using the
server's clock. Invalid ``access_control`` produces an import error with both
+current and supported older FAB releases. Permission backend failures remain
retryable
+server errors. Empty imports still publish a receipt and clear the file's
earlier import
+errors. A bundle refresh, including one forced by a parse request, waits for
that bundle's
+active imports instead of discarding them; only initializing a bundle again
invalidates and
+requeues older in-flight results from that manager.
+
+Inventories report discovered definitions and sources whose discovery failed
separately.
+Failed sources preserve their existing definitions and diagnostics, including
archive members;
+other sources still undergo deletion reconciliation. Each inventory carries
the previously
+observed server revision; conflicting refreshes must
+rediscover before retrying. Imports carry the accepted revision, so an older
snapshot cannot
+publish after a changed inventory. Only accepted publications establish
missing Dags within
+a file; rejected requests and incomplete discovery do not establish absence.
Every
+``[scheduler] parsing_cleanup_interval``, the scheduler marks the Dags of
inactive bundles stale,
+removes stale warnings, and ends expired processor Jobs, in bounded batches
and in
+both modes.
+
+Callback and priority requests remain stored while claimed. A lost claim
response or
+acknowledgment can be retried. Claims use increasing sequences per Job and
work type; the
+client waits for a claim response before sending the next sequence for that
type.
+Retrying an empty claim remains empty, even if work arrives
+afterward, and retrying an older claim cannot consume newly requested work.
Requests from a
+retired Job become available to another processor. A callback acknowledgment
confirms
+delivery, not successful execution of user code. As with other at-least-once
delivery, a crash
+after a callback's side effect and before acknowledgment can repeat that side
effect. Before the
+processor's Job ends, a graceful shutdown acknowledges callbacks whose
processes have already exited
+cleanly; callbacks that are still running or that the shutdown signal stopped
are delivered again.
+When a callback's bundle
+version cannot be prepared, the processor retries with backoff for up to
+``[dag_processor] job_heartbeat_timeout``
+and then acknowledges the request without running the callback. The same retry
applies after an
+Execution API request fails, but only when none of the callbacks processed
together for a file
+found its Dag or task, so none of them ran. When one of them did, a callback
whose Dag or task did
+not load is acknowledged without running. Priority requests are assigned
+to an import started after their requested refresh. Only that import's
publication or terminal
+failure acknowledges the request; an older in-flight import cannot consume it.
Discovery after
+the requested refresh can also acknowledge a request if it confirms the
definition is absent.
+A new request for a file already being imported remains pending for a
subsequent import,
+even if the earlier import has not yet been acknowledged.
+
+The manager retains a bounded queue of pending publications. Each loop it
sends queued
+publications and acknowledgments for up to a second, stopping at the first
failure that may
+be transient. Connecting is limited to a second, but a publication or
inventory waits up to
+``[workers] execution_api_timeout`` for its response, so a slow write can
still be
+acknowledged. Retries keep the same payload and identity while
+heartbeats and subprocess supervision continue between attempts. Pending
results
+count against parsing capacity. Failed publication attempts throttle reparsing
but
+do not advance the timestamps used to detect stale Dags.
+
+The returned ``dag_parse`` token binds these identifiers to the processor's
session and
+Job. It authorizes parse-time operations within that bundle and has no
Job-management or
+token-exchange access. A bundle header cannot override its signed bundle. The
token expires
+no later than the management credential; ending or replacing the Job
invalidates it.
+The supervisor caches it for the attempt and attempts renewal at 80% of its
lifetime.
+Transient early-renewal failures leave the existing token usable until expiry.
+
+Parse-time requests use ``[workers] execution_api_timeout`` (five seconds by
default),
+unless the request specifies a timeout. Required parsing-token exchanges and
expired
+Job-token renewals use the same timeout. Each makes one attempt without retry
sleeps in
+the manager loop; optional early token renewal caps each HTTP timeout phase at
one second.
+
+Heartbeats use ``[scheduler] job_heartbeat_sec`` and are also sent between
bundle
+operations and before each callback's bundle is prepared. Transient API
failures are
+retried on the next heartbeat interval. Once a Job has gone
+``[dag_processor] job_heartbeat_timeout`` without a recorded heartbeat, the
scheduler's next cleanup (every
+``[scheduler] parsing_cleanup_interval``) ends it and releases its claimed
work, and a processor
+that could not reach the API for that long stops. Set the timeout above the
longest single
+bundle initialization, refresh, or listing, and set it for the scheduler as
well as the
+processor: the scheduler's value decides when the Job ends. The processor does
not update its
+local health file during one of those operations either, so
``health_check_threshold`` and the
+liveness probe's failure budget must also cover it. Likewise, the API server
reads
+``[dag_processor] health_check_threshold`` to decide whether a session's
previous Job is still
+alive, so keep it the same for the API server and the processor. A restart
request, replaced Job, or retired
+registration also stops parsing and terminates its children. The deployment's
process
+supervisor must restart the command, which creates a new registration. After a
crash,
+registration retries ``job_running`` with backoff for up to the health-check
threshold
+plus one heartbeat interval, retaining the same registration ID while waiting
for the old
+Job to expire.
+Failed runs attempt to record failure without masking the original error;
completion
+failures after a successful run are reported to the caller.
+
+Use ``airflow dag-processor --check-health`` for local liveness and
``--check-ready`` for
+local liveness plus a recent API heartbeat. Both read ``[dag_processor]
health_check_file``
+without querying the database. Give each processor its own writable health
file. A paused
+manager fails liveness even while the API remains healthy; a transient API
outage affects
+readiness without immediately failing liveness.
Review Comment:
The `loop` field is written once per manager iteration, but
`_service_processor_sockets` runs twice per iteration and answers each child's
API request in turn, each waiting up to `execution_api_timeout` (5s) with no
retry. During an outage with 4 parsing processes whose Dag files read a
Variable at import, one iteration can take 2 x 4 x 5s, past the default
`health_check_threshold` of 30s, so `--check-health` fails as well. Could the
manager refresh the health file between socket handlers, or could this sentence
give that bound?
##########
airflow-core/tests/unit/dag_processing/test_manager.py:
##########
@@ -7601,14 +8050,374 @@ def
test_oversized_inventory_byte_boundary_preserves_auth_failures(auth_status):
manager._publish_bundle_inventory("testing", {})
[email protected]("status_code", [403, 404, 409, 503])
+def test_requested_work_ack_retains_only_retryable_failures(status_code):
Review Comment:
This only exercises `ExecutionAPIWorkSource.flush`, with no manager code on
the path, and `test_flush_discards_permanent_refusal_and_continues` in
`test_api_requested_work.py` already covers the refusals. Could it move there
as a 404 case? `_make_work` and `_make_status_error` are also copied between
these test modules, and `requested_work_helpers.py` looks like the place for
them.
##########
airflow-core/src/airflow/dag_processing/api_requested_work.py:
##########
@@ -0,0 +1,378 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Claim and acknowledgment lifecycle for Dag processor work from the
Execution API."""
+
+from __future__ import annotations
+
+import time
+from collections import Counter
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import attrs
+import httpx2
+
+from airflow.callbacks.callback_requests import CallbackRequest,
DagCallbackRequest
+from airflow.configuration import conf
+from airflow.dag_processing.collection import is_source_discovery_failed
+from airflow.dag_processing.requested_work import RequestedWorkSource
+from airflow.utils.file import find_enclosing_file
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+if TYPE_CHECKING:
+ from collections.abc import Collection
+ from uuid import UUID
+
+ from airflow.api_fastapi.execution_api.datamodels.dag_parsing import
ProcessorWorkItem
+ from airflow.dag_processing.api_client import DagProcessorAPIClient
+ from airflow.dag_processing.bundles.base import BaseDagBundle
+ from airflow.dag_processing.manager import DagFileInfo
+
+_API_SEND_BUDGET_SECONDS = 1.0
+_MAX_CALLBACK_RETRY_DELAY = 60.0
+
+
+def _get_status_code(error: httpx2.HTTPError) -> int | None:
+ return error.response.status_code if isinstance(error,
httpx2.HTTPStatusError) else None
+
+
[email protected]
+class DeferredCallback:
+ """A claimed callback that could not be prepared or run, retried until it
expires."""
+
+ expires_at: float
+ attempts: int = 0
+ next_attempt_time: float = 0.0
+
+
[email protected]
+class PriorityParseClaim:
+ """A priority request served only after an inventory captured following
the claim."""
+
+ work: ProcessorWorkItem
+ inventory_sequence: int
+ attempt_id: UUID | None = None
+
+
[email protected]
+class ExecutionAPIWorkSource(RequestedWorkSource, LoggingMixin):
+ """Retain claimed work until its terminal outcome can be acknowledged by
the API."""
+
+ client: DagProcessorAPIClient
+ _callback_claims: dict[int, ProcessorWorkItem] = attrs.field(factory=dict,
init=False)
+ _deferred_api_callbacks: list[CallbackRequest] = attrs.field(factory=list,
init=False)
+ _deferred_callback_retries: dict[int, DeferredCallback] =
attrs.field(factory=dict, init=False)
+ _interrupted_callback_retries: dict[str, DeferredCallback] =
attrs.field(factory=dict, init=False)
+ """Retries of callbacks an API failure interrupted, by work ID, so later
attempts share one expiry."""
+ _callback_attempts: dict[UUID, list[ProcessorWorkItem]] =
attrs.field(factory=dict, init=False)
+ _priority_claims: dict[tuple[str, UUID], PriorityParseClaim] =
attrs.field(factory=dict, init=False)
+ _pending_work_acks: list[tuple[Literal["callbacks", "priority"],
ProcessorWorkItem]] = attrs.field(
+ factory=list, init=False
+ )
+
+ requires_acknowledgment = True
+
+ def claim_priority_files(
+ self,
+ bundles: Collection[BaseDagBundle],
+ *,
+ capacity: int,
+ unavailable_bundles: Collection[str],
+ active_attempts: Collection[UUID],
+ inventory_sequence: int,
+ ) -> list[DagFileInfo]:
+ """Claim priority files within the capacity left by active and
retained work."""
+ from airflow.dag_processing.manager import DagFileInfo
+
+ bundles_by_name = {bundle.name: bundle for bundle in bundles}
+ unavailable = set(unavailable_bundles)
+ available = (
+ capacity
+ - sum(
+ claim.attempt_id in active_attempts or claim.work.bundle_name
not in unavailable
+ for claim in self._priority_claims.values()
+ )
+ - sum(kind == "priority" for kind, _ in self._pending_work_acks)
+ )
+ if available <= 0:
+ return []
+ # Retain failed-bundle claims for recovery without letting them
consume healthy bundles' capacity.
+ # Excluding those bundles from new claims also bounds the suspended
backlog.
+ eligible_bundles = [name for name in bundles_by_name if name not in
unavailable]
+ try:
+ work = self.client.claim_work("priority", eligible_bundles,
available)
+ except httpx2.HTTPError as error:
+ self.log.warning(
+ "Unable to claim priority parses; retrying on a later loop",
+ error_type=type(error).__name__,
+ status_code=_get_status_code(error),
+ )
+ return []
+ files = []
+ for item in work:
+ self._priority_claims[item.id, item.claim_id] = PriorityParseClaim(
+ work=item, inventory_sequence=inventory_sequence
+ )
+ bundle = bundles_by_name[item.bundle_name]
+ try:
+ path = find_enclosing_file(bundle.path / item.relative_fileloc)
+ except OSError as error:
+ # The next inventory acknowledges the claim if the bundle
listing cannot see the file either.
+ self.log.warning(
+ "Unable to resolve priority parse %s in bundle %s: %s",
+ item.relative_fileloc,
+ bundle.name,
+ error,
+ )
+ path = None
+ files.append(
+ DagFileInfo(
+ rel_path=path.relative_to(bundle.path) if path else
Path(item.relative_fileloc),
+ bundle_name=bundle.name,
+ bundle_path=bundle.path,
+ )
+ )
+ return files
+
+ def fetch_callbacks(self, bundles: Collection[BaseDagBundle], *, limit:
int) -> list[CallbackRequest]:
+ """Claim callbacks for ready bundles and include requests awaiting
bundle preparation."""
+ # Callbacks waiting for a retry take no capacity from other bundles.
Their bundle stops claiming
+ # only once it holds a full batch, so one unavailable version does not
block the rest.
+ deferred = Counter(
+ request.bundle_name
+ for request in self._deferred_api_callbacks
+ if id(request) in self._deferred_callback_retries
+ )
+ available = (
+ limit
+ - (len(self._callback_claims) - deferred.total())
+ - sum(len(work) for work in self._callback_attempts.values())
+ - sum(kind == "callbacks" for kind, _ in self._pending_work_acks)
+ )
+ if available <= 0:
+ return list(self._deferred_api_callbacks)
+ eligible_bundles = [
+ b.name
+ for b in bundles
+ if (not b.supports_versioning or b.is_initialized) and
deferred[b.name] < limit
+ ]
+ try:
+ work = self.client.claim_work("callbacks", eligible_bundles,
available)
+ except httpx2.HTTPError as error:
+ self.log.warning(
+ "Unable to claim callbacks; retrying on a later loop",
+ error_type=type(error).__name__,
+ status_code=_get_status_code(error),
+ )
+ return list(self._deferred_api_callbacks)
+ for item in work:
+ self._retain_callback_claim(item)
+ return list(self._deferred_api_callbacks)
+
+ def _retain_callback_claim(self, item: ProcessorWorkItem) ->
CallbackRequest:
+ if item.callback is None:
+ raise ValueError("Callback claim returned no request")
+ # Each claim needs its own identity, even when a client reuses an
equal callback object.
+ request = item.callback.model_copy(deep=True)
+ self._callback_claims[id(request)] = item
+ self._deferred_api_callbacks.append(request)
+ return request
+
+ def callbacks_started(self, attempt_id: UUID, callbacks:
Collection[CallbackRequest]) -> None:
+ """Associate a callback process with the claims it will execute."""
+ claims = [
+ self._callback_claims.pop(id(request))
+ for request in callbacks
+ if id(request) in self._callback_claims
+ ]
+ if claims:
+ self._callback_attempts[attempt_id] = claims
+
+ def parse_started(self, file: DagFileInfo, attempt_id: UUID) -> None:
+ """Associate an ordinary parse with the priority requests it will
satisfy."""
+ for claim in self._get_priority_claims(file):
+ claim.attempt_id = attempt_id
+
+ def attempt_finished(self, attempt_id: UUID, *, had_callbacks: bool) ->
None:
+ """Queue acknowledgments only for claims assigned to this terminal
attempt."""
+ if had_callbacks:
+ for item in self._callback_attempts.pop(attempt_id, []):
+ self._interrupted_callback_retries.pop(item.id, None)
+ self._pending_work_acks.append(("callbacks", item))
+ return
+ for key, claim in list(self._priority_claims.items()):
+ if claim.attempt_id == attempt_id:
+ self._pending_work_acks.append(("priority", claim.work))
+ del self._priority_claims[key]
+
+ def callbacks_interrupted(self, attempt_id: UUID) -> None:
+ """Retry an attempt's callbacks with backoff until the expiry set by
their first failure."""
+ for item in self._callback_attempts.pop(attempt_id, []):
+ request = self._retain_callback_claim(item)
+ if retry := self._interrupted_callback_retries.get(item.id):
+ self._deferred_callback_retries[id(request)] = retry
+ self.callback_deferred(request)
+ if retry := self._deferred_callback_retries.get(id(request)):
+ self._interrupted_callback_retries[item.id] = retry
+
+ def flush(self, *, deadline: float | None = None) -> None:
+ """Send queued acknowledgments in order until the deadline, keeping
claims after a failure."""
+ # Runtime API-client imports stay local so managers using the
in-process API do not load it.
+ from airflow.dag_processing.api_client import is_transient_api_error
+
+ if deadline is None:
+ deadline = time.monotonic() + _API_SEND_BUDGET_SECONDS
+ while self._pending_work_acks and time.monotonic() < deadline:
+ kind, item = self._pending_work_acks[0]
+ try:
+ self.client.acknowledge_work(kind, item)
+ except httpx2.HTTPError as error:
+ status_code = (
+ error.response.status_code if isinstance(error,
httpx2.HTTPStatusError) else None
Review Comment:
This is `_get_status_code(error)` from the top of this module (the runner's
heartbeat warning and `manager.py:1085` inline it too). Likewise,
`callback_deferred` rebuilds the target string that
`DagFileProcessorManager._get_callback_target` already builds.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
