kaxil commented on code in PR #74794:
URL: https://github.com/apache/airflow/pull/74794#discussion_r4336375725
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -787,9 +818,36 @@ def claim_priority_files(self) -> list[DagFileInfo]:
Default implementation reads from the metadata DB; override to source
requests from an API.
"""
- if self.api_client is not None:
+ if self.api_client is None:
+ return self._claim_priority_files()
+ available = (
Review Comment:
This budget is shared by every bundle, but a claim for a bundle whose
`initialize()` or `refresh()` keeps raising is never released. The bundle stays
in `_force_refresh_bundles`, its files wait in `_files_awaiting_source`, and
none of the release paths (`_finish_priority_claims`, the inventory publish,
`_discard_unusable_bundle_source`) run. With the default `parsing_processes` of
2, two reparse requests against a bundle with broken Git credentials stop
priority parsing for every other bundle. DB mode doesn't have this because
`_claim_priority_files` deletes each request as it queues the file. Could the
claim leave out bundles whose last refresh failed (the way `fetch_callbacks`
leaves out uninitialized ones), or only count claims whose bundle has a
dispatch source?
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -2071,13 +2303,18 @@ def _start_new_processes(self):
self._readmit_files_with_source()
bundle_to_team = self._get_team_names({file.bundle_name for file in
self._file_queue})
- while (
- self._parallelism > len(self._processors) +
len(self._pending_publications) and self._file_queue
- ):
+ # Requested work skipped in this pass goes back to the front once the
pass is over.
+ requested_later: dict[DagFileInfo, None] = {}
+ for _ in range(len(self._file_queue)):
+ if len(self._processors) + len(self._pending_publications) >=
self._parallelism:
+ break
file, _ = self._file_queue.popitem(last=False)
+ priority_claims = self._get_priority_claims(file)
# Stop creating duplicate processor i.e. processor with the same
filepath
if file in self._processors or file in self._pending_publications:
self._priority_files.discard(file)
+ if priority_claims:
Review Comment:
Callback files don't get this treatment. If a second callback for the same
file is queued while the first callback-only processor is still running, the
file is popped here and never put back: the request stays in
`_callback_to_execute[file]`, and for a versioned bundle `prepare_file_queue`
never lists a `DagFileInfo` carrying that `bundle_version`, so nothing
re-queues it. In API mode the claim also stays in `_callback_claims`
unacknowledged, so it keeps counting against `max_callbacks_per_loop` in
`fetch_callbacks` until the processor restarts. Two tasks of one run failing a
few seconds apart is enough to hit it. Would `if priority_claims or file in
self._callback_to_execute:` work here?
##########
airflow-core/src/airflow/jobs/dag_processor_job_runner.py:
##########
@@ -49,12 +61,92 @@ def __init__(
):
super().__init__(job)
self.processor = processor
+ self._last_api_heartbeat = 0.0
+ self._next_api_heartbeat = 0.0
self.processor.heartbeat = lambda: perform_heartbeat(
job=self.job,
heartbeat_callback=self.heartbeat_callback,
only_if_necessary=True,
)
+ def run_with_api(self, client: DagProcessorAPIClient) -> int | None:
+ """Let the API own the Job row while retaining the normal processor
cleanup and listeners."""
+ # The client loads versioned routes which import this runner.
+ from airflow.api_fastapi.execution_api.datamodels.job import
TerminalJobState
+
+ self.processor.api_client = client
+ self.job.id = self._register_api_job(client)
+ self.job.state = JobState.RUNNING
+ self.job.start_date = self.job.latest_heartbeat = timezone.utcnow()
+ self._last_api_heartbeat = monotonic()
+ self._next_api_heartbeat = self._last_api_heartbeat +
self.job.heartrate
+ self.processor.heartbeat = lambda: self._heartbeat_api(client)
+ stats.incr("job_start", 1, 1)
+ failed = True
+ try:
+ result = execute_job(self.job, execute_callable=self._execute)
+ failed = False
+ return result
+ finally:
+ self.job.end_date = timezone.utcnow()
+ self.job.state = JobState.FAILED if failed else JobState.SUCCESS
+ try:
+ get_listener_manager().hook.before_stopping(component=self.job)
+ except Exception:
+ self.log.exception("Error calling Dag processor shutdown
listener")
+ try:
+ client.complete_job(TerminalJobState(self.job.state))
+ except Exception:
+ if not failed:
+ raise
+ self.log.exception("Unable to record the failed Dag processor
Job")
+ finally:
+ stats.incr("job_end", 1, 1)
+
+ def _register_api_job(self, client: DagProcessorAPIClient) -> int:
+ from airflow.dag_processing.api_client import
DagProcessorJobAlreadyRunning
+
+ deadline = monotonic() + conf.getint("dag_processor",
"health_check_threshold") + self.job.heartrate
+ delay = 1.0
+ while True:
+ try:
+ return client.register_job()
+ except DagProcessorJobAlreadyRunning:
+ remaining = deadline - monotonic()
+ if remaining <= 0:
+ raise
+ self.log.info("Waiting for the previous Dag processor Job
heartbeat to expire")
+ sleep(min(delay, remaining))
+ delay = min(delay * 2, 5.0)
+
+ def _heartbeat_api(self, client: DagProcessorAPIClient) -> None:
+ from airflow.dag_processing.api_client import
DagProcessorRegistrationRetired
+
+ write_processor_health(self._last_api_heartbeat)
+ if client.restart_required:
+ raise DagProcessorRegistrationRetired("The Dag processor
registration requires a restart")
+ now = monotonic()
+ if now < self._next_api_heartbeat:
+ return
+ self._next_api_heartbeat = now + self.job.heartrate
+ try:
+ state = client.heartbeat()
+ except httpx2.HTTPError as error:
+ if isinstance(error, httpx2.HTTPStatusError) and
error.response.status_code < 500:
+ raise
+ stats.incr("dag_processor_heartbeat_failure", 1, 1)
Review Comment:
In DB mode a failed heartbeat increments `job_heartbeat_failure` (from
`Job.heartbeat`), and `run_with_api` keeps the DB-mode `job_start`/`job_end`
names. Should this be `job_heartbeat_failure` as well, so existing alerts keep
firing after a processor moves to API mode? If the new name is intentional, the
API-mode docs could say it replaces the old one.
##########
airflow-core/src/airflow/jobs/dag_processor_job_runner.py:
##########
@@ -49,12 +61,92 @@ def __init__(
):
super().__init__(job)
self.processor = processor
+ self._last_api_heartbeat = 0.0
+ self._next_api_heartbeat = 0.0
self.processor.heartbeat = lambda: perform_heartbeat(
job=self.job,
heartbeat_callback=self.heartbeat_callback,
only_if_necessary=True,
)
+ def run_with_api(self, client: DagProcessorAPIClient) -> int | None:
+ """Let the API own the Job row while retaining the normal processor
cleanup and listeners."""
+ # The client loads versioned routes which import this runner.
+ from airflow.api_fastapi.execution_api.datamodels.job import
TerminalJobState
+
+ self.processor.api_client = client
+ self.job.id = self._register_api_job(client)
+ self.job.state = JobState.RUNNING
+ self.job.start_date = self.job.latest_heartbeat = timezone.utcnow()
+ self._last_api_heartbeat = monotonic()
+ self._next_api_heartbeat = self._last_api_heartbeat +
self.job.heartrate
+ self.processor.heartbeat = lambda: self._heartbeat_api(client)
+ stats.incr("job_start", 1, 1)
+ failed = True
+ try:
+ result = execute_job(self.job, execute_callable=self._execute)
+ failed = False
+ return result
+ finally:
+ self.job.end_date = timezone.utcnow()
+ self.job.state = JobState.FAILED if failed else JobState.SUCCESS
+ try:
+ get_listener_manager().hook.before_stopping(component=self.job)
+ except Exception:
+ self.log.exception("Error calling Dag processor shutdown
listener")
+ try:
+ client.complete_job(TerminalJobState(self.job.state))
+ except Exception:
+ if not failed:
+ raise
+ self.log.exception("Unable to record the failed Dag processor
Job")
+ finally:
+ stats.incr("job_end", 1, 1)
+
+ def _register_api_job(self, client: DagProcessorAPIClient) -> int:
+ from airflow.dag_processing.api_client import
DagProcessorJobAlreadyRunning
+
+ deadline = monotonic() + conf.getint("dag_processor",
"health_check_threshold") + self.job.heartrate
+ delay = 1.0
+ while True:
+ try:
+ return client.register_job()
+ except DagProcessorJobAlreadyRunning:
+ remaining = deadline - monotonic()
+ if remaining <= 0:
+ raise
+ self.log.info("Waiting for the previous Dag processor Job
heartbeat to expire")
+ sleep(min(delay, remaining))
+ delay = min(delay * 2, 5.0)
+
+ def _heartbeat_api(self, client: DagProcessorAPIClient) -> None:
+ from airflow.dag_processing.api_client import
DagProcessorRegistrationRetired
+
+ write_processor_health(self._last_api_heartbeat)
+ if client.restart_required:
+ raise DagProcessorRegistrationRetired("The Dag processor
registration requires a restart")
+ now = monotonic()
+ if now < self._next_api_heartbeat:
+ return
+ self._next_api_heartbeat = now + self.job.heartrate
+ try:
+ state = client.heartbeat()
+ except httpx2.HTTPError as error:
+ if isinstance(error, httpx2.HTTPStatusError) and
error.response.status_code < 500:
+ raise
+ stats.incr("dag_processor_heartbeat_failure", 1, 1)
+ if monotonic() - self._last_api_heartbeat >= conf.getint(
+ "dag_processor", "job_heartbeat_timeout"
+ ):
+ raise DagProcessorHeartbeatTimeout("Dag processor API
heartbeat timed out") from error
+ self.log.warning("Unable to heartbeat the Dag processor Job;
retrying on the next interval")
Review Comment:
During an outage this logs the same line every heartrate for up to
`job_heartbeat_timeout` with nothing about the cause, so a reader can't tell a
DNS failure from a proxy 503. Could it include `type(error).__name__` (or the
status code) and the elapsed time against `[dag_processor]
job_heartbeat_timeout`? The `DagProcessorHeartbeatTimeout` message above would
help with the same two values too, since that's what the operator sees when the
processor stops.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -846,9 +904,80 @@ def fetch_callbacks(self) -> list[CallbackRequest]:
Default implementation reads from the metadata DB; override to source
callbacks from an API.
"""
- if self.api_client is not None:
+ if self.api_client is None:
+ return self._fetch_callbacks_from_db()
+ available = (
+ self.max_callbacks_per_loop
+ - len(self._callback_claims)
+ - 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)
+ bundles = [b.name for b in self._dag_bundles if not
b.supports_versioning or b.is_initialized]
+ try:
+ work = self.api_client.claim_work("callbacks", bundles, available)
+ except httpx2.HTTPError:
+ self.log.warning("Unable to claim callbacks; retrying on a later
loop")
Review Comment:
This branch and the matching one in `claim_priority_files` are what keep a
503 or connect error from `claim_work` out of the manager loop, but no test
makes `claim_work` raise: every mock of it in `test_manager.py` returns work.
Could one test, parametrized over callbacks and priority, raise
`httpx2.ConnectError` and a 503 and check that the deferred callbacks come back
unchanged?
##########
airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py:
##########
@@ -0,0 +1,275 @@
+# 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.
+from __future__ import annotations
+
+from unittest import mock
+
+import httpx2
+import pytest
+
+from airflow.api_fastapi.execution_api.datamodels.job import TerminalJobState
+from airflow.dag_processing.api_client import (
+ DagProcessorAPIClient,
+ DagProcessorJobAlreadyRunning,
+ DagProcessorRegistrationRetired,
+)
+from airflow.dag_processing.manager import DagFileProcessorManager
+from airflow.jobs.dag_processor_job_runner import
DagProcessorHeartbeatTimeout, DagProcessorJobRunner
+from airflow.jobs.job import Job, JobState
+
+from tests_common.test_utils.config import conf_vars
+
+
[email protected]
+def client():
+ client = mock.create_autospec(DagProcessorAPIClient, instance=True)
+ client.register_job.return_value = 42
+ client.restart_required = False
+ client.heartbeat.return_value = JobState.RUNNING
+ return client
+
+
[email protected]
+def runner():
Review Comment:
`_heartbeat_api` writes the health file on every call, and nothing here
points `health_check_file` away from its default, so the heartbeat tests write
`{AIRFLOW_HOME}/dag_processor_health.json` with the test's pid and leave it
behind. Could this fixture (or an autouse one) set `health_check_file` under
`tmp_path`? The subprocess test in `test_dag_processor_client.py` passes
`os.environ` through without it as well.
##########
airflow-core/src/airflow/jobs/dag_processor_job_runner.py:
##########
@@ -49,12 +61,92 @@ def __init__(
):
super().__init__(job)
self.processor = processor
+ self._last_api_heartbeat = 0.0
+ self._next_api_heartbeat = 0.0
self.processor.heartbeat = lambda: perform_heartbeat(
job=self.job,
heartbeat_callback=self.heartbeat_callback,
only_if_necessary=True,
)
+ def run_with_api(self, client: DagProcessorAPIClient) -> int | None:
+ """Let the API own the Job row while retaining the normal processor
cleanup and listeners."""
+ # The client loads versioned routes which import this runner.
Review Comment:
`datamodels/job.py` only imports `core_api.base` and `airflow.jobs.job`, so
this import can't cycle and can go to the top of the module. The cycle this
comment describes is real for the `api_client` imports in `_register_api_job`
and `_heartbeat_api`, which have no comment.
##########
airflow-core/tests/unit/dag_processing/test_manager.py:
##########
@@ -478,19 +482,21 @@ def
test_find_files_in_bundle_parses_zip_once_and_observes_its_members(self, tmp
"test_zip.zip/broken_dag.py",
}
+ @pytest.mark.parametrize("api_mode", [False, True], ids=["direct", "api"])
Review Comment:
`_find_files_in_bundle` never reads `self.api_client`, so the `api` cases
run exactly the same code as `direct`. Could this parametrize go, or move to a
test where API mode actually branches?
##########
airflow-core/docs/security/jwt_token_authentication.rst:
##########
@@ -557,14 +577,154 @@ 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
checkpoint
+and the metadata writes commit together. Replaying the latest identical
request returns
+the saved receipt; conflicting payloads and superseded sequences are rejected.
Checkpoints
+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. 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. A graceful
+shutdown acknowledges finished work before the processor's Job ends. 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. 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.
+
+Set ``CONNECTION_CHECK_MAX_COUNT=0`` in the processor container too: the
standard image
+entrypoint otherwise checks the database before starting the command. Do not
set
+``_AIRFLOW_DB_MIGRATE`` or ``_AIRFLOW_WWW_USER_CREATE`` in that container.
+
+Isolated deployment examples
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The following examples run only the processor, alongside an existing trusted
deployment.
Review Comment:
Should this say to stop the trusted deployment's own `dag-processor` for the
bundles this one serves, or to run it with `--bundle-name` for the others? The
standalone compose serves `dags-folder` from its own `./dags`, while the
quick-start keeps a DB-mode processor on the same bundle. Both reconcile that
bundle's inventory, so each marks the other's Dags stale when the files differ,
and the DB-credentialed processor keeps running Dag code, which is what this
section is trying to avoid.
##########
airflow-core/src/airflow/jobs/dag_processor_job_runner.py:
##########
@@ -49,12 +61,92 @@ def __init__(
):
super().__init__(job)
self.processor = processor
+ self._last_api_heartbeat = 0.0
+ self._next_api_heartbeat = 0.0
self.processor.heartbeat = lambda: perform_heartbeat(
job=self.job,
heartbeat_callback=self.heartbeat_callback,
only_if_necessary=True,
)
+ def run_with_api(self, client: DagProcessorAPIClient) -> int | None:
Review Comment:
This re-implements `run_job` / `Job.prepare_for_execution` /
`Job.complete_execution`, and the two copies already differ:
`complete_execution` calls `before_stopping` before setting `end_date`, while
this sets `end_date` and the final state first. The next change to the DB
lifecycle (a listener hook, a metric) will silently skip API mode. Would
something like `run_job(job, execute_callable,
lifecycle=ApiJobLifecycle(client))`, with DB and API lifecycle objects behind
`Job`, keep metric names, listener order and state rules in one place? If not,
a note here that this mirrors `run_job` would help.
##########
airflow-core/tests/unit/dag_processing/test_manager.py:
##########
@@ -5457,6 +5524,34 @@ def
test_unusable_inventory_preserves_version_and_recovers_after_refresh(self, o
assert manager._bundle_versions[bundle.name] is None
assert not manager._pending_inventories
+ def test_refresh_dag_bundles_heartbeats_before_each_bundle(self):
+ events = []
+ manager = DagFileProcessorManager(max_runs=1)
+ manager.heartbeat = lambda: events.append("heartbeat")
+ manager._dag_bundles = []
+ for name in ("first", "second"):
+ bundle = self._make_refresh_bundle()
+ bundle.name = name
+ bundle.refresh.side_effect = lambda name=name: events.append(name)
+ manager._dag_bundles.append(bundle)
+ manager._force_refresh_bundles = set()
+
+ with (
Review Comment:
This repeats the eight patches from `_refresh_with_mocked_state`, none with
`autospec=True`, so a renamed or re-signatured method still passes. Could the
helper take a list of bundles so this test can reuse it, with autospec added
there?
##########
airflow-core/src/airflow/dag_processing/health.py:
##########
@@ -0,0 +1,56 @@
+# 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.
+
+"""Local probes for a processor that does not have metadata database
credentials."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from time import monotonic
+
+from airflow.configuration import conf
+
+
+def write_processor_health(last_api_heartbeat: float) -> None:
+ path = Path(conf.get("dag_processor", "health_check_file"))
+ path.parent.mkdir(parents=True, exist_ok=True)
+ temporary = path.with_name(f".{path.name}.{os.getpid()}")
+ temporary.write_text(
+ json.dumps(
+ {
+ "pid": os.getpid(),
+ "loop": monotonic(),
+ "api": last_api_heartbeat,
+ }
+ )
+ )
+ temporary.replace(path)
+
+
+def check_processor_health(*, readiness: bool = False) -> None:
+ try:
+ record = json.loads(Path(conf.get("dag_processor",
"health_check_file")).read_text())
+ os.kill(record["pid"], 0)
+ now = monotonic()
+ maximum_age = conf.getint("dag_processor", "health_check_threshold")
+ fields = ("loop", "api") if readiness else ("loop",)
+ if any(not 0 <= now - record[field] <= maximum_age for field in
fields):
+ raise ValueError("Processor heartbeat has expired")
+ except (OSError, ValueError, KeyError, TypeError) as error:
+ raise SystemExit("Dag processor is not healthy") from error
Review Comment:
An uncaught `SystemExit` prints only its message, so every probe failure
(missing file, `os.kill` from the wrong user, stale loop, stale API heartbeat)
shows up in `docker inspect` or a kubelet event as the same "Dag processor is
not healthy". Could the message carry the cause, e.g. `f"Dag processor is not
healthy: {error}"`, with the expiry naming the field, its age and
`health_check_threshold`?
##########
airflow-core/docs/security/security_model.rst:
##########
@@ -499,6 +499,14 @@ potentially still executes with direct database access in
the Dag File Processor
the database directly, including the JWT signing key configuration if it is
available in the
process environment. If a Dag author obtains the JWT signing key, they
could forge arbitrary tokens.
+ Since Airflow 3.5 the Dag File Processor can instead run without metadata
database access, by
+ setting ``[dag_processor] execution_api_token_file``. It then uses
short-lived Execution API
+ credentials limited to its granted bundles, and each parsing attempt gets a
bundle-scoped token
+ carrying its file and attempt identity. This mode removes the processor's
need for database
Review Comment:
Could this say that the file and attempt identity aren't used for access
checks yet, as `jwt_token_authentication.rst` does? Read alone, "limited to its
granted bundles" plus a per-file token suggests Dag code is confined to its
file, while it runs as the processor's user and can read the session token the
client re-reads.
##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -939,22 +1077,61 @@ def _add_callback_to_queue(self, request:
CallbackRequest) -> None:
self.log.debug("Queuing %s CallbackRequest: %s",
type(request).__name__, request)
if get_claiming_importer(request.filepath, request.bundle_name) is not
None:
self._log_dropped_lang_sdk_callback(request)
+ if item := self._callback_claims.pop(id(request), None):
+ self._pending_work_acks.append(("callbacks", item))
+ self._remove_deferred_callback(request)
return
+ retry = self._deferred_callback_retries.get(id(request))
+ if retry is not None and time.monotonic() < retry.next_attempt_time:
+ return
+ self.heartbeat()
bundle = self.prepare_callback_bundle(request)
if bundle is None:
+ if id(request) in self._callback_claims:
+ self._defer_claimed_callback(request)
return
-
file_info = DagFileInfo(
rel_path=Path(request.filepath),
bundle_path=bundle.path,
bundle_name=request.bundle_name,
bundle_version=request.bundle_version,
)
self._callback_to_execute[file_info].append(request)
+ if id(request) in self._callback_claims:
+ self._remove_deferred_callback(request)
self._add_files_to_queue([file_info], mode="front")
team_name = self._get_team_name(file_info.bundle_name)
stats.incr("dag_processing.other_callback_count",
tags=prune_dict({"team_name": team_name}))
+ def _remove_deferred_callback(self, request: CallbackRequest) -> None:
+ # Equal callback payloads can belong to different claims.
+ self._deferred_api_callbacks = [
+ deferred for deferred in self._deferred_api_callbacks if deferred
is not request
+ ]
+ self._deferred_callback_retries.pop(id(request), None)
+
+ def _defer_claimed_callback(self, request: CallbackRequest) -> None:
+ """Retry preparing a claimed callback's bundle with backoff, then drop
the unpreparable request."""
+ now = time.monotonic()
+ timeout = conf.getint("dag_processor", "job_heartbeat_timeout")
+ retry = self._deferred_callback_retries.setdefault(id(request),
DeferredCallback(now + timeout))
+ if now >= retry.expires_at:
+ self.log.error(
+ "Dropping %s for %s in bundle %s: its bundle version %s could
not be prepared within %s seconds",
Review Comment:
This is the only record that a claimed callback was acknowledged without
running, and it names neither the setting behind the deadline (`[dag_processor]
job_heartbeat_timeout`, which doesn't read as a callback setting) nor the Dag
run whose callback was lost. Could it include both, the way
`_log_dropped_lang_sdk_callback` builds its target? The retries before this
also go through `prepare_callback_bundle`'s `"..., skipping"` with a full
traceback each time, which reads as if the callback was dropped on the first
failure.
##########
airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py:
##########
@@ -371,3 +408,118 @@ def test_replaced_job_signals_restart(clock, api_url,
provision_token, session):
assert replaced.restart_required
assert replacement.heartbeat() == JobState.RUNNING
+
+
+class SecretReadingLocalBundle(LocalDagBundle):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+ if SDKVariable.get("processor-client-key") != "from-api":
+ raise ValueError("Bundle constructor did not resolve its Variable")
+
+
[email protected]("multi_team", [False, True])
[email protected]("database_access", [False, True])
+def test_normal_processor_command_uses_authenticated_api(
+ api_url,
+ api_requests,
+ provision_token,
+ tmp_path,
+ session,
+ multi_team,
+ database_access,
+):
+ clear_db_dags()
+ clear_db_dag_bundles()
+ clear_db_teams()
+ team_name = "processor-team" if multi_team else None
+ if team_name:
+ session.add(Team(name=team_name))
+ session.commit()
+ dag_file = tmp_path / "api_dag.py"
+ dag_file.write_text(
+ "from airflow.sdk import DAG, Variable\n"
+ 'dag = DAG("authenticated_processor", schedule=None,
description=Variable.get("processor-client-key"))\n'
+ )
+ bundle_config = [
+ {
+ "name": "bundle-a",
+ "classpath": f"{__name__}.SecretReadingLocalBundle",
+ "kwargs": {"path": str(tmp_path)},
+ **({"team_name": team_name} if team_name else {}),
+ }
+ ]
+ config = {
+ ("dag_processor", "execution_api_token_file"): str(provision_token()),
+ ("core", "execution_api_server_url"): api_url,
+ ("core", "load_examples"): "False",
+ ("core", "multi_team"): str(multi_team),
+ ("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config),
+ ("scheduler", "job_heartbeat_sec"): "0",
+ ("database", "sql_alchemy_conn"): conf.get("database",
"sql_alchemy_conn"),
+ }
+ try:
+ with conf_vars(config):
+ Variable.set("processor-client-key", "from-api",
team_name=team_name)
+ # Provision the row directly: constructing the bundle here would
read its secret without API auth.
+ bundle = DagBundleModel(name="bundle-a")
+ if team_name:
+ bundle.teams = [session.scalars(select(Team).where(Team.name
== team_name)).one()]
+ session.add(bundle)
+ session.commit()
+ processor_config = dict(config)
+ if not database_access:
+ processor_config.update(
+ {
+ (
+ "database",
+ "sql_alchemy_conn",
+ ):
"postgresql+psycopg://blocked:[email protected]:1/blocked",
Review Comment:
This literal hardcodes the psycopg 3 driver, and the LowestDeps core job on
this head fails both `database_access=False` cases with `ModuleNotFoundError:
No module named 'psycopg'`. `airflow dag-processor` still builds an engine at
`import airflow`, so the driver in the placeholder URL has to be importable.
Could the blocked URL be derived from the backend under test instead, e.g.
`make_url(conf.get("database", "sql_alchemy_conn")).set(username="blocked",
password="blocked", port=1, database="blocked")`? `test_api_isolation.py:148`
uses the same literal.
##########
airflow-core/src/airflow/jobs/dag_processor_job_runner.py:
##########
@@ -49,12 +61,92 @@ def __init__(
):
super().__init__(job)
self.processor = processor
+ self._last_api_heartbeat = 0.0
+ self._next_api_heartbeat = 0.0
self.processor.heartbeat = lambda: perform_heartbeat(
job=self.job,
heartbeat_callback=self.heartbeat_callback,
only_if_necessary=True,
)
+ def run_with_api(self, client: DagProcessorAPIClient) -> int | None:
+ """Let the API own the Job row while retaining the normal processor
cleanup and listeners."""
+ # The client loads versioned routes which import this runner.
+ from airflow.api_fastapi.execution_api.datamodels.job import
TerminalJobState
+
+ self.processor.api_client = client
+ self.job.id = self._register_api_job(client)
+ self.job.state = JobState.RUNNING
+ self.job.start_date = self.job.latest_heartbeat = timezone.utcnow()
+ self._last_api_heartbeat = monotonic()
+ self._next_api_heartbeat = self._last_api_heartbeat +
self.job.heartrate
+ self.processor.heartbeat = lambda: self._heartbeat_api(client)
+ stats.incr("job_start", 1, 1)
+ failed = True
+ try:
+ result = execute_job(self.job, execute_callable=self._execute)
+ failed = False
+ return result
+ finally:
+ self.job.end_date = timezone.utcnow()
+ self.job.state = JobState.FAILED if failed else JobState.SUCCESS
+ try:
+ get_listener_manager().hook.before_stopping(component=self.job)
+ except Exception:
+ self.log.exception("Error calling Dag processor shutdown
listener")
+ try:
+ client.complete_job(TerminalJobState(self.job.state))
+ except Exception:
+ if not failed:
+ raise
+ self.log.exception("Unable to record the failed Dag processor
Job")
+ finally:
+ stats.incr("job_end", 1, 1)
+
+ def _register_api_job(self, client: DagProcessorAPIClient) -> int:
+ from airflow.dag_processing.api_client import
DagProcessorJobAlreadyRunning
+
+ deadline = monotonic() + conf.getint("dag_processor",
"health_check_threshold") + self.job.heartrate
+ delay = 1.0
+ while True:
+ try:
+ return client.register_job()
+ except DagProcessorJobAlreadyRunning:
+ remaining = deadline - monotonic()
+ if remaining <= 0:
+ raise
+ self.log.info("Waiting for the previous Dag processor Job
heartbeat to expire")
+ sleep(min(delay, remaining))
+ delay = min(delay * 2, 5.0)
+
+ def _heartbeat_api(self, client: DagProcessorAPIClient) -> None:
+ from airflow.dag_processing.api_client import
DagProcessorRegistrationRetired
+
+ write_processor_health(self._last_api_heartbeat)
Review Comment:
Nothing pins what this writes for the readiness field. The job-runner tests
never read the health file, `test_health.py` calls `write_processor_health`
directly, and the integration test only checks `--check-ready` on the healthy
path. As far as I can tell, replacing `self._last_api_heartbeat` with
`monotonic()` here would keep a processor that lost the API reporting ready,
and every test in the PR would still pass. A runner test that drives heartbeat
failures past `health_check_threshold` and asserts
`check_processor_health(readiness=True)` exits while the liveness check passes
would cover it. The dead-pid check in `health.py` is unpinned too, since the
`{}` case fails on the missing `loop` key first.
##########
airflow-core/docs/security/security_model.rst:
##########
@@ -775,6 +783,10 @@ model — Airflow does not enforce these natively.
value itself. However, access to the file it points to should be restricted
to the Scheduler
(which generates tokens) and the API Server (which validates them).
+ A Dag File Processor running with ``[dag_processor]
execution_api_token_file`` needs none of the
Review Comment:
The table above lists `AIRFLOW__API_AUTH__JWT_SECRET` for the API Server and
Scheduler, not the Dag File Processor, so "listed for it above" doesn't match.
`AIRFLOW__CORE__ASSET_MANAGER_KWARGS` is the one listed only for the processor,
and in this mode it moves the other way: the API server runs
`update_dag_parsing_results_in_db`, which calls `asset_manager.create_assets`.
Could the note say the processor doesn't need it and the API server does?
##########
airflow-core/src/airflow/cli/cli_config.py:
##########
@@ -2304,6 +2304,21 @@ class GroupCommand(NamedTuple):
help="Start a dag processor instance",
func=lazy_load_command("airflow.cli.commands.dag_processor_command.dag_processor"),
args=(
+ Arg(
+ ("--check-health",),
+ action="store_true",
+ help="Check the API-mode processor's local loop heartbeat and
exit",
+ ),
+ Arg(
+ ("--check-ready",),
+ action="store_true",
+ help="Check the API-mode processor's loop and API heartbeat
and exit",
+ ),
+ Arg(
+ ("--sync-bundles-only",),
Review Comment:
I'm not sure a flag on `dag-processor` is the right home for this. In API
mode `dag-processor` is the least-trusted principal, while
`--sync-bundles-only` needs the database and signing keys, behaves the same
whether or not `execution_api_token_file` is set, and sits beside `airflow
dag-processor-token` in the docs' trusted-host steps. Would a separate command
next to `dag-processor-token` make that trust boundary visible in the name?
--
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]
