kaxil commented on code in PR #74794:
URL: https://github.com/apache/airflow/pull/74794#discussion_r4468989876


##########
airflow-core/src/airflow/dag_processing/manager.py:
##########
@@ -2161,6 +2237,8 @@ def _collect_results(self):
 
         for file in finished:
             processor = self._processors.pop(file)
+            if processor.had_callbacks:
+                self.work_source.attempt_finished(processor.id, 
had_callbacks=True)

Review Comment:
   This acknowledges a callback attempt whatever happened in the child. A 503 
on an import-time `Variable.get` sets `api_request_error` and the Dag never 
loads, but `_execute_callbacks` swallows the lookup error and the child exits 
normally, so the callback is never delivered again. The timeout kill further 
down acknowledges after SIGKILL the same way, while `terminate()` only does it 
when `exited_cleanly`. Could both paths require `exited_cleanly` and no 
`api_request_error`, and requeue under the existing claim otherwise?



##########
airflow-core/src/airflow/dag_processing/api_requested_work.py:
##########
@@ -0,0 +1,329 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Claim and acknowledgment lifecycle for Dag processor work from the 
Execution API."""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import attrs
+import httpx2
+
+from airflow.callbacks.callback_requests import CallbackRequest, 
DagCallbackRequest
+from airflow.configuration import conf
+from airflow.dag_processing.collection import is_source_discovery_failed
+from airflow.dag_processing.requested_work import RequestedWorkSource
+from airflow.utils.file import find_enclosing_file
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+if TYPE_CHECKING:
+    from collections.abc import Collection
+    from uuid import UUID
+
+    from airflow.api_fastapi.execution_api.datamodels.dag_parsing import 
ProcessorWorkItem
+    from airflow.dag_processing.api_client import DagProcessorAPIClient
+    from airflow.dag_processing.bundles.base import BaseDagBundle
+    from airflow.dag_processing.manager import DagFileInfo
+
+_API_SEND_BUDGET_SECONDS = 1.0
+_MAX_CALLBACK_RETRY_DELAY = 60.0
+
+
[email protected]
+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
+
+
[email protected]
+class PriorityParseClaim:
+    """A priority request served only after an inventory captured following 
the claim."""
+
+    work: ProcessorWorkItem
+    inventory_sequence: int
+    attempt_id: UUID | None = None
+
+
[email protected]
+class ExecutionAPIWorkSource(RequestedWorkSource, LoggingMixin):
+    """Retain claimed work until its terminal outcome can be acknowledged by 
the API."""
+
+    client: DagProcessorAPIClient
+    _callback_claims: dict[int, ProcessorWorkItem] = attrs.field(factory=dict, 
init=False)
+    _deferred_api_callbacks: list[CallbackRequest] = attrs.field(factory=list, 
init=False)
+    _deferred_callback_retries: dict[int, DeferredCallback] = 
attrs.field(factory=dict, init=False)
+    _callback_attempts: dict[UUID, list[ProcessorWorkItem]] = 
attrs.field(factory=dict, init=False)
+    _priority_claims: dict[tuple[str, UUID], PriorityParseClaim] = 
attrs.field(factory=dict, init=False)
+    _pending_work_acks: list[tuple[Literal["callbacks", "priority"], 
ProcessorWorkItem]] = attrs.field(
+        factory=list, init=False
+    )
+
+    requires_acknowledgment = True
+
+    def claim_priority_files(
+        self,
+        bundles: Collection[BaseDagBundle],
+        *,
+        capacity: int,
+        unavailable_bundles: Collection[str],
+        active_attempts: Collection[UUID],
+        inventory_sequence: int,
+    ) -> list[DagFileInfo]:
+        """Claim priority files within the capacity left by active and 
retained work."""
+        from airflow.dag_processing.manager import DagFileInfo
+
+        bundles_by_name = {bundle.name: bundle for bundle in bundles}
+        unavailable = set(unavailable_bundles)
+        available = (
+            capacity
+            - sum(
+                claim.attempt_id in active_attempts or claim.work.bundle_name 
not in unavailable
+                for claim in self._priority_claims.values()
+            )
+            - sum(kind == "priority" for kind, _ in self._pending_work_acks)
+        )
+        if available <= 0:
+            return []
+        # Retain failed-bundle claims for recovery without letting them 
consume healthy bundles' capacity.
+        # Excluding those bundles from new claims also bounds the suspended 
backlog.
+        eligible_bundles = [name for name in bundles_by_name if name not in 
unavailable]
+        try:
+            work = self.client.claim_work("priority", eligible_bundles, 
available)
+        except httpx2.HTTPError:
+            self.log.warning("Unable to claim priority parses; retrying on a 
later loop")

Review Comment:
   This warning and the ones at :154, :201 and :212 don't say what failed, so a 
DNS error and a gateway 503 read the same. Could they carry the error type and 
status code the way the heartbeat warning does now?



##########
airflow-core/src/airflow/dag_processing/api_requested_work.py:
##########
@@ -0,0 +1,329 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Claim and acknowledgment lifecycle for Dag processor work from the 
Execution API."""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal
+
+import attrs
+import httpx2
+
+from airflow.callbacks.callback_requests import CallbackRequest, 
DagCallbackRequest
+from airflow.configuration import conf
+from airflow.dag_processing.collection import is_source_discovery_failed
+from airflow.dag_processing.requested_work import RequestedWorkSource
+from airflow.utils.file import find_enclosing_file
+from airflow.utils.log.logging_mixin import LoggingMixin
+
+if TYPE_CHECKING:
+    from collections.abc import Collection
+    from uuid import UUID
+
+    from airflow.api_fastapi.execution_api.datamodels.dag_parsing import 
ProcessorWorkItem
+    from airflow.dag_processing.api_client import DagProcessorAPIClient
+    from airflow.dag_processing.bundles.base import BaseDagBundle
+    from airflow.dag_processing.manager import DagFileInfo
+
+_API_SEND_BUDGET_SECONDS = 1.0
+_MAX_CALLBACK_RETRY_DELAY = 60.0
+
+
[email protected]
+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
+
+
[email protected]
+class PriorityParseClaim:
+    """A priority request served only after an inventory captured following 
the claim."""
+
+    work: ProcessorWorkItem
+    inventory_sequence: int
+    attempt_id: UUID | None = None
+
+
[email protected]
+class ExecutionAPIWorkSource(RequestedWorkSource, LoggingMixin):
+    """Retain claimed work until its terminal outcome can be acknowledged by 
the API."""
+
+    client: DagProcessorAPIClient
+    _callback_claims: dict[int, ProcessorWorkItem] = attrs.field(factory=dict, 
init=False)
+    _deferred_api_callbacks: list[CallbackRequest] = attrs.field(factory=list, 
init=False)
+    _deferred_callback_retries: dict[int, DeferredCallback] = 
attrs.field(factory=dict, init=False)
+    _callback_attempts: dict[UUID, list[ProcessorWorkItem]] = 
attrs.field(factory=dict, init=False)
+    _priority_claims: dict[tuple[str, UUID], PriorityParseClaim] = 
attrs.field(factory=dict, init=False)
+    _pending_work_acks: list[tuple[Literal["callbacks", "priority"], 
ProcessorWorkItem]] = attrs.field(
+        factory=list, init=False
+    )
+
+    requires_acknowledgment = True
+
+    def claim_priority_files(
+        self,
+        bundles: Collection[BaseDagBundle],
+        *,
+        capacity: int,
+        unavailable_bundles: Collection[str],
+        active_attempts: Collection[UUID],
+        inventory_sequence: int,
+    ) -> list[DagFileInfo]:
+        """Claim priority files within the capacity left by active and 
retained work."""
+        from airflow.dag_processing.manager import DagFileInfo
+
+        bundles_by_name = {bundle.name: bundle for bundle in bundles}
+        unavailable = set(unavailable_bundles)
+        available = (
+            capacity
+            - sum(
+                claim.attempt_id in active_attempts or claim.work.bundle_name 
not in unavailable
+                for claim in self._priority_claims.values()
+            )
+            - sum(kind == "priority" for kind, _ in self._pending_work_acks)
+        )
+        if available <= 0:
+            return []
+        # Retain failed-bundle claims for recovery without letting them 
consume healthy bundles' capacity.
+        # Excluding those bundles from new claims also bounds the suspended 
backlog.
+        eligible_bundles = [name for name in bundles_by_name if name not in 
unavailable]
+        try:
+            work = self.client.claim_work("priority", eligible_bundles, 
available)
+        except httpx2.HTTPError:
+            self.log.warning("Unable to claim priority parses; retrying on a 
later loop")
+            return []
+        files = []
+        for item in work:
+            self._priority_claims[item.id, item.claim_id] = PriorityParseClaim(
+                work=item, inventory_sequence=inventory_sequence
+            )
+            bundle = bundles_by_name[item.bundle_name]
+            try:
+                path = find_enclosing_file(bundle.path / item.relative_fileloc)
+            except OSError as error:
+                # The next inventory acknowledges the claim if the bundle 
listing cannot see the file either.
+                self.log.warning(
+                    "Unable to resolve priority parse %s in bundle %s: %s",
+                    item.relative_fileloc,
+                    bundle.name,
+                    error,
+                )
+                path = None
+            files.append(
+                DagFileInfo(
+                    rel_path=path.relative_to(bundle.path) if path else 
Path(item.relative_fileloc),
+                    bundle_name=bundle.name,
+                    bundle_path=bundle.path,
+                )
+            )
+        return files
+
+    def fetch_callbacks(self, bundles: Collection[BaseDagBundle], *, limit: 
int) -> list[CallbackRequest]:
+        """Claim callbacks for ready bundles and include requests awaiting 
bundle preparation."""
+        available = (
+            limit
+            - len(self._callback_claims)

Review Comment:
   A callback whose bundle version can't be prepared stays in 
`_callback_claims` for up to `job_heartbeat_timeout` (300s by default), and its 
bundle is still in `eligible_bundles`. So a burst of callbacks for a bundle 
with an unreachable remote fills `max_callbacks_per_loop` (20), and no other 
bundle gets a callback claimed until they expire. Could deferred claims stop 
counting against the limit, the way `claim_priority_files` handles unavailable 
bundles?



-- 
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]

Reply via email to