From 3d0e8d94a1ce678497c24411fc663da200555555 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 15 Oct 2026 11:23:24 +0100 Subject: [PATCH 1/5] Run Dag processors without metadata database access Standalone processors need to resolve bundle credentials and publish parsed metadata without inheriting database or signing credentials. API-managed Jobs also need lifecycle and requested-work recovery so transient outages do not lose callbacks or priority parses. --- .../dag-processor-api/docker-compose.yaml | 53 + .../security/jwt_token_authentication.rst | 196 +++- airflow-core/docs/security/security_model.rst | 15 +- airflow-core/src/airflow/cli/cli_config.py | 15 + .../cli/commands/dag_processor_command.py | 66 +- .../src/airflow/config_templates/config.yml | 40 +- .../src/airflow/dag_processing/api_client.py | 44 +- .../src/airflow/dag_processing/health.py | 56 ++ .../src/airflow/dag_processing/manager.py | 316 +++++- .../src/airflow/dag_processing/processor.py | 7 +- .../airflow/jobs/dag_processor_job_runner.py | 94 +- airflow-core/src/airflow/utils/cli.py | 16 +- .../dag_processing/test_api_isolation.py | 306 ++++++ .../test_dag_processor_client.py | 154 ++- .../commands/test_dag_processor_command.py | 136 +++ .../tests/unit/dag_processing/test_health.py | 64 ++ .../tests/unit/dag_processing/test_manager.py | 951 +++++++++++++++++- .../jobs/test_dag_processor_job_runner.py | 275 +++++ .../metrics/metrics_template.yaml | 36 +- 19 files changed, 2749 insertions(+), 91 deletions(-) create mode 100644 airflow-core/docs/security/dag-processor-api/docker-compose.yaml create mode 100644 airflow-core/src/airflow/dag_processing/health.py create mode 100644 airflow-core/tests/integration/dag_processing/test_api_isolation.py create mode 100644 airflow-core/tests/unit/dag_processing/test_health.py create mode 100644 airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py diff --git a/airflow-core/docs/security/dag-processor-api/docker-compose.yaml b/airflow-core/docs/security/dag-processor-api/docker-compose.yaml new file mode 100644 index 0000000000000..a73b114243f2f --- /dev/null +++ b/airflow-core/docs/security/dag-processor-api/docker-compose.yaml @@ -0,0 +1,53 @@ +# 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. +--- +# Run independently of the quick-start Compose file so its credentials are not inherited. +services: + dag-processor: + image: ${AIRFLOW_IMAGE_NAME:?Set an Airflow image with processor API support} + command: [dag-processor] + user: "${AIRFLOW_UID:-50000}:0" + environment: + USER: airflow + AIRFLOW_CONFIG: /tmp/airflow.cfg + AIRFLOW__CORE__LOAD_EXAMPLES: "False" + AIRFLOW__CORE__EXECUTION_API_SERVER_URL: http://airflow-apiserver:8080/execution/ + AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE: /run/dag-processor/processor.jwt + AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_STORAGE_PATH: /tmp/dag-bundles + AIRFLOW__DAG_PROCESSOR__HEALTH_CHECK_FILE: /tmp/dag-processor-health.json + CONNECTION_CHECK_MAX_COUNT: "0" + volumes: + - ./tokens:/run/dag-processor:ro + - ./dags:/opt/airflow/dags:ro + - ./logs:/opt/airflow/logs + read_only: true + tmpfs: + - /tmp + cap_drop: [ALL] + security_opt: [no-new-privileges:true] + healthcheck: + test: [CMD, airflow, dag-processor, --check-health] + interval: 30s + timeout: 10s + start_period: 30s + retries: 5 + networks: [processor-api] + restart: unless-stopped +networks: + processor-api: + name: ${DAG_PROCESSOR_API_NETWORK:-dag-processor-api} + external: true diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst index e72f9d3faa6a5..e978e054bf23b 100644 --- a/airflow-core/docs/security/jwt_token_authentication.rst +++ b/airflow-core/docs/security/jwt_token_authentication.rst @@ -471,8 +471,10 @@ Dag File Processor and Triggerer --------------------------------- The **Dag File Processor** and **Triggerer** are internal Airflow components that also -interact with the Execution API, but they do so via an **in-process** transport -(``InProcessExecutionAPI``) rather than over the network. This in-process API: +interact with the Execution API, but by default they do so via an **in-process** transport +(``InProcessExecutionAPI``) rather than over the network. A Dag File Processor can instead use +authenticated HTTP without database access; see `Dag processor HTTP client credentials`_. This +in-process API: - Runs the Execution API application directly within the same process, using an ASGI/WSGI bridge. @@ -514,16 +516,18 @@ guidance, and the planned strategic and tactical improvements. Dag processor HTTP client credentials ------------------------------------- -The core ``DagProcessorAPIClient`` lets processor integrations authenticate to the Execution -API without holding its signing key. The standard ``airflow dag-processor`` command does -not yet use this client. Provisioning a token file alone does not switch that command to -HTTP or remove its need for database access. +Set ``[dag_processor] execution_api_token_file`` to run ``airflow dag-processor`` with +authenticated HTTP requests. Job registration, heartbeats, completion, bundle secret +lookups, bundle metadata, result publication, and callback and priority delivery then use +the Execution API. The API client is created +in the processor process, after any daemon or hot-reload fork. Leaving the setting +unset preserves the in-process API and database-backed Job lifecycle. -On a trusted host with the API signing key, provision a session token for the bundles -the client may access:: +On a trusted host with database access and the server keys, synchronize the configured +bundle catalog and provision a separate token file for each processor:: - airflow dag-processor-token --token-file /run/airflow/processor.jwt \ - --bundle-name dags-folder --rotate + airflow dag-processor --sync-bundles-only + airflow dag-processor-token --token-file /run/airflow/processor.jwt --bundle-name dags-folder --rotate The command writes the token atomically. A new file is readable only by its owner, and replacing an existing file keeps its mode and group, so grant the processor read access once. ACLs are not @@ -532,14 +536,30 @@ token replacement stays visible. Each processor process needs its own session. R provisioner with the same token file and bundles keeps that session, so running processors keep their Jobs. +Configure ``[core] execution_api_server_url`` to reach the API server, and start the +processor:: + + AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE=/run/airflow/processor.jwt \ + airflow dag-processor --bundle-name dags-folder + +The processor does not need metadata database credentials, the Fernet key, or API signing +keys. Remove them from its environment and mounted configuration, and block its database +network path. Enabling the setting alone does not remove existing credentials or isolate +code from other secrets on the processor host. Catalog synchronization is a trusted +deployment operation: repeat it when bundle configuration changes. It loads provider +classes for bundle URLs; the API server does not construct bundles or access their sources. +The processor does not use the shared SDK secret cache in this mode, even when +``[secrets] use_cache`` is enabled, because SDK cache lookups do not include the bundle +identity that the API uses for authorization. + The three credentials have different issuers and permissions: - ``dag_processor_session`` is issued by trusted provisioning. It can only register a Job or renew that Job's credential through ``POST /jobs``. - ``dag_processor`` is returned by Job registration. It can heartbeat and complete that - Job, read Connections and Variables for a granted bundle, and exchange a parsing - credential through ``POST /jobs/{job_id}/parse-token``. It cannot outlive the session - token used to obtain it. + Job, read Connections and Variables for a granted bundle, publish parsing results and + bundle inventories, claim requested work, and exchange a parsing credential through + ``POST /jobs/{job_id}/parse-token``. It cannot outlive the session token used to obtain it. - ``dag_parse`` is returned by the parsing-token exchange for one attempt, bundle, and relative file location. It permits parse-time requests, including callback-context reads, but cannot manage Jobs or exchange tokens. It cannot outlive its parent Job credential. @@ -557,14 +577,148 @@ 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 or unassigned +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, and 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 callback as failed. 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. + +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. +Use an image containing processor API support; this feature targets Airflow 3.5. +Synchronize the bundle catalog and run token provisioning in the trusted deployment as +described above. Keep the token provisioner running, and mount its directory rather than +a single file so atomic token replacement remains visible. + +For Compose, use the :download:`standalone processor definition ` +as its own project, rather than merging it with the quick-start file. It has no shared +``env_file``, server configuration mount, or dependency on a database/migration service. +Create its ``tokens``, ``dags``, and ``logs`` directories with permissions for the configured +container user; provision ``tokens/processor.jwt``. Set ``AIRFLOW_IMAGE_NAME`` to the image +with API support. Create an internal network and attach only the API server and processor:: + + docker network create --internal dag-processor-api + docker network connect --alias airflow-apiserver dag-processor-api + docker compose -p dag-processor-api -f docker-compose.yaml up -d + +The API server keeps its existing database network. The processor joins only +``dag-processor-api``. Verify readiness with ``docker compose -p dag-processor-api +-f docker-compose.yaml exec dag-processor airflow dag-processor --check-ready``. +The examples do not grant provider credentials through environment variables; bundle +Connection lookups use the API. Review any image customizations, added mounts, or environment +settings for server credentials before deployment. + +The token carries file and attempt identity, but the API does not currently use them +for access checks or request attribution. Connection and Variable permissions remain +bundle/team based. Archives retain one file identity under the current processor model. +This exchange does not isolate hostile code from credentials accessible on the processor +host. The process and deployment limitations described above still apply. Workload Isolation and Current Limitations diff --git a/airflow-core/docs/security/security_model.rst b/airflow-core/docs/security/security_model.rst index ec5c340851dfc..ace40cfc59ae6 100644 --- a/airflow-core/docs/security/security_model.rst +++ b/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 + credentials, the Fernet key and signing keys, but does not remove them: Deployment Managers must + take them out of the processor's environment and block its database network path. The Triggerer + still needs database access. See :doc:`/security/jwt_token_authentication`. + **Dag File Processor and Triggerer are shared across teams** In the default deployment, a **single Dag File Processor instance** parses all Dag files and a **single Triggerer instance** handles all triggers — regardless of team assignment. There is no @@ -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 + metadata database variables, the Fernet key or the JWT signing secret listed for it above; it + authenticates with the token file instead. + .. START AUTOGENERATED PROVIDER SENSITIVE VARS **Provider-specific sensitive configuration variables:** @@ -848,7 +860,8 @@ model — Airflow does not enforce these natively. the metadata database exclusively through the API server (similar to how workers use the Execution API today). This would eliminate the need for these components to have database credentials at all, providing security by design rather than relying on deployment-level - measures. + measures. Airflow 3.5 offers this for the Dag File Processor as an opt-in mode + (``[dag_processor] execution_api_token_file``); the Triggerer still uses the database. * **Tactical (shorter-term)**: Native support for Unix user impersonation in the Dag File Processor and Triggerer child processes, so that Dag author code runs as a different, low- privilege user that cannot access the parent's credentials or the database. diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py index 961a295ddee0e..c75f4bd4e149e 100644 --- a/airflow-core/src/airflow/cli/cli_config.py +++ b/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",), + action="store_true", + help="Provision the configured bundle catalog using trusted database and signing credentials, then exit", + ), ARG_PID, ARG_DAEMON, ARG_BUNDLE_NAME, diff --git a/airflow-core/src/airflow/cli/commands/dag_processor_command.py b/airflow-core/src/airflow/cli/commands/dag_processor_command.py index b8195bead2895..b49a4d6662393 100644 --- a/airflow-core/src/airflow/cli/commands/dag_processor_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_processor_command.py @@ -19,19 +19,28 @@ from __future__ import annotations import logging -from typing import Any +from typing import TYPE_CHECKING, Any from airflow.cli.commands.daemon_utils import run_command_with_daemon_option from airflow.configuration import conf -from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names +from airflow.dag_processing.bundles.manager import ( + _load_bundle_config_snapshot, + get_configured_bundle_team_names, +) +from airflow.dag_processing.health import check_processor_health from airflow.dag_processing.manager import DagFileProcessorManager +from airflow.executors.base_executor import get_execution_api_server_url from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner from airflow.jobs.job import Job, run_job from airflow.utils import cli as cli_utils from airflow.utils.memray_utils import MemrayTraceComponents, enable_memray_trace +from airflow.utils.net import get_hostname from airflow.utils.process_utils import set_component_mp_start_method from airflow.utils.providers_configuration_loader import providers_configuration_loaded +if TYPE_CHECKING: + from collections.abc import Callable + log = logging.getLogger(__name__) @@ -60,7 +69,12 @@ def _get_team_names(bundle_names: list[str] | None) -> list[str]: def _create_dag_processor_job_runner(args: Any) -> DagProcessorJobRunner: """Create DagFileProcessorProcess instance.""" if args.bundle_name: - cli_utils.validate_dag_bundle_arg(args.bundle_name) + if conf.get("dag_processor", "execution_api_token_file", fallback=None): + # Bundle constructors can read secrets; defer them until the API client is registered. + if unknown := set(args.bundle_name) - _load_bundle_config_snapshot().names: + raise SystemExit(f"Bundles not found: {', '.join(sorted(unknown))}") + else: + cli_utils.validate_dag_bundle_arg(args.bundle_name) return DagProcessorJobRunner( job=Job(bundle_names=args.bundle_name, team_names=_get_team_names(args.bundle_name)), processor=DagFileProcessorManager( @@ -72,13 +86,55 @@ 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. + from airflow.dag_processing.api_client import DagProcessorAPIClient + + with DagProcessorAPIClient( + base_url=get_execution_api_server_url(), + token_file=token_file, + hostname=get_hostname(), + unixname=job_runner.job.unixname, + bundle_names=job_runner.processor.bundle_names_to_parse + or sorted(_load_bundle_config_snapshot().names), + ) as client: + job_runner.run_with_api(client) -@cli_utils.action_cli @providers_configuration_loaded def dag_processor(args): """Start Airflow Dag Processor Job.""" + if getattr(args, "check_health", False) or getattr(args, "check_ready", False): + if not conf.get("dag_processor", "execution_api_token_file", fallback=None): + raise SystemExit("Local health probes require Dag processor API mode") + return check_processor_health(readiness=getattr(args, "check_ready", False)) + if getattr(args, "sync_bundles_only", False): + return _run_with_action_log(_sync_bundle_catalog, args) + if not conf.get("dag_processor", "execution_api_token_file", fallback=None): + return _run_with_action_log(_start_dag_processor, args) + if getattr(args, "verbose", False): + cli_utils.set_verbose_logging() + return _start_dag_processor(args) + + +def _run_with_action_log(func: Callable[[Any], Any], args: Any) -> Any: + # The action log names its event after the decorated function; keep the command's own name. + def dag_processor(args): + return func(args) + + return cli_utils.action_cli(dag_processor)(args) + + +def _sync_bundle_catalog(args): + if args.bundle_name: + cli_utils.validate_dag_bundle_arg(args.bundle_name) + DagFileProcessorManager(max_runs=0, bundle_names_to_parse=args.bundle_name).sync_bundles() + + +def _start_dag_processor(args): set_component_mp_start_method("dag_processor") job_runner = _create_dag_processor_job_runner(args) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 4a46f287088bd..8106e448eaa9f 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3313,6 +3313,32 @@ dag_processor: - how often to refresh DAGs from those sources - how many files to parse concurrently options: + execution_api_token_file: + description: | + Path to a session token written by ``airflow dag-processor-token``. When set, + ``airflow dag-processor`` uses the authenticated Execution API for its Job lifecycle, + bundle metadata and secrets, parse-time requests, result publication, and callback + and priority-request delivery. The API URL comes from + ``[core] execution_api_server_url``. Each processor needs its own session token. + + Provision the bundle catalog on a trusted host before starting the processor. + The processor does not need metadata database credentials, the Fernet key, or + API signing keys. Remove those credentials from its deployment and block database + network access to enforce isolation; setting this option does not remove them. + Leave unset to use the in-process API and database-backed Job lifecycle. + version_added: 3.5.0 + type: string + default: ~ + example: "/run/airflow/dag-processor.jwt" + health_check_file: + description: | + Local heartbeat file written by a Dag processor using the Execution API. + Use a separate writable file for each processor. ``airflow dag-processor --check-health`` + checks its loop heartbeat; ``--check-ready`` also checks the last successful API heartbeat. + version_added: 3.5.0 + type: string + default: "{AIRFLOW_HOME}/dag_processor_health.json" + example: ~ dag_bundle_storage_path: description: | String path to folder where Airflow bundles can store files locally. Not templated. @@ -3476,6 +3502,10 @@ dag_processor: dag processor is considered unhealthy. This is used by the health check in the **/health** endpoint and in ``airflow jobs check`` CLI for DagProcessorJob. + For a Dag processor using the Execution API, the API server also reads this value to decide + whether a session's previous Job is still alive, and the processor reads it for its local + ``--check-health`` and ``--check-ready`` probes and to bound how long registration waits for + a previous Job to expire. Set the same value for the API server and the processor. version_added: 3.2.0 type: integer example: ~ @@ -3484,7 +3514,15 @@ dag_processor: description: | Seconds a Dag processor using the Execution API may go without a recorded Job heartbeat. After this, the scheduler's next cleanup (every ``[scheduler] parsing_cleanup_interval``) - ends the Job and releases the callbacks and parse requests it claimed. + ends the Job and releases the callbacks and parse requests it claimed, and a processor that + could not reach the API stops. Processors do not heartbeat while initializing, refreshing, + or listing a single bundle, so set this above the longest of those operations; the local + health check must also tolerate them, through ``health_check_threshold`` and the probe's + failure budget. It also bounds how long a processor retries preparing the bundle + version of a claimed callback before acknowledging the callback as failed. + The scheduler reads this value to end Jobs and the processor reads it to stop, so set the + same value for both; changing it only in the processor's configuration has no effect on + when the scheduler ends the Job. A callback that ran before its acknowledgment was recorded can run again after the Job times out, repeating its side effects. Callbacks should tolerate repeated execution. version_added: 3.5.0 diff --git a/airflow-core/src/airflow/dag_processing/api_client.py b/airflow-core/src/airflow/dag_processing/api_client.py index 964b0aead80ab..14f6e6db4c750 100644 --- a/airflow-core/src/airflow/dag_processing/api_client.py +++ b/airflow-core/src/airflow/dag_processing/api_client.py @@ -25,7 +25,7 @@ from dataclasses import dataclass, field from pathlib import Path from time import monotonic -from typing import Any +from typing import Any, Literal from urllib.parse import quote from uuid import UUID @@ -39,6 +39,9 @@ DagBundleStateResponse, DagParseResultBody, DagParseResultResponse, + ProcessorWorkAckBody, + ProcessorWorkClaimBody, + ProcessorWorkItem, ) from airflow.api_fastapi.execution_api.datamodels.job import ( DagParseTokenBody, @@ -142,6 +145,7 @@ def __init__( self._job_id: int | None = None self._completion_state: TerminalJobState | None = None self._completed = False + self._pending_claims: dict[str, ProcessorWorkClaimBody] = {} super().__init__(base_url=base_url, token="", **kwargs) @property @@ -398,6 +402,44 @@ def get_bundles(self) -> 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(), + 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(), + ) + work = [ProcessorWorkItem.model_validate(item) for item in response.json()] + del self._pending_claims[kind] + return work + + def acknowledge_work( + self, kind: Literal["callbacks", "priority"], work: ProcessorWorkItem, *, failed: bool = False + ) -> None: + body = ProcessorWorkAckBody(claim_id=work.claim_id, state="failed" if failed else "success") + self.request( + "POST", + f"jobs/{self._require_job_id()}/requested-work/{kind}/{work.id}/ack", + content=body.model_dump_json(), + headers=_JOB_API_HEADERS, + retry=False, + timeout=self._get_bounded_timeout(), + ) + def publish_inventory(self, bundle_name: str, body: DagBundleInventoryBody) -> DagBundleInventoryResponse: response = self.request( "POST", diff --git a/airflow-core/src/airflow/dag_processing/health.py b/airflow-core/src/airflow/dag_processing/health.py new file mode 100644 index 0000000000000..640724a85ef00 --- /dev/null +++ b/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 diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 818b5d6538af8..2ae733beaf57e 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -23,6 +23,7 @@ import gc import inspect import logging +import math import os import random import selectors @@ -42,7 +43,7 @@ import attrs import httpx2 import structlog -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from sqlalchemy import select, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import load_only @@ -58,9 +59,10 @@ DagParseResultBody, ParseSourceCode, ParseWarning, + ProcessorWorkItem, ) from airflow.api_fastapi.execution_api.in_process import InProcessExecutionAPI -from airflow.callbacks.callback_requests import DagCallbackRequest +from airflow.callbacks.callback_requests import CallbackRequest, DagCallbackRequest from airflow.configuration import conf from airflow.dag_processing.bundles.base import ( BundleUsageTrackingManager, @@ -113,7 +115,6 @@ from sqlalchemy.orm import Session from sqlalchemy.sql import Select - from airflow.callbacks.callback_requests import CallbackRequest from airflow.dag_processing.api_client import DagProcessorAPIClient from airflow.dag_processing.bundles.base import BaseDagBundle from airflow.sdk.api.client import Client @@ -183,6 +184,27 @@ class PendingBundleInventory: next_attempt_time: float = 0.0 +@attrs.define +class DeferredCallback: + """A claimed callback whose bundle could not be prepared, retried until it expires.""" + + expires_at: float + attempts: int = 0 + next_attempt_time: float = 0.0 + + +@attrs.define +class PriorityParseClaim: + """A priority request served only by an import after its requested refresh.""" + + work: ProcessorWorkItem + refresh_generation: int + attempt_id: UUID | None = None + + +_MAX_CALLBACK_RETRY_DELAY = 60.0 + + @dataclass(frozen=True) class DagFileInfo: """Information about a DAG file.""" @@ -342,6 +364,14 @@ class DagFileProcessorManager(LoggingMixin): _pending_inventories: dict[str, PendingBundleInventory] = attrs.field(factory=dict, init=False) _parsing_start_time: float | None = attrs.field(default=None, init=False) + _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) + _callback_attempts: dict[UUID, list[ProcessorWorkItem]] = attrs.field(factory=dict, init=False) + _priority_claims: dict[str, PriorityParseClaim] = attrs.field(factory=dict, init=False) + _pending_work_acks: list[tuple[Literal["callbacks", "priority"], ProcessorWorkItem, bool]] = attrs.field( + factory=list, init=False + ) _num_run: int = attrs.field(default=0, init=False) _callback_to_execute: dict[DagFileInfo, list[CallbackRequest]] = attrs.field( @@ -737,6 +767,7 @@ def _run_parsing_loop(self): self._collect_results() self._publish_pending_results() + self._acknowledge_requested_work() for callback in self.fetch_callbacks(): self._add_callback_to_queue(callback) self._scan_stale_dags() @@ -797,9 +828,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 = ( + self._parallelism + - len(self._priority_claims) + - sum(kind == "priority" for kind, _, _ in self._pending_work_acks) + ) + if available <= 0: return [] - return self._claim_priority_files() + bundles = {bundle.name: bundle for bundle in self._dag_bundles} + try: + work = self.api_client.claim_work("priority", list(bundles), available) + except httpx2.HTTPError: + self.log.warning("Unable to claim priority parses; retrying on a later loop") + return [] + files = [] + for item in work: + self._priority_claims[item.id] = PriorityParseClaim( + work=item, refresh_generation=self._bundle_refresh_generations[item.bundle_name] + 1 + ) + bundle = bundles[item.bundle_name] + path = find_enclosing_file(bundle.path / item.relative_fileloc) + 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 request_bundle_refresh(self, bundle_names: str | Iterable[str]) -> None: """ @@ -856,9 +914,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") + return list(self._deferred_api_callbacks) + for item in work: + if item.callback is None: + raise ValueError("Callback claim returned no request") + request: CallbackRequest = TypeAdapter(CallbackRequest).validate_json(item.callback) + self._callback_claims[id(request)] = item + self._deferred_api_callbacks.append(request) + return list(self._deferred_api_callbacks) + + def _acknowledge_requested_work(self, *, deadline: float | None = None) -> None: + """Send queued acknowledgments in order until the deadline, keeping claims after a failure.""" + if self.api_client is None: + return + if deadline is None: + deadline = time.monotonic() + _API_SEND_BUDGET_SECONDS + while self._pending_work_acks and time.monotonic() < deadline: + kind, item, failed = self._pending_work_acks[0] + try: + self.api_client.acknowledge_work(kind, item, failed=failed) + except httpx2.HTTPStatusError as error: + if error.response.status_code >= 500: + self.log.warning("Unable to acknowledge requested work; retaining the claim") + return + if error.response.status_code != 409: + # Retrying cannot succeed; the server releases the claim when this Job ends. + self.log.warning( + "The API refused to acknowledge %s work %s (HTTP %s)", + kind, + item.id, + error.response.status_code, + ) + except httpx2.RequestError: + self.log.warning("Unable to acknowledge requested work; retaining the claim") + return + self._pending_work_acks.pop(0) + + def _finish_callback_attempt(self, proc: BaseDagFileProcessorProcess) -> None: + for item in self._callback_attempts.pop(proc.id, []): + self._pending_work_acks.append(("callbacks", item, proc._exit_code != 0)) + + def _finish_priority_claims(self, attempt_id: UUID, *, failed: bool) -> None: + """Acknowledge only requests assigned to this import, never newer requests for its file.""" + for key, claim in list(self._priority_claims.items()): + if claim.attempt_id == attempt_id: + self._pending_work_acks.append(("priority", claim.work, failed)) + del self._priority_claims[key] + + def _get_priority_claims(self, file: DagFileInfo) -> list[PriorityParseClaim]: + if file.bundle_version is not None: return [] - return self._fetch_callbacks_from_db() + return [ + claim + for claim in self._priority_claims.values() + if claim.work.bundle_name == file.bundle_name + and ( + claim.work.relative_fileloc == str(file.rel_path) + or claim.work.relative_fileloc.startswith(str(file.rel_path) + "/") + ) + ] @provide_session @retry_db_transaction @@ -936,7 +1065,16 @@ def prepare_callback_bundle(self, request: CallbackRequest) -> BaseDagBundle | N ) if bundle.supports_versioning: bundle.initialize() - except Exception: + except Exception as error: + if self.api_client is not None: + from airflow.dag_processing.api_client import DagProcessorRegistrationRetired + + if isinstance(error, DagProcessorRegistrationRetired): + raise + if self.api_client.restart_required: + raise DagProcessorRegistrationRetired( + "The Dag processor registration has ended; restart required" + ) from error self.log.exception( "Error preparing bundle %s version %s for callback, skipping", request.bundle_name, @@ -949,11 +1087,19 @@ 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, True)) + 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, @@ -961,10 +1107,41 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None: 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 acknowledge it as failed.""" + 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", + type(request).__name__, + request.filepath, + request.bundle_name, + request.bundle_version, + timeout, + ) + self._pending_work_acks.append(("callbacks", self._callback_claims.pop(id(request)), True)) + self._remove_deferred_callback(request) + return + retry.attempts += 1 + retry.next_attempt_time = min( + now + min(2 ** (retry.attempts - 1), _MAX_CALLBACK_RETRY_DELAY), retry.expires_at + ) + def _log_dropped_lang_sdk_callback(self, request: CallbackRequest) -> None: if isinstance(request, DagCallbackRequest): target = f"dag_id={request.dag_id} run_id={request.run_id}" @@ -1059,12 +1236,35 @@ def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set self._bundle_versions[bundle_name] = pending.source.bundle_version self._bundle_failed_sources[bundle_name] = frozenset(pending.body.failed_sources) known_files[bundle_name] = pending.found_files + failed_sources = set(pending.body.failed_sources) + for key, claim in list(self._priority_claims.items()): + item = claim.work + if ( + item.bundle_name == bundle_name + and pending.source.refresh_generation >= claim.refresh_generation + and item.relative_fileloc not in pending.body.files + and not is_source_discovery_failed(item.relative_fileloc, failed_sources) + ): + self._pending_work_acks.append(("priority", item, True)) + del self._priority_claims[key] def _discard_unusable_bundle_source(self, bundle_name: str) -> None: self._pending_inventories.pop(bundle_name, None) self._bundle_parse_sources.pop(bundle_name, None) self._force_refresh_bundles.discard(bundle_name) self._bundles_waiting_for_refresh.discard(bundle_name) + active_attempts = { + processor.id for file, processor in self._processors.items() if file.bundle_name == bundle_name + } + active_attempts.update( + pending.body.attempt_id + for file, pending in self._pending_publications.items() + if file.bundle_name == bundle_name + ) + for key, claim in list(self._priority_claims.items()): + if claim.work.bundle_name == bundle_name and claim.attempt_id not in active_attempts: + self._pending_work_acks.append(("priority", claim.work, True)) + del self._priority_claims[key] def _invalidate_bundle_parse_source(self, bundle_name: str) -> None: # Initialization or refresh may change files even when the provider raises afterwards. @@ -1116,6 +1316,8 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): for bundle in self._dag_bundles: if not periodic_check_due and bundle.name not in urgent_bundles: continue + # Bundle operations can block for minutes; an API-mode Job ends after job_heartbeat_timeout. + self.heartbeat() with self._use_bundle(bundle.name): if catalog is not None and bundle.name not in catalog: self._discard_unusable_bundle_source(bundle.name) @@ -1594,11 +1796,15 @@ def handle_removed_files(self, known_files: dict[str, set[DagFileInfo]]): def purge_removed_files_from_queue(self, present: set[DagFileInfo]): """Remove from queue any files no longer observed locally.""" present_keys = {file.presence_key for file in present} - self._file_queue = OrderedDict((x, None) for x in self._file_queue if x.presence_key in present_keys) - self._files_awaiting_source = { - file: None for file in self._files_awaiting_source if file.presence_key in present_keys - } - self._priority_files = {file for file in self._priority_files if file.presence_key in present_keys} + + def is_retained(file: DagFileInfo) -> bool: + return file.presence_key in present_keys or bool( + self.api_client is not None and self._callback_to_execute.get(file) + ) + + self._file_queue = OrderedDict((x, None) for x in self._file_queue if is_retained(x)) + self._files_awaiting_source = {x: None for x in self._files_awaiting_source if is_retained(x)} + self._priority_files = {x for x in self._priority_files if is_retained(x)} stats.gauge("dag_processing.file_path_queue_size", len(self._file_queue)) self._record_files_awaiting_source() @@ -1619,6 +1825,8 @@ def terminate_orphan_processes(self, present: set[DagFileInfo]): bundle_to_team = self._get_team_names({file.bundle_name for file in self._processors}) for file in list(self._processors.keys()): + if self.api_client is not None and self._processors[file].id in self._callback_attempts: + continue if file.presence_key not in present_keys: processor = self._processors.pop(file, None) if not processor: @@ -1728,8 +1936,11 @@ def handle_parsing_result( file.bundle_name, ) self._record_failed_publication(file, next_stat) + self._finish_priority_claims(proc.id, failed=True) return + if proc.parsing_result is None and not is_callback_only: + self._finish_priority_claims(proc.id, failed=True) self._file_stats[file] = next_stat def _record_failed_publication(self, file: DagFileInfo, stat: DagFileStat) -> None: @@ -1796,6 +2007,7 @@ def _publish_pending_results(self, *, deadline: float | None = None) -> None: detail += f": {message}" self._pending_publications[file] = self._build_rejection_report( file, + attempt_id=pending.body.attempt_id, bundle_version=pending.body.bundle_version, bundle_revision=pending.body.bundle_revision, parse_duration=pending.body.parse_duration, @@ -1807,8 +2019,10 @@ def _publish_pending_results(self, *, deadline: float | None = None) -> None: continue self.log.exception("Failed to publish parsing result for %s", file) self._record_failed_publication(file, pending.stat) + self._finish_priority_claims(pending.body.attempt_id, failed=True) else: self._file_stats[file] = pending.stat + self._finish_priority_claims(pending.body.attempt_id, failed=False) del self._pending_publications[file] def _build_pending_publication( @@ -1832,6 +2046,7 @@ def _build_pending_publication( reason = f"it exceeds the {MAX_PUBLICATION_BYTES // (1024 * 1024)} MiB publication limit" return self._build_rejection_report( file, + attempt_id=proc.id, bundle_version=source.bundle_version, bundle_revision=source.bundle_revision, parse_duration=run_duration, @@ -1844,6 +2059,7 @@ def _build_rejection_report( self, file: DagFileInfo, *, + attempt_id: UUID, bundle_version: str | None, bundle_revision: UUID | None, parse_duration: float, @@ -1851,7 +2067,12 @@ def _build_rejection_report( refresh_generation: int, reason: str, ) -> PendingDagPublication: - """Replace a result that cannot be published with an import error users can see.""" + """ + Replace a result that cannot be published with an import error users can see. + + :param attempt_id: The import whose result is replaced; its priority requests are released. + """ + self._finish_priority_claims(attempt_id, failed=True) self._dispatch_sequence += 1 if len(reason) > _MAX_REJECTION_REASON_LENGTH: reason = reason[:_MAX_REJECTION_REASON_LENGTH] + "..." @@ -1972,6 +2193,7 @@ def _collect_results(self): for file in finished: processor = self._processors.pop(file) + self._finish_callback_attempt(processor) processor.close() def _get_log_dir(self) -> str: @@ -2037,17 +2259,24 @@ def client(self) -> Client: def _create_process( self, dag_file: DagFileInfo, *, source: DagParseSource | None = None ) -> BaseDagFileProcessorProcess: - id = uuid7() + process_id = uuid7() if source is None: source = self._get_dispatch_source(dag_file) or DagParseSource() callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, []) + claims = [ + self._callback_claims.pop(id(request)) + for request in callback_to_execute_for_file + if id(request) in self._callback_claims + ] + if claims: + self._callback_attempts[process_id] = claims logger, logger_filehandle = self._get_logger_for_dag_file(dag_file) subprocess_logs_to_stdout = conf.get("logging", "dag_processor_log_target") == "stdout" if get_claiming_importer(dag_file.absolute_path, dag_file.bundle_name) is not None: return LangSDKDagFileProcessorProcess.start( - id=id, + id=process_id, path=dag_file.absolute_path, bundle_path=cast("Path", dag_file.bundle_path), bundle_name=dag_file.bundle_name, @@ -2061,7 +2290,7 @@ def _create_process( ) return DagFileProcessorProcess.start( - id=id, + id=process_id, path=dag_file.absolute_path, bundle_path=cast("Path", dag_file.bundle_path), bundle_name=dag_file.bundle_name, @@ -2089,6 +2318,10 @@ def _get_dispatch_source(self, file: DagFileInfo) -> DagParseSource | None: return None return self._bundle_parse_sources.get(file.bundle_name) + def _has_priority_request(self, file: DagFileInfo) -> bool: + # A claim outlives the queue marker when its file is requeued behind a running or stale import. + return file in self._priority_files or bool(self._get_priority_claims(file)) + def _readmit_files_with_source(self) -> None: ready = [file for file in self._files_awaiting_source if self._get_dispatch_source(file) is not None] if not ready: @@ -2102,12 +2335,12 @@ def _readmit_files_with_source(self) -> None: priority = [ file for file in queued_first - if file in self._priority_files and file not in self._callback_to_execute + if self._has_priority_request(file) and file not in self._callback_to_execute ] regular = [ file for file in restored_first - if file not in self._priority_files and file not in self._callback_to_execute + if not self._has_priority_request(file) and file not in self._callback_to_execute ] self._file_queue.clear() self._add_files_to_queue(callbacks + priority + regular, mode="back") @@ -2117,14 +2350,22 @@ 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: - if file in self._priority_files: + if priority_claims: + self._priority_files.discard(file) + requested_later[file] = None + elif self.api_client is None and file in self._priority_files: self._files_awaiting_source[file] = None + else: + self._priority_files.discard(file) continue source = self._get_dispatch_source(file) @@ -2133,13 +2374,19 @@ def _start_new_processes(self): continue has_callbacks = file in self._callback_to_execute processor = self._create_process(file, source=source) - if has_callbacks and file in self._priority_files: + if self.api_client is None and has_callbacks and file in self._priority_files: # Callback execution cannot satisfy a request for a fresh ordinary parse. self._files_awaiting_source[file] = None else: self._priority_files.discard(file) self._dispatch_sequence += 1 processor.dispatch_sequence = self._dispatch_sequence + if file.bundle_version is None: + for claim in priority_claims: + if processor.had_callbacks or source.refresh_generation < claim.refresh_generation: + requested_later[file] = None + else: + claim.attempt_id = processor.id stats.incr( "dag_processing.processes", tags=prune_dict( @@ -2153,6 +2400,8 @@ def _start_new_processes(self): ) self._processors[file] = processor + if requested_later: + self._add_files_to_queue(list(reversed(requested_later)), mode="front") stats.gauge("dag_processing.file_path_queue_size", len(self._file_queue)) self._record_files_awaiting_source() @@ -2381,10 +2630,13 @@ def _kill_timed_out_processors(self): last_num_of_db_queries=0, ) self._file_stats[file] = stat + if not processor.had_callbacks: + self._finish_priority_claims(processor.id, failed=True) # Clean up `self._processors` after iterating over it for proc in processors_to_remove: processor = self._processors.pop(proc) + self._finish_callback_attempt(processor) processor.close() def _add_files_to_queue( @@ -2437,6 +2689,13 @@ def max_runs_reached(self): return False if self._pending_inventories: return False + if ( + self._pending_work_acks + or self._callback_claims + or self._callback_attempts + or self._priority_claims + ): + return False if any( file not in self._callback_to_execute for file in [*self._file_queue, *self._files_awaiting_source] @@ -2464,10 +2723,15 @@ def terminate(self): processor.kill(signal.SIGTERM, escalation_delay=5.0) def end(self): - """Kill all child processes on exit since we don't want to leave them as orphaned.""" + """Kill child processes so none are orphaned, then acknowledge work that already finished.""" pids_to_kill = [p.pid for p in self._processors.values()] if pids_to_kill: kill_child_processes_by_pids(pids_to_kill) + # Claims the server releases after this Job ends are delivered again, repeating their callbacks. + try: + self._acknowledge_requested_work(deadline=math.inf) + except Exception: + self.log.exception("Unable to acknowledge requested work before exiting") def emit_metrics(*, parse_time: float, dag_file_stats: Sequence[DagFileStat]): diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 405ad379f4c0c..90307c9f14407 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -40,6 +40,7 @@ ) from airflow.configuration import conf from airflow.dag_processing.bundles.base import BundleVersionLock +from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names from airflow.dag_processing.dagbag import BundleDagBag, DagBag, LangSDKSerializedDAG from airflow.models.dag import DagModel from airflow.sdk.exceptions import TaskNotFound @@ -434,7 +435,11 @@ def _execute_dag_callbacks(dagbag: DagBag, request: DagCallbackRequest, log: Fil { "dag_id": request.dag_id, "team_name": ( - DagModel.get_team_name(request.dag_id) + ( + get_configured_bundle_team_names().get(request.bundle_name) + if conf.get("dag_processor", "execution_api_token_file", fallback=None) + else DagModel.get_team_name(request.dag_id) + ) if conf.getboolean("core", "multi_team") else None ), diff --git a/airflow-core/src/airflow/jobs/dag_processor_job_runner.py b/airflow-core/src/airflow/jobs/dag_processor_job_runner.py index 3b25d32d4a52c..3af40ae3f91be 100644 --- a/airflow-core/src/airflow/jobs/dag_processor_job_runner.py +++ b/airflow-core/src/airflow/jobs/dag_processor_job_runner.py @@ -17,19 +17,31 @@ from __future__ import annotations +from time import monotonic, sleep from typing import TYPE_CHECKING +import httpx2 + from airflow._shared.observability.metrics import stats +from airflow._shared.timezones import timezone +from airflow.configuration import conf +from airflow.dag_processing.health import write_processor_health from airflow.jobs.base_job_runner import BaseJobRunner -from airflow.jobs.job import Job, perform_heartbeat +from airflow.jobs.job import Job, JobState, execute_job, perform_heartbeat +from airflow.listeners.listener import get_listener_manager from airflow.utils.log.logging_mixin import LoggingMixin if TYPE_CHECKING: from sqlalchemy.orm import Session + from airflow.dag_processing.api_client import DagProcessorAPIClient from airflow.dag_processing.manager import DagFileProcessorManager +class DagProcessorHeartbeatTimeout(RuntimeError): + """The processor could not confirm its Job heartbeat before the scheduler may end the Job.""" + + class DagProcessorJobRunner(BaseJobRunner, LoggingMixin): """ DagProcessorJobRunner is a job runner that runs a DagFileProcessorManager processor. @@ -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") + return + if client.restart_required or state != JobState.RUNNING: + raise DagProcessorRegistrationRetired("The API requested a Dag processor restart") + self._last_api_heartbeat = monotonic() + self.job.latest_heartbeat = timezone.utcnow() + self.heartbeat_callback() + def _execute(self) -> int | None: self.log.info("Starting the Dag Processor Job") try: diff --git a/airflow-core/src/airflow/utils/cli.py b/airflow-core/src/airflow/utils/cli.py index 165c26c8b7297..85ad48d156773 100644 --- a/airflow-core/src/airflow/utils/cli.py +++ b/airflow-core/src/airflow/utils/cli.py @@ -60,6 +60,14 @@ def _check_cli_args(args): ) +def set_verbose_logging() -> None: + """Enable debug output for the root logger and its handlers.""" + root_logger = logging.getLogger() + root_logger.setLevel(logging.DEBUG) + for handler in root_logger.handlers: + handler.setLevel(logging.DEBUG) + + def action_cli(func=None, check_db=True): def action_logging(f: T) -> T: """ @@ -97,12 +105,8 @@ def wrapper(*args, **kwargs): _check_cli_args(args) metrics = _build_metrics(f.__name__, args[0]) cli_action_loggers.on_pre_execution(**metrics) - verbose = getattr(args[0], "verbose", False) - root_logger = logging.getLogger() - if verbose: - root_logger.setLevel(logging.DEBUG) - for handler in root_logger.handlers: - handler.setLevel(logging.DEBUG) + if getattr(args[0], "verbose", False): + set_verbose_logging() try: # Check and run migrations if necessary if check_db: diff --git a/airflow-core/tests/integration/dag_processing/test_api_isolation.py b/airflow-core/tests/integration/dag_processing/test_api_isolation.py new file mode 100644 index 0000000000000..6701c7f4a26c9 --- /dev/null +++ b/airflow-core/tests/integration/dag_processing/test_api_isolation.py @@ -0,0 +1,306 @@ +# 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. + +"""Run with Breeze and PostgreSQL; the processor joins an API-only internal network.""" + +from __future__ import annotations + +import json +import secrets +import shutil +import socket +import subprocess +import time +from pathlib import Path +from urllib.parse import urlsplit + +import pytest +import yaml +from sqlalchemy import delete, select +from sqlalchemy.engine import make_url +from uuid6 import uuid7 + +from airflow.callbacks.callback_requests import DagCallbackRequest +from airflow.configuration import conf +from airflow.jobs.job import Job +from airflow.models.callback import DagProcessorCallback +from airflow.models.dag import DagModel +from airflow.models.dagbag import DagPriorityParsingRequest +from airflow.models.dagbundle import DagBundleModel +from airflow.models.dagcode import DagCode +from airflow.models.variable import Variable + +from tests_common.test_utils.db import clear_db_callbacks, clear_db_dag_bundles, clear_db_dags +from unit.api_fastapi.execution_api.conftest import async_db_engine as async_db_engine +from unit.api_fastapi.execution_api.test_dag_processor_client import ( + api_requests as api_requests, + api_url as api_url, + clean_db as clean_db, + freeze_time as freeze_time, + provision_token as provision_token, +) + +pytestmark = [ + pytest.mark.db_test, + pytest.mark.skipif( + not Path("/var/run/docker.sock").exists() or not shutil.which("docker"), + reason="Run in Breeze with its Docker socket mounted", + ), +] + + +@pytest.fixture +def api_bind_host(): + return "0.0.0.0" + + +@pytest.fixture +def api_secret(): + return secrets.token_hex(32) + + +def run_docker(*args, input=None, check=True, timeout=30): + return subprocess.run( + ["docker", *args], + input=input, + capture_output=True, + text=True, + check=check, + timeout=timeout, + ) + + +@pytest.mark.parametrize("shutdown", ["signal", "api_restart"]) +def test_processor_without_database_network_or_server_credentials( + api_url, + provision_token, + session, + tmp_path, + shutdown, +): + database = make_url(conf.get("database", "sql_alchemy_conn")) + if not database.host: + pytest.skip("Requires PostgreSQL or MySQL in Breeze") + database_ip = socket.gethostbyname(database.host) + database_port = database.port or (5432 if database.drivername.startswith("postgres") else 3306) + current = json.loads(run_docker("inspect", socket.gethostname()).stdout)[0] + mounts = {mount["Destination"]: mount["Source"] for mount in current["Mounts"]} + suffix = uuid7().hex + network = f"processor-isolation-{suffix}" + worker = f"processor-isolation-worker-{suffix}" + clear_db_callbacks() + clear_db_dags() + clear_db_dag_bundles() + session.execute(delete(DagPriorityParsingRequest)) + session.add(DagBundleModel(name="bundle-a")) + session.add(DagPriorityParsingRequest(bundle_name="bundle-a", relative_fileloc="a.py")) + callback = DagProcessorCallback( + priority_weight=1, + callback=DagCallbackRequest( + filepath="a.py", + bundle_name="bundle-a", + bundle_version=None, + dag_id="isolated", + run_id="run", + ), + ) + session.add(callback) + session.commit() + callback_id = callback.id + Variable.set("isolation-value", "from-api") + Variable.delete("isolation-callback") + source_code = ( + "from airflow.sdk import DAG, Variable\n" + "def callback(context):\n" + " Variable.set('isolation-callback', 'done')\n" + "dag = DAG('isolated', schedule=None, description=Variable.get('isolation-value'), " + "on_failure_callback=callback)\n" + ) + run_docker("network", "create", "--internal", network) + try: + run_docker("network", "connect", network, current["Id"]) + inspected = json.loads(run_docker("inspect", current["Id"]).stdout)[0] + api_ip = inspected["NetworkSettings"]["Networks"][network]["IPAddress"] + worker_api_url = f"http://{api_ip}:{urlsplit(api_url).port}/execution/" + example = Path(__file__).resolve().parents[3] / "docs/security/dag-processor-api/docker-compose.yaml" + processor_environment = yaml.safe_load(example.read_text())["services"]["dag-processor"][ + "environment" + ] + environment = { + "USER": processor_environment["USER"], + "AIRFLOW_USER_HOME_DIR": "/tmp/airflow-user", + "AIRFLOW_HOME": "/tmp/isolated-airflow", + "AIRFLOW_CONFIG": "/tmp/isolated-airflow/airflow.cfg", + "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN": f"postgresql+psycopg://blocked:blocked@{database_ip}:{database_port}/blocked", + "AIRFLOW__CORE__FERNET_KEY": "", + "AIRFLOW__API_AUTH__JWT_SECRET": "", + "AIRFLOW__API_AUTH__JWT_PRIVATE_KEY_PATH": "", + "AIRFLOW__CORE__LOAD_EXAMPLES": "False", + "AIRFLOW__CORE__EXECUTION_API_SERVER_URL": worker_api_url, + "AIRFLOW__DAG_PROCESSOR__EXECUTION_API_TOKEN_FILE": "/tmp/processor.jwt", + "AIRFLOW__DAG_PROCESSOR__HEALTH_CHECK_THRESHOLD": "5", + "AIRFLOW__SCHEDULER__JOB_HEARTBEAT_SEC": "1", + "CONNECTION_CHECK_MAX_COUNT": "0", + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps( + [ + { + "name": "bundle-a", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/tmp/worker-only-dags"}, + } + ] + ), + } + arguments = [ + "run", + "--detach", + "--name", + worker, + "--network", + network, + "--cap-drop", + "ALL", + "--security-opt", + "no-new-privileges", + "--user", + "54321:0", + "--read-only", + "--tmpfs", + "/tmp:rw,exec,size=256m", + "--entrypoint", + "sleep", + ] + for package in ("airflow-core", "task-sdk", "shared", "providers", "devel-common"): + arguments += [ + "--mount", + f"type=bind,source={mounts[f'/opt/airflow/{package}']},target=/opt/airflow/{package},readonly", + ] + for key, value in environment.items(): + arguments += ["--env", f"{key}={value}"] + arguments += [current["Image"], "600"] + run_docker(*arguments) + run_docker( + "exec", + "-i", + worker, + "sh", + "-c", + "umask 077; cat > /tmp/processor.jwt", + input=provision_token().read_text(), + ) + run_docker( + "exec", + "-i", + worker, + "sh", + "-c", + "mkdir -p /tmp/worker-only-dags; cat > /tmp/worker-only-dags/a.py", + input=source_code, + ) + entrypoint = Path(__file__).resolve().parents[4] / "scripts/docker/entrypoint_prod.sh" + run_docker( + "exec", + "-i", + worker, + "sh", + "-c", + "cat > /tmp/processor-entrypoint.sh", + input=entrypoint.read_text(), + ) + assert not Path("/tmp/worker-only-dags/a.py").exists() + denied = run_docker( + "exec", + worker, + "python", + "-c", + ( + "import socket,sys\n" + f"try: socket.create_connection(({database_ip!r}, {database_port}), timeout=2)\n" + "except OSError: sys.exit(0)\n" + "sys.exit(1)\n" + ), + ) + assert denied.returncode == 0 + with (tmp_path / "processor.log").open("w+") as output: + process = subprocess.Popen( + [ + "docker", + "exec", + worker, + "bash", + "/tmp/processor-entrypoint.sh", + "airflow", + "dag-processor", + "--bundle-name", + "bundle-a", + ], + stdout=output, + stderr=subprocess.STDOUT, + ) + deadline = time.monotonic() + 90 + while time.monotonic() < deadline: + session.expire_all() + dag = session.get(DagModel, "isolated") + delivered = session.get(DagProcessorCallback, callback_id) + if dag and delivered is None: + break + if process.poll() is not None: + output.seek(0) + pytest.fail(output.read()) + time.sleep(0.2) + else: + output.seek(0) + pytest.fail(output.read()) + assert Variable.get("isolation-callback") == "done" + assert session.scalar(select(DagPriorityParsingRequest)) is None + assert ( + session.scalar(select(DagCode.source_code).where(DagCode.dag_id == "isolated")) == source_code + ) + run_docker("exec", worker, "airflow", "dag-processor", "--check-ready") + pause = ( + "import json,os,signal; " + "pid=json.load(open('/tmp/isolated-airflow/dag_processor_health.json'))['pid']; " + "os.kill(pid, signal.SIGSTOP)" + ) + run_docker("exec", worker, "python", "-c", pause) + time.sleep(6) + assert ( + run_docker( + "exec", worker, "airflow", "dag-processor", "--check-health", check=False + ).returncode + != 0 + ) + run_docker("exec", worker, "python", "-c", pause.replace("SIGSTOP", "SIGCONT")) + if shutdown == "signal": + run_docker("exec", worker, "python", "-c", pause.replace("SIGSTOP", "SIGTERM")) + else: + session.scalar(select(Job)).state = "restarting" + session.commit() + assert (process.wait(timeout=30) == 0) == (shutdown == "signal") + session.expire_all() + assert session.scalar(select(Job.state)) == ("success" if shutdown == "signal" else "failed") + finally: + run_docker("rm", "--force", worker, check=False) + run_docker("network", "disconnect", network, current["Id"], check=False) + run_docker("network", "rm", network, check=False) + clear_db_callbacks() + clear_db_dags() + clear_db_dag_bundles() + session.execute(delete(DagPriorityParsingRequest)) + session.commit() + Variable.delete("isolation-value") + Variable.delete("isolation-callback") diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py index 0c7484acffd15..087633dcd1cf4 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py @@ -16,9 +16,13 @@ # under the License. from __future__ import annotations +import json +import os import socket +import subprocess import threading from datetime import UTC, datetime +from pathlib import Path from unittest import mock import httpx2 @@ -32,20 +36,31 @@ from airflow.api_fastapi.app import create_app from airflow.api_fastapi.auth.tokens import JWTGenerator, JWTValidator from airflow.api_fastapi.execution_api.app import lifespan +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import DagParseResultBody from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody, JobState, TerminalJobState +from airflow.configuration import conf from airflow.dag_processing.api_client import ( DagParseContext, DagProcessorAPIClient, DagProcessorRegistrationRetired, ) +from airflow.dag_processing.bundles.local import LocalDagBundle from airflow.jobs.job import Job +from airflow.models.dag import DagModel +from airflow.models.dag_parse_checkpoint import DagParseCheckpoint from airflow.models.dagbundle import DagBundleModel +from airflow.models.serialized_dag import SerializedDagModel +from airflow.models.team import Team from airflow.models.variable import Variable +from airflow.sdk import Variable as SDKVariable from airflow.sdk.api.client import Client from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import ( + clear_db_dag_bundles, + clear_db_dags, clear_db_jobs, + clear_db_teams, clear_db_variables, ) @@ -218,7 +233,7 @@ def __init__(self): def handle_request(self, request): response = self.transport.handle_request(request) if ( - request.url.path.endswith(("/jobs", "/complete", "/parse-token")) + request.url.path.endswith(("/jobs", "/complete", "/parse-token", "/parse-results")) and request.url.path not in self.lost and response.is_success ): @@ -247,16 +262,38 @@ def test_lost_registration_and_completion_responses_over_http(api_url, provision request=DagParseTokenBody(attempt_id=uuid7(), bundle_name="bundle-a", relative_fileloc="dag.py") ) with processor.use_parse(context): + with pytest.raises(httpx2.ReadError): + processor.variables.get("parse-token-key") assert processor.variables.get("parse-token-key").value == "value" with pytest.raises(httpx2.HTTPStatusError) as error: processor.post(f"jobs/{job_id}/heartbeat") assert error.value.response.status_code == 403 + with pytest.raises(httpx2.HTTPStatusError) as publication_error: + processor.post(f"jobs/{job_id}/parse-results", json={}) + assert publication_error.value.response.status_code == 403 + body = DagParseResultBody( + attempt_id=uuid7(), + dispatch_sequence=1, + bundle_name="bundle-a", + relative_fileloc="dag.py", + parse_duration=0.1, + serialized_dags=[], + source_codes={}, + ) + with pytest.raises(httpx2.ReadError): + processor.publish_parse_result(body) + processor.heartbeat() + receipt = processor.publish_parse_result(body) + assert receipt.attempt_id == body.attempt_id + assert session.scalar(select(DagParseCheckpoint)).attempt_id == body.attempt_id + assert processor.publish_parse_result(body) == receipt processor.complete_job(TerminalJobState.SUCCESS) assert session.get(Job, job_id).state == JobState.SUCCESS.value assert transport.lost == { "/execution/jobs", f"/execution/jobs/{job_id}/parse-token", f"/execution/jobs/{job_id}/complete", + f"/execution/jobs/{job_id}/parse-results", } @@ -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") + + +@pytest.mark.parametrize("multi_team", [False, True]) +@pytest.mark.parametrize("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:blocked@127.0.0.1:1/blocked", + ("core", "fernet_key"): "", + ("api_auth", "jwt_secret"): "", + ("api_auth", "jwt_private_key_path"): "", + } + ) + result = subprocess.run( + ["airflow", "dag-processor", "--num-runs", "1", "--bundle-name", "bundle-a"], + env={ + **os.environ, + **{ + f"AIRFLOW__{section.upper()}__{key.upper()}": value + for (section, key), value in processor_config.items() + }, + "PYTHONPATH": f"{Path(__file__).parents[3]}:{os.environ.get('PYTHONPATH', '')}", + }, + capture_output=True, + text=True, + timeout=90, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + + session.expire_all() + job = session.scalars(select(Job)).one() + assert job.state == JobState.SUCCESS + assert job.end_date is not None + assert job.team_names == ([team_name] if team_name else []) + assert session.get(DagModel, "authenticated_processor").description == "from-api" + assert ( + session.scalar( + select(SerializedDagModel).where(SerializedDagModel.dag_id == "authenticated_processor") + ) + is not None + ) + requests = {(path, claims.get("scope")) for path, claims in api_requests} + assert ("/execution/jobs", "dag_processor_session") in requests + assert (f"/execution/jobs/{job.id}/heartbeat", "dag_processor") in requests + assert (f"/execution/jobs/{job.id}/complete", "dag_processor") in requests + assert (f"/execution/jobs/{job.id}/parse-results", "dag_processor") in requests + assert ("/execution/variables/processor-client-key", "dag_processor") in requests + assert ("/execution/variables/processor-client-key", "dag_parse") in requests + parsing = [claims for _, claims in api_requests if claims.get("scope") == "dag_parse"] + assert all(claims["relative_fileloc"] == "api_dag.py" for claims in parsing) + finally: + clear_db_jobs() + clear_db_dags() + clear_db_dag_bundles() + clear_db_variables() + clear_db_teams() diff --git a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py index 77de6f928d33e..b0a69dd9378f5 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +import logging import os from unittest import mock @@ -24,6 +25,9 @@ from airflow.cli import cli_parser from airflow.cli.commands import dag_processor_command +from airflow.dag_processing.manager import DagFileProcessorManager +from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner +from airflow.jobs.job import Job from tests_common.test_utils.config import conf_vars @@ -166,3 +170,135 @@ def test_memray_traces_the_job_and_not_the_parent_process(self, mock_runner, moc mock_daemon_option.call_args.kwargs["callback"]() memray.Tracker.assert_called_once() + + @conf_vars({("dag_processor", "execution_api_token_file"): "/run/processor.jwt"}) + @mock.patch("airflow.cli.commands.dag_processor_command.run_command_with_daemon_option", autospec=True) + @mock.patch("airflow.dag_processing.api_client.DagProcessorAPIClient", autospec=True) + @mock.patch.object(DagProcessorJobRunner, "run_with_api", autospec=True) + @mock.patch("airflow.cli.commands.dag_processor_command.run_job", autospec=True) + def test_api_client_is_created_in_the_daemon_callback( + self, run_job, run_with_api, client_class, daemon, configure_testing_dag_bundle + ): + with configure_testing_dag_bundle(os.devnull): + args = self.parser.parse_args(["dag-processor", "--bundle-name", "testing"]) + dag_processor_command.dag_processor(args) + client_class.assert_not_called() + daemon.call_args.kwargs["callback"]() + assert client_class.call_args.kwargs["token_file"] == "/run/processor.jwt" + assert client_class.call_args.kwargs["bundle_names"] == ["testing"] + run_with_api.assert_called_once() + assert run_with_api.call_args.args[1] is client_class.return_value.__enter__.return_value + client_class.return_value.__exit__.assert_called_once() + run_job.assert_not_called() + + @conf_vars({("dag_processor", "execution_api_token_file"): "/run/processor.jwt"}) + @mock.patch("airflow.utils.cli.validate_dag_bundle_arg", autospec=True) + def test_api_bundle_validation_does_not_construct_bundles(self, validate, configure_testing_dag_bundle): + with configure_testing_dag_bundle(os.devnull): + runner = dag_processor_command._create_dag_processor_job_runner( + self.parser.parse_args(["dag-processor", "--bundle-name", "testing"]) + ) + assert runner.processor.bundle_names_to_parse == ["testing"] + with pytest.raises(SystemExit, match="Bundles not found: missing"): + dag_processor_command._create_dag_processor_job_runner( + self.parser.parse_args(["dag-processor", "--bundle-name", "missing"]) + ) + validate.assert_not_called() + + @conf_vars( + { + ("dag_processor", "execution_api_token_file"): "/run/processor.jwt", + ("core", "load_examples"): "False", + } + ) + @mock.patch("airflow.dag_processing.api_client.DagProcessorAPIClient", autospec=True) + def test_api_client_closes_on_failure(self, client_class, configure_testing_dag_bundle): + runner = mock.create_autospec(DagProcessorJobRunner, instance=True) + runner.job = Job() + runner.processor = mock.create_autospec(DagFileProcessorManager, instance=True) + runner.processor.bundle_names_to_parse = None + runner.run_with_api.side_effect = RuntimeError("failed") + with configure_testing_dag_bundle(os.devnull), pytest.raises(RuntimeError, match="failed"): + dag_processor_command._run_dag_processor_job(runner) + assert client_class.call_args.kwargs["bundle_names"] == ["testing"] + client_class.return_value.__exit__.assert_called_once() + + +@pytest.mark.parametrize("api_mode", [False, True]) +@pytest.mark.parametrize("verbose", [False, True]) +@mock.patch("airflow.cli.commands.dag_processor_command._start_dag_processor", autospec=True) +@mock.patch("airflow.utils.cli.cli_action_loggers.on_pre_execution", autospec=True) +@mock.patch("airflow.utils.cli.cli_action_loggers.on_post_execution", autospec=True) +@mock.patch("airflow.utils.db.check_and_run_migrations", autospec=True) +@mock.patch("airflow.utils.db.synchronize_log_template", autospec=True) +def test_verbose_logging_in_both_modes( + sync_log_template, + check_migrations, + on_post_execution, + on_pre_execution, + start, + api_mode, + verbose, + monkeypatch, +): + root_logger = logging.getLogger() + original_level = root_logger.level + handler = logging.NullHandler(level=logging.INFO) + monkeypatch.setattr(root_logger, "handlers", [handler]) + root_logger.setLevel(logging.INFO) + try: + with conf_vars( + {("dag_processor", "execution_api_token_file"): "/run/processor.jwt" if api_mode else ""} + ): + args = cli_parser.get_parser().parse_args(["dag-processor", *(["--verbose"] if verbose else [])]) + dag_processor_command.dag_processor(args) + + expected_level = logging.DEBUG if verbose else logging.INFO + assert root_logger.level == expected_level + assert handler.level == expected_level + start.assert_called_once_with(args) + if api_mode: + on_pre_execution.assert_not_called() + on_post_execution.assert_not_called() + check_migrations.assert_not_called() + sync_log_template.assert_not_called() + else: + on_pre_execution.assert_called_once() + on_post_execution.assert_called_once() + finally: + root_logger.setLevel(original_level) + + +@mock.patch.object(DagFileProcessorManager, "sync_bundles", autospec=True) +@mock.patch.object(DagProcessorJobRunner, "_execute", autospec=True) +def test_catalog_provisioning_does_not_start_a_processor(execute, sync): + args = cli_parser.get_parser().parse_args(["dag-processor", "--sync-bundles-only"]) + dag_processor_command.dag_processor(args) + sync.assert_called_once() + execute.assert_not_called() + + +@pytest.mark.parametrize("command", [["dag-processor"], ["dag-processor", "--sync-bundles-only"]]) +@mock.patch("airflow.utils.cli.cli_action_loggers.on_pre_execution", autospec=True) +@mock.patch.object(DagFileProcessorManager, "sync_bundles", autospec=True) +@mock.patch("airflow.cli.commands.dag_processor_command.DagProcessorJobRunner") +def test_action_log_keeps_the_command_name(mock_runner, sync, on_pre_execution, command): + mock_runner.return_value.job_type = "DagProcessorJob" + dag_processor_command.dag_processor(cli_parser.get_parser().parse_args(command)) + assert on_pre_execution.call_args.kwargs["sub_command"] == "dag_processor" + + +@mock.patch.object(DagFileProcessorManager, "sync_bundles", autospec=True) +def test_catalog_provisioning_for_named_bundles_keeps_other_bundles(sync, configure_testing_dag_bundle): + parser = cli_parser.get_parser() + with configure_testing_dag_bundle(os.devnull): + dag_processor_command.dag_processor( + parser.parse_args(["dag-processor", "--sync-bundles-only", "--bundle-name", "testing"]) + ) + with pytest.raises(SystemExit, match="Bundles not found: unknown"): + dag_processor_command.dag_processor( + parser.parse_args(["dag-processor", "--sync-bundles-only", "--bundle-name", "unknown"]) + ) + manager = sync.call_args.args[0] + assert manager.bundle_names_to_parse == ["testing"] + sync.assert_called_once() diff --git a/airflow-core/tests/unit/dag_processing/test_health.py b/airflow-core/tests/unit/dag_processing/test_health.py new file mode 100644 index 0000000000000..ffd118fdb31e7 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_health.py @@ -0,0 +1,64 @@ +# 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 pytest + +from airflow.dag_processing.health import check_processor_health, write_processor_health + +from tests_common.test_utils.config import conf_vars + + +@pytest.mark.parametrize( + ("loop_age", "api_age", "readiness", "healthy"), + [ + (0, 0, False, True), + (31, 0, False, False), + (0, 31, False, True), + (0, 31, True, False), + (0, 0, True, True), + (-1, 0, False, False), + ], +) +@mock.patch("airflow.dag_processing.health.monotonic", autospec=True) +def test_local_health_requires_a_fresh_loop(clock, tmp_path, loop_age, api_age, readiness, healthy): + with conf_vars( + { + ("dag_processor", "health_check_file"): str(tmp_path / "health.json"), + ("dag_processor", "health_check_threshold"): "30", + } + ): + clock.return_value = 100 - loop_age + write_processor_health(100 - api_age) + clock.return_value = 100 + if healthy: + check_processor_health(readiness=readiness) + else: + with pytest.raises(SystemExit): + check_processor_health(readiness=readiness) + + +@pytest.mark.parametrize("contents", [None, "broken", "{}"]) +def test_missing_or_invalid_health_file_fails_closed(tmp_path, contents): + path = tmp_path / "health.json" + if contents is not None: + path.write_text(contents) + with conf_vars({("dag_processor", "health_check_file"): str(path)}), pytest.raises(SystemExit): + check_processor_health() diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index d958c942f2d7b..61dc6f3b34df3 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -59,11 +59,13 @@ DagBundleStateResponse, DagParseResultBody, ParseSourceCode, + ProcessorWorkItem, ) from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance as TIDataModel from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.dag_processing.api_client import ( DagProcessorAPIClient, + DagProcessorRegistrationRetired, DagProcessorSecretsComms, ) from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion, DagBundleRefreshError @@ -79,8 +81,10 @@ DagFileInfo, DagFileProcessorManager, DagFileStat, + DeferredCallback, PendingBundleInventory, PendingDagPublication, + PriorityParseClaim, ) from airflow.dag_processing.processor import ( DagFileParseRequest, @@ -499,10 +503,11 @@ 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"]) @pytest.mark.parametrize("absolute", [False, True], ids=["relative", "absolute"]) @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) def test_find_files_in_bundle_makes_discovery_errors_bundle_relative( - self, mock_get_registry, tmp_path, absolute + self, mock_get_registry, tmp_path, absolute, api_mode ): (tmp_path / "broken.dag").write_text("") source_reference = os.fspath(tmp_path / "broken.dag") if absolute else "broken.dag" @@ -510,8 +515,9 @@ def test_find_files_in_bundle_makes_discovery_errors_bundle_relative( (mock.sentinel.importer, DagImportError(source_reference=source_reference, message="boom")) ] bundle = _make_bundle(tmp_path) + client = mock.create_autospec(DagProcessorAPIClient, instance=True) if api_mode else None - found_files = DagFileProcessorManager(max_runs=1)._find_files_in_bundle(bundle) + found_files = DagFileProcessorManager(max_runs=1, api_client=client)._find_files_in_bundle(bundle) assert [file.definition_locs for file in found_files.files] == [frozenset({"broken.dag"})] assert found_files.failed_sources == frozenset({"broken.dag"}) @@ -4026,6 +4032,7 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): with mock.patch.object( DagFileProcessorProcess, "start", side_effect=lambda *args, **kwargs: self.mock_processor()[0] ) as start: + manager._bundle_parse_sources["testing"] = DagParseSource() manager._start_new_processes() # Callbacks passed to processor assert start.call_args_list == [ @@ -5564,7 +5571,7 @@ def test_unpack_bundle_version_with_string_non_versioned_bundle(self): @pytest.mark.usefixtures("testing_dag_bundle") def test_bundle_version_data_stored_after_refresh(self, session): - """Capture version metadata alongside the refreshed source.""" + """Keep the local source metadata independently of its persisted bundle state.""" from airflow.dag_processing.bundles.base import BundleVersion manager = DagFileProcessorManager(max_runs=1) @@ -5971,9 +5978,59 @@ def test_pending_publications_bound_admission_and_prevent_reimports(self, create read_end.close() proc.close() + def test_requested_work_is_readmitted_ahead_of_queued_files(self): + manager = DagFileProcessorManager(max_runs=1) + requested = DagFileInfo(bundle_name="testing", rel_path=Path("requested.py")) + parked = DagFileInfo(bundle_name="testing", rel_path=Path("parked.py")) + queued = DagFileInfo(bundle_name="other", rel_path=Path("queued.py")) + _claim_priority_parse(manager, requested) + manager._files_awaiting_source = dict.fromkeys([parked, requested]) + manager._file_queue[queued] = None + manager._bundle_parse_sources["testing"] = DagParseSource() + + manager._readmit_files_with_source() + + assert list(manager._file_queue) == [requested, parked, queued] + + @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) + def test_requested_file_still_parsing_stays_ahead_of_the_queue(self, create): + manager = DagFileProcessorManager(max_runs=1) + manager._parallelism = 2 + requested, other, later = _get_file_infos(["requested.py", "other.py", "later.py"]) + _claim_priority_parse(manager, requested) + manager._bundle_parse_sources["testing"] = DagParseSource() + manager._processors[requested] = self.mock_processor()[0] + manager._file_queue = OrderedDict.fromkeys([requested, other, later]) + create.return_value = self.mock_processor()[0] + + manager._start_new_processes() + + assert other in manager._processors + assert list(manager._file_queue) == [requested, later] + + @mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True) + def test_queuing_a_callback_heartbeats_before_preparing_its_bundle(self, mock_prepare_callback_bundle): + events = [] + manager = DagFileProcessorManager(max_runs=1) + manager.heartbeat = lambda: events.append("heartbeat") + mock_prepare_callback_bundle.side_effect = lambda manager, request: events.append("prepare") + request = DagCallbackRequest( + filepath="file1.py", + dag_id="dag1", + run_id="run1", + is_failure_callback=False, + bundle_name="testing", + bundle_version="some_commit_hash", + msg=None, + ) + + manager._add_callback_to_queue(request) + + assert events == ["heartbeat", "prepare"] + @pytest.mark.parametrize("failure", ["invalid_path", "too_many_files"]) @mock.patch.object(DagFileProcessorManager, "_get_observed_filelocs", autospec=True) - def test_unusable_inventory_preserves_version_and_recovers_after_refresh(self, observed, failure): + def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed, failure): observed.return_value = ( {r"bad\dag.py"} if failure == "invalid_path" else {f"{index}.py" for index in range(100001)} ) @@ -5982,6 +6039,10 @@ def test_unusable_inventory_preserves_version_and_recovers_after_refresh(self, o max_runs=1, api_client=client, parallelism=1, bundle_refresh_check_interval=0 ) bundle = self._make_refresh_bundle() + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py" + ) + manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1) manager._bundles_waiting_for_refresh.add(bundle.name) manager._bundle_versions[bundle.name] = "previous" @@ -5996,6 +6057,12 @@ def test_unusable_inventory_preserves_version_and_recovers_after_refresh(self, o assert manager._bundle_versions[bundle.name] == "previous" assert not manager._force_refresh_bundles assert not manager._bundles_waiting_for_refresh + assert not manager._priority_claims + assert manager._pending_work_acks == [("priority", work, True)] + manager._acknowledge_requested_work() + client.claim_work.return_value = [] + manager.claim_priority_files() + client.claim_work.assert_called_once_with("priority", [bundle.name], 1) observed.return_value = set() client.publish_inventory.return_value.revision = uuid7() @@ -6004,6 +6071,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 ( + mock.patch.object( + manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None) + ), + mock.patch.object(manager, "update_bundle_state"), + mock.patch.object(manager, "_find_files_in_bundle", return_value=set()), + mock.patch.object(manager, "deactivate_deleted_dags"), + mock.patch.object(manager, "clear_orphaned_import_errors"), + mock.patch.object(manager, "handle_removed_files"), + mock.patch.object(manager, "_resort_file_queue"), + mock.patch.object(manager, "_add_new_files_to_queue"), + ): + manager._refresh_dag_bundles({}) + + assert events == ["heartbeat", "first", "heartbeat", "second"] + class TestMultiTeamMetrics: """Tests for team_name tag on dag processing metrics in multi-team mode.""" @@ -6114,13 +6209,13 @@ def test_log_file_processing_stats_omits_team_name_when_not_multi_team( @mock.patch("airflow.dag_processing.manager.stats.incr") def test_start_new_processes_includes_team_name(self, mock_incr, mock_get_team_name): manager = DagFileProcessorManager(max_runs=1) + manager._bundle_parse_sources["testing"] = DagParseSource() dag_file = DagFileInfo( bundle_name="testing", rel_path=Path("dag_file.py"), bundle_path=TEST_DAGS_FOLDER, ) manager._file_queue[dag_file] = None - manager._bundle_parse_sources["testing"] = DagParseSource() manager._create_process = MagicMock(return_value=self.mock_processor()) manager._start_new_processes() @@ -6616,6 +6711,247 @@ def create(_, name, version=None, version_data=None): assert make_manager().prepare_callback_bundle(request) is not None assert self._get_connection_bundles(api_requests) == ["bundle_b"] + @mock.patch("airflow.sdk.api.client.ConnectionOperations.get", autospec=True) + @mock.patch.object(DagBundlesManager, "get_bundle", autospec=True) + def test_git_callback_connection_outage_is_deferred(self, get_bundle, get_connection, make_manager): + git_bundle = pytest.importorskip("airflow.providers.git.bundles.git").GitDagBundle + get_bundle.side_effect = lambda _, **kwargs: git_bundle(tracking_ref="main", **kwargs) + get_connection.side_effect = httpx2.ConnectError("API unavailable") + manager = make_manager() + request = DagCallbackRequest( + filepath="a.py", bundle_name="bundle_b", bundle_version="v1", dag_id="dag", run_id="run" + ) + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="bundle_b", relative_fileloc="a.py" + ) + manager._callback_claims[id(request)] = work + manager._deferred_api_callbacks.append(request) + + manager._add_callback_to_queue(request) + + get_connection.assert_called_once() + assert id(request) in manager._deferred_callback_retries + assert not manager._pending_work_acks + + +@pytest.mark.parametrize( + ("failure", "stage"), + [ + ("transport", "construct"), + ("server", "construct"), + ("retired", "construct"), + ("wrapped_retired", "construct"), + ("retired", "initialize"), + ], +) +@mock.patch.object(DagFileProcessorManager, "_create_bundle_manager", autospec=True) +def test_callback_construction_failure_uses_bounded_retry(create, failure, stage): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.restart_required = failure == "wrapped_retired" + manager = DagFileProcessorManager(max_runs=1, api_client=client) + request = DagCallbackRequest( + filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run" + ) + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py" + ) + manager._callback_claims[id(request)] = work + manager._deferred_api_callbacks.append(request) + response = httpx2.Response(503, request=httpx2.Request("GET", "http://api/connections/git")) + errors = { + "transport": httpx2.ConnectError("API unavailable"), + "server": httpx2.HTTPStatusError("API unavailable", request=response.request, response=response), + "retired": DagProcessorRegistrationRetired("retired"), + "wrapped_retired": RuntimeError("provider wrapped a retired registration"), + } + if stage == "construct": + create.return_value.get_bundle.side_effect = errors[failure] + else: + create.return_value.get_bundle.return_value.supports_versioning = True + create.return_value.get_bundle.return_value.initialize.side_effect = errors[failure] + if "retired" in failure: + with pytest.raises(DagProcessorRegistrationRetired): + manager._add_callback_to_queue(request) + assert not manager._deferred_callback_retries + return + + manager._add_callback_to_queue(request) + assert id(request) in manager._deferred_callback_retries + assert not manager._pending_work_acks + manager._deferred_callback_retries[id(request)].next_attempt_time = 0 + create.return_value.get_bundle.side_effect = None + create.return_value.get_bundle.return_value = _make_bundle(TEST_DAGS_FOLDER) + manager._add_callback_to_queue(request) + assert not manager._deferred_callback_retries + assert not manager._deferred_api_callbacks + assert list(manager._callback_to_execute.values()) == [[request]] + + +@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"}) +@mock.patch("airflow.dag_processing.manager.time", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_create_bundle_manager", autospec=True) +def test_callback_constructor_outage_allows_other_work_until_retry_expires( + create_bundle_manager, create_process, mock_time, tmp_path +): + mock_time.monotonic.return_value = 0 + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.restart_required = False + manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=2) + bundle = _make_bundle(tmp_path) + bundle.is_initialized = True + manager._dag_bundles = [bundle] + manager._bundle_parse_sources[bundle.name] = DagParseSource() + + def get_bundle(*, name, version, version_data): + if version == "unavailable": + raise httpx2.ConnectError("Bundle connection lookup unavailable") + return bundle + + create_bundle_manager.return_value.get_bundle.side_effect = get_bundle + work = [] + for version in ("unavailable", "available"): + request = DagCallbackRequest( + filepath=f"{version}.py", + bundle_name=bundle.name, + bundle_version=version, + dag_id=version, + run_id="run", + ) + work.append( + ProcessorWorkItem( + id=str(uuid7()), + claim_id=uuid7(), + bundle_name=bundle.name, + relative_fileloc=request.filepath, + callback=request.model_dump_json(), + ) + ) + client.claim_work.return_value = work + for request in manager.fetch_callbacks(): + manager._add_callback_to_queue(request) + assert len(manager._deferred_api_callbacks) == 1 + normal = DagFileInfo(bundle_name=bundle.name, rel_path=Path("normal.py"), bundle_path=tmp_path) + manager._add_files_to_queue([normal], mode="back") + + def create(manager, file): + proc = mock.create_autospec(DagFileProcessorProcess, instance=True) + proc.id = uuid7() + proc.had_callbacks = file.bundle_version is not None + return proc + + create_process.side_effect = create + manager._start_new_processes() + assert {file.rel_path for file in manager._processors} == {Path("available.py"), Path("normal.py")} + assert not manager._pending_work_acks + + mock_time.monotonic.return_value = 10 + for request in list(manager._deferred_api_callbacks): + manager._add_callback_to_queue(request) + assert not manager._deferred_api_callbacks + assert not manager._deferred_callback_retries + assert manager._pending_work_acks == [("callbacks", work[0], True)] + manager._acknowledge_requested_work() + manager._acknowledge_requested_work() + client.acknowledge_work.assert_called_once_with("callbacks", work[0], failed=True) + + +@pytest.mark.parametrize("outcome", ["queued", "expired", "unsupported"]) +@mock.patch("airflow.dag_processing.manager.time", autospec=True) +@mock.patch("airflow.dag_processing.manager.get_claiming_importer", autospec=True) +@mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True) +def test_equal_claimed_callbacks_are_removed_by_identity(prepare, importer, clock, outcome): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + first = DagCallbackRequest( + filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run" + ) + second = first.model_copy(deep=True) + work = [ + ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py") + for _ in range(2) + ] + manager._callback_claims = {id(first): work[0], id(second): work[1]} + manager._deferred_api_callbacks = [first, second] + manager._deferred_callback_retries = { + id(first): DeferredCallback(expires_at=20, next_attempt_time=5), + id(second): DeferredCallback(expires_at=0), + } + clock.monotonic.return_value = 0 + importer.return_value = object() if outcome == "unsupported" else None + prepare.return_value = None if outcome == "expired" else _make_bundle(TEST_DAGS_FOLDER) + + manager._add_callback_to_queue(second) + + assert len(manager._deferred_api_callbacks) == 1 + assert manager._deferred_api_callbacks[0] is first + assert set(manager._deferred_callback_retries) == {id(first)} + assert manager._callback_claims[id(first)] is work[0] + if outcome == "queued": + assert next(iter(manager._callback_to_execute.values()))[0] is second + assert manager._callback_claims[id(second)] is work[1] + assert not manager._pending_work_acks + else: + assert id(second) not in manager._callback_claims + assert manager._pending_work_acks == [("callbacks", work[1], True)] + + +@mock.patch("airflow.dag_processing.manager.get_claiming_importer", autospec=True, return_value=None) +@mock.patch.object(DagFileProcessorProcess, "start", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True) +@mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True) +def test_equal_claimed_callbacks_dispatch_once_each_after_preparation_recovers( + prepare, logger, start, importer +): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.claim_work.return_value = [] + manager = DagFileProcessorManager(max_runs=1, api_client=client) + first = DagCallbackRequest( + filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run" + ) + second = first.model_copy(deep=True) + work = [ + ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py") + for _ in range(2) + ] + manager._callback_claims = {id(first): work[0], id(second): work[1]} + manager._deferred_api_callbacks = [first, second] + bundle = _make_bundle(TEST_DAGS_FOLDER) + prepare.side_effect = [None, bundle, bundle] + logger.return_value = (mock.MagicMock(spec=structlog.stdlib.BoundLogger), mock.mock_open()()) + + def create_process(**kwargs): + process = mock.create_autospec(DagFileProcessorProcess, instance=True) + process.id = kwargs["id"] + process._exit_code = 0 + return process + + start.side_effect = create_process + for request in manager.fetch_callbacks(): + manager._add_callback_to_queue(request) + + for expected in (second, first): + file = next(iter(manager._callback_to_execute)) + process = manager._create_process(file) + assert start.call_args.kwargs["callbacks"][0] is expected + manager._finish_callback_attempt(process) + manager._acknowledge_requested_work() + if expected is second: + manager._deferred_callback_retries[id(first)].next_attempt_time = 0 + for request in manager.fetch_callbacks(): + manager._add_callback_to_queue(request) + + assert start.call_count == 2 + assert manager.fetch_callbacks() == [] + assert not manager._callback_claims + assert not manager._callback_attempts + assert not manager._deferred_callback_retries + assert not manager._callback_to_execute + assert client.acknowledge_work.call_args_list == [ + mock.call("callbacks", work[1], failed=False), + mock.call("callbacks", work[0], failed=False), + ] + @mock.patch.object( DagFileProcessorManager, @@ -6679,7 +7015,7 @@ def test_refresh_reads_one_fresh_catalog_per_pass(find_files, tmp_path): @pytest.mark.parametrize("has_pending_inventory", [False, True]) @mock.patch("airflow.dag_processing.manager.time", autospec=True) @mock.patch.object(DagFileProcessorManager, "_find_files_in_bundle", autospec=True) -def test_catalog_omission_releases_source_and_allows_recovery( +def test_catalog_omission_releases_priority_capacity_and_allows_recovery( find_files, clock, catalog_failed, has_pending_inventory, tmp_path ): clock.monotonic.return_value = 100 @@ -6688,6 +7024,7 @@ def test_catalog_omission_releases_source_and_allows_recovery( manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=10) bundles = {} files = {} + claims = {} for name in ("missing", "ready"): bundle = _make_bundle(tmp_path) bundle.name = name @@ -6699,6 +7036,9 @@ def test_catalog_omission_releases_source_and_allows_recovery( files[name] = {file} manager._bundle_versions[name] = None manager._bundle_parse_sources[name] = DagParseSource() + work = ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="a.py") + claims[name] = PriorityParseClaim(work=work, refresh_generation=1) + manager._priority_claims[work.id] = claims[name] manager._dag_bundles = list(bundles.values()) previous_files = files.copy() previous_versions = manager._bundle_versions.copy() @@ -6729,7 +7069,10 @@ def test_catalog_omission_releases_source_and_allows_recovery( assert files == previous_files assert manager._bundle_versions == previous_versions assert manager._bundle_parse_sources["ready"] is ready_source + assert manager._priority_claims[claims["ready"].work.id] is claims["ready"] if catalog_failed: + assert manager._priority_claims[claims["missing"].work.id] is claims["missing"] + assert not manager._pending_work_acks assert manager._force_refresh_bundles == {"missing"} assert manager._bundles_waiting_for_refresh == {"missing"} if has_pending_inventory: @@ -6741,12 +7084,18 @@ def test_catalog_omission_releases_source_and_allows_recovery( return client.publish_inventory.assert_not_called() + assert claims["missing"].work.id not in manager._priority_claims + assert manager._pending_work_acks == [("priority", claims["missing"].work, True)] assert not manager._pending_inventories assert not manager._force_refresh_bundles assert not manager._bundles_waiting_for_refresh assert manager._get_dispatch_source(next(iter(files["missing"]))) is None manager._refresh_dag_bundles(files) client.get_bundles.assert_called_once_with() + manager._acknowledge_requested_work() + client.claim_work.return_value = [] + manager.claim_priority_files() + client.claim_work.assert_called_once_with("priority", ["missing", "ready"], 1) clock.monotonic.return_value = 110 client.get_bundles.return_value = [ready_state, ready_state.model_copy(update={"name": "missing"})] @@ -6936,6 +7285,108 @@ def publish(bundle_name, body): assert client.get_bundles.call_count == calls +@pytest.mark.parametrize("rejection", ["local_limit", "server_limit"]) +@pytest.mark.parametrize("retained_attempt", [None, "processor", "publication"]) +@mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 1024) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_oversized_inventory_releases_unassigned_priority_capacity( + create, retained_attempt, rejection, tmp_path +): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=3 if retained_attempt else 1) + for name in ("oversized", "ready"): + bundle = _make_bundle(tmp_path) + bundle.name = name + manager._dag_bundles.append(bundle) + (tmp_path / "a.py").touch() + rejected_work, ready_work = [ + ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="a.py") + for name in ("oversized", "ready") + ] + retained_claims = {} + if retained_attempt: + for name, attempt_id in (("oversized", uuid7()), ("ready", None)): + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="b.py" + ) + retained_claims[work.id] = PriorityParseClaim( + work=work, refresh_generation=1, attempt_id=attempt_id + ) + if name == "oversized": + active_file = DagFileInfo(bundle_name=name, rel_path=Path("b.py"), bundle_path=tmp_path) + if retained_attempt == "processor": + process = mock.create_autospec(DagFileProcessorProcess, instance=True) + process.id = attempt_id + manager._processors[active_file] = process + else: + pending_body = mock.create_autospec(DagParseResultBody, instance=True) + pending_body.attempt_id = attempt_id + manager._pending_publications[active_file] = PendingDagPublication( + body=pending_body, stat=DagFileStat(), refresh_generation=1 + ) + manager._priority_claims.update(retained_claims) + client.claim_work.side_effect = [[rejected_work], [ready_work]] + manager._queue_requested_files_for_parsing() + assert manager._priority_claims[rejected_work.id].attempt_id is None + body = DagBundleInventoryBody( + attempt_id=uuid7(), + dispatch_sequence=1, + expected_revision=None, + files=[f"{index}/{'é' * 80}.py" for index in range(6)] if rejection == "local_limit" else [], + ) + manager._pending_inventories["oversized"] = PendingBundleInventory( + body, set(), DagParseSource(refresh_generation=1) + ) + client.publish_inventory.side_effect = _make_status_error(413) + + manager._publish_bundle_inventory("oversized", {}) + + assert manager._priority_claims == retained_claims + assert manager._pending_work_acks == [("priority", rejected_work, True)] + if rejection == "local_limit": + client.publish_inventory.assert_not_called() + else: + client.publish_inventory.assert_called_once_with("oversized", body) + assert manager.claim_priority_files() == [] + client.claim_work.assert_called_once_with("priority", ["oversized", "ready"], 1) + manager._acknowledge_requested_work() + client.acknowledge_work.assert_called_once_with("priority", rejected_work, failed=True) + manager._queue_requested_files_for_parsing() + assert client.claim_work.call_args_list == [ + mock.call("priority", ["oversized", "ready"], 1), + mock.call("priority", ["oversized", "ready"], 1), + ] + manager._force_refresh_bundles.remove("ready") + manager._invalidate_bundle_parse_source("ready") + ready_file = DagFileInfo(bundle_name="ready", rel_path=Path("a.py"), bundle_path=tmp_path) + body = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=2, expected_revision=None, files=["a.py", "b.py"] + ) + manager._pending_inventories["ready"] = PendingBundleInventory( + body, + {ready_file}, + DagParseSource(refresh_generation=1), + ) + client.publish_inventory.side_effect = None + client.publish_inventory.return_value = DagBundleInventoryResponse( + attempt_id=body.attempt_id, accepted_at=timezone.utcnow(), revision=uuid7() + ) + manager._publish_bundle_inventory("ready", {}) + process = mock.create_autospec(DagFileProcessorProcess, instance=True) + process.id = uuid7() + process.had_callbacks = False + create.return_value = process + + manager._start_new_processes() + + create.assert_called_once_with(manager, ready_file) + assert manager._priority_claims[ready_work.id].attempt_id == process.id + assert {key: claim for key, claim in manager._priority_claims.items() if key != ready_work.id} == ( + retained_claims + ) + assert not manager._pending_work_acks + + @pytest.mark.parametrize("auth_status", [401, 403]) def test_oversized_inventory_byte_boundary_preserves_auth_failures(auth_status): client = mock.create_autospec(DagProcessorAPIClient, instance=True) @@ -6964,14 +7415,320 @@ def test_oversized_inventory_byte_boundary_preserves_auth_failures(auth_status): manager._publish_bundle_inventory("testing", {}) +@pytest.mark.parametrize("status_code", [403, 404, 409, 503]) +def test_requested_work_ack_retains_only_retryable_failures(status_code): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py" + ) + manager._pending_work_acks.append(("callbacks", work, False)) + response = httpx2.Response(status_code, request=httpx2.Request("POST", "http://api/ack")) + client.acknowledge_work.side_effect = httpx2.HTTPStatusError( + "Unavailable", request=response.request, response=response + ) + manager._acknowledge_requested_work() + assert bool(manager._pending_work_acks) == (status_code == 503) + if status_code == 503: + client.acknowledge_work.side_effect = None + manager._acknowledge_requested_work() + assert not manager._pending_work_acks + assert client.acknowledge_work.call_args.args == ("callbacks", work) + + +def test_api_inventory_cleanup_preserves_claimed_callbacks(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["removed.py"])[0] + request = DagCallbackRequest( + filepath="removed.py", bundle_name="testing", bundle_version="older", dag_id="dag", run_id="run" + ) + manager._callback_to_execute[file].append(request) + manager._file_queue[file] = None + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = uuid7() + manager._processors[file] = processor + manager._callback_attempts[processor.id] = [ + ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="removed.py" + ) + ] + manager.purge_removed_files_from_queue(set()) + manager.terminate_orphan_processes(set()) + assert list(manager._file_queue) == [file] + processor.kill.assert_not_called() + processor.close.assert_not_called() + + +@pytest.mark.parametrize("generation", [0, 1]) +@pytest.mark.parametrize("discovery_failed", [False, True]) +def test_inventory_only_fails_priority_requests_after_authoritative_refresh(generation, discovery_failed): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.publish_inventory.return_value.revision = uuid7() + manager = DagFileProcessorManager(max_runs=1, api_client=client) + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="archive.zip/a.py" + ) + manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1) + body = DagBundleInventoryBody( + attempt_id=uuid7(), + dispatch_sequence=1, + expected_revision=None, + files=["archive.zip"], + failed_sources=["archive.zip"] if discovery_failed else [], + ) + manager._pending_inventories["testing"] = PendingBundleInventory( + body, set(), DagParseSource(refresh_generation=generation) + ) + + manager._publish_bundle_inventory("testing", {}) + + failed = generation == 1 and not discovery_failed + assert manager._pending_work_acks == ([("priority", work, True)] if failed else []) + assert (work.id in manager._priority_claims) is not failed + + +@pytest.mark.parametrize("kind", ["callbacks", "priority"]) +def test_pending_requested_work_does_not_block_other_available_capacity(kind): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.claim_work.return_value = [] + manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=2, max_callbacks_per_loop=2) + bundle = mock.create_autospec(BaseDagBundle, instance=True) + bundle.name = "testing" + bundle.supports_versioning = False + manager._dag_bundles = [bundle] + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="missing.py" + ) + if kind == "priority": + manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1) + manager.claim_priority_files() + else: + manager._callback_claims[1] = work + manager.fetch_callbacks() + client.claim_work.assert_called_once_with(kind, ["testing"], 1) + + +def _claim_priority_parse(manager: DagFileProcessorManager, file: DagFileInfo) -> ProcessorWorkItem: + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name=file.bundle_name, relative_fileloc=str(file.rel_path) + ) + manager._priority_claims[work.id] = PriorityParseClaim( + work=work, refresh_generation=0, attempt_id=uuid7() + ) + manager._file_stats[file] = DagFileStat() + return work + + +@pytest.mark.parametrize("parallelism", [1, 2]) +@pytest.mark.parametrize("older_outcome", ["published", "completed", "timeout"]) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_priority_request_waits_for_its_own_refreshed_import(create, tmp_path, parallelism, older_outcome): + (tmp_path / "a.py").write_text("from airflow.sdk import DAG\ndag = DAG('test')\n") + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager( + max_runs=1, + api_client=client, + parallelism=parallelism, + processor_timeout=1, + bundle_refresh_check_interval=0, + ) + bundle = _make_bundle(tmp_path) + bundle.is_initialized = True + bundle.supports_versioning = False + bundle.refresh_interval = 300 + manager._dag_bundles = [bundle] + file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("a.py"), bundle_path=tmp_path) + manager._bundle_parse_sources[bundle.name] = DagParseSource() + old = mock.create_autospec(DagFileProcessorProcess, instance=True) + old.id = uuid7() + old.start_time = time.monotonic() - 5 + old.had_callbacks = False + old.parse_source = DagParseSource() + old.dispatch_sequence = 1 + old.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) + client.get_bundles.return_value = [ + DagBundleStateResponse( + name=bundle.name, version=None, last_refreshed=None, revision=None, team_name=None + ) + ] + client.publish_inventory.side_effect = lambda name, body: DagBundleInventoryResponse( + attempt_id=body.attempt_id, revision=uuid7(), accepted_at=timezone.utcnow() + ) + known_files = {bundle.name: {file}} + if older_outcome == "published": + body = mock.create_autospec(DagParseResultBody, instance=True) + body.attempt_id = old.id + manager._pending_publications[file] = PendingDagPublication( + body=body, stat=DagFileStat(), refresh_generation=0 + ) + else: + manager._processors[file] = old + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py" + ) + client.claim_work.return_value = [work] + manager._queue_requested_files_for_parsing() + manager._refresh_dag_bundles(known_files) + bundle.refresh.assert_not_called() + manager._start_new_processes() + if older_outcome == "published": + manager._publish_pending_results() + elif older_outcome == "completed": + manager._processors.pop(file) + manager.handle_parsing_result(file, old) + manager._publish_pending_results() + else: + manager._kill_timed_out_processors() + assert not manager._pending_work_acks + assert work.id in manager._priority_claims + assert file in manager._file_queue + create.assert_not_called() + + manager._start_new_processes() + create.assert_not_called() + manager._refresh_dag_bundles(known_files) + bundle.refresh.assert_called_once_with() + assert manager._bundle_parse_sources[bundle.name].refresh_generation == 1 + new = mock.create_autospec(DagFileProcessorProcess, instance=True) + new.id = uuid7() + new.start_time = time.monotonic() + new.had_callbacks = False + new.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) + create.return_value = new + manager._start_new_processes() + create.assert_called_once_with(manager, file) + + later = work.model_copy(update={"id": str(uuid7()), "claim_id": uuid7()}) + manager._parallelism = 2 + client.claim_work.return_value = [later] + manager._queue_requested_files_for_parsing() + manager.handle_parsing_result(file, new) + manager._publish_pending_results() + assert manager._pending_work_acks == [("priority", work, False)] + assert later.id in manager._priority_claims + assert file in manager._file_queue + + +@pytest.mark.parametrize("failure", ["response_lost", "server_error", "source_changed"]) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_priority_request_during_publication_retry_completes_after_refresh(create, tmp_path, failure): + (tmp_path / "a.py").write_text("from airflow.sdk import DAG\ndag = DAG('test')\n") + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_refresh_check_interval=0) + bundle = _make_bundle(tmp_path) + bundle.is_initialized = True + bundle.supports_versioning = False + bundle.refresh_interval = 300 + manager._dag_bundles = [bundle] + file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("a.py"), bundle_path=tmp_path) + client.get_bundles.return_value = [ + DagBundleStateResponse( + name=bundle.name, version=None, last_refreshed=None, revision=None, team_name=None + ) + ] + client.publish_inventory.side_effect = lambda name, body: DagBundleInventoryResponse( + attempt_id=body.attempt_id, revision=uuid7(), accepted_at=timezone.utcnow() + ) + known_files = {bundle.name: {file}} + first = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py" + ) + client.claim_work.return_value = [first] + manager._queue_requested_files_for_parsing() + manager._refresh_dag_bundles(known_files) + proc = mock.create_autospec(DagFileProcessorProcess, instance=True) + proc.id = uuid7() + proc.start_time = time.monotonic() + proc.had_callbacks = False + proc.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) + create.return_value = proc + manager._start_new_processes() + manager._processors.pop(file) + manager.handle_parsing_result(file, proc) + if failure == "response_lost": + error = httpx2.ReadError("Lost publication acknowledgment") + else: + response = httpx2.Response( + 409 if failure == "source_changed" else 503, + json={"detail": {"reason": failure}}, + request=httpx2.Request("POST", "http://api/parse-results"), + ) + error = httpx2.HTTPStatusError("Publication failed", request=response.request, response=response) + client.publish_parse_result.side_effect = error + manager._publish_pending_results() + assert not manager._pending_work_acks + + later = first.model_copy(update={"id": str(uuid7()), "claim_id": uuid7()}) + client.claim_work.return_value = [later] + manager._queue_requested_files_for_parsing() + client.publish_parse_result.side_effect = None + if failure != "source_changed": + manager._refresh_dag_bundles(known_files) + assert bundle.refresh.call_count == 1 + manager._start_new_processes() + assert create.call_count == 1 + manager._pending_publications[file].next_attempt_time = 0 + manager._publish_pending_results() + assert manager._pending_work_acks == [("priority", first, False)] + assert later.id in manager._priority_claims + manager._refresh_dag_bundles(known_files) + assert bundle.refresh.call_count == 2 + replacement = mock.create_autospec(DagFileProcessorProcess, instance=True) + replacement.id = uuid7() + replacement.start_time = time.monotonic() + replacement.had_callbacks = False + replacement.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) + create.return_value = replacement + manager._start_new_processes() + assert manager._processors[file] is replacement + manager._processors.pop(file) + manager.handle_parsing_result(file, replacement) + manager._publish_pending_results() + assert manager._pending_work_acks == [("priority", first, False), ("priority", later, False)] + assert not manager._priority_claims + assert not manager._pending_publications + manager._acknowledge_requested_work() + manager._acknowledge_requested_work() + assert client.acknowledge_work.call_args_list == [ + mock.call("priority", first, failed=False), + mock.call("priority", later, failed=False), + ] + + +@pytest.mark.parametrize( + ("status_code", "failed"), [(None, False), (409, True)], ids=["published", "rejected"] +) +def test_publication_outcome_acknowledges_priority_claims(status_code, failed): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) + manager._pending_publications[file] = PendingDagPublication( + body=mock.create_autospec(DagParseResultBody, instance=True), stat=DagFileStat(), refresh_generation=0 + ) + manager._pending_publications[file].body.attempt_id = manager._priority_claims[work.id].attempt_id + if status_code: + response = httpx2.Response(status_code, request=httpx2.Request("POST", "http://api/parse-results")) + client.publish_parse_result.side_effect = httpx2.HTTPStatusError( + "Rejected", request=response.request, response=response + ) + + manager._publish_pending_results() + + assert manager._pending_work_acks == [("priority", work, failed)] + assert not manager._priority_claims + + def test_publication_for_a_changed_source_refreshes_and_reparses(): client = mock.create_autospec(DagProcessorAPIClient, instance=True) manager = DagFileProcessorManager(max_runs=1, api_client=client) file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) manager._pending_publications[file] = PendingDagPublication( body=mock.create_autospec(DagParseResultBody, instance=True), stat=DagFileStat(), refresh_generation=0 ) - manager._pending_publications[file].body.attempt_id = uuid7() + manager._pending_publications[file].body.attempt_id = manager._priority_claims[work.id].attempt_id response = httpx2.Response( 409, json={"detail": {"reason": "source_changed"}}, @@ -6986,6 +7743,46 @@ def test_publication_for_a_changed_source_refreshes_and_reparses(): assert not manager._pending_publications assert manager._force_refresh_bundles == {"testing"} assert list(manager._file_queue) == [file] + assert manager._priority_claims[work.id].work == work + assert not manager._pending_work_acks + + +def test_rejected_inventory_releases_priority_attempt_discarded_after_source_change(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, parallelism=1) + file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) + body = mock.create_autospec(DagParseResultBody, instance=True) + body.attempt_id = manager._priority_claims[work.id].attempt_id + manager._pending_publications[file] = PendingDagPublication( + body=body, stat=DagFileStat(), refresh_generation=0 + ) + client.publish_parse_result.side_effect = _make_status_error(409, "source_changed") + + manager._publish_pending_results() + + assert not manager._pending_publications + assert manager._priority_claims[work.id].attempt_id is not None + inventory = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=2, expected_revision=None, files=["a.py"] + ) + manager._pending_inventories[file.bundle_name] = PendingBundleInventory( + inventory, + {file}, + DagParseSource(refresh_generation=1), + ) + client.publish_inventory.side_effect = _make_status_error(413) + + manager._publish_bundle_inventory(file.bundle_name, {}) + + assert not manager._priority_claims + assert manager._pending_work_acks == [("priority", work, True)] + assert not manager._force_refresh_bundles + manager._acknowledge_requested_work() + client.acknowledge_work.assert_called_once_with("priority", work, failed=True) + client.claim_work.return_value = [] + manager.claim_priority_files() + client.claim_work.assert_called_once_with("priority", [], 1) @conf_vars({("core", "multi_team"): "True"}) @@ -7005,6 +7802,91 @@ def test_api_team_lookup_failure_is_retried_without_stopping_the_processor(): assert manager._get_team_names({"testing"}) == {"testing": "team_a"} +@pytest.mark.parametrize( + ("outcome", "acknowledged"), + [("no_result", True), ("build_failed", True), ("callback_only", False)], +) +def test_parse_without_publication_acknowledges_priority_claims(outcome, acknowledged): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) + proc = mock.create_autospec(DagFileProcessorProcess, instance=True) + proc.id = manager._priority_claims[work.id].attempt_id + proc.start_time = time.monotonic() + proc.had_callbacks = outcome == "callback_only" + proc.parse_source = DagParseSource() + proc.parsing_result = ( + DagFileParsingResult(fileloc="a.py", serialized_dags=[]) if outcome == "build_failed" else None + ) + + with mock.patch.object(manager, "_build_parse_result", side_effect=ValueError("invalid result")): + manager.handle_parsing_result(file, proc) + + assert manager._pending_work_acks == ([("priority", work, True)] if acknowledged else []) + assert bool(manager._priority_claims) is not acknowledged + + +@pytest.mark.parametrize("had_callbacks", [False, True]) +def test_timed_out_parse_acknowledges_priority_claims(had_callbacks): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client, processor_timeout=1) + file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = manager._priority_claims[work.id].attempt_id + processor.start_time = time.monotonic() - 5 + processor.had_callbacks = had_callbacks + manager._processors[file] = processor + + manager._kill_timed_out_processors() + + processor.kill.assert_called_once_with(signal.SIGKILL) + assert manager._pending_work_acks == ([] if had_callbacks else [("priority", work, True)]) + + +@pytest.mark.parametrize("recovers", [False, True], ids=["dropped", "recovered"]) +@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"}) +@mock.patch("airflow.dag_processing.manager.time", autospec=True) +@mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True, return_value=None) +def test_unpreparable_callback_is_retried_with_backoff_until_the_job_heartbeat_timeout( + prepare, mock_time, recovers +): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + request = DagCallbackRequest( + filepath="a.py", bundle_name="testing", bundle_version="v1", dag_id="dag", run_id="run" + ) + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py" + ) + manager._callback_claims[id(request)] = work + manager._deferred_api_callbacks.append(request) + + for now in (0, 0.5, 1, 2): + mock_time.monotonic.return_value = now + manager._add_callback_to_queue(request) + assert prepare.call_count == 2 + + if recovers: + bundle = mock.create_autospec(BaseDagBundle, instance=True) + bundle.path = TEST_DAGS_FOLDER + prepare.return_value = bundle + mock_time.monotonic.return_value = 10 + manager._add_callback_to_queue(request) + + assert prepare.call_count == 3 + assert not manager._deferred_callback_retries + assert not manager._deferred_api_callbacks + if recovers: + assert manager._callback_claims == {id(request): work} + assert list(manager._callback_to_execute.values()) == [[request]] + assert not manager._pending_work_acks + else: + assert not manager._callback_claims + assert manager._pending_work_acks == [("callbacks", work, True)] + + def _make_pending_publication( file: DagFileInfo, *, attempt_id: UUID | None = None, import_errors: dict[str, str] | None = None ) -> PendingDagPublication: @@ -7075,9 +7957,10 @@ def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message): manager = DagFileProcessorManager(max_runs=1, api_client=client) manager._dispatch_sequence = 7 file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) original = _make_pending_publication( file, - attempt_id=uuid7(), + attempt_id=manager._priority_claims[work.id].attempt_id, import_errors={"a.py/member.py": "member failed"}, ) manager._pending_publications[file] = original @@ -7095,6 +7978,7 @@ def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message): assert not manager._pending_publications assert manager._file_stats[file].num_dags == 0 assert manager._file_stats[file].import_errors == 1 + assert manager._pending_work_acks == [("priority", work, True)] @conf_vars({("workers", "execution_api_retries"): "1"}) @@ -7103,9 +7987,10 @@ def test_server_failure_after_publication_retries_preserves_previous_result(): client.publish_parse_result.side_effect = [_make_status_error(500), None] manager = DagFileProcessorManager(max_runs=1, api_client=client) file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) previous = DagFileStat(num_dags=2, last_finish_time=DEFAULT_DATE) manager._file_stats[file] = previous - pending = _make_pending_publication(file, attempt_id=uuid7()) + pending = _make_pending_publication(file, attempt_id=manager._priority_claims[work.id].attempt_id) manager._pending_publications[file] = pending manager._publish_pending_results() @@ -7115,6 +8000,7 @@ def test_server_failure_after_publication_retries_preserves_previous_result(): assert manager._file_stats[file].num_dags == previous.num_dags assert manager._file_stats[file].last_finish_time == previous.last_finish_time assert manager._file_stats[file].last_attempt_time is not None + assert manager._pending_work_acks == [("priority", work, True)] @mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 32 * 1024) @@ -7221,11 +8107,12 @@ def test_result_over_the_publication_limit_is_reported_without_sending_it(): client = mock.create_autospec(DagProcessorAPIClient, instance=True) manager = DagFileProcessorManager(max_runs=1, api_client=client) file = _get_file_infos(["a.py"])[0] + work = _claim_priority_parse(manager, file) dag = SdkDAG("large", schedule=None, doc_md="x" * 1024 * 1024) dag.fileloc = "/worker/a.py" dag.relative_fileloc = "a.py" proc = mock.create_autospec(DagFileProcessorProcess, instance=True) - proc.id = uuid7() + proc.id = manager._priority_claims[work.id].attempt_id proc.start_time = time.monotonic() proc.had_callbacks = False proc.dispatch_sequence = 1 @@ -7242,17 +8129,45 @@ def test_result_over_the_publication_limit_is_reported_without_sending_it(): assert pending.body.import_errors == { "a.py": "The parsing result was not published: it exceeds the 1 MiB publication limit." } + assert manager._pending_work_acks == [("priority", work, True)] client.publish_parse_result.assert_not_called() -@pytest.mark.parametrize("method", ["claim_priority_files", "fetch_callbacks"]) -def test_internal_api_manager_leaves_requested_work_for_activation(method, mocker): +def test_acknowledgments_are_sent_within_the_loop_budget(): client = mock.create_autospec(DagProcessorAPIClient, instance=True) manager = DagFileProcessorManager(max_runs=1, api_client=client) - database_method = ( - "_claim_priority_files" if method == "claim_priority_files" else "_fetch_callbacks_from_db" - ) - read_database = mocker.patch.object(manager, database_method, autospec=True) + work = [ + ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py") + for _ in range(3) + ] + manager._pending_work_acks = [("priority", item, False) for item in work] - assert getattr(manager, method)() == [] - read_database.assert_not_called() + manager._acknowledge_requested_work(deadline=time.monotonic()) + client.acknowledge_work.assert_not_called() + + manager._acknowledge_requested_work() + assert client.acknowledge_work.call_args_list == [ + mock.call("priority", item, failed=False) for item in work + ] + assert not manager._pending_work_acks + + +@pytest.mark.parametrize("failure", [None, httpx2.ConnectError("unavailable"), RuntimeError("closed")]) +def test_shutdown_acknowledges_finished_work(failure): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.acknowledge_work.side_effect = failure + manager = DagFileProcessorManager(max_runs=1, api_client=client) + work = [ + ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py") + for _ in range(2) + ] + manager._pending_work_acks = [("callbacks", item, False) for item in work] + + manager.end() + + if failure is None: + assert client.acknowledge_work.call_count == 2 + assert not manager._pending_work_acks + else: + client.acknowledge_work.assert_called_once() + assert len(manager._pending_work_acks) == 2 diff --git a/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py b/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py new file mode 100644 index 0000000000000..f55ed3ca5242c --- /dev/null +++ b/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 + + +@pytest.fixture +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 + + +@pytest.fixture +def runner(): + return DagProcessorJobRunner( + job=Job(heartrate=5), processor=mock.create_autospec(DagFileProcessorManager, instance=True) + ) + + +@pytest.fixture +def clock(): + with mock.patch( + "airflow.jobs.dag_processor_job_runner.monotonic", autospec=True, return_value=0 + ) as clock: + yield clock + + +@pytest.mark.parametrize( + ("error", "outcome"), + [ + (None, TerminalJobState.SUCCESS), + (SystemExit(0), TerminalJobState.SUCCESS), + (ValueError("parse failed"), TerminalJobState.FAILED), + (KeyboardInterrupt(), TerminalJobState.FAILED), + ], +) +@mock.patch.object(Job, "prepare_for_execution", autospec=True) +@mock.patch.object(Job, "complete_execution", autospec=True) +def test_api_lifecycle_cleans_up_before_completion(complete, prepare, runner, client, error, outcome): + events = [] + runner.processor.run.side_effect = error + runner.processor.terminate.side_effect = lambda: events.append("terminate") + runner.processor.end.side_effect = lambda: events.append("end") + client.complete_job.side_effect = lambda state: events.append(state) + + if outcome == TerminalJobState.FAILED: + with pytest.raises(type(error)): + runner.run_with_api(client) + else: + runner.run_with_api(client) + + assert runner.job.id == 42 + assert runner.job.state == outcome.value + assert runner.processor.api_client is client + runner.processor.sync_bundles.assert_not_called() + client.register_job.assert_called_once_with() + assert events == ["terminate", "end", outcome] + prepare.assert_not_called() + complete.assert_not_called() + + +def test_failed_registration_never_starts_parsing(runner, client): + client.register_job.side_effect = httpx2.ConnectError("unavailable") + with pytest.raises(httpx2.ConnectError): + runner.run_with_api(client) + runner.processor.run.assert_not_called() + client.complete_job.assert_not_called() + + +@pytest.mark.parametrize("available_at", [12, 15, None]) +@conf_vars({("dag_processor", "health_check_threshold"): "10"}) +@mock.patch("airflow.jobs.dag_processor_job_runner.sleep", autospec=True) +def test_registration_waits_for_a_crashed_job_with_a_bounded_backoff( + sleep, runner, client, clock, available_at +): + def wait(seconds): + clock.return_value += seconds + + def register(): + runner.processor.run.assert_not_called() + if available_at is None or clock.return_value < available_at: + raise DagProcessorJobAlreadyRunning("Still alive") + return 42 + + sleep.side_effect = wait + client.register_job.side_effect = register + + if available_at is None: + with pytest.raises(DagProcessorJobAlreadyRunning): + runner.run_with_api(client) + runner.processor.run.assert_not_called() + client.complete_job.assert_not_called() + else: + runner.run_with_api(client) + runner.processor.run.assert_called_once_with() + assert runner.job.id == 42 + + assert sleep.call_args_list == [mock.call(1), mock.call(2), mock.call(4), mock.call(5)] + ( + [mock.call(3)] if available_at != 12 else [] + ) + assert clock.return_value == (available_at or 15) + assert client.register_job.call_count == sleep.call_count + 1 + + +@pytest.mark.parametrize("reason", ["registration_conflict", "registration_retired", None]) +@mock.patch("airflow.jobs.dag_processor_job_runner.sleep", autospec=True) +def test_registration_does_not_wait_on_other_conflicts(sleep, runner, client, reason): + error = httpx2.HTTPStatusError( + "conflict", + request=httpx2.Request("POST", "http://api/jobs"), + response=httpx2.Response(409, json={"detail": {"reason": reason}}), + ) + client.register_job.side_effect = error + + with pytest.raises(httpx2.HTTPStatusError): + runner.run_with_api(client) + + sleep.assert_not_called() + client.register_job.assert_called_once_with() + runner.processor.run.assert_not_called() + + +@mock.patch("airflow.jobs.dag_processor_job_runner.get_listener_manager", autospec=True) +def test_shutdown_listener_failure_does_not_prevent_completion(listeners, runner, client): + listeners.return_value.hook.before_stopping.side_effect = ValueError("listener failed") + runner.run_with_api(client) + listeners.return_value.hook.before_stopping.assert_called_once_with(component=runner.job) + client.complete_job.assert_called_once_with(TerminalJobState.SUCCESS) + + +@pytest.mark.parametrize("parse_error", [None, ValueError("original error")]) +def test_completion_failure_preserves_the_original_error(runner, client, parse_error): + runner.processor.run.side_effect = parse_error + client.complete_job.side_effect = httpx2.ConnectError("completion unavailable") + with pytest.raises(type(parse_error) if parse_error else httpx2.ConnectError) as raised: + runner.run_with_api(client) + if parse_error: + assert raised.value is parse_error + + +def test_heartbeats_are_throttled_and_use_the_api(runner, client, clock): + def run(): + runner.processor.heartbeat() + client.heartbeat.assert_not_called() + clock.return_value = 5 + runner.processor.heartbeat() + runner.processor.heartbeat() + client.heartbeat.assert_called_once_with() + clock.return_value = 10 + runner.processor.heartbeat() + assert client.heartbeat.call_count == 2 + + runner.processor.run.side_effect = run + runner.run_with_api(client) + + +@pytest.mark.parametrize("restart", ["already_retired", "renewal_retired", "server_restart", "job_closed"]) +def test_restart_stops_parsing_and_releases_children(runner, client, clock, restart): + def heartbeat(): + if restart == "job_closed": + raise DagProcessorRegistrationRetired("replaced") + if restart == "renewal_retired": + client.restart_required = True + return JobState.RESTARTING if restart == "server_restart" else JobState.RUNNING + + def run(): + client.restart_required = restart == "already_retired" + clock.return_value = 1 if client.restart_required else 5 + runner.processor.heartbeat() + pytest.fail("Parsing continued after a restart request") + + client.heartbeat.side_effect = heartbeat + runner.processor.run.side_effect = run + with pytest.raises(DagProcessorRegistrationRetired): + runner.run_with_api(client) + runner.processor.terminate.assert_called_once_with() + runner.processor.end.assert_called_once_with() + client.complete_job.assert_called_once_with(TerminalJobState.FAILED) + + +@pytest.mark.parametrize("failure", ["transport", 503, 403]) +@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"}) +def test_heartbeat_failure_is_bounded(runner, client, clock, failure): + error = ( + httpx2.ConnectError("unavailable") + if failure == "transport" + else httpx2.HTTPStatusError( + "refused", + request=httpx2.Request("POST", "http://api/jobs/42/heartbeat"), + response=httpx2.Response(failure), + ) + ) + client.heartbeat.side_effect = error + + def run(): + clock.return_value = 5 + runner.processor.heartbeat() + clock.return_value = 10 + runner.processor.heartbeat() + + runner.processor.run.side_effect = run + with pytest.raises(httpx2.HTTPStatusError if failure == 403 else DagProcessorHeartbeatTimeout): + runner.run_with_api(client) + assert client.heartbeat.call_count == (1 if failure == 403 else 2) + + +@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"}) +def test_recovered_heartbeat_resets_the_lease(runner, client, clock): + client.heartbeat.side_effect = [ + httpx2.ConnectError("offline"), + JobState.RUNNING, + httpx2.ConnectError("offline"), + ] + + def run(): + for tick in (5, 10, 15): + clock.return_value = tick + runner.processor.heartbeat() + + runner.processor.run.side_effect = run + runner.run_with_api(client) + client.complete_job.assert_called_once_with(TerminalJobState.SUCCESS) + + +@conf_vars( + { + ("dag_processor", "health_check_threshold"): "30", + ("dag_processor", "job_heartbeat_timeout"): "300", + } +) +def test_a_stall_past_the_health_threshold_keeps_the_lease(runner, client, clock): + client.heartbeat.side_effect = [httpx2.ConnectError("offline"), JobState.RUNNING] + + def run(): + for tick in (120, 125): + clock.return_value = tick + runner.processor.heartbeat() + + runner.processor.run.side_effect = run + runner.run_with_api(client) + assert client.heartbeat.call_count == 2 + client.complete_job.assert_called_once_with(TerminalJobState.SUCCESS) diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index 93978e902280a..cf054af9091c5 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -39,6 +39,24 @@ metrics: legacy_name: "-" name_variables: ["job_name"] + - name: "job_start" + description: "Number of started jobs." + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "job_end" + description: "Number of ended jobs." + type: "counter" + legacy_name: "-" + name_variables: [] + + - name: "dag_processor_heartbeat_failure" + description: "Number of failed Dag processor heartbeats through the Execution API." + type: "counter" + legacy_name: "-" + name_variables: [] + - name: "operator_failures" description: "Operator ``{operator_name}`` failures." type: "counter" @@ -122,15 +140,6 @@ metrics: legacy_name: "-" name_variables: [] - - name: "dag_processing.results_discarded_on_refresh" - description: "Parse attempts discarded because their bundle source generation was invalidated. - Metric with bundle_name tagging, and team_name tagging in multi-team mode. - This defensive counter should normally remain zero because refreshes wait for active imports. - Use dag_processing.files_awaiting_source to monitor normal refresh delays." - type: "counter" - legacy_name: "-" - name_variables: [] - - name: "dag_processing.processes" description: "Relative number of currently running Dag parsing processes (ie this delta is negative when, since the last metric was sent, processes have completed). Metric with file_path, bundle_name, and action @@ -152,6 +161,15 @@ metrics: legacy_name: "-" name_variables: [] + - name: "dag_processing.results_discarded_on_refresh" + description: "Parse attempts discarded because their bundle source generation was invalidated. + Metric with bundle_name tagging, and team_name tagging in multi-team mode. + This defensive counter should normally remain zero because refreshes wait for active imports. + Use dag_processing.files_awaiting_source to monitor normal refresh delays." + type: "counter" + legacy_name: "-" + name_variables: [] + - name: "dag_processing.callback_only_count" description: "Number of DAG file processing runs that processed callbacks only, without full DAG parsing" type: "counter" From 573f161b13983493e5dc1ab35b3004e0a4b369c2 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Mon, 19 Oct 2026 12:21:45 +0100 Subject: [PATCH 2/5] 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. --- .../security/jwt_token_authentication.rst | 30 ++-- .../src/airflow/config_templates/config.yml | 6 +- .../src/airflow/dag_processing/api_client.py | 9 +- .../src/airflow/dag_processing/manager.py | 41 +++--- .../unit/dag_processing/test_api_client.py | 57 ++++++++ .../tests/unit/dag_processing/test_manager.py | 135 +++++++++--------- 6 files changed, 174 insertions(+), 104 deletions(-) diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst index e978e054bf23b..405eed4c19f32 100644 --- a/airflow-core/docs/security/jwt_token_authentication.rst +++ b/airflow-core/docs/security/jwt_token_authentication.rst @@ -620,21 +620,27 @@ 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 or unassigned -bundles stale, removes stale warnings, and ends expired processor Jobs, in bounded batches and in +``[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, and 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 callback as failed. 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. +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 diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 8106e448eaa9f..706e3f2be795d 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3519,12 +3519,12 @@ dag_processor: or listing a single bundle, so set this above the longest of those operations; the local health check must also tolerate them, through ``health_check_threshold`` and the probe's failure budget. It also bounds how long a processor retries preparing the bundle - version of a claimed callback before acknowledging the callback as failed. + version of a claimed callback before acknowledging it without running it. + A callback that ran before its acknowledgment was recorded can run again after the Job + times out, repeating its side effects. Callbacks should tolerate repeated execution. The scheduler reads this value to end Jobs and the processor reads it to stop, so set the same value for both; changing it only in the processor's configuration has no effect on when the scheduler ends the Job. - A callback that ran before its acknowledgment was recorded can run again after the Job - times out, repeating its side effects. Callbacks should tolerate repeated execution. version_added: 3.5.0 type: integer example: ~ diff --git a/airflow-core/src/airflow/dag_processing/api_client.py b/airflow-core/src/airflow/dag_processing/api_client.py index 14f6e6db4c750..b1db01ba87bc4 100644 --- a/airflow-core/src/airflow/dag_processing/api_client.py +++ b/airflow-core/src/airflow/dag_processing/api_client.py @@ -146,6 +146,7 @@ def __init__( self._completion_state: TerminalJobState | None = None self._completed = False self._pending_claims: dict[str, ProcessorWorkClaimBody] = {} + self._claim_sequences: dict[str, int] = {} super().__init__(base_url=base_url, token="", **kwargs) @property @@ -411,6 +412,7 @@ def claim_work( kind, ProcessorWorkClaimBody( claim_id=uuid7(), + claim_sequence=self._claim_sequences.get(kind, 0) + 1, bundle_names=bundle_names, limit=min(limit, 100), ), @@ -424,13 +426,12 @@ def claim_work( timeout=self._get_bounded_timeout(), ) work = [ProcessorWorkItem.model_validate(item) for item in response.json()] + self._claim_sequences[kind] = body.claim_sequence del self._pending_claims[kind] return work - def acknowledge_work( - self, kind: Literal["callbacks", "priority"], work: ProcessorWorkItem, *, failed: bool = False - ) -> None: - body = ProcessorWorkAckBody(claim_id=work.claim_id, state="failed" if failed else "success") + def acknowledge_work(self, kind: Literal["callbacks", "priority"], work: ProcessorWorkItem) -> None: + body = ProcessorWorkAckBody(claim_id=work.claim_id) self.request( "POST", f"jobs/{self._require_job_id()}/requested-work/{kind}/{work.id}/ack", diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 2ae733beaf57e..115f8c12c85ba 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -368,8 +368,8 @@ class DagFileProcessorManager(LoggingMixin): _deferred_api_callbacks: list[CallbackRequest] = attrs.field(factory=list, init=False) _deferred_callback_retries: dict[int, DeferredCallback] = attrs.field(factory=dict, init=False) _callback_attempts: dict[UUID, list[ProcessorWorkItem]] = attrs.field(factory=dict, init=False) - _priority_claims: dict[str, PriorityParseClaim] = attrs.field(factory=dict, init=False) - _pending_work_acks: list[tuple[Literal["callbacks", "priority"], ProcessorWorkItem, bool]] = attrs.field( + _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 ) _num_run: int = attrs.field(default=0, init=False) @@ -833,7 +833,7 @@ def claim_priority_files(self) -> list[DagFileInfo]: available = ( self._parallelism - len(self._priority_claims) - - sum(kind == "priority" for kind, _, _ in self._pending_work_acks) + - sum(kind == "priority" for kind, _ in self._pending_work_acks) ) if available <= 0: return [] @@ -845,7 +845,7 @@ def claim_priority_files(self) -> list[DagFileInfo]: return [] files = [] for item in work: - self._priority_claims[item.id] = PriorityParseClaim( + self._priority_claims[item.id, item.claim_id] = PriorityParseClaim( work=item, refresh_generation=self._bundle_refresh_generations[item.bundle_name] + 1 ) bundle = bundles[item.bundle_name] @@ -920,7 +920,7 @@ def fetch_callbacks(self) -> list[CallbackRequest]: 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) + - sum(kind == "callbacks" for kind, _ in self._pending_work_acks) ) if available <= 0: return list(self._deferred_api_callbacks) @@ -945,9 +945,9 @@ def _acknowledge_requested_work(self, *, deadline: float | None = None) -> None: if deadline is None: deadline = time.monotonic() + _API_SEND_BUDGET_SECONDS while self._pending_work_acks and time.monotonic() < deadline: - kind, item, failed = self._pending_work_acks[0] + kind, item = self._pending_work_acks[0] try: - self.api_client.acknowledge_work(kind, item, failed=failed) + self.api_client.acknowledge_work(kind, item) except httpx2.HTTPStatusError as error: if error.response.status_code >= 500: self.log.warning("Unable to acknowledge requested work; retaining the claim") @@ -967,13 +967,13 @@ def _acknowledge_requested_work(self, *, deadline: float | None = None) -> None: def _finish_callback_attempt(self, proc: BaseDagFileProcessorProcess) -> None: for item in self._callback_attempts.pop(proc.id, []): - self._pending_work_acks.append(("callbacks", item, proc._exit_code != 0)) + self._pending_work_acks.append(("callbacks", item)) - def _finish_priority_claims(self, attempt_id: UUID, *, failed: bool) -> None: + def _finish_priority_claims(self, attempt_id: UUID) -> None: """Acknowledge only requests assigned to this import, never newer requests for its file.""" for key, claim in list(self._priority_claims.items()): if claim.attempt_id == attempt_id: - self._pending_work_acks.append(("priority", claim.work, failed)) + self._pending_work_acks.append(("priority", claim.work)) del self._priority_claims[key] def _get_priority_claims(self, file: DagFileInfo) -> list[PriorityParseClaim]: @@ -1088,7 +1088,7 @@ def _add_callback_to_queue(self, request: CallbackRequest) -> None: 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, True)) + self._pending_work_acks.append(("callbacks", item)) self._remove_deferred_callback(request) return retry = self._deferred_callback_retries.get(id(request)) @@ -1121,7 +1121,7 @@ def _remove_deferred_callback(self, request: CallbackRequest) -> None: 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 acknowledge it as failed.""" + """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)) @@ -1134,7 +1134,7 @@ def _defer_claimed_callback(self, request: CallbackRequest) -> None: request.bundle_version, timeout, ) - self._pending_work_acks.append(("callbacks", self._callback_claims.pop(id(request)), True)) + self._pending_work_acks.append(("callbacks", self._callback_claims.pop(id(request)))) self._remove_deferred_callback(request) return retry.attempts += 1 @@ -1245,7 +1245,7 @@ def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set and item.relative_fileloc not in pending.body.files and not is_source_discovery_failed(item.relative_fileloc, failed_sources) ): - self._pending_work_acks.append(("priority", item, True)) + self._pending_work_acks.append(("priority", item)) del self._priority_claims[key] def _discard_unusable_bundle_source(self, bundle_name: str) -> None: @@ -1263,7 +1263,7 @@ def _discard_unusable_bundle_source(self, bundle_name: str) -> None: ) for key, claim in list(self._priority_claims.items()): if claim.work.bundle_name == bundle_name and claim.attempt_id not in active_attempts: - self._pending_work_acks.append(("priority", claim.work, True)) + self._pending_work_acks.append(("priority", claim.work)) del self._priority_claims[key] def _invalidate_bundle_parse_source(self, bundle_name: str) -> None: @@ -1936,11 +1936,11 @@ def handle_parsing_result( file.bundle_name, ) self._record_failed_publication(file, next_stat) - self._finish_priority_claims(proc.id, failed=True) + self._finish_priority_claims(proc.id) return if proc.parsing_result is None and not is_callback_only: - self._finish_priority_claims(proc.id, failed=True) + self._finish_priority_claims(proc.id) self._file_stats[file] = next_stat def _record_failed_publication(self, file: DagFileInfo, stat: DagFileStat) -> None: @@ -2019,10 +2019,9 @@ def _publish_pending_results(self, *, deadline: float | None = None) -> None: continue self.log.exception("Failed to publish parsing result for %s", file) self._record_failed_publication(file, pending.stat) - self._finish_priority_claims(pending.body.attempt_id, failed=True) else: self._file_stats[file] = pending.stat - self._finish_priority_claims(pending.body.attempt_id, failed=False) + self._finish_priority_claims(pending.body.attempt_id) del self._pending_publications[file] def _build_pending_publication( @@ -2072,7 +2071,7 @@ def _build_rejection_report( :param attempt_id: The import whose result is replaced; its priority requests are released. """ - self._finish_priority_claims(attempt_id, failed=True) + self._finish_priority_claims(attempt_id) self._dispatch_sequence += 1 if len(reason) > _MAX_REJECTION_REASON_LENGTH: reason = reason[:_MAX_REJECTION_REASON_LENGTH] + "..." @@ -2631,7 +2630,7 @@ def _kill_timed_out_processors(self): ) self._file_stats[file] = stat if not processor.had_callbacks: - self._finish_priority_claims(processor.id, failed=True) + self._finish_priority_claims(processor.id) # Clean up `self._processors` after iterating over it for proc in processors_to_remove: diff --git a/airflow-core/tests/unit/dag_processing/test_api_client.py b/airflow-core/tests/unit/dag_processing/test_api_client.py index 9e685a38463af..01fc98b648abf 100644 --- a/airflow-core/tests/unit/dag_processing/test_api_client.py +++ b/airflow-core/tests/unit/dag_processing/test_api_client.py @@ -30,6 +30,7 @@ from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( DagBundleInventoryBody, DagParseResultBody, + ProcessorWorkItem, ) from airflow.api_fastapi.execution_api.datamodels.job import DagParseTokenBody, JobState, TerminalJobState from airflow.dag_processing import api_client @@ -1079,3 +1080,59 @@ def test_inventory_waits_for_a_slow_write_but_connects_quickly(make_client): timeout = requests[-1].extensions["timeout"] assert timeout["connect"] <= 1 assert timeout["read"] == client.timeout.read > 1 + + +@pytest.mark.parametrize("kind", ["callbacks", "priority"]) +@pytest.mark.parametrize( + "failure", + [ + httpx2.ReadError("Lost claim response"), + httpx2.Response(503), + httpx2.Response(200, json=[{"id": "incomplete"}]), + ], + ids=["lost_response", "server_error", "invalid_response"], +) +def test_claim_sequence_survives_retries_and_renewal_then_advances_per_kind(make_client, kind, failure): + client, requests = make_client( + make_registration_response(), + failure, + make_registration_response(), + httpx2.Response(200, json=[]), + httpx2.Response(200, json=[]), + httpx2.Response(200, json=[]), + ) + client.register_job() + with pytest.raises((httpx2.HTTPError, ValueError)): + client.claim_work(kind, ["bundle-a"], 2) + + client.register_job() + assert client.claim_work(kind, ["bundle-b"], 5) == [] + other_kind = "callbacks" if kind == "priority" else "priority" + assert client.claim_work(other_kind, ["bundle-a"], 2) == [] + assert client.claim_work(kind, ["bundle-b"], 5) == [] + + assert requests[1].content == requests[3].content + original = json.loads(requests[1].content) + assert original["claim_sequence"] == 1 + assert original["bundle_names"] == ["bundle-a"] + assert original["limit"] == 2 + assert json.loads(requests[4].content)["claim_sequence"] == 1 + subsequent = json.loads(requests[5].content) + assert subsequent["claim_sequence"] == 2 + assert subsequent["claim_id"] != original["claim_id"] + assert subsequent["bundle_names"] == ["bundle-b"] + assert subsequent["limit"] == 5 + + +@pytest.mark.parametrize("kind", ["callbacks", "priority"]) +def test_work_acknowledgment_sends_only_the_claim_identity(make_client, kind): + client, requests = make_client(make_registration_response(), httpx2.Response(204)) + client.register_job() + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name="bundle-a", relative_fileloc="dag.py" + ) + + client.acknowledge_work(kind, work) + + assert requests[-1].url.path == f"/execution/jobs/1/requested-work/{kind}/{work.id}/ack" + assert json.loads(requests[-1].content) == {"claim_id": str(work.claim_id)} diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 61dc6f3b34df3..d81f04ac13889 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -6042,7 +6042,7 @@ def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed work = ProcessorWorkItem( id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py" ) - manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1) + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, refresh_generation=1) manager._bundles_waiting_for_refresh.add(bundle.name) manager._bundle_versions[bundle.name] = "previous" @@ -6058,7 +6058,7 @@ def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed assert not manager._force_refresh_bundles assert not manager._bundles_waiting_for_refresh assert not manager._priority_claims - assert manager._pending_work_acks == [("priority", work, True)] + assert manager._pending_work_acks == [("priority", work)] manager._acknowledge_requested_work() client.claim_work.return_value = [] manager.claim_priority_files() @@ -6850,10 +6850,10 @@ def create(manager, file): manager._add_callback_to_queue(request) assert not manager._deferred_api_callbacks assert not manager._deferred_callback_retries - assert manager._pending_work_acks == [("callbacks", work[0], True)] + assert manager._pending_work_acks == [("callbacks", work[0])] manager._acknowledge_requested_work() manager._acknowledge_requested_work() - client.acknowledge_work.assert_called_once_with("callbacks", work[0], failed=True) + client.acknowledge_work.assert_called_once_with("callbacks", work[0]) @pytest.mark.parametrize("outcome", ["queued", "expired", "unsupported"]) @@ -6893,7 +6893,7 @@ def test_equal_claimed_callbacks_are_removed_by_identity(prepare, importer, cloc assert not manager._pending_work_acks else: assert id(second) not in manager._callback_claims - assert manager._pending_work_acks == [("callbacks", work[1], True)] + assert manager._pending_work_acks == [("callbacks", work[1])] @mock.patch("airflow.dag_processing.manager.get_claiming_importer", autospec=True, return_value=None) @@ -6948,8 +6948,8 @@ def create_process(**kwargs): assert not manager._deferred_callback_retries assert not manager._callback_to_execute assert client.acknowledge_work.call_args_list == [ - mock.call("callbacks", work[1], failed=False), - mock.call("callbacks", work[0], failed=False), + mock.call("callbacks", work[1]), + mock.call("callbacks", work[0]), ] @@ -7038,7 +7038,7 @@ def test_catalog_omission_releases_priority_capacity_and_allows_recovery( manager._bundle_parse_sources[name] = DagParseSource() work = ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="a.py") claims[name] = PriorityParseClaim(work=work, refresh_generation=1) - manager._priority_claims[work.id] = claims[name] + manager._priority_claims[work.id, work.claim_id] = claims[name] manager._dag_bundles = list(bundles.values()) previous_files = files.copy() previous_versions = manager._bundle_versions.copy() @@ -7069,9 +7069,12 @@ def test_catalog_omission_releases_priority_capacity_and_allows_recovery( assert files == previous_files assert manager._bundle_versions == previous_versions assert manager._bundle_parse_sources["ready"] is ready_source - assert manager._priority_claims[claims["ready"].work.id] is claims["ready"] + assert manager._priority_claims[claims["ready"].work.id, claims["ready"].work.claim_id] is claims["ready"] if catalog_failed: - assert manager._priority_claims[claims["missing"].work.id] is claims["missing"] + assert ( + manager._priority_claims[claims["missing"].work.id, claims["missing"].work.claim_id] + is claims["missing"] + ) assert not manager._pending_work_acks assert manager._force_refresh_bundles == {"missing"} assert manager._bundles_waiting_for_refresh == {"missing"} @@ -7084,8 +7087,8 @@ def test_catalog_omission_releases_priority_capacity_and_allows_recovery( return client.publish_inventory.assert_not_called() - assert claims["missing"].work.id not in manager._priority_claims - assert manager._pending_work_acks == [("priority", claims["missing"].work, True)] + assert (claims["missing"].work.id, claims["missing"].work.claim_id) not in manager._priority_claims + assert manager._pending_work_acks == [("priority", claims["missing"].work)] assert not manager._pending_inventories assert not manager._force_refresh_bundles assert not manager._bundles_waiting_for_refresh @@ -7309,7 +7312,7 @@ def test_oversized_inventory_releases_unassigned_priority_capacity( work = ProcessorWorkItem( id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="b.py" ) - retained_claims[work.id] = PriorityParseClaim( + retained_claims[work.id, work.claim_id] = PriorityParseClaim( work=work, refresh_generation=1, attempt_id=attempt_id ) if name == "oversized": @@ -7327,7 +7330,7 @@ def test_oversized_inventory_releases_unassigned_priority_capacity( manager._priority_claims.update(retained_claims) client.claim_work.side_effect = [[rejected_work], [ready_work]] manager._queue_requested_files_for_parsing() - assert manager._priority_claims[rejected_work.id].attempt_id is None + assert manager._priority_claims[rejected_work.id, rejected_work.claim_id].attempt_id is None body = DagBundleInventoryBody( attempt_id=uuid7(), dispatch_sequence=1, @@ -7342,7 +7345,7 @@ def test_oversized_inventory_releases_unassigned_priority_capacity( manager._publish_bundle_inventory("oversized", {}) assert manager._priority_claims == retained_claims - assert manager._pending_work_acks == [("priority", rejected_work, True)] + assert manager._pending_work_acks == [("priority", rejected_work)] if rejection == "local_limit": client.publish_inventory.assert_not_called() else: @@ -7350,7 +7353,7 @@ def test_oversized_inventory_releases_unassigned_priority_capacity( assert manager.claim_priority_files() == [] client.claim_work.assert_called_once_with("priority", ["oversized", "ready"], 1) manager._acknowledge_requested_work() - client.acknowledge_work.assert_called_once_with("priority", rejected_work, failed=True) + client.acknowledge_work.assert_called_once_with("priority", rejected_work) manager._queue_requested_files_for_parsing() assert client.claim_work.call_args_list == [ mock.call("priority", ["oversized", "ready"], 1), @@ -7380,10 +7383,12 @@ def test_oversized_inventory_releases_unassigned_priority_capacity( manager._start_new_processes() create.assert_called_once_with(manager, ready_file) - assert manager._priority_claims[ready_work.id].attempt_id == process.id - assert {key: claim for key, claim in manager._priority_claims.items() if key != ready_work.id} == ( - retained_claims - ) + assert manager._priority_claims[ready_work.id, ready_work.claim_id].attempt_id == process.id + assert { + key: claim + for key, claim in manager._priority_claims.items() + if key != (ready_work.id, ready_work.claim_id) + } == (retained_claims) assert not manager._pending_work_acks @@ -7422,7 +7427,7 @@ def test_requested_work_ack_retains_only_retryable_failures(status_code): work = ProcessorWorkItem( id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py" ) - manager._pending_work_acks.append(("callbacks", work, False)) + manager._pending_work_acks.append(("callbacks", work)) response = httpx2.Response(status_code, request=httpx2.Request("POST", "http://api/ack")) client.acknowledge_work.side_effect = httpx2.HTTPStatusError( "Unavailable", request=response.request, response=response @@ -7469,7 +7474,7 @@ def test_inventory_only_fails_priority_requests_after_authoritative_refresh(gene work = ProcessorWorkItem( id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="archive.zip/a.py" ) - manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1) + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, refresh_generation=1) body = DagBundleInventoryBody( attempt_id=uuid7(), dispatch_sequence=1, @@ -7484,8 +7489,8 @@ def test_inventory_only_fails_priority_requests_after_authoritative_refresh(gene manager._publish_bundle_inventory("testing", {}) failed = generation == 1 and not discovery_failed - assert manager._pending_work_acks == ([("priority", work, True)] if failed else []) - assert (work.id in manager._priority_claims) is not failed + assert manager._pending_work_acks == ([("priority", work)] if failed else []) + assert ((work.id, work.claim_id) in manager._priority_claims) is not failed @pytest.mark.parametrize("kind", ["callbacks", "priority"]) @@ -7501,7 +7506,7 @@ def test_pending_requested_work_does_not_block_other_available_capacity(kind): id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="missing.py" ) if kind == "priority": - manager._priority_claims[work.id] = PriorityParseClaim(work=work, refresh_generation=1) + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, refresh_generation=1) manager.claim_priority_files() else: manager._callback_claims[1] = work @@ -7513,7 +7518,7 @@ def _claim_priority_parse(manager: DagFileProcessorManager, file: DagFileInfo) - work = ProcessorWorkItem( id=str(uuid7()), claim_id=uuid7(), bundle_name=file.bundle_name, relative_fileloc=str(file.rel_path) ) - manager._priority_claims[work.id] = PriorityParseClaim( + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim( work=work, refresh_generation=0, attempt_id=uuid7() ) manager._file_stats[file] = DagFileStat() @@ -7581,7 +7586,7 @@ def test_priority_request_waits_for_its_own_refreshed_import(create, tmp_path, p else: manager._kill_timed_out_processors() assert not manager._pending_work_acks - assert work.id in manager._priority_claims + assert (work.id, work.claim_id) in manager._priority_claims assert file in manager._file_queue create.assert_not_called() @@ -7599,14 +7604,14 @@ def test_priority_request_waits_for_its_own_refreshed_import(create, tmp_path, p manager._start_new_processes() create.assert_called_once_with(manager, file) - later = work.model_copy(update={"id": str(uuid7()), "claim_id": uuid7()}) + later = work.model_copy(update={"claim_id": uuid7()}) manager._parallelism = 2 client.claim_work.return_value = [later] manager._queue_requested_files_for_parsing() manager.handle_parsing_result(file, new) manager._publish_pending_results() - assert manager._pending_work_acks == [("priority", work, False)] - assert later.id in manager._priority_claims + assert manager._pending_work_acks == [("priority", work)] + assert (later.id, later.claim_id) in manager._priority_claims assert file in manager._file_queue @@ -7659,7 +7664,7 @@ def test_priority_request_during_publication_retry_completes_after_refresh(creat manager._publish_pending_results() assert not manager._pending_work_acks - later = first.model_copy(update={"id": str(uuid7()), "claim_id": uuid7()}) + later = first.model_copy(update={"claim_id": uuid7()}) client.claim_work.return_value = [later] manager._queue_requested_files_for_parsing() client.publish_parse_result.side_effect = None @@ -7670,8 +7675,8 @@ def test_priority_request_during_publication_retry_completes_after_refresh(creat assert create.call_count == 1 manager._pending_publications[file].next_attempt_time = 0 manager._publish_pending_results() - assert manager._pending_work_acks == [("priority", first, False)] - assert later.id in manager._priority_claims + assert manager._pending_work_acks == [("priority", first)] + assert (later.id, later.claim_id) in manager._priority_claims manager._refresh_dag_bundles(known_files) assert bundle.refresh.call_count == 2 replacement = mock.create_autospec(DagFileProcessorProcess, instance=True) @@ -7685,21 +7690,19 @@ def test_priority_request_during_publication_retry_completes_after_refresh(creat manager._processors.pop(file) manager.handle_parsing_result(file, replacement) manager._publish_pending_results() - assert manager._pending_work_acks == [("priority", first, False), ("priority", later, False)] + assert manager._pending_work_acks == [("priority", first), ("priority", later)] assert not manager._priority_claims assert not manager._pending_publications manager._acknowledge_requested_work() manager._acknowledge_requested_work() assert client.acknowledge_work.call_args_list == [ - mock.call("priority", first, failed=False), - mock.call("priority", later, failed=False), + mock.call("priority", first), + mock.call("priority", later), ] -@pytest.mark.parametrize( - ("status_code", "failed"), [(None, False), (409, True)], ids=["published", "rejected"] -) -def test_publication_outcome_acknowledges_priority_claims(status_code, failed): +@pytest.mark.parametrize("status_code", [None, 409], ids=["published", "rejected"]) +def test_publication_outcome_acknowledges_priority_claims(status_code): client = mock.create_autospec(DagProcessorAPIClient, instance=True) manager = DagFileProcessorManager(max_runs=1, api_client=client) file = _get_file_infos(["a.py"])[0] @@ -7707,7 +7710,9 @@ def test_publication_outcome_acknowledges_priority_claims(status_code, failed): manager._pending_publications[file] = PendingDagPublication( body=mock.create_autospec(DagParseResultBody, instance=True), stat=DagFileStat(), refresh_generation=0 ) - manager._pending_publications[file].body.attempt_id = manager._priority_claims[work.id].attempt_id + manager._pending_publications[file].body.attempt_id = manager._priority_claims[ + work.id, work.claim_id + ].attempt_id if status_code: response = httpx2.Response(status_code, request=httpx2.Request("POST", "http://api/parse-results")) client.publish_parse_result.side_effect = httpx2.HTTPStatusError( @@ -7716,7 +7721,7 @@ def test_publication_outcome_acknowledges_priority_claims(status_code, failed): manager._publish_pending_results() - assert manager._pending_work_acks == [("priority", work, failed)] + assert manager._pending_work_acks == [("priority", work)] assert not manager._priority_claims @@ -7728,7 +7733,9 @@ def test_publication_for_a_changed_source_refreshes_and_reparses(): manager._pending_publications[file] = PendingDagPublication( body=mock.create_autospec(DagParseResultBody, instance=True), stat=DagFileStat(), refresh_generation=0 ) - manager._pending_publications[file].body.attempt_id = manager._priority_claims[work.id].attempt_id + manager._pending_publications[file].body.attempt_id = manager._priority_claims[ + work.id, work.claim_id + ].attempt_id response = httpx2.Response( 409, json={"detail": {"reason": "source_changed"}}, @@ -7743,7 +7750,7 @@ def test_publication_for_a_changed_source_refreshes_and_reparses(): assert not manager._pending_publications assert manager._force_refresh_bundles == {"testing"} assert list(manager._file_queue) == [file] - assert manager._priority_claims[work.id].work == work + assert manager._priority_claims[work.id, work.claim_id].work == work assert not manager._pending_work_acks @@ -7753,7 +7760,7 @@ def test_rejected_inventory_releases_priority_attempt_discarded_after_source_cha file = _get_file_infos(["a.py"])[0] work = _claim_priority_parse(manager, file) body = mock.create_autospec(DagParseResultBody, instance=True) - body.attempt_id = manager._priority_claims[work.id].attempt_id + body.attempt_id = manager._priority_claims[work.id, work.claim_id].attempt_id manager._pending_publications[file] = PendingDagPublication( body=body, stat=DagFileStat(), refresh_generation=0 ) @@ -7762,7 +7769,7 @@ def test_rejected_inventory_releases_priority_attempt_discarded_after_source_cha manager._publish_pending_results() assert not manager._pending_publications - assert manager._priority_claims[work.id].attempt_id is not None + assert manager._priority_claims[work.id, work.claim_id].attempt_id is not None inventory = DagBundleInventoryBody( attempt_id=uuid7(), dispatch_sequence=2, expected_revision=None, files=["a.py"] ) @@ -7776,10 +7783,10 @@ def test_rejected_inventory_releases_priority_attempt_discarded_after_source_cha manager._publish_bundle_inventory(file.bundle_name, {}) assert not manager._priority_claims - assert manager._pending_work_acks == [("priority", work, True)] + assert manager._pending_work_acks == [("priority", work)] assert not manager._force_refresh_bundles manager._acknowledge_requested_work() - client.acknowledge_work.assert_called_once_with("priority", work, failed=True) + client.acknowledge_work.assert_called_once_with("priority", work) client.claim_work.return_value = [] manager.claim_priority_files() client.claim_work.assert_called_once_with("priority", [], 1) @@ -7812,7 +7819,7 @@ def test_parse_without_publication_acknowledges_priority_claims(outcome, acknowl file = _get_file_infos(["a.py"])[0] work = _claim_priority_parse(manager, file) proc = mock.create_autospec(DagFileProcessorProcess, instance=True) - proc.id = manager._priority_claims[work.id].attempt_id + proc.id = manager._priority_claims[work.id, work.claim_id].attempt_id proc.start_time = time.monotonic() proc.had_callbacks = outcome == "callback_only" proc.parse_source = DagParseSource() @@ -7823,7 +7830,7 @@ def test_parse_without_publication_acknowledges_priority_claims(outcome, acknowl with mock.patch.object(manager, "_build_parse_result", side_effect=ValueError("invalid result")): manager.handle_parsing_result(file, proc) - assert manager._pending_work_acks == ([("priority", work, True)] if acknowledged else []) + assert manager._pending_work_acks == ([("priority", work)] if acknowledged else []) assert bool(manager._priority_claims) is not acknowledged @@ -7834,7 +7841,7 @@ def test_timed_out_parse_acknowledges_priority_claims(had_callbacks): file = _get_file_infos(["a.py"])[0] work = _claim_priority_parse(manager, file) processor = mock.create_autospec(DagFileProcessorProcess, instance=True) - processor.id = manager._priority_claims[work.id].attempt_id + processor.id = manager._priority_claims[work.id, work.claim_id].attempt_id processor.start_time = time.monotonic() - 5 processor.had_callbacks = had_callbacks manager._processors[file] = processor @@ -7842,7 +7849,7 @@ def test_timed_out_parse_acknowledges_priority_claims(had_callbacks): manager._kill_timed_out_processors() processor.kill.assert_called_once_with(signal.SIGKILL) - assert manager._pending_work_acks == ([] if had_callbacks else [("priority", work, True)]) + assert manager._pending_work_acks == ([] if had_callbacks else [("priority", work)]) @pytest.mark.parametrize("recovers", [False, True], ids=["dropped", "recovered"]) @@ -7884,7 +7891,7 @@ def test_unpreparable_callback_is_retried_with_backoff_until_the_job_heartbeat_t assert not manager._pending_work_acks else: assert not manager._callback_claims - assert manager._pending_work_acks == [("callbacks", work, True)] + assert manager._pending_work_acks == [("callbacks", work)] def _make_pending_publication( @@ -7960,7 +7967,7 @@ def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message): work = _claim_priority_parse(manager, file) original = _make_pending_publication( file, - attempt_id=manager._priority_claims[work.id].attempt_id, + attempt_id=manager._priority_claims[work.id, work.claim_id].attempt_id, import_errors={"a.py/member.py": "member failed"}, ) manager._pending_publications[file] = original @@ -7978,7 +7985,7 @@ def test_result_the_api_refuses_is_replaced_by_an_import_error(error, message): assert not manager._pending_publications assert manager._file_stats[file].num_dags == 0 assert manager._file_stats[file].import_errors == 1 - assert manager._pending_work_acks == [("priority", work, True)] + assert manager._pending_work_acks == [("priority", work)] @conf_vars({("workers", "execution_api_retries"): "1"}) @@ -7990,7 +7997,9 @@ def test_server_failure_after_publication_retries_preserves_previous_result(): work = _claim_priority_parse(manager, file) previous = DagFileStat(num_dags=2, last_finish_time=DEFAULT_DATE) manager._file_stats[file] = previous - pending = _make_pending_publication(file, attempt_id=manager._priority_claims[work.id].attempt_id) + pending = _make_pending_publication( + file, attempt_id=manager._priority_claims[work.id, work.claim_id].attempt_id + ) manager._pending_publications[file] = pending manager._publish_pending_results() @@ -8000,7 +8009,7 @@ def test_server_failure_after_publication_retries_preserves_previous_result(): assert manager._file_stats[file].num_dags == previous.num_dags assert manager._file_stats[file].last_finish_time == previous.last_finish_time assert manager._file_stats[file].last_attempt_time is not None - assert manager._pending_work_acks == [("priority", work, True)] + assert manager._pending_work_acks == [("priority", work)] @mock.patch("airflow.dag_processing.manager.MAX_PUBLICATION_BYTES", 32 * 1024) @@ -8112,7 +8121,7 @@ def test_result_over_the_publication_limit_is_reported_without_sending_it(): dag.fileloc = "/worker/a.py" dag.relative_fileloc = "a.py" proc = mock.create_autospec(DagFileProcessorProcess, instance=True) - proc.id = manager._priority_claims[work.id].attempt_id + proc.id = manager._priority_claims[work.id, work.claim_id].attempt_id proc.start_time = time.monotonic() proc.had_callbacks = False proc.dispatch_sequence = 1 @@ -8129,7 +8138,7 @@ def test_result_over_the_publication_limit_is_reported_without_sending_it(): assert pending.body.import_errors == { "a.py": "The parsing result was not published: it exceeds the 1 MiB publication limit." } - assert manager._pending_work_acks == [("priority", work, True)] + assert manager._pending_work_acks == [("priority", work)] client.publish_parse_result.assert_not_called() @@ -8140,15 +8149,13 @@ def test_acknowledgments_are_sent_within_the_loop_budget(): ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py") for _ in range(3) ] - manager._pending_work_acks = [("priority", item, False) for item in work] + manager._pending_work_acks = [("priority", item) for item in work] manager._acknowledge_requested_work(deadline=time.monotonic()) client.acknowledge_work.assert_not_called() manager._acknowledge_requested_work() - assert client.acknowledge_work.call_args_list == [ - mock.call("priority", item, failed=False) for item in work - ] + assert client.acknowledge_work.call_args_list == [mock.call("priority", item) for item in work] assert not manager._pending_work_acks @@ -8161,7 +8168,7 @@ def test_shutdown_acknowledges_finished_work(failure): ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="a.py") for _ in range(2) ] - manager._pending_work_acks = [("callbacks", item, False) for item in work] + manager._pending_work_acks = [("callbacks", item) for item in work] manager.end() From 8ad77a6194e669186c4cd8ae1164ce85d8a270ae Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 22 Oct 2026 12:46:35 +0100 Subject: [PATCH 3/5] Wait for fresh inventories before running priority parses Safe refresh failures should not stop ordinary Dag parsing or consume a reparse request before fresh source is available. Callbacks and finite processing runs must preserve that distinction, and metrics must remain meaningful while files wait for a usable source. Known limitation: the Git provider's Airflow 3.4 version gate still assumes DagBundleRefreshError exists in the installed core. Compatibility with unpatched 3.4 core requires a follow-up before publication. --- .../src/airflow/dag_processing/manager.py | 29 +- .../src/airflow/dag_processing/processor.py | 1 + .../tests/unit/dag_processing/test_manager.py | 90 +++--- .../test_priority_inventory_sequence.py | 306 ++++++++++++++++++ 4 files changed, 379 insertions(+), 47 deletions(-) create mode 100644 airflow-core/tests/unit/dag_processing/test_priority_inventory_sequence.py diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 115f8c12c85ba..ceeb25922188f 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -195,10 +195,10 @@ class DeferredCallback: @attrs.define class PriorityParseClaim: - """A priority request served only by an import after its requested refresh.""" + """A priority request served only after an inventory captured following the claim.""" work: ProcessorWorkItem - refresh_generation: int + inventory_sequence: int attempt_id: UUID | None = None @@ -846,7 +846,7 @@ def claim_priority_files(self) -> list[DagFileInfo]: files = [] for item in work: self._priority_claims[item.id, item.claim_id] = PriorityParseClaim( - work=item, refresh_generation=self._bundle_refresh_generations[item.bundle_name] + 1 + work=item, inventory_sequence=self._dispatch_sequence + 1 ) bundle = bundles[item.bundle_name] path = find_enclosing_file(bundle.path / item.relative_fileloc) @@ -1231,17 +1231,19 @@ def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set return del self._pending_inventories[bundle_name] self._bundle_parse_sources[bundle_name] = attrs.evolve( - pending.source, bundle_revision=receipt.revision + pending.source, + bundle_revision=receipt.revision, + inventory_sequence=pending.body.dispatch_sequence, ) self._bundle_versions[bundle_name] = pending.source.bundle_version - self._bundle_failed_sources[bundle_name] = frozenset(pending.body.failed_sources) known_files[bundle_name] = pending.found_files + self._bundle_failed_sources[bundle_name] = frozenset(pending.body.failed_sources) failed_sources = set(pending.body.failed_sources) for key, claim in list(self._priority_claims.items()): item = claim.work if ( item.bundle_name == bundle_name - and pending.source.refresh_generation >= claim.refresh_generation + and pending.body.dispatch_sequence >= claim.inventory_sequence and item.relative_fileloc not in pending.body.files and not is_source_discovery_failed(item.relative_fileloc, failed_sources) ): @@ -1473,6 +1475,7 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): continue found_files = discovery.files + if self.api_client is not None: self._dispatch_sequence += 1 try: @@ -1506,8 +1509,8 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): else: self._bundle_versions[bundle.name] = version_after_refresh - self._bundle_failed_sources[bundle.name] = discovery.failed_sources known_files[bundle.name] = found_files + self._bundle_failed_sources[bundle.name] = discovery.failed_sources self.deactivate_deleted_dags(bundle_name=bundle.name, present=found_files) self.clear_orphaned_import_errors( @@ -2315,7 +2318,12 @@ def _get_dispatch_source(self, file: DagFileInfo) -> DagParseSource | None: or (file in self._priority_files and file.bundle_name in self._bundles_needing_refresh) ): return None - return self._bundle_parse_sources.get(file.bundle_name) + source = self._bundle_parse_sources.get(file.bundle_name) + if source is not None and any( + source.inventory_sequence < claim.inventory_sequence for claim in self._get_priority_claims(file) + ): + return None + return source def _has_priority_request(self, file: DagFileInfo) -> bool: # A claim outlives the queue marker when its file is requeued behind a running or stale import. @@ -2382,7 +2390,7 @@ def _start_new_processes(self): processor.dispatch_sequence = self._dispatch_sequence if file.bundle_version is None: for claim in priority_claims: - if processor.had_callbacks or source.refresh_generation < claim.refresh_generation: + if processor.had_callbacks: requested_later[file] = None else: claim.attempt_id = processor.id @@ -2676,7 +2684,8 @@ def max_runs_reached(self): A file that is still queued, waiting for its bundle source, or being parsed is not done, even if its stats already show *max_runs*. - Callback-only work is ignored because it does not count as a parse. + Callbacks do not count as parses, but API mode waits for claimed callback + work and its pending acknowledgments before exiting. """ if self.max_runs == -1: # Unlimited runs. return False diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 90307c9f14407..45c92addcedb7 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -600,6 +600,7 @@ class DagParseSource: version_data: Mapping[str, Any] | None = None refresh_generation: int = 0 bundle_revision: UUID | None = None + inventory_sequence: int = 0 @classmethod def for_callback(cls, bundle_version: str | None) -> Self: diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index d81f04ac13889..2340fb9c190b2 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -36,6 +36,7 @@ from datetime import datetime, timedelta from pathlib import Path from socket import socket, socketpair +from types import MappingProxyType from unittest import mock from unittest.mock import MagicMock from uuid import UUID @@ -2349,17 +2350,13 @@ def test_forced_refresh_blocks_dispatch_with_a_valid_source(self, create): assert not manager._bundles_waiting_for_refresh assert list(manager._files_awaiting_source) == [file] - @pytest.mark.parametrize("api_mode", [False, True]) @mock.patch("airflow.dag_processing.manager.emit_metrics", autospec=True) @mock.patch("airflow.dag_processing.manager.time.perf_counter", autospec=True) @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) @mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) @conf_vars({("dag_processor", "file_parsing_sort_mode"): "alphabetical"}) - def test_parse_time_spans_source_deferral_until_dispatch( - self, create, persist, perf_counter, emit, api_mode - ): - client = mock.create_autospec(DagProcessorAPIClient, instance=True) if api_mode else None - manager = DagFileProcessorManager(max_runs=1, api_client=client) + def test_parse_time_spans_source_deferral_until_dispatch(self, create, persist, perf_counter, emit): + manager = DagFileProcessorManager(max_runs=1) manager._parallelism = 2 blocked = DagFileInfo(bundle_name="blocked", rel_path=Path("blocked.py")) healthy = DagFileInfo(bundle_name="healthy", rel_path=Path("healthy.py")) @@ -2401,7 +2398,6 @@ def test_parse_time_spans_source_deferral_until_dispatch( manager.prepare_file_queue(known_files=known_files) assert emit.call_count == 1 - @pytest.mark.parametrize("api_mode", [False, True]) @mock.patch("airflow.dag_processing.manager.emit_metrics", autospec=True) @mock.patch("airflow.dag_processing.manager.time.perf_counter", autospec=True) @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) @@ -2413,19 +2409,17 @@ def test_parse_time_spans_source_deferral_until_dispatch( } ) def test_parse_time_emits_when_processors_finish_in_alternating_cycles( - self, create, persist, perf_counter, emit, api_mode + self, create, persist, perf_counter, emit ): - client = mock.create_autospec(DagProcessorAPIClient, instance=True) if api_mode else None - manager = DagFileProcessorManager(max_runs=-1, api_client=client) + manager = DagFileProcessorManager(max_runs=-1) manager._parallelism = 2 first, second = _get_file_infos(["first.py", "second.py"]) manager._bundle_parse_sources[first.bundle_name] = DagParseSource() known_files = {first.bundle_name: {first, second}} processors = [mock.create_autospec(DagFileProcessorProcess, instance=True) for _ in range(5)] for processor in processors: - processor.id = uuid7() - processor.had_callbacks = False processor.start_time = 0 + processor.had_callbacks = False processor.is_ready = False processor.parsing_result = DagFileParsingResult(fileloc="file.py", serialized_dags=[]) self._configure_mock_processes(create, *processors) @@ -2436,8 +2430,6 @@ def test_parse_time_emits_when_processors_finish_in_alternating_cycles( for count, completed in enumerate((first, second, first), start=1): manager._processors[completed].is_ready = True manager._collect_results() - if api_mode: - manager._publish_pending_results() assert len(manager._processors) == 1 perf_counter.return_value = 10 * (count + 1) @@ -5756,13 +5748,11 @@ def test_api_publication_uses_dispatch_identity_without_local_persistence( proc, read_end = self.mock_processor() proc.had_callbacks = False proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) - proc.parse_source = manager._bundle_parse_sources["testing"] - create.return_value = proc + self._configure_mock_processes(create, proc) if publication_fails: client.publish_parse_result.side_effect = RuntimeError("publication failed") try: manager._start_new_processes() - create.assert_called_once_with(manager, file, source=proc.parse_source) manager.handle_parsing_result(file, proc, session=session) client.publish_parse_result.assert_not_called() manager._publish_pending_results() @@ -5812,6 +5802,24 @@ def test_api_publication_reports_a_result_the_api_would_refuse(self, session): assert manager._file_stats[file].import_errors == 1 assert not manager._pending_publications + @pytest.mark.parametrize( + "version_data", + [None, {}, {"revision": "v1"}, MappingProxyType({}), MappingProxyType({"revision": "v1"})], + ) + def test_api_publication_preserves_version_data(self, version_data): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + manager = DagFileProcessorManager(max_runs=1, api_client=client) + file = _get_file_infos(["dag.py"])[0] + proc = mock.create_autospec(DagFileProcessorProcess, instance=True) + proc.id = uuid7() + proc.dispatch_sequence = 1 + proc.parse_source = DagParseSource(version_data=version_data) + proc.parsing_result = DagFileParsingResult(fileloc="dag.py", serialized_dags=[]) + + body = manager._build_parse_result(file, proc, 0.5) + + assert body.version_data == version_data + @pytest.mark.parametrize("source_available", [True, False]) def test_api_publication_carries_diagnostics_and_explicit_source_availability(self, source_available): client = mock.create_autospec(DagProcessorAPIClient, instance=True) @@ -6001,7 +6009,7 @@ def test_requested_file_still_parsing_stays_ahead_of_the_queue(self, create): manager._bundle_parse_sources["testing"] = DagParseSource() manager._processors[requested] = self.mock_processor()[0] manager._file_queue = OrderedDict.fromkeys([requested, other, later]) - create.return_value = self.mock_processor()[0] + self._configure_mock_processes(create, self.mock_processor()[0]) manager._start_new_processes() @@ -6042,7 +6050,7 @@ def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed work = ProcessorWorkItem( id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py" ) - manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, refresh_generation=1) + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, inventory_sequence=1) manager._bundles_waiting_for_refresh.add(bundle.name) manager._bundle_versions[bundle.name] = "previous" @@ -6088,7 +6096,11 @@ def test_refresh_dag_bundles_heartbeats_before_each_bundle(self): manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None) ), mock.patch.object(manager, "update_bundle_state"), - mock.patch.object(manager, "_find_files_in_bundle", return_value=set()), + mock.patch.object( + manager, + "_find_files_in_bundle", + return_value=DagBundleDiscoveryResult(files=set(), failed_sources=frozenset()), + ), mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -6834,10 +6846,11 @@ def get_bundle(*, name, version, version_data): normal = DagFileInfo(bundle_name=bundle.name, rel_path=Path("normal.py"), bundle_path=tmp_path) manager._add_files_to_queue([normal], mode="back") - def create(manager, file): + def create(manager, file, *, source): proc = mock.create_autospec(DagFileProcessorProcess, instance=True) proc.id = uuid7() proc.had_callbacks = file.bundle_version is not None + proc.parse_source = source return proc create_process.side_effect = create @@ -6957,7 +6970,7 @@ def create_process(**kwargs): DagFileProcessorManager, "_find_files_in_bundle", autospec=True, - return_value=DagBundleDiscoveryResult(files=set()), + return_value=DagBundleDiscoveryResult(files=set(), failed_sources=frozenset()), ) @conf_vars({("core", "multi_team"): "True"}) def test_refresh_reads_one_fresh_catalog_per_pass(find_files, tmp_path): @@ -7037,7 +7050,7 @@ def test_catalog_omission_releases_priority_capacity_and_allows_recovery( manager._bundle_versions[name] = None manager._bundle_parse_sources[name] = DagParseSource() work = ProcessorWorkItem(id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="a.py") - claims[name] = PriorityParseClaim(work=work, refresh_generation=1) + claims[name] = PriorityParseClaim(work=work, inventory_sequence=1) manager._priority_claims[work.id, work.claim_id] = claims[name] manager._dag_bundles = list(bundles.values()) previous_files = files.copy() @@ -7105,7 +7118,7 @@ def test_catalog_omission_releases_priority_capacity_and_allows_recovery( client.publish_inventory.side_effect = None revision = uuid7() client.publish_inventory.return_value.revision = revision - find_files.return_value = DagBundleDiscoveryResult(files=files["missing"]) + find_files.return_value = DagBundleDiscoveryResult(files=files["missing"], failed_sources=frozenset()) manager._refresh_dag_bundles(files) bundles["missing"].refresh.assert_called_once() @@ -7313,7 +7326,7 @@ def test_oversized_inventory_releases_unassigned_priority_capacity( id=str(uuid7()), claim_id=uuid7(), bundle_name=name, relative_fileloc="b.py" ) retained_claims[work.id, work.claim_id] = PriorityParseClaim( - work=work, refresh_generation=1, attempt_id=attempt_id + work=work, inventory_sequence=1, attempt_id=attempt_id ) if name == "oversized": active_file = DagFileInfo(bundle_name=name, rel_path=Path("b.py"), bundle_path=tmp_path) @@ -7382,7 +7395,7 @@ def test_oversized_inventory_releases_unassigned_priority_capacity( manager._start_new_processes() - create.assert_called_once_with(manager, ready_file) + create.assert_called_once_with(manager, ready_file, source=manager._bundle_parse_sources["ready"]) assert manager._priority_claims[ready_work.id, ready_work.claim_id].attempt_id == process.id assert { key: claim @@ -7465,30 +7478,30 @@ def test_api_inventory_cleanup_preserves_claimed_callbacks(): processor.close.assert_not_called() -@pytest.mark.parametrize("generation", [0, 1]) +@pytest.mark.parametrize("inventory_sequence", [1, 2]) @pytest.mark.parametrize("discovery_failed", [False, True]) -def test_inventory_only_fails_priority_requests_after_authoritative_refresh(generation, discovery_failed): +def test_inventory_only_fails_priority_requests_after_authoritative_refresh( + inventory_sequence, discovery_failed +): client = mock.create_autospec(DagProcessorAPIClient, instance=True) client.publish_inventory.return_value.revision = uuid7() manager = DagFileProcessorManager(max_runs=1, api_client=client) work = ProcessorWorkItem( id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="archive.zip/a.py" ) - manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, refresh_generation=1) + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, inventory_sequence=2) body = DagBundleInventoryBody( attempt_id=uuid7(), - dispatch_sequence=1, + dispatch_sequence=inventory_sequence, expected_revision=None, files=["archive.zip"], failed_sources=["archive.zip"] if discovery_failed else [], ) - manager._pending_inventories["testing"] = PendingBundleInventory( - body, set(), DagParseSource(refresh_generation=generation) - ) + manager._pending_inventories["testing"] = PendingBundleInventory(body, set(), DagParseSource()) manager._publish_bundle_inventory("testing", {}) - failed = generation == 1 and not discovery_failed + failed = inventory_sequence == 2 and not discovery_failed assert manager._pending_work_acks == ([("priority", work)] if failed else []) assert ((work.id, work.claim_id) in manager._priority_claims) is not failed @@ -7506,7 +7519,7 @@ def test_pending_requested_work_does_not_block_other_available_capacity(kind): id=str(uuid7()), claim_id=uuid7(), bundle_name="testing", relative_fileloc="missing.py" ) if kind == "priority": - manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, refresh_generation=1) + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim(work=work, inventory_sequence=1) manager.claim_priority_files() else: manager._callback_claims[1] = work @@ -7519,7 +7532,7 @@ def _claim_priority_parse(manager: DagFileProcessorManager, file: DagFileInfo) - id=str(uuid7()), claim_id=uuid7(), bundle_name=file.bundle_name, relative_fileloc=str(file.rel_path) ) manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim( - work=work, refresh_generation=0, attempt_id=uuid7() + work=work, inventory_sequence=0, attempt_id=uuid7() ) manager._file_stats[file] = DagFileStat() return work @@ -7599,10 +7612,11 @@ def test_priority_request_waits_for_its_own_refreshed_import(create, tmp_path, p new.id = uuid7() new.start_time = time.monotonic() new.had_callbacks = False + new.parse_source = manager._bundle_parse_sources[bundle.name] new.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) create.return_value = new manager._start_new_processes() - create.assert_called_once_with(manager, file) + create.assert_called_once_with(manager, file, source=manager._bundle_parse_sources[file.bundle_name]) later = work.model_copy(update={"claim_id": uuid7()}) manager._parallelism = 2 @@ -7646,6 +7660,7 @@ def test_priority_request_during_publication_retry_completes_after_refresh(creat proc.id = uuid7() proc.start_time = time.monotonic() proc.had_callbacks = False + proc.parse_source = manager._bundle_parse_sources[bundle.name] proc.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) create.return_value = proc manager._start_new_processes() @@ -7683,6 +7698,7 @@ def test_priority_request_during_publication_retry_completes_after_refresh(creat replacement.id = uuid7() replacement.start_time = time.monotonic() replacement.had_callbacks = False + replacement.parse_source = manager._bundle_parse_sources[bundle.name] replacement.parsing_result = DagFileParsingResult(fileloc="a.py", serialized_dags=[]) create.return_value = replacement manager._start_new_processes() diff --git a/airflow-core/tests/unit/dag_processing/test_priority_inventory_sequence.py b/airflow-core/tests/unit/dag_processing/test_priority_inventory_sequence.py new file mode 100644 index 0000000000000..48500e6c61267 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_priority_inventory_sequence.py @@ -0,0 +1,306 @@ +# 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 dataclasses import replace +from pathlib import Path +from unittest import mock + +import pytest +from uuid6 import uuid7 + +from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryBody, + DagBundleInventoryResponse, + DagBundleStateResponse, + ProcessorWorkItem, +) +from airflow.callbacks.callback_requests import DagCallbackRequest +from airflow.dag_processing.api_client import DagProcessorAPIClient +from airflow.dag_processing.bundles.base import BaseDagBundle, DagBundleRefreshError +from airflow.dag_processing.bundles.local import LocalDagBundle +from airflow.dag_processing.manager import DagFileInfo, DagFileProcessorManager, PendingBundleInventory +from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess, DagParseSource + + +@pytest.fixture +def api_manager(tmp_path): + (tmp_path / "a.py").write_text("from airflow.sdk import DAG\ndag = DAG('test')\n") + bundle = LocalDagBundle(name="testing", path=str(tmp_path), refresh_interval=0) + bundle.is_initialized = True + file = DagFileInfo(bundle_name=bundle.name, rel_path=Path("a.py"), bundle_path=tmp_path) + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + revision = uuid7() + client.get_bundles.return_value = [ + DagBundleStateResponse( + name=bundle.name, version=None, last_refreshed=None, revision=revision, team_name=None + ) + ] + client.publish_inventory.side_effect = lambda name, body: DagBundleInventoryResponse( + attempt_id=body.attempt_id, revision=revision, accepted_at=timezone.utcnow() + ) + manager = DagFileProcessorManager( + max_runs=1, api_client=client, parallelism=2, bundle_refresh_check_interval=0 + ) + manager._dag_bundles = [bundle] + manager._bundle_versions[bundle.name] = None + manager._bundle_refresh_generations[bundle.name] = 3 + manager._dispatch_sequence = 7 + manager._bundle_parse_sources[bundle.name] = DagParseSource( + refresh_generation=3, bundle_revision=revision, inventory_sequence=7 + ) + work = ProcessorWorkItem( + id=str(uuid7()), claim_id=uuid7(), bundle_name=bundle.name, relative_fileloc="a.py" + ) + client.claim_work.return_value = [work] + return manager, client, bundle, file, work + + +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_local_refresh_attaches_claim_after_new_inventory_without_invalidating_active_parse( + create, api_manager +): + manager, client, bundle, file, work = api_manager + old = mock.create_autospec(DagFileProcessorProcess, instance=True) + old.id = uuid7() + old.had_callbacks = False + old.parse_source = manager._bundle_parse_sources[bundle.name] + manager._processors[file] = old + + manager._queue_requested_files_for_parsing() + manager._refresh_dag_bundles({bundle.name: {file}}) + + client.publish_inventory.assert_called_once() + source = manager._bundle_parse_sources[bundle.name] + assert source.refresh_generation == old.parse_source.refresh_generation == 3 + assert source.inventory_sequence > old.parse_source.inventory_sequence + claim = manager._priority_claims[work.id, work.claim_id] + assert claim.attempt_id is None + manager._start_new_processes() + create.assert_not_called() + + manager._processors.pop(file) + new = mock.create_autospec(DagFileProcessorProcess, instance=True) + new.id = uuid7() + new.had_callbacks = False + create.return_value = new + manager._start_new_processes() + + create.assert_called_once_with(manager, file, source=source) + assert claim.attempt_id == new.id + + +def test_inventory_created_before_claim_cannot_acknowledge_new_missing_file_request(api_manager): + manager, _, bundle, _, work = api_manager + source = manager._bundle_parse_sources[bundle.name] + old_inventory = DagBundleInventoryBody( + attempt_id=uuid7(), dispatch_sequence=7, expected_revision=source.bundle_revision, files=[] + ) + manager._pending_inventories[bundle.name] = PendingBundleInventory(old_inventory, set(), source) + manager._queue_requested_files_for_parsing() + + manager._publish_bundle_inventory(bundle.name, {}) + + assert (work.id, work.claim_id) in manager._priority_claims + assert not manager._pending_work_acks + + new_inventory = old_inventory.model_copy(update={"attempt_id": uuid7(), "dispatch_sequence": 8}) + manager._pending_inventories[bundle.name] = PendingBundleInventory(new_inventory, set(), source) + manager._publish_bundle_inventory(bundle.name, {}) + + assert (work.id, work.claim_id) not in manager._priority_claims + assert manager._pending_work_acks == [("priority", work)] + + +@pytest.mark.parametrize("callback_version", [None, "v1"]) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_callbacks_run_while_priority_parse_waits_for_new_inventory(create, api_manager, callback_version): + manager, _, _, file, work = api_manager + manager._parallelism = 1 + manager._queue_requested_files_for_parsing() + manager._force_refresh_bundles.clear() + callback_file = replace(file, bundle_version=callback_version) + request = DagCallbackRequest( + filepath=str(file.rel_path), + bundle_name=file.bundle_name, + bundle_version=callback_version, + dag_id="test", + run_id="run", + ) + manager._callback_to_execute[callback_file].append(request) + manager._add_files_to_queue([callback_file], mode="front") + + def create_process(_, file, *, source): + manager._callback_to_execute.pop(file) + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = uuid7() + processor.had_callbacks = True + return processor + + create.side_effect = create_process + manager._start_new_processes() + + create.assert_called_once_with( + manager, callback_file, source=DagParseSource.for_callback(callback_version) + ) + assert manager._priority_claims[work.id, work.claim_id].attempt_id is None + manager._processors.pop(callback_file) + manager._start_new_processes() + + assert create.call_count == 1 + assert file in manager._files_awaiting_source + assert file not in manager._file_queue + assert not manager._pending_work_acks + + +@pytest.mark.parametrize("callback_already_running", [False, True]) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_callback_for_removed_file_does_not_restore_acknowledged_priority_request( + create, api_manager, callback_already_running +): + manager, client, bundle, file, work = api_manager + manager._queue_requested_files_for_parsing() + request = DagCallbackRequest( + filepath=str(file.rel_path), bundle_name=bundle.name, bundle_version=None, dag_id="test", run_id="run" + ) + manager._callback_to_execute[file].append(request) + file.absolute_path.unlink() + + manager._refresh_dag_bundles({bundle.name: {file}}) + + assert client.publish_inventory.call_args.args[1].files == [] + assert not manager._priority_claims + assert manager._pending_work_acks == [("priority", work)] + assert file in manager._file_queue + + def create_process(_, file, *, source): + manager._callback_to_execute.pop(file) + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = uuid7() + processor.had_callbacks = True + processor.is_ready = True + processor.start_time = 0 + processor.parse_source = source + processor.parsing_result = None + return processor + + create.side_effect = create_process + if callback_already_running: + manager._processors[file] = manager._create_process(file, source=DagParseSource.for_callback(None)) + manager._start_new_processes() + assert not manager._priority_files + assert not manager._files_awaiting_source + assert not manager._file_queue + manager._collect_results() + manager._acknowledge_requested_work() + manager._start_new_processes() + + create.assert_called_once_with(manager, file, source=DagParseSource.for_callback(None)) + client.acknowledge_work.assert_called_once_with("priority", work) + assert not manager._priority_files + assert not manager._files_awaiting_source + assert not manager._file_queue + assert not manager._processors + + +@pytest.mark.parametrize("file_process_interval", [0, 300]) +@pytest.mark.parametrize("max_runs", [-1, 1]) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_failed_fetch_parks_priority_work_until_new_inventory_without_starving_regular_files( + create, api_manager, max_runs, file_process_interval +): + manager, client, local_bundle, file, work = api_manager + bundle = mock.create_autospec(BaseDagBundle, instance=True) + bundle.name = local_bundle.name + bundle.path = local_bundle.path + bundle.is_initialized = True + bundle.supports_versioning = True + bundle.refresh_changes_source = True + bundle.refresh_interval = 300 + bundle.get_current_version.return_value = "v1" + bundle.refresh.side_effect = DagBundleRefreshError("fetch failed", source_unchanged=True) + manager._dag_bundles = [bundle] + manager._bundle_versions[bundle.name] = "v1" + source = manager._bundle_parse_sources[bundle.name] + manager._bundle_parse_sources[bundle.name] = DagParseSource( + bundle_version="v1", + refresh_generation=source.refresh_generation, + bundle_revision=source.bundle_revision, + inventory_sequence=source.inventory_sequence, + ) + client.get_bundles.return_value[0].version = "v1" + manager.max_runs = max_runs + manager._file_process_interval = file_process_interval + manager._parallelism = 1 + manager._num_run = 1 + healthy_file = replace(file, rel_path=Path("healthy.py")) + (bundle.path / healthy_file.rel_path).write_text("from airflow.sdk import DAG\ndag = DAG('healthy')\n") + known_files = {bundle.name: {file, healthy_file}} + manager._queue_requested_files_for_parsing() + manager._refresh_dag_bundles(known_files) + manager._add_files_to_queue([healthy_file], mode="back") + + def create_process(_, file, *, source): + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = uuid7() + processor.had_callbacks = False + processor.is_ready = True + processor.start_time = 0 + processor.parse_source = source + processor.parsing_result = DagFileParsingResult(fileloc=str(file.rel_path), serialized_dags=[]) + return processor + + create.side_effect = create_process + + manager._start_new_processes() + + create.assert_called_once_with(manager, healthy_file, source=manager._bundle_parse_sources[bundle.name]) + assert file in manager._files_awaiting_source + assert file not in manager._file_queue + assert manager._bundle_parse_sources[bundle.name].refresh_generation == 4 + assert manager._bundle_parse_sources[bundle.name].inventory_sequence == 7 + claim = manager._priority_claims[work.id, work.claim_id] + assert claim.attempt_id is None + client.publish_inventory.assert_not_called() + + manager._collect_results() + manager._publish_pending_results() + assert manager._file_stats[healthy_file].run_count == 1 + assert not manager.max_runs_reached() + for _ in range(3): + manager.prepare_file_queue(known_files) + manager._start_new_processes() + assert file not in manager._processors + assert file in manager._files_awaiting_source + assert file not in manager._file_queue + manager._collect_results() + manager._publish_pending_results() + assert all(call.args[1] == healthy_file for call in create.call_args_list) + assert not manager._pending_work_acks + + bundle.refresh.side_effect = None + manager._refresh_dag_bundles(known_files) + manager._start_new_processes() + + assert claim.attempt_id == manager._processors[file].id + assert file not in manager._files_awaiting_source + assert manager._bundle_parse_sources[bundle.name].inventory_sequence >= claim.inventory_sequence + manager._collect_results() + manager._publish_pending_results() + assert manager._file_stats[file].run_count == 1 + assert manager._pending_work_acks == [("priority", work)] From 09c26e05e79aebcdd45707e27c705b69bb1e9e86 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 22 Oct 2026 13:46:30 +0100 Subject: [PATCH 4/5] Preserve requested Dag processing work during bundle failures Callbacks queued behind an active import must still run, and an unavailable bundle must not starve priority work from healthy bundles. Operators also need consistent heartbeat signals and clear credential requirements when moving processors to the Execution API. --- .../security/jwt_token_authentication.rst | 15 +- airflow-core/docs/security/security_model.rst | 17 +- airflow-core/src/airflow/cli/cli_config.py | 5 +- .../src/airflow/dag_processing/health.py | 11 +- .../src/airflow/dag_processing/manager.py | 70 ++- .../airflow/jobs/dag_processor_job_runner.py | 33 +- .../dag_processing/test_api_isolation.py | 12 +- .../test_dag_processor_client.py | 17 +- .../tests/unit/dag_processing/test_health.py | 49 ++- .../tests/unit/dag_processing/test_manager.py | 155 +++---- .../test_requested_work_recovery.py | 404 ++++++++++++++++++ .../jobs/test_dag_processor_job_runner.py | 82 +++- .../metrics/metrics_template.yaml | 4 +- 13 files changed, 714 insertions(+), 160 deletions(-) create mode 100644 airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py diff --git a/airflow-core/docs/security/jwt_token_authentication.rst b/airflow-core/docs/security/jwt_token_authentication.rst index 405eed4c19f32..f050fcd3fc8e7 100644 --- a/airflow-core/docs/security/jwt_token_authentication.rst +++ b/airflow-core/docs/security/jwt_token_authentication.rst @@ -523,8 +523,11 @@ the Execution API. The API client is created in the processor process, after any daemon or hot-reload fork. Leaving the setting unset preserves the in-process API and database-backed Job lifecycle. -On a trusted host with database access and the server keys, synchronize the configured -bundle catalog and provision a separate token file for each processor:: +Run catalog synchronization and token provisioning on a trusted host. Catalog synchronization +requires database access and the deployment's Fernet key when it signs bundle URLs. Token +provisioning requires the API's JWT signing keys and bundle-name configuration, but does not +access the metadata database. Synchronize the configured bundle catalog and provision a +separate token file for each processor:: airflow dag-processor --sync-bundles-only airflow dag-processor-token --token-file /run/airflow/processor.jwt --bundle-name dags-folder --rotate @@ -702,9 +705,17 @@ Synchronize the bundle catalog and run token provisioning in the trusted deploym described above. Keep the token provisioner running, and mount its directory rather than a single file so atomic token replacement remains visible. +Before starting the isolated processor, stop the trusted deployment's processor for the bundles +being moved. If that processor must continue serving other bundles, restrict it with +``--bundle-name`` to those other bundles. This prevents conflicting inventory reconciliation +and keeps Dag code for the moved bundles out of the database-credentialed processor. + For Compose, use the :download:`standalone processor definition ` as its own project, rather than merging it with the quick-start file. It has no shared ``env_file``, server configuration mount, or dependency on a database/migration service. +It serves ``dags-folder`` from its own ``./dags`` directory. Disable the quick-start deployment's +``airflow-dag-processor`` service unless it is restricted to other bundles: processors that +reconcile different file sets for the same bundle mark each other's missing Dags stale. Create its ``tokens``, ``dags``, and ``logs`` directories with permissions for the configured container user; provision ``tokens/processor.jwt``. Set ``AIRFLOW_IMAGE_NAME`` to the image with API support. Create an internal network and attach only the API server and processor:: diff --git a/airflow-core/docs/security/security_model.rst b/airflow-core/docs/security/security_model.rst index ace40cfc59ae6..ad06fa064a999 100644 --- a/airflow-core/docs/security/security_model.rst +++ b/airflow-core/docs/security/security_model.rst @@ -502,10 +502,13 @@ potentially still executes with direct database access in the Dag File Processor 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 - credentials, the Fernet key and signing keys, but does not remove them: Deployment Managers must - take them out of the processor's environment and block its database network path. The Triggerer - still needs database access. See :doc:`/security/jwt_token_authentication`. + carrying its file and attempt identity. The API does not currently use those identities for + access checks. This does not provide per-file isolation: Dag code runs as the processor's Unix + user and can read its session token file and other credentials available to that user. + This mode removes the processor's need for database credentials, the Fernet key and signing keys, + but does not remove them: Deployment Managers must take them out of the processor's environment + and block its database network path. The Triggerer still needs database access. + See :doc:`/security/jwt_token_authentication`. **Dag File Processor and Triggerer are shared across teams** In the default deployment, a **single Dag File Processor instance** parses all Dag files and a @@ -784,8 +787,10 @@ model — Airflow does not enforce these natively. (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 - metadata database variables, the Fernet key or the JWT signing secret listed for it above; it - authenticates with the token file instead. + metadata database variables, the Fernet key or JWT signing keys; it authenticates with the token + file instead. It also no longer needs asset manager credentials for publishing parsing results. + In this mode, configure ``[core] asset_manager_class`` and ``[core] asset_manager_kwargs`` on the + API Server, which creates assets when it persists those results. .. START AUTOGENERATED PROVIDER SENSITIVE VARS diff --git a/airflow-core/src/airflow/cli/cli_config.py b/airflow-core/src/airflow/cli/cli_config.py index c75f4bd4e149e..b0b2e80a94524 100644 --- a/airflow-core/src/airflow/cli/cli_config.py +++ b/airflow-core/src/airflow/cli/cli_config.py @@ -2317,7 +2317,10 @@ class GroupCommand(NamedTuple): Arg( ("--sync-bundles-only",), action="store_true", - help="Provision the configured bundle catalog using trusted database and signing credentials, then exit", + help=( + "Provision the bundle catalog on a trusted host using database access " + "and the Fernet key for signing bundle URLs, then exit" + ), ), ARG_PID, ARG_DAEMON, diff --git a/airflow-core/src/airflow/dag_processing/health.py b/airflow-core/src/airflow/dag_processing/health.py index 640724a85ef00..618d6c0651674 100644 --- a/airflow-core/src/airflow/dag_processing/health.py +++ b/airflow-core/src/airflow/dag_processing/health.py @@ -50,7 +50,12 @@ def check_processor_health(*, readiness: bool = False) -> None: 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") + for field in fields: + age = now - record[field] + if not 0 <= age <= maximum_age: + raise ValueError( + f"Processor {field} heartbeat age {age:.1f}s is outside 0..{maximum_age}s " + "([dag_processor] health_check_threshold)" + ) except (OSError, ValueError, KeyError, TypeError) as error: - raise SystemExit("Dag processor is not healthy") from error + raise SystemExit(f"Dag processor is not healthy: {error}") from error diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index ceeb25922188f..4995600eb2f5f 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -830,16 +830,27 @@ def claim_priority_files(self) -> list[DagFileInfo]: """ if self.api_client is None: return self._claim_priority_files() + bundles = {bundle.name: bundle for bundle in self._dag_bundles} + unavailable = self._bundles_needing_refresh | { + bundle.name for bundle in bundles.values() if not bundle.is_initialized + } + active_attempts = {process.id for process in self._processors.values()} + active_attempts.update(pending.body.attempt_id for pending in self._pending_publications.values()) available = ( self._parallelism - - len(self._priority_claims) + - 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 [] - bundles = {bundle.name: bundle for bundle in self._dag_bundles} + # 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 if name not in unavailable] try: - work = self.api_client.claim_work("priority", list(bundles), available) + work = self.api_client.claim_work("priority", eligible_bundles, available) except httpx2.HTTPError: self.log.warning("Unable to claim priority parses; retrying on a later loop") return [] @@ -1075,11 +1086,20 @@ def prepare_callback_bundle(self, request: CallbackRequest) -> BaseDagBundle | N raise DagProcessorRegistrationRetired( "The Dag processor registration has ended; restart required" ) from error - self.log.exception( - "Error preparing bundle %s version %s for callback, skipping", - request.bundle_name, - request.bundle_version, - ) + if id(request) in self._callback_claims: + self.log.warning( + "Unable to prepare callback bundle; retrying", + bundle_name=request.bundle_name, + bundle_version=request.bundle_version, + target=self._get_callback_target(request), + error=str(error), + ) + else: + self.log.exception( + "Error preparing bundle %s version %s for callback, skipping", + request.bundle_name, + request.bundle_version, + ) return None return bundle @@ -1127,12 +1147,14 @@ def _defer_claimed_callback(self, request: CallbackRequest) -> None: 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", - type(request).__name__, - request.filepath, - request.bundle_name, - request.bundle_version, - timeout, + "Dropping callback after bundle preparation timed out", + callback_type=type(request).__name__, + filepath=request.filepath, + bundle_name=request.bundle_name, + bundle_version=request.bundle_version, + target=self._get_callback_target(request), + timeout=timeout, + timeout_setting="[dag_processor] job_heartbeat_timeout", ) self._pending_work_acks.append(("callbacks", self._callback_claims.pop(id(request)))) self._remove_deferred_callback(request) @@ -1142,17 +1164,19 @@ def _defer_claimed_callback(self, request: CallbackRequest) -> None: now + min(2 ** (retry.attempts - 1), _MAX_CALLBACK_RETRY_DELAY), retry.expires_at ) - def _log_dropped_lang_sdk_callback(self, request: CallbackRequest) -> None: + @staticmethod + def _get_callback_target(request: CallbackRequest) -> str: if isinstance(request, DagCallbackRequest): - target = f"dag_id={request.dag_id} run_id={request.run_id}" - else: - ti = request.ti - target = f"dag_id={ti.dag_id} run_id={ti.run_id} task_id={ti.task_id}" + return f"dag_id={request.dag_id} run_id={request.run_id}" + ti = request.ti + return f"dag_id={ti.dag_id} run_id={ti.run_id} task_id={ti.task_id}" + + def _log_dropped_lang_sdk_callback(self, request: CallbackRequest) -> None: self.log.warning( "Dropping %s for %s (%s): Lang-SDK runtimes do not run callbacks", type(request).__name__, request.filepath, - target, + self._get_callback_target(request), ) @provide_session @@ -1341,6 +1365,7 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): any_refreshed = True except Exception as e: self.log.exception("Error initializing bundle %s: %s", bundle.name, e) + self._bundles_needing_refresh.add(bundle.name) continue try: bundle_state = ( @@ -2366,8 +2391,9 @@ def _start_new_processes(self): 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: - if priority_claims: - self._priority_files.discard(file) + if priority_claims or file in self._callback_to_execute: + if self.api_client is not None: + self._priority_files.discard(file) requested_later[file] = None elif self.api_client is None and file in self._priority_files: self._files_awaiting_source[file] = None diff --git a/airflow-core/src/airflow/jobs/dag_processor_job_runner.py b/airflow-core/src/airflow/jobs/dag_processor_job_runner.py index 3af40ae3f91be..c3382e130551a 100644 --- a/airflow-core/src/airflow/jobs/dag_processor_job_runner.py +++ b/airflow-core/src/airflow/jobs/dag_processor_job_runner.py @@ -24,6 +24,7 @@ from airflow._shared.observability.metrics import stats from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.job import TerminalJobState from airflow.configuration import conf from airflow.dag_processing.health import write_processor_health from airflow.jobs.base_job_runner import BaseJobRunner @@ -70,10 +71,7 @@ def __init__( ) 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 - + """Mirror ``run_job`` metrics and listener ordering while the API owns Job persistence.""" self.processor.api_client = client self.job.id = self._register_api_job(client) self.job.state = JobState.RUNNING @@ -88,12 +86,12 @@ def run_with_api(self, client: DagProcessorAPIClient) -> int | None: 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") + self.job.end_date = timezone.utcnow() try: client.complete_job(TerminalJobState(self.job.state)) except Exception: @@ -104,6 +102,7 @@ def run_with_api(self, client: DagProcessorAPIClient) -> int | None: stats.incr("job_end", 1, 1) def _register_api_job(self, client: DagProcessorAPIClient) -> int: + # Client imports stay local: its versioned routes import this runner. from airflow.dag_processing.api_client import DagProcessorJobAlreadyRunning deadline = monotonic() + conf.getint("dag_processor", "health_check_threshold") + self.job.heartrate @@ -134,12 +133,24 @@ def _heartbeat_api(self, client: DagProcessorAPIClient) -> None: 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") + stats.incr("job_heartbeat_failure", 1, 1) + elapsed = monotonic() - self._last_api_heartbeat + timeout = conf.getint("dag_processor", "job_heartbeat_timeout") + error_type = type(error).__name__ + status_code = error.response.status_code if isinstance(error, httpx2.HTTPStatusError) else None + if elapsed >= timeout: + reason = f"{error_type} (HTTP {status_code})" if status_code is not None else error_type + raise DagProcessorHeartbeatTimeout( + f"Dag processor API heartbeat timed out after {elapsed:.1f}s without a successful heartbeat " + f"([dag_processor] job_heartbeat_timeout={timeout}s; {reason})" + ) from error + self.log.warning( + "Unable to heartbeat the Dag processor Job; retrying on the next interval", + error_type=error_type, + status_code=status_code, + elapsed_seconds=elapsed, + timeout_seconds=timeout, + ) return if client.restart_required or state != JobState.RUNNING: raise DagProcessorRegistrationRetired("The API requested a Dag processor restart") diff --git a/airflow-core/tests/integration/dag_processing/test_api_isolation.py b/airflow-core/tests/integration/dag_processing/test_api_isolation.py index 6701c7f4a26c9..fbb4224e7ea40 100644 --- a/airflow-core/tests/integration/dag_processing/test_api_isolation.py +++ b/airflow-core/tests/integration/dag_processing/test_api_isolation.py @@ -145,7 +145,14 @@ def test_processor_without_database_network_or_server_credentials( "AIRFLOW_USER_HOME_DIR": "/tmp/airflow-user", "AIRFLOW_HOME": "/tmp/isolated-airflow", "AIRFLOW_CONFIG": "/tmp/isolated-airflow/airflow.cfg", - "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN": f"postgresql+psycopg://blocked:blocked@{database_ip}:{database_port}/blocked", + "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN": database.set( + username="blocked", + password="blocked", + host=database_ip, + port=database_port, + database="blocked", + query={}, + ).render_as_string(hide_password=False), "AIRFLOW__CORE__FERNET_KEY": "", "AIRFLOW__API_AUTH__JWT_SECRET": "", "AIRFLOW__API_AUTH__JWT_PRIVATE_KEY_PATH": "", @@ -256,7 +263,8 @@ def test_processor_without_database_network_or_server_credentials( session.expire_all() dag = session.get(DagModel, "isolated") delivered = session.get(DagProcessorCallback, callback_id) - if dag and delivered is None: + priority = session.scalar(select(DagPriorityParsingRequest)) + if dag and delivered is None and priority is None: break if process.poll() is not None: output.seek(0) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py b/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py index 087633dcd1cf4..3f2e3992fb519 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/test_dag_processor_client.py @@ -30,6 +30,7 @@ import pytest import uvicorn from sqlalchemy import select, update +from sqlalchemy.engine import make_url from tenacity import wait_none from uuid6 import uuid7 @@ -450,6 +451,7 @@ def test_normal_processor_command_uses_authenticated_api( ] config = { ("dag_processor", "execution_api_token_file"): str(provision_token()), + ("dag_processor", "health_check_file"): str(tmp_path / "processor-health.json"), ("core", "execution_api_server_url"): api_url, ("core", "load_examples"): "False", ("core", "multi_team"): str(multi_team), @@ -468,12 +470,24 @@ def test_normal_processor_command_uses_authenticated_api( session.commit() processor_config = dict(config) if not database_access: + database = make_url(conf.get("database", "sql_alchemy_conn")) + if database.get_backend_name() == "sqlite": + blocked_database = database.set(database=str(tmp_path / "unavailable" / "metadata.db")) + else: + blocked_database = database.set( + username="blocked", + password="blocked", + host="127.0.0.1", + port=1, + database="blocked", + query={}, + ) processor_config.update( { ( "database", "sql_alchemy_conn", - ): "postgresql+psycopg://blocked:blocked@127.0.0.1:1/blocked", + ): blocked_database.render_as_string(hide_password=False), ("core", "fernet_key"): "", ("api_auth", "jwt_secret"): "", ("api_auth", "jwt_private_key_path"): "", @@ -495,6 +509,7 @@ def test_normal_processor_command_uses_authenticated_api( check=False, ) assert result.returncode == 0, result.stdout + result.stderr + assert (tmp_path / "processor-health.json").is_file() session.expire_all() job = session.scalars(select(Job)).one() diff --git a/airflow-core/tests/unit/dag_processing/test_health.py b/airflow-core/tests/unit/dag_processing/test_health.py index ffd118fdb31e7..aaf04aa69ec12 100644 --- a/airflow-core/tests/unit/dag_processing/test_health.py +++ b/airflow-core/tests/unit/dag_processing/test_health.py @@ -17,6 +17,9 @@ from __future__ import annotations +import json +import subprocess +import sys from unittest import mock import pytest @@ -51,14 +54,52 @@ def test_local_health_requires_a_fresh_loop(clock, tmp_path, loop_age, api_age, if healthy: check_processor_health(readiness=readiness) else: - with pytest.raises(SystemExit): + field, age = ("loop", loop_age) if not 0 <= loop_age <= 30 else ("api", api_age) + with pytest.raises( + SystemExit, + match=rf"Processor {field} heartbeat age {age:.1f}s.*health_check_threshold", + ): check_processor_health(readiness=readiness) -@pytest.mark.parametrize("contents", [None, "broken", "{}"]) -def test_missing_or_invalid_health_file_fails_closed(tmp_path, contents): +@pytest.mark.parametrize( + ("contents", "cause", "detail"), + [ + (None, FileNotFoundError, "health.json"), + ("broken", json.JSONDecodeError, "Expecting value"), + ("{}", KeyError, "pid"), + ], +) +def test_missing_or_invalid_health_file_fails_closed(tmp_path, contents, cause, detail): path = tmp_path / "health.json" if contents is not None: path.write_text(contents) - with conf_vars({("dag_processor", "health_check_file"): str(path)}), pytest.raises(SystemExit): + with ( + conf_vars({("dag_processor", "health_check_file"): str(path)}), + pytest.raises(SystemExit, match=f"Dag processor is not healthy: .*{detail}") as raised, + ): check_processor_health() + assert isinstance(raised.value.__cause__, cause) + + +def test_dead_processor_fails_health_with_a_fresh_heartbeat(tmp_path): + path = tmp_path / "health.json" + with subprocess.Popen([sys.executable, "-c", "pass"]) as process: + process.wait(timeout=10) + with conf_vars({("dag_processor", "health_check_file"): str(path)}): + write_processor_health(0) + record = json.loads(path.read_text()) + record["pid"] = process.pid + path.write_text(json.dumps(record)) + with pytest.raises(SystemExit, match="Dag processor is not healthy: .*No such process") as raised: + check_processor_health() + assert isinstance(raised.value.__cause__, ProcessLookupError) + + +@mock.patch("airflow.dag_processing.health.os.kill", autospec=True) +def test_inaccessible_processor_reports_the_probe_error(kill, tmp_path): + kill.side_effect = PermissionError("Operation not permitted") + with conf_vars({("dag_processor", "health_check_file"): str(tmp_path / "health.json")}): + write_processor_health(0) + with pytest.raises(SystemExit, match="Dag processor is not healthy: Operation not permitted"): + check_processor_health() diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 2340fb9c190b2..1e462f14e11f9 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -504,11 +504,10 @@ 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"]) @pytest.mark.parametrize("absolute", [False, True], ids=["relative", "absolute"]) @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) def test_find_files_in_bundle_makes_discovery_errors_bundle_relative( - self, mock_get_registry, tmp_path, absolute, api_mode + self, mock_get_registry, tmp_path, absolute ): (tmp_path / "broken.dag").write_text("") source_reference = os.fspath(tmp_path / "broken.dag") if absolute else "broken.dag" @@ -516,9 +515,7 @@ def test_find_files_in_bundle_makes_discovery_errors_bundle_relative( (mock.sentinel.importer, DagImportError(source_reference=source_reference, message="boom")) ] bundle = _make_bundle(tmp_path) - client = mock.create_autospec(DagProcessorAPIClient, instance=True) if api_mode else None - - found_files = DagFileProcessorManager(max_runs=1, api_client=client)._find_files_in_bundle(bundle) + found_files = DagFileProcessorManager(max_runs=1)._find_files_in_bundle(bundle) assert [file.definition_locs for file in found_files.files] == [frozenset({"broken.dag"})] assert found_files.failed_sources == frozenset({"broken.dag"}) @@ -2162,7 +2159,7 @@ def test_forced_refresh_waits_for_in_flight_imports(self, create, persist, curre manager._start_new_processes() self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version="v1"), force=True + manager, [bundle], BundleState(last_refreshed=None, version="v1"), force=True ) bundle.refresh.assert_not_called() assert manager._bundles_waiting_for_refresh == {bundle.name} @@ -2172,7 +2169,7 @@ def test_forced_refresh_waits_for_in_flight_imports(self, create, persist, curre assert not manager._file_queue self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version="v1"), force=True + manager, [bundle], BundleState(last_refreshed=None, version="v1"), force=True ) bundle.refresh.assert_called_once_with() @@ -2206,7 +2203,7 @@ def test_reinitialized_bundle_discards_only_affected_results_without_advancing_f bundle.is_initialized = False self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version="v1"), force=True + manager, [bundle], BundleState(last_refreshed=None, version="v1"), force=True ) manager._collect_results() @@ -2666,13 +2663,13 @@ def test_failed_refresh_blocks_dispatch_until_the_source_recovers(self, create, self._configure_mock_processes(create) state = BundleState(last_refreshed=timezone.utcnow(), version="v1") - self._refresh_with_mocked_state(manager, bundle, state, force=True) + self._refresh_with_mocked_state(manager, [bundle], state, force=True) manager._start_new_processes() create.assert_not_called() assert not manager._file_queue assert list(manager._files_awaiting_source) == [file] - self._refresh_with_mocked_state(manager, bundle, state) + self._refresh_with_mocked_state(manager, [bundle], state) manager._start_new_processes() create.assert_called_once() assert create.return_value.parse_source.refresh_generation == 2 @@ -2745,7 +2742,7 @@ def test_refresh_failure_preserves_dispatch_only_with_an_explicit_unchanged_sour self._configure_mock_processes(create) self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version="v1"), force=force + manager, [bundle], BundleState(last_refreshed=None, version="v1"), force=force ) manager._start_new_processes() @@ -2776,16 +2773,16 @@ def test_unchanged_source_failure_retries_at_the_next_periodic_check(self, monot bundle.refresh.side_effect = [failure, failure, None] monotonic.return_value = 30 - self._refresh_with_mocked_state(manager, bundle, state, force=True) + self._refresh_with_mocked_state(manager, [bundle], state, force=True) assert manager._bundles_needing_refresh == {bundle.name} assert not manager._force_refresh_bundles monotonic.return_value = 31 - self._refresh_with_mocked_state(manager, bundle, state) + self._refresh_with_mocked_state(manager, [bundle], state) assert bundle.refresh.call_count == 1 monotonic.return_value = 60 - self._refresh_with_mocked_state(manager, bundle, state) + self._refresh_with_mocked_state(manager, [bundle], state) assert bundle.refresh.call_count == 2 warnings = [ event @@ -2797,7 +2794,7 @@ def test_unchanged_source_failure_retries_at_the_next_periodic_check(self, monot assert warnings[0]["log_level"] == "warning" monotonic.return_value = 90 - self._refresh_with_mocked_state(manager, bundle, state) + self._refresh_with_mocked_state(manager, [bundle], state) assert bundle.refresh.call_count == 3 assert not manager._bundles_needing_refresh @@ -2973,7 +2970,7 @@ def test_periodic_refresh_drains_imports_before_admitting_more( manager._file_queue[queued] = None self._configure_mock_processes(create) - self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version="v1")) + self._refresh_with_mocked_state(manager, [bundle], BundleState(last_refreshed=None, version="v1")) manager._start_new_processes() bundle.refresh.assert_not_called() create.assert_not_called() @@ -2983,7 +2980,7 @@ def test_periodic_refresh_drains_imports_before_admitting_more( manager._collect_results() if api_mode: persist.assert_not_called() - self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version="v1")) + self._refresh_with_mocked_state(manager, [bundle], BundleState(last_refreshed=None, version="v1")) bundle.refresh.assert_not_called() manager._publish_pending_results() assert client.publish_parse_result.call_args.args[0].bundle_version == "v1" @@ -2991,7 +2988,7 @@ def test_periodic_refresh_drains_imports_before_admitting_more( persist.assert_called_once() assert persist.call_args.kwargs["bundle_version"] == "v1" self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=timezone.utcnow(), version="v1") + manager, [bundle], BundleState(last_refreshed=timezone.utcnow(), version="v1") ) manager._start_new_processes() @@ -3026,7 +3023,7 @@ def test_refresh_drains_unversioned_callbacks_before_admitting_more(self, create manager._processors = {finished: finished_processor, active: active_processor} state = BundleState(last_refreshed=timezone.utcnow() if force else None, version="v1") - self._refresh_with_mocked_state(manager, bundle, state, force=force) + self._refresh_with_mocked_state(manager, [bundle], state, force=force) manager._collect_results() assert list(manager._processors) == [active] @@ -3035,7 +3032,7 @@ def test_refresh_drains_unversioned_callbacks_before_admitting_more(self, create manager._file_queue = OrderedDict.fromkeys([regular, callback, pinned]) self._configure_mock_processes(create, pinned_processor, callback_processor, regular_processor) - self._refresh_with_mocked_state(manager, bundle, state, force=force) + self._refresh_with_mocked_state(manager, [bundle], state, force=force) manager._start_new_processes() bundle.refresh.assert_not_called() @@ -3047,7 +3044,7 @@ def test_refresh_drains_unversioned_callbacks_before_admitting_more(self, create active_processor._open_sockets.clear() manager._collect_results() assert list(manager._processors) == [pinned] - self._refresh_with_mocked_state(manager, bundle, state, force=force) + self._refresh_with_mocked_state(manager, [bundle], state, force=force) manager._start_new_processes() bundle.refresh.assert_called_once_with() @@ -5179,32 +5176,25 @@ def _make_refresh_bundle(self, *, supports_versioning=False, current_version=Non bundle.get_current_version.return_value = current_version return bundle - def _refresh_with_mocked_state(self, manager, bundle, initial_state, *, force=False): - """Run _refresh_dag_bundles with get/update_bundle_state mocked out. - - Returns the two MagicMock objects for post-call assertions. MagicMock retains its - call records after the ``with`` block exits (un-patching only restores the original - attribute; it does not clear the mock's recorded calls), so callers can assert on - them normally after this method returns. - """ - manager._dag_bundles = [bundle] - manager._force_refresh_bundles = {bundle.name} if force else set() - mock_get = mock.patch.object(manager, "get_bundle_state", return_value=initial_state) - mock_update = mock.patch.object(manager, "update_bundle_state") - mock_catalog = ( - mock.patch.object( - manager.api_client, - "get_bundles", - return_value=[ - DagBundleStateResponse( - name=bundle.name, - last_refreshed=initial_state.last_refreshed, - version=initial_state.version, - revision=initial_state.revision, - team_name=None, - ) - ], + def _refresh_with_mocked_state(self, manager, bundles, initial_state, *, force=False): + """Refresh bundles with mocked state and return the get/update mocks for assertions.""" + manager._dag_bundles = bundles + manager._force_refresh_bundles = {bundle.name for bundle in bundles} if force else set() + mock_get = mock.patch.object(manager, "get_bundle_state", autospec=True, return_value=initial_state) + mock_update = mock.patch.object(manager, "update_bundle_state", autospec=True) + get_catalog = mock.create_autospec(DagProcessorAPIClient, instance=True).get_bundles + get_catalog.return_value = [ + DagBundleStateResponse( + name=bundle.name, + last_refreshed=initial_state.last_refreshed, + version=initial_state.version, + revision=initial_state.revision, + team_name=None, ) + for bundle in bundles + ] + mock_catalog = ( + mock.patch.object(manager.api_client, "get_bundles", new=get_catalog) if manager.api_client is not None else nullcontext() ) @@ -5215,13 +5205,14 @@ def _refresh_with_mocked_state(self, manager, bundle, initial_state, *, force=Fa mock.patch.object( manager, "_find_files_in_bundle", + autospec=True, return_value=DagBundleDiscoveryResult(files=set(), failed_sources=frozenset()), ), - mock.patch.object(manager, "deactivate_deleted_dags"), - mock.patch.object(manager, "clear_orphaned_import_errors"), - mock.patch.object(manager, "handle_removed_files"), - mock.patch.object(manager, "_resort_file_queue"), - mock.patch.object(manager, "_add_new_files_to_queue"), + mock.patch.object(manager, "deactivate_deleted_dags", autospec=True), + mock.patch.object(manager, "clear_orphaned_import_errors", autospec=True), + mock.patch.object(manager, "handle_removed_files", autospec=True), + mock.patch.object(manager, "_resort_file_queue", autospec=True), + mock.patch.object(manager, "_add_new_files_to_queue", autospec=True), ): manager._refresh_dag_bundles({}) return patched_get, patched_update @@ -5232,7 +5223,7 @@ def test_refresh_dag_bundles_non_versioned_calls_update_bundle_state(self): bundle = self._make_refresh_bundle(supports_versioning=False) mock_get, mock_update = self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version=None) + manager, [bundle], BundleState(last_refreshed=None, version=None) ) mock_get.assert_called_once_with("mock_bundle") @@ -5244,7 +5235,7 @@ def test_refresh_dag_bundles_clears_team_name_cache(self): manager._bundle_name_to_team_name = {"stale_bundle": "old_team"} bundle = self._make_refresh_bundle(supports_versioning=False) - self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version=None)) + self._refresh_with_mocked_state(manager, [bundle], BundleState(last_refreshed=None, version=None)) assert manager._bundle_name_to_team_name == {} @@ -5256,7 +5247,7 @@ def test_refresh_dag_bundles_versioned_version_changed_calls_update_bundle_state manager._bundle_versions["mock_bundle"] = "v1" mock_get, mock_update = self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version="v1") + manager, [bundle], BundleState(last_refreshed=None, version="v1") ) mock_get.assert_called_once_with("mock_bundle") @@ -5271,7 +5262,7 @@ def test_refresh_dag_bundles_versioned_version_unchanged_calls_update_bundle_sta manager._bundle_versions["mock_bundle"] = "v1" mock_get, mock_update = self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version="v1") + manager, [bundle], BundleState(last_refreshed=None, version="v1") ) mock_get.assert_called_once_with("mock_bundle") @@ -5327,7 +5318,7 @@ def test_refresh_dag_bundles_versioned_first_seen_skips_short_circuit(self): # _bundle_versions is empty → previously_seen=False mock_get, mock_update = self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version="v1") + manager, [bundle], BundleState(last_refreshed=None, version="v1") ) mock_get.assert_called_once_with("mock_bundle") @@ -5361,25 +5352,9 @@ def test_refresh_dag_bundles_initialize_non_airflow_exception_skips_bundle(self) healthy_bundle.name = "healthy_bundle" manager._dag_bundles = [failing_bundle, healthy_bundle] - manager._force_refresh_bundles = set() - - with ( - mock.patch.object( - manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None) - ), - mock.patch.object(manager, "update_bundle_state"), - mock.patch.object( - manager, - "_find_files_in_bundle", - return_value=DagBundleDiscoveryResult(files=set(), failed_sources=frozenset()), - ), - mock.patch.object(manager, "deactivate_deleted_dags"), - mock.patch.object(manager, "clear_orphaned_import_errors"), - mock.patch.object(manager, "handle_removed_files"), - mock.patch.object(manager, "_resort_file_queue"), - mock.patch.object(manager, "_add_new_files_to_queue"), - ): - manager._refresh_dag_bundles({}) + self._refresh_with_mocked_state( + manager, manager._dag_bundles, BundleState(last_refreshed=None, version=None) + ) healthy_bundle.refresh.assert_called_once() @@ -5576,7 +5551,9 @@ def test_bundle_version_data_stored_after_refresh(self, session): # Pre-populate so previously_seen=True and version differs manager._bundle_versions["mock_bundle"] = "oldhash" - self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version="oldhash")) + self._refresh_with_mocked_state( + manager, [bundle], BundleState(last_refreshed=None, version="oldhash") + ) test_data["files"]["dag.py"] = "changed" assert manager._bundle_versions["mock_bundle"] == "newhash" @@ -6055,7 +6032,7 @@ def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed manager._bundle_versions[bundle.name] = "previous" self._refresh_with_mocked_state( - manager, bundle, BundleState(last_refreshed=None, version=None), force=True + manager, [bundle], BundleState(last_refreshed=None, version=None), force=True ) bundle.refresh.assert_called_once() @@ -6074,7 +6051,7 @@ def test_listing_the_api_cannot_accept_releases_priority_capacity(self, observed observed.return_value = set() client.publish_inventory.return_value.revision = uuid7() - self._refresh_with_mocked_state(manager, bundle, BundleState(last_refreshed=None, version=None)) + self._refresh_with_mocked_state(manager, [bundle], BundleState(last_refreshed=None, version=None)) assert bundle.name in manager._bundle_parse_sources assert manager._bundle_versions[bundle.name] is None assert not manager._pending_inventories @@ -6089,25 +6066,9 @@ def test_refresh_dag_bundles_heartbeats_before_each_bundle(self): bundle.name = name bundle.refresh.side_effect = lambda name=name: events.append(name) manager._dag_bundles.append(bundle) - manager._force_refresh_bundles = set() - - with ( - mock.patch.object( - manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None) - ), - mock.patch.object(manager, "update_bundle_state"), - mock.patch.object( - manager, - "_find_files_in_bundle", - return_value=DagBundleDiscoveryResult(files=set(), failed_sources=frozenset()), - ), - mock.patch.object(manager, "deactivate_deleted_dags"), - mock.patch.object(manager, "clear_orphaned_import_errors"), - mock.patch.object(manager, "handle_removed_files"), - mock.patch.object(manager, "_resort_file_queue"), - mock.patch.object(manager, "_add_new_files_to_queue"), - ): - manager._refresh_dag_bundles({}) + self._refresh_with_mocked_state( + manager, manager._dag_bundles, BundleState(last_refreshed=None, version=None) + ) assert events == ["heartbeat", "first", "heartbeat", "second"] diff --git a/airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py b/airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py new file mode 100644 index 0000000000000..08f4e1897df81 --- /dev/null +++ b/airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py @@ -0,0 +1,404 @@ +# 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 pathlib import Path +from unittest import mock + +import httpx2 +import pytest +from uuid6 import uuid7 + +from airflow._shared.timezones import timezone +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleInventoryResponse, + DagBundleStateResponse, + DagParseResultBody, + ProcessorWorkItem, +) +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TaskInstance +from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest +from airflow.dag_processing.api_client import DagProcessorAPIClient +from airflow.dag_processing.bundles.base import BaseDagBundle, DagBundleRefreshError +from airflow.dag_processing.bundles.local import LocalDagBundle +from airflow.dag_processing.manager import ( + DagFileInfo, + DagFileProcessorManager, + DeferredCallback, + PendingDagPublication, + PriorityParseClaim, +) +from airflow.dag_processing.processor import DagFileProcessorProcess, DagParseSource + +from tests_common.test_utils.config import conf_vars + + +def _make_work(bundle_name, relative_fileloc="a.py", *, callback=None): + return ProcessorWorkItem( + id=str(uuid7()), + claim_id=uuid7(), + bundle_name=bundle_name, + relative_fileloc=relative_fileloc, + callback=callback.model_dump_json() if callback is not None else None, + ) + + +@pytest.fixture +def api_manager(tmp_path): + ready_path = tmp_path / "ready" + ready_path.mkdir() + bundle = LocalDagBundle(name="ready", path=str(ready_path), refresh_interval=0) + bundle.is_initialized = True + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.restart_required = False + client.claim_work.return_value = [] + client.get_bundles.return_value = [ + DagBundleStateResponse( + name=bundle.name, version=None, last_refreshed=None, revision=None, team_name=None + ) + ] + client.publish_inventory.side_effect = lambda name, body: DagBundleInventoryResponse( + attempt_id=body.attempt_id, revision=uuid7(), accepted_at=timezone.utcnow() + ) + manager = DagFileProcessorManager( + max_runs=1, api_client=client, parallelism=2, bundle_refresh_check_interval=0 + ) + manager._dag_bundles = [bundle] + manager._bundle_parse_sources[bundle.name] = DagParseSource() + return manager, client, bundle + + +@pytest.mark.parametrize("bundle_version", [None, "v1"]) +@mock.patch.object(DagFileProcessorProcess, "start", autospec=True) +@mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True) +@mock.patch.object(DagFileProcessorManager, "prepare_callback_bundle", autospec=True) +def test_callback_queued_behind_running_callback_dispatches_and_acknowledges( + prepare, get_logger, start, api_manager, bundle_version +): + manager, client, bundle = api_manager + prepare.return_value = bundle + get_logger.return_value = (mock.Mock(spec=[]), mock.Mock(spec=[])) + + def start_process(**kwargs): + process = mock.create_autospec(DagFileProcessorProcess, instance=True) + process.id = kwargs["id"] + process.had_callbacks = True + process.parse_source = kwargs["parse_source"] + return process + + start.side_effect = start_process + requests = [ + DagCallbackRequest( + filepath="a.py", bundle_name=bundle.name, bundle_version=bundle_version, dag_id="dag", run_id=run + ) + for run in ("first", "second") + ] + work = [_make_work(bundle.name, callback=request) for request in requests] + for item in work: + client.claim_work.return_value = [item] + for request in manager.fetch_callbacks(): + manager._add_callback_to_queue(request) + manager._start_new_processes() + + file, first = next(iter(manager._processors.items())) + assert start.call_count == 1 + assert list(manager._file_queue) == [file] + assert [request.run_id for request in manager._callback_to_execute[file]] == ["second"] + assert list(manager._callback_claims.values()) == [work[1]] + + manager._processors.pop(file) + manager._finish_callback_attempt(first) + manager._start_new_processes() + second = manager._processors.pop(file) + manager._finish_callback_attempt(second) + manager._acknowledge_requested_work() + + assert [call.kwargs["callbacks"][0].run_id for call in start.call_args_list] == ["first", "second"] + assert client.acknowledge_work.call_args_list == [mock.call("callbacks", item) for item in work] + assert not manager._callback_claims + assert not manager._callback_attempts + assert not manager._callback_to_execute + assert not manager._file_queue + + +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_queued_callback_preserves_db_priority_request_behind_running_parse(create, api_manager): + manager, _, bundle = api_manager + manager.api_client = None + file = DagFileInfo(bundle_name=bundle.name, bundle_path=bundle.path, rel_path=Path("a.py")) + running = mock.create_autospec(DagFileProcessorProcess, instance=True) + manager._processors[file] = running + manager._add_files_to_queue([file], mode="frontprio") + request = DagCallbackRequest( + filepath="a.py", bundle_name=bundle.name, bundle_version=None, dag_id="dag", run_id="run" + ) + manager._add_callback_to_queue(request) + + def create_process(_, file, *, source): + process = mock.create_autospec(DagFileProcessorProcess, instance=True) + process.id = uuid7() + process.had_callbacks = bool(manager._callback_to_execute.pop(file, [])) + process.parse_source = source + return process + + create.side_effect = create_process + manager._start_new_processes() + create.assert_not_called() + assert file in manager._file_queue + assert file in manager._priority_files + + del manager._processors[file] + manager._start_new_processes() + assert manager._processors[file].had_callbacks + assert file in manager._priority_files + assert file in manager._files_awaiting_source + + del manager._processors[file] + manager._start_new_processes() + assert not manager._processors[file].had_callbacks + assert create.call_count == 2 + assert file not in manager._priority_files + assert file not in manager._files_awaiting_source + + +@pytest.mark.parametrize("failure", ["initialize", "refresh_changed", "refresh_unchanged"]) +@mock.patch.object(DagFileProcessorManager, "_create_process", autospec=True) +def test_failed_bundle_claims_allow_healthy_priority_work_then_resume(create, api_manager, tmp_path, failure): + manager, client, ready = api_manager + broken = mock.create_autospec(BaseDagBundle, instance=True) + broken.name = "broken" + broken.path = tmp_path / "broken" + broken.path.mkdir() + broken.is_initialized = True + broken.supports_versioning = False + broken.refresh_changes_source = True + broken.refresh_interval = 0 + manager._dag_bundles.append(broken) + client.get_bundles.return_value.append( + DagBundleStateResponse( + name=broken.name, version=None, last_refreshed=None, revision=None, team_name=None + ) + ) + manager._bundle_parse_sources[broken.name] = DagParseSource() + for filename in ("a.py", "b.py"): + (broken.path / filename).write_text("from airflow.sdk import DAG\ndag = DAG('broken')\n") + (ready.path / "ready.py").write_text("from airflow.sdk import DAG\ndag = DAG('ready')\n") + work = [_make_work(broken.name, filename) for filename in ("a.py", "b.py")] + client.claim_work.return_value = work + manager._queue_requested_files_for_parsing() + if failure == "initialize": + broken.is_initialized = False + broken.initialize.side_effect = OSError("unavailable") + else: + broken.refresh.side_effect = DagBundleRefreshError( + "unavailable", source_unchanged=failure == "refresh_unchanged" + ) + known_files = {} + manager._refresh_dag_bundles(known_files) + + assert manager._bundles_needing_refresh == {broken.name} + assert len(manager._priority_claims) == 2 + assert not manager._pending_work_acks + healthy_work = _make_work(ready.name, "ready.py") + client.claim_work.return_value = [healthy_work] + manager._queue_requested_files_for_parsing() + client.claim_work.assert_called_with("priority", [ready.name], 2) + client.claim_work.return_value = [] + for _ in range(3): + manager.claim_priority_files() + client.claim_work.assert_called_with("priority", [ready.name], 1) + assert len(manager._priority_claims) == 3 + + def create_process(_, file, *, source): + process = mock.create_autospec(DagFileProcessorProcess, instance=True) + process.id = uuid7() + process.had_callbacks = False + process.parse_source = source + return process + + create.side_effect = create_process + manager._start_new_processes() + create.assert_not_called() + manager._refresh_dag_bundles(known_files) + manager._start_new_processes() + assert {file.bundle_name for file in manager._processors} == {ready.name} + _, process = manager._processors.popitem() + manager._finish_priority_claims(process.id) + manager._acknowledge_requested_work() + client.acknowledge_work.assert_called_once_with("priority", healthy_work) + assert len(manager._priority_claims) == 2 + + broken.initialize.side_effect = lambda: setattr(broken, "is_initialized", True) + broken.refresh.side_effect = None + manager._refresh_dag_bundles(known_files) + manager._start_new_processes() + + assert not manager._bundles_needing_refresh + assert len(manager._processors) == 2 + assert {file.bundle_name for file in manager._processors} == {broken.name} + for file, process in list(manager._processors.items()): + claim = next(claim for claim in manager._priority_claims.values() if claim.attempt_id == process.id) + assert process.parse_source.inventory_sequence >= claim.inventory_sequence + manager._finish_priority_claims(process.id) + del manager._processors[file] + manager._acknowledge_requested_work() + assert not manager._priority_claims + assert client.acknowledge_work.call_args_list == [ + mock.call("priority", item) for item in [healthy_work, *work] + ] + + +def test_priority_claims_waiting_for_fresh_inventory_keep_the_budget_bounded(api_manager): + manager, client, bundle = api_manager + client.claim_work.return_value = [_make_work(bundle.name, filename) for filename in ("a.py", "b.py")] + manager._queue_requested_files_for_parsing() + + assert manager.claim_priority_files() == [] + + client.claim_work.assert_called_once_with("priority", [bundle.name], 2) + assert len(manager._priority_claims) == 2 + + +@pytest.mark.parametrize("stage", ["running", "publishing", "acknowledging", "stale"]) +def test_failed_bundle_only_suspends_priority_claims_without_active_attempts(api_manager, stage): + manager, client, bundle = api_manager + manager._parallelism = 1 + manager._bundles_needing_refresh.add(bundle.name) + file = DagFileInfo(bundle_name=bundle.name, bundle_path=bundle.path, rel_path=Path("a.py")) + work = _make_work(bundle.name) + attempt_id = uuid7() + if stage == "acknowledging": + manager._pending_work_acks.append(("priority", work)) + else: + manager._priority_claims[work.id, work.claim_id] = PriorityParseClaim( + work=work, inventory_sequence=1, attempt_id=attempt_id + ) + if stage == "running": + process = mock.create_autospec(DagFileProcessorProcess, instance=True) + process.id = attempt_id + manager._processors[file] = process + elif stage == "publishing": + pending = mock.create_autospec(PendingDagPublication, instance=True) + pending.body = mock.create_autospec(DagParseResultBody, instance=True) + pending.body.attempt_id = attempt_id + manager._pending_publications[file] = pending + + assert manager.claim_priority_files() == [] + + if stage == "stale": + client.claim_work.assert_called_once_with("priority", [], 1) + assert manager._priority_claims[work.id, work.claim_id].attempt_id == attempt_id + else: + client.claim_work.assert_not_called() + + +def test_uninitialized_bundle_is_not_claimed_until_initialization_recovers(api_manager): + manager, client, bundle = api_manager + bundle.is_initialized = False + + manager.claim_priority_files() + + client.claim_work.assert_called_once_with("priority", [], 2) + manager._refresh_dag_bundles({}) + assert bundle.is_initialized + manager.claim_priority_files() + client.claim_work.assert_called_with("priority", [bundle.name], 2) + + +@pytest.mark.parametrize("kind", ["callbacks", "priority"]) +@pytest.mark.parametrize("failure", ["connection", "503"]) +def test_claim_failure_preserves_deferred_callbacks_and_claims(api_manager, kind, failure): + manager, client, bundle = api_manager + request = DagCallbackRequest( + filepath="a.py", bundle_name=bundle.name, bundle_version="v1", dag_id="dag", run_id="run" + ) + deferred = [request, request.model_copy(deep=True)] + manager._deferred_api_callbacks = deferred.copy() + claims = {id(callback): _make_work(bundle.name, callback=callback) for callback in deferred} + manager._callback_claims = claims.copy() + retries = {id(request): DeferredCallback(expires_at=100, next_attempt_time=10, attempts=2)} + manager._deferred_callback_retries = retries.copy() + if failure == "connection": + client.claim_work.side_effect = httpx2.ConnectError("unavailable") + else: + response = httpx2.Response(503, request=httpx2.Request("POST", "http://api/claim")) + client.claim_work.side_effect = httpx2.HTTPStatusError( + "unavailable", request=response.request, response=response + ) + + result = manager.fetch_callbacks() if kind == "callbacks" else manager.claim_priority_files() + + assert result == (deferred if kind == "callbacks" else []) + if kind == "callbacks": + assert all(actual is expected for actual, expected in zip(result, deferred)) + assert manager._callback_claims == claims + assert manager._deferred_callback_retries == retries + assert all(actual is expected for actual, expected in zip(manager._deferred_api_callbacks, deferred)) + assert not manager._pending_work_acks + + +@pytest.mark.parametrize("callback_kind", ["dag", "task"]) +@conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"}) +@mock.patch.object(DagFileProcessorManager, "_create_bundle_manager", autospec=True) +def test_callback_preparation_logs_retry_and_identifies_expired_work( + create_bundle_manager, api_manager, caplog, callback_kind +): + manager, client, bundle = api_manager + fields = dict(filepath="a.py", bundle_name=bundle.name, bundle_version="v1") + if callback_kind == "dag": + request = DagCallbackRequest(**fields, dag_id="dag", run_id="run") + target = "dag_id=dag run_id=run" + else: + request = TaskCallbackRequest( + **fields, + ti=TaskInstance( + id=uuid7(), dag_id="dag", run_id="run", task_id="task", try_number=1, dag_version_id=uuid7() + ), + ) + target = "dag_id=dag run_id=run task_id=task" + work = _make_work(bundle.name, callback=request) + manager._callback_claims[id(request)] = work + manager._deferred_api_callbacks.append(request) + create_bundle_manager.return_value.get_bundle.side_effect = OSError("version unavailable") + + manager._add_callback_to_queue(request) + + assert { + "event": "Unable to prepare callback bundle; retrying", + "bundle_name": bundle.name, + "bundle_version": "v1", + "target": target, + "error": "version unavailable", + } in caplog + assert not manager._pending_work_acks + retry = manager._deferred_callback_retries[id(request)] + retry.expires_at = retry.next_attempt_time = 0 + manager._add_callback_to_queue(request) + + assert { + "event": "Dropping callback after bundle preparation timed out", + "callback_type": type(request).__name__, + "filepath": "a.py", + "bundle_name": bundle.name, + "bundle_version": "v1", + "target": target, + "timeout": 10, + "timeout_setting": "[dag_processor] job_heartbeat_timeout", + } in caplog + assert manager._pending_work_acks == [("callbacks", work)] + assert not manager._callback_claims + assert not manager._deferred_api_callbacks diff --git a/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py b/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py index f55ed3ca5242c..c9f5601fc7877 100644 --- a/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py +++ b/airflow-core/tests/unit/jobs/test_dag_processor_job_runner.py @@ -20,6 +20,7 @@ import httpx2 import pytest +from sqlalchemy.orm import Session from airflow.api_fastapi.execution_api.datamodels.job import TerminalJobState from airflow.dag_processing.api_client import ( @@ -27,9 +28,10 @@ DagProcessorJobAlreadyRunning, DagProcessorRegistrationRetired, ) +from airflow.dag_processing.health import check_processor_health 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 airflow.jobs.job import Job, JobState, run_job from tests_common.test_utils.config import conf_vars @@ -44,10 +46,11 @@ def client(): @pytest.fixture -def runner(): - return DagProcessorJobRunner( - job=Job(heartrate=5), processor=mock.create_autospec(DagFileProcessorManager, instance=True) - ) +def runner(tmp_path): + with conf_vars({("dag_processor", "health_check_file"): str(tmp_path / "health.json")}): + yield DagProcessorJobRunner( + job=Job(heartrate=5), processor=mock.create_autospec(DagFileProcessorManager, instance=True) + ) @pytest.fixture @@ -161,6 +164,41 @@ def test_shutdown_listener_failure_does_not_prevent_completion(listeners, runner client.complete_job.assert_called_once_with(TerminalJobState.SUCCESS) +@pytest.mark.parametrize("use_api", [False, True], ids=["database", "api"]) +@pytest.mark.parametrize("error", [None, ValueError("parse failed")], ids=["success", "failed"]) +@mock.patch("airflow.jobs.job.get_listener_manager", autospec=True) +@mock.patch("airflow.jobs.dag_processor_job_runner.get_listener_manager", autospec=True) +@mock.patch.object(Job, "prepare_for_execution", autospec=True) +def test_shutdown_listener_observes_terminal_state_before_end_date( + prepare, api_listeners, db_listeners, runner, client, use_api, error +): + db_listeners.return_value = api_listeners.return_value + observed = [] + api_listeners.return_value.hook.before_stopping.side_effect = lambda component: observed.append( + (component.state, component.end_date) + ) + runner.processor.run.side_effect = error + + def run(): + if use_api: + runner.run_with_api(client) + else: + run_job( + runner.job, + execute_callable=runner._execute, + session=mock.create_autospec(Session, instance=True), + ) + + if error is None: + run() + else: + with pytest.raises(ValueError, match="parse failed"): + run() + + assert observed == [(JobState.FAILED if error else JobState.SUCCESS, None)] + assert runner.job.end_date is not None + + @pytest.mark.parametrize("parse_error", [None, ValueError("original error")]) def test_completion_failure_preserves_the_original_error(runner, client, parse_error): runner.processor.run.side_effect = parse_error @@ -213,7 +251,8 @@ def run(): @pytest.mark.parametrize("failure", ["transport", 503, 403]) @conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"}) -def test_heartbeat_failure_is_bounded(runner, client, clock, failure): +@mock.patch("airflow.jobs.dag_processor_job_runner.stats.incr", autospec=True) +def test_heartbeat_failure_is_bounded(incr, runner, client, clock, failure, caplog): error = ( httpx2.ConnectError("unavailable") if failure == "transport" @@ -232,9 +271,26 @@ def run(): runner.processor.heartbeat() runner.processor.run.side_effect = run - with pytest.raises(httpx2.HTTPStatusError if failure == 403 else DagProcessorHeartbeatTimeout): + with pytest.raises(httpx2.HTTPStatusError if failure == 403 else DagProcessorHeartbeatTimeout) as raised: runner.run_with_api(client) assert client.heartbeat.call_count == (1 if failure == 403 else 2) + failure_metrics = [call for call in incr.call_args_list if call.args[0].endswith("_heartbeat_failure")] + assert failure_metrics == [mock.call("job_heartbeat_failure", 1, 1)] * (0 if failure == 403 else 2) + if failure != 403: + assert { + "event": "Unable to heartbeat the Dag processor Job; retrying on the next interval", + "error_type": type(error).__name__, + "status_code": None if failure == "transport" else failure, + "elapsed_seconds": 5, + "timeout_seconds": 10, + } in caplog + message = str(raised.value) + assert "after 10.0s" in message + assert "[dag_processor] job_heartbeat_timeout=10s" in message + assert type(error).__name__ in message + if failure == 503: + assert "HTTP 503" in message + assert raised.value.__cause__ is error @conf_vars({("dag_processor", "job_heartbeat_timeout"): "10"}) @@ -261,13 +317,21 @@ def run(): ("dag_processor", "job_heartbeat_timeout"): "300", } ) -def test_a_stall_past_the_health_threshold_keeps_the_lease(runner, client, clock): +@mock.patch("airflow.dag_processing.health.monotonic", autospec=True) +def test_api_outage_fails_readiness_until_heartbeat_recovers(health_clock, runner, client, clock): client.heartbeat.side_effect = [httpx2.ConnectError("offline"), JobState.RUNNING] + health_clock.side_effect = lambda: clock.return_value def run(): - for tick in (120, 125): + clock.return_value = 120 + runner.processor.heartbeat() + check_processor_health() + with pytest.raises(SystemExit, match="api heartbeat age 120.0s"): + check_processor_health(readiness=True) + for tick in (125, 126): clock.return_value = tick runner.processor.heartbeat() + check_processor_health(readiness=True) runner.processor.run.side_effect = run runner.run_with_api(client) diff --git a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml index cf054af9091c5..dfceef659bead 100644 --- a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml +++ b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml @@ -51,8 +51,8 @@ metrics: legacy_name: "-" name_variables: [] - - name: "dag_processor_heartbeat_failure" - description: "Number of failed Dag processor heartbeats through the Execution API." + - name: "job_heartbeat_failure" + description: "Number of failed job heartbeats through the metadata database or Execution API." type: "counter" legacy_name: "-" name_variables: [] From ff4ff4540262f15459bf336561ee7eb5fe3ce207 Mon Sep 17 00:00:00 2001 From: Ephraim Anierobi Date: Thu, 22 Oct 2026 16:05:47 +0100 Subject: [PATCH 5/5] Allow independent priority parses to finish in either order in recovery tests Recovered files can finish independently after the healthy bundle's work. The regression must require every claim to be acknowledged exactly once without imposing an ordering between separate priority parses. --- .../unit/dag_processing/test_requested_work_recovery.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py b/airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py index 08f4e1897df81..074843f15e647 100644 --- a/airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py +++ b/airflow-core/tests/unit/dag_processing/test_requested_work_recovery.py @@ -257,9 +257,8 @@ def create_process(_, file, *, source): del manager._processors[file] manager._acknowledge_requested_work() assert not manager._priority_claims - assert client.acknowledge_work.call_args_list == [ - mock.call("priority", item) for item in [healthy_work, *work] - ] + assert client.acknowledge_work.call_count == 1 + len(work) + client.acknowledge_work.assert_has_calls([mock.call("priority", item) for item in work], any_order=True) def test_priority_claims_waiting_for_fresh_inventory_keep_the_budget_bounded(api_manager):