ephraimbuddy commented on code in PR #74793: URL: https://github.com/apache/airflow/pull/74793#discussion_r4490967959
########## airflow-core/tests/unit/dag_processing/test_manager_api_failures.py: ########## @@ -0,0 +1,333 @@ +# 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 + +import signal +from pathlib import Path +from unittest import mock + +import httpx2 +import pytest +from sqlalchemy import select +from uuid6 import uuid7 + +from airflow.api_fastapi.execution_api.datamodels.dag_parsing import ( + DagBundleStateResponse, + DagParseResultBody, + ParseSourceCode, +) +from airflow.api_fastapi.execution_api.datamodels.token import DagProcessorClaims, DagProcessorToken +from airflow.api_fastapi.execution_api.routes.dag_parsing import publish_parse_result +from airflow.dag_processing.api_client import DagProcessorAPIClient, DagProcessorRegistrationRetired +from airflow.dag_processing.manager import ( + DagFileInfo, + DagFileProcessorManager, + DagFileStat, + PendingDagPublication, +) +from airflow.dag_processing.processor import DagFileProcessorProcess, DagParseSource +from airflow.jobs.job import Job, JobState +from airflow.models.dag import DagModel +from airflow.models.dag_processor_publication_receipt import DagProcessorPublicationReceipt +from airflow.models.dagbundle import DagBundleModel +from airflow.models.errors import ParseImportError +from airflow.sdk import DAG, timezone +from airflow.serialization.serialized_objects import LazyDeserializedDAG + +from tests_common.test_utils.config import conf_vars + + [email protected]_test [email protected]("outcome", ["no_result", "timeout"]) [email protected]("lose_acknowledgment", [False, True]) +@conf_vars({("workers", "execution_api_retries"): "2"}) [email protected]("airflow.dag_processing.manager.time", autospec=True) +def test_failed_parse_publishes_an_error_and_ages_out_prior_dags( + clock, session, time_machine, outcome, lose_acknowledgment +): + time_machine.move_to("2026-10-05T12:00:00Z", tick=False) + clock.monotonic.return_value = 100.0 + suffix = uuid7().hex + bundle_name = f"failed_parse_{suffix}" + dag_id = f"failed_parse_{suffix}" + session_id = uuid7() + revision = uuid7() + job = Job(job_type="DagProcessorJob", state=JobState.RUNNING) + job.session_id = session_id + job.registration_id = uuid7() + job.bundle_names = [bundle_name] + bundle = DagBundleModel(name=bundle_name, version="v1") + bundle.parse_revision = revision + session.add_all([job, bundle]) + session.flush() + token = DagProcessorToken( + id=session_id, + claims=DagProcessorClaims(exp=1, job_id=job.id, dag_bundles=frozenset({bundle_name})), + ) + dag = DAG(dag_id, schedule=None) + dag.fileloc = "/worker/a.py" + dag.relative_fileloc = "a.py" + body = DagParseResultBody( + attempt_id=uuid7(), + dispatch_sequence=1, + bundle_name=bundle_name, + relative_fileloc="a.py", + bundle_version="v1", + bundle_revision=revision, + parse_duration=0.1, + serialized_dags=[LazyDeserializedDAG.from_dag(dag).data], + source_codes=[ParseSourceCode(source_code="# captured source", language="python")], + dag_source_indexes={dag_id: 0}, + ) + publish_parse_result(job.id, body, session=session, token=token) + assert session.get(DagModel, dag_id).is_stale is False + previous_finish = timezone.utcnow() + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.return_value = [] + + def publish(body): + receipt = publish_parse_result(job.id, body, session=session, token=token) + if lose_acknowledgment and client.publish_parse_result.call_count == 1: + raise httpx2.ReadTimeout("Lost publication acknowledgment") + return receipt + + client.publish_parse_result.side_effect = publish + manager = DagFileProcessorManager(max_runs=1, api_client=client, processor_timeout=10) + manager._dispatch_sequence = 2 + time_machine.shift(manager.stale_dag_threshold + 1) + file = DagFileInfo(bundle_name=bundle_name, rel_path=Path("a.py")) + previous_stat = DagFileStat(num_dags=1, last_finish_time=previous_finish) + manager._file_stats[file] = previous_stat + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = uuid7() + processor.pid = 123 + processor.dispatch_sequence = 2 + processor.start_time = 89 if outcome == "timeout" else 99 + processor.is_ready = True + processor.had_callbacks = False + processor.api_request_error = None + processor.parsing_result = None + processor.parse_source = DagParseSource(bundle_version="v1", bundle_revision=revision) + manager._processors[file] = processor + + if outcome == "timeout": + manager._kill_timed_out_processors() + processor.kill.assert_called_once_with(signal.SIGKILL) + else: + manager._collect_results() + processor.close.assert_called_once_with() + assert not manager._processors + pending = manager._pending_publications[file] + assert pending.body.attempt_id == processor.id + assert pending.body.dispatch_sequence == processor.dispatch_sequence + assert manager._dispatch_sequence == processor.dispatch_sequence + assert pending.body.bundle_name == bundle_name + assert pending.body.relative_fileloc == "a.py" + assert pending.body.bundle_version == "v1" + assert pending.body.bundle_revision == revision + assert pending.body.parse_duration == (11 if outcome == "timeout" else 1) + assert pending.body.serialized_dags == [] + assert pending.body.import_errors == {"a.py": mock.ANY} + assert 0 < len(pending.body.import_errors["a.py"]) < 2048 + assert manager._file_stats[file] is previous_stat + + manager._publish_pending_results() + session.expire_all() + receipt = session.scalar( + select(DagProcessorPublicationReceipt).where(DagProcessorPublicationReceipt.job_id == job.id) + ) + accepted_at = receipt.accepted_at + if lose_acknowledgment: + assert manager._pending_publications[file] is pending + assert manager._file_stats[file] is previous_stat + clock.monotonic.return_value = pending.next_attempt_time + time_machine.shift(10) + manager._publish_pending_results() + session.expire_all() + assert receipt.accepted_at == accepted_at + + assert not manager._pending_publications + assert manager._file_stats[file].num_dags == 0 + assert manager._file_stats[file].import_errors == 1 + assert manager._file_stats[file].run_count == 1 + assert session.get(DagModel, dag_id).is_stale is True + error = session.scalar(select(ParseImportError).where(ParseImportError.bundle_name == bundle_name)) + assert error is not None + assert error.filename == "a.py" + assert error.stacktrace == pending.body.import_errors["a.py"] + assert receipt.attempt_id == processor.id + assert receipt.dispatch_sequence == processor.dispatch_sequence + assert client.publish_parse_result.call_args_list == [mock.call(pending.body)] * ( + 2 if lose_acknowledgment else 1 + ) + + [email protected]("outcome", ["callback", "refreshed"]) [email protected]("airflow.settings.Session", autospec=True, side_effect=AssertionError("Unexpected DB session")) [email protected]("airflow.dag_processing.manager.time", autospec=True) +def test_timeout_does_not_publish_for_callbacks_or_obsolete_sources(clock, session_factory, outcome): + clock.monotonic.return_value = 100.0 + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.return_value = [] + manager = DagFileProcessorManager(max_runs=1, api_client=client, processor_timeout=10) + file = DagFileInfo(bundle_name="testing", rel_path=Path("a.py")) + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = uuid7() + processor.pid = 123 + processor.dispatch_sequence = 1 + processor.start_time = 89 + processor.had_callbacks = outcome == "callback" + processor.api_request_error = None + processor.parsing_result = None + processor.parse_source = DagParseSource(bundle_revision=uuid7()) + manager._processors[file] = processor + manager._bundle_refresh_generations[file.bundle_name] = 1 if outcome == "refreshed" else 0 + + manager._kill_timed_out_processors() + + assert not manager._pending_publications + manager._publish_pending_results() + processor.kill.assert_called_once_with(signal.SIGKILL) + processor.close.assert_called_once_with() + assert not manager._processors + client.publish_parse_result.assert_not_called() + session_factory.assert_not_called() + assert list(manager._file_queue) == ([file] if outcome == "refreshed" else []) + + [email protected]("outcome", ["no_result", "timeout"]) [email protected]("airflow.dag_processing.manager.time", autospec=True) +def test_failed_parse_with_unrepresentable_diagnostic_does_not_stop_other_files(clock, time_machine, outcome): + time_machine.move_to("2026-10-05T12:00:00Z", tick=False) + clock.monotonic.return_value = 100.0 + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.return_value = [] + manager = DagFileProcessorManager(max_runs=1, api_client=client, processor_timeout=10) + time_machine.shift(manager.stale_dag_threshold + 1) + long_file, short_file = [ + DagFileInfo(bundle_name="testing", rel_path=Path(path)) + for path in ("directory/" * 110 + "a.py", "b.py") + ] + manager._file_stats[long_file] = DagFileStat(num_dags=1) + for file in (long_file, short_file): + processor = mock.create_autospec(DagFileProcessorProcess, instance=True) + processor.id = uuid7() + processor.pid = 123 + processor.dispatch_sequence = 1 + processor.start_time = 89 + processor.is_ready = True + processor.had_callbacks = False + processor.api_request_error = None + processor.parsing_result = None + processor.parse_source = DagParseSource() + manager._processors[file] = processor + processes = list(manager._processors.values()) + + if outcome == "timeout": + manager._kill_timed_out_processors() + else: + manager._collect_results() + manager._publish_pending_results() + + assert not manager._processors + assert not manager._pending_publications + assert manager._file_stats[long_file].num_dags == 1 + assert manager._file_stats[long_file].last_finish_time is None + assert manager._file_stats[long_file].last_attempt_time is not None + assert manager._file_stats[short_file].import_errors == 1 + client.publish_parse_result.assert_called_once() + assert client.publish_parse_result.call_args.args[0].relative_fileloc == "b.py" + for processor in processes: + processor.close.assert_called_once_with() + + +def test_startup_catalog_read_waits_out_an_unavailable_api(): + client = mock.create_autospec(DagProcessorAPIClient, instance=True) + client.get_bundles.return_value = [ + DagBundleStateResponse(name="a", version=None, last_refreshed=None, revision=None, team_name=None) + ] + manager = DagFileProcessorManager(max_runs=1, api_client=client, bundle_names_to_parse=["a"]) + + manager.sync_bundles() + + client.get_bundles.assert_called_once_with(retry=True) + + [email protected](DagFileProcessorManager, "after_run", autospec=True) [email protected](DagFileProcessorManager, "_run_parsing_loop", autospec=True) [email protected](DagFileProcessorManager, "before_run", autospec=True) +def test_only_the_parsing_loop_retries_without_waiting(before_run, run_parsing_loop, after_run, tmp_path): + with DagProcessorAPIClient( + base_url="http://unused/";, token_file=tmp_path / "processor.jwt", hostname="processor-1" + ) as client: + retries_at_once = [] + for phase in (before_run, run_parsing_loop, after_run): + phase.side_effect = lambda _: retries_at_once.append(client._retry_without_waiting.get()) + + DagFileProcessorManager(max_runs=1, api_client=client).run() + + assert retries_at_once == [False, True, False] + + [email protected]("outcome", ["no_result", "timeout"]) [email protected]("airflow.dag_processing.manager.time", autospec=True) +def test_failed_parse_is_reported_at_once_as_an_unavailable_result(clock, time_machine, outcome): Review Comment: Merged them into `test_failed_parse_is_reported_at_once_as_an_unavailable_result`, parametrized over the outcome and `api_request_error`. Each case now asserts the full import-error text for its reason. Fixed in 5c9861c736. ########## airflow-core/src/airflow/dag_processing/manager.py: ########## @@ -1442,41 +1813,307 @@ def handle_parsing_result( team_name=team_name, ) - if proc.parsing_result is not None: + try: + if self.api_client is not None and not is_callback_only: + if proc.parsing_result is not None and proc.api_request_error is None: + self._pending_publications[file] = self._build_pending_publication( + file, proc, run_duration, next_stat + ) + else: + self._queue_failed_parse_result( + file, + proc, + run_duration, + next_stat, + reason="the parsing process exited without a result", + ) + return + if proc.parsing_result is not None: + with create_session() if session is None else nullcontext(session) as session: + self.persist_parsing_result( + bundle_name=file.bundle_name, + bundle_version=proc.parse_source.bundle_version, + version_data=( + dict(proc.parse_source.version_data) + if proc.parse_source.version_data is not None + else None + ), + parsing_result=proc.parsing_result, + run_duration=run_duration, + relative_fileloc=str(file.rel_path), + session=session, + ) + except Exception: + self.log.exception( + "Failed to persist parsing result for %s in bundle %s; " + "keeping previous persisted stats while throttling retries. " + "Other files in this cycle are still processed.", + str(file.rel_path), + file.bundle_name, + ) + self._record_failed_publication(file, next_stat) + return + + self._file_stats[file] = next_stat + + def _queue_failed_parse_result( + self, + file: DagFileInfo, + proc: BaseDagFileProcessorProcess, + run_duration: float, + stat: DagFileStat, + *, + reason: str, + ) -> None: + source = proc.parse_source + if source.refresh_generation != self._bundle_refresh_generations.get(file.bundle_name, 0): + self._add_files_to_queue([file], mode="front") + return + if proc.api_request_error is not None: + # The Task SDK answers a failed request as a missing value, so the Dags may hold a fallback. + reason = f"an Execution API request failed while parsing: {proc.api_request_error}" + try: + self._pending_publications[file] = self._build_rejection_report( + file, + bundle_version=source.bundle_version, + bundle_revision=source.bundle_revision, + parse_duration=run_duration, + stat=stat, + refresh_generation=source.refresh_generation, + reason=reason, + attempt_id=proc.id, + dispatch_sequence=proc.dispatch_sequence, + ) + except ValidationError: + self.log.exception("Unable to report failed parse for %s; preserving previous result", file) + self._record_failed_publication(file, stat) + + def _record_failed_publication(self, file: DagFileInfo, stat: DagFileStat) -> None: + self._file_stats[file] = attrs.evolve( + self._file_stats[file], + last_attempt_time=timezone.utcnow(), + last_duration=stat.last_duration, + run_count=stat.run_count, + ) + + def _publish_pending_results(self, *, deadline: float | None = None) -> None: + """Publish queued results until the deadline, stopping early while the API is unavailable.""" + if self.api_client is None: + return + from airflow.dag_processing.api_client import DagProcessorRegistrationRetired, get_error_reason Review Comment: Done: both imports carry the same comment as `after_run`'s, and so does the new local import in `_publish_bundle_inventory`. Fixed in 5c9861c736. ########## airflow-core/src/airflow/dag_processing/manager.py: ########## @@ -874,6 +1034,101 @@ def update_bundle_state( values["version"] = version session.execute(update(DagBundleModel).where(DagBundleModel.name == bundle_name).values(**values)) + def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set[DagFileInfo]]) -> None: + if self.api_client is None: + raise ValueError("Inventory publication requires an API client") + pending = self._pending_inventories[bundle_name] + if len(pending.body.model_dump_json().encode()) > MAX_PUBLICATION_BYTES: + self.log.error( + "Bundle %s inventory exceeds the %s-byte publication limit; retaining the previous inventory", + bundle_name, + MAX_PUBLICATION_BYTES, + ) + self._defer_unusable_bundle_inventory(bundle_name) + return + try: + receipt = self.api_client.publish_inventory(bundle_name, pending.body) + except httpx2.HTTPError as error: + status_code = error.response.status_code if isinstance(error, httpx2.HTTPStatusError) else None + if status_code is None or status_code in (408, 429) or status_code >= 500: + pending.attempts += 1 + pending.next_attempt_time = time.monotonic() + _compute_retry_delay(pending.attempts) + self.log.warning( + "Unable to publish inventory for %s; retrying on a later refresh", bundle_name + ) + return + if status_code == 409: + del self._pending_inventories[bundle_name] + self.request_bundle_refresh(bundle_name) + elif status_code == 413: + self._defer_unusable_bundle_inventory(bundle_name) + else: + # The Job, its bundle grants or the API contract no longer match this processor, so a + # resend cannot succeed; restarting registers again and reads the catalog again. + raise + self.log.warning( + "Inventory publication was rejected for %s", bundle_name, status_code=status_code + ) + return + del self._pending_inventories[bundle_name] + self._bundle_inventory_retry_times.pop(bundle_name, None) + self._bundle_parse_sources[bundle_name] = attrs.evolve( + pending.source, bundle_revision=receipt.revision + ) + 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 + + def _discard_unusable_bundle_source(self, bundle_name: str) -> None: + self._pending_inventories.pop(bundle_name, None) + self._bundle_inventory_retry_times.pop(bundle_name, None) + self._bundle_refresh_retry_at.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) + + def _defer_unusable_bundle_inventory(self, bundle_name: str) -> None: + refresh_interval = next( + (bundle.refresh_interval for bundle in self._dag_bundles if bundle.name == bundle_name), + self.bundle_refresh_check_interval, + ) + self._discard_unusable_bundle_source(bundle_name) + self._bundle_inventory_retry_times[bundle_name] = time.monotonic() + max( + refresh_interval, self.bundle_refresh_check_interval, _compute_retry_delay(1) + ) + + def _delay_bundle_listing_retry(self, bundle_name: str) -> None: + """Retry a failed listing soon, then double the wait up to the refresh interval while it fails.""" + refresh_interval = next( + (bundle.refresh_interval for bundle in self._dag_bundles if bundle.name == bundle_name), + self.bundle_refresh_check_interval, + ) + first_delay = max(self.bundle_refresh_check_interval, _compute_retry_delay(1)) + previous_delay = self._bundle_listing_retry_delays.get(bundle_name) + delay = ( + first_delay + if previous_delay is None + else min(max(refresh_interval, first_delay), previous_delay * 2) + ) + self._bundle_listing_retry_delays[bundle_name] = delay + self._bundle_inventory_retry_times[bundle_name] = time.monotonic() + delay Review Comment: Agreed. The deadline was read after `bundle.refresh()` and the failed listing, while the next periodic check counts from the check's start, so any refresh time pushed the first retry to the check after. `_delay_bundle_listing_retry` now takes the check's `now_seconds` as `check_started_at`, and the test's refresh advances the clock so it sees the gap. If a refresh outlasts the current wait, the next check retries at once, and the doubling still spaces out later attempts. Fixed in 5c9861c736. ########## airflow-core/docs/administration-and-deployment/dag-bundles.rst: ########## @@ -275,6 +275,10 @@ Some caveats: If listing the bundle's Dag files or recording the refresh fails after a successful refresh, the processor keeps parsing the files it listed before, from the refreshed source, and retries the refresh no later than the bundle's :ref:`config:dag_processor__refresh_interval`. + When the Dag processor publishes its results through the Execution API, a failed listing instead Review Comment: Agreed. The listing's parse source waits in the pending inventory until the API accepts it. The paragraph now says parsing stays paused until then: a transient failure is resent with a backoff of up to `[workers] execution_api_retry_wait_max`, and a 413, or a listing too large to send, waits for the bundle's `refresh_interval`. Fixed in 5c9861c736. ########## airflow-core/src/airflow/dag_processing/manager.py: ########## @@ -874,6 +1034,101 @@ def update_bundle_state( values["version"] = version session.execute(update(DagBundleModel).where(DagBundleModel.name == bundle_name).values(**values)) + def _publish_bundle_inventory(self, bundle_name: str, known_files: dict[str, set[DagFileInfo]]) -> None: + if self.api_client is None: + raise ValueError("Inventory publication requires an API client") + pending = self._pending_inventories[bundle_name] + if len(pending.body.model_dump_json().encode()) > MAX_PUBLICATION_BYTES: + self.log.error( + "Bundle %s inventory exceeds the %s-byte publication limit; retaining the previous inventory", + bundle_name, + MAX_PUBLICATION_BYTES, + ) + self._defer_unusable_bundle_inventory(bundle_name) + return + try: + receipt = self.api_client.publish_inventory(bundle_name, pending.body) + except httpx2.HTTPError as error: + status_code = error.response.status_code if isinstance(error, httpx2.HTTPStatusError) else None + if status_code is None or status_code in (408, 429) or status_code >= 500: Review Comment: Agreed. Added `is_transient_api_error` and used it in all three places. It sits in api_client.py next to `get_error_reason` rather than next to `_is_rejected_content`: `_get_parse_token` needs it too, and the client shouldn't import the manager. The manager imports it locally, like its other API-client imports. Behaviour is unchanged, and a parametrized test pins the rule. Fixed in 5c9861c736. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
