kaxil commented on code in PR #74791:
URL: https://github.com/apache/airflow/pull/74791#discussion_r4401074656
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py:
##########
@@ -411,3 +421,132 @@ def is_removed(relative_fileloc: str | None) -> bool:
accepted_at=checkpoint.accepted_at,
revision=bundle.parse_revision,
)
+
+
[email protected](
+ "/{job_id}/requested-work/{kind}/claim",
+ responses=create_openapi_http_exception_doc(
+ [status.HTTP_403_FORBIDDEN, status.HTTP_404_NOT_FOUND,
status.HTTP_409_CONFLICT]
+ ),
+ dependencies=[Security(require_auth, scopes=["token:dag_processor"])],
+)
+def claim_work(
+ job_id: int,
+ kind: Literal["callbacks", "priority"],
+ body: ProcessorWorkClaimBody,
+ *,
+ session: SessionDep,
+ token=CurrentExecutionToken,
+) -> list[ProcessorWorkItem]:
+ """Replay outstanding work without letting an old or empty claim take new
requests."""
+ if not set(body.bundle_names) <= token.claims.dag_bundles:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"bundle_not_granted"})
+ _validate_job(job_id, token, session=session)
+ model = DagProcessorCallback if kind == "callbacks" else
DagPriorityParsingRequest
+ receipt = session.get(DagProcessorWorkClaim, (job_id, kind))
+ payload_hash = _compute_payload_hash(body)
+ replaying = receipt is not None and body.claim_sequence <=
receipt.claim_sequence
+ if receipt is not None and replaying:
+ if body.claim_sequence < receipt.claim_sequence:
+ return []
+ if receipt.claim_id != body.claim_id or receipt.payload_hash !=
payload_hash:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"claim_conflict"})
+ query = select(model).where(
+ model.processor_job_id == job_id, model.processor_claim_id ==
body.claim_id
+ )
+ else:
+ if receipt is None:
+ receipt = DagProcessorWorkClaim(job_id=job_id, kind=kind)
+ session.add(receipt)
+ receipt.claim_id = body.claim_id
+ receipt.claim_sequence = body.claim_sequence
+ receipt.payload_hash = payload_hash
+ ready =
select(DagBundleModel.name).where(DagBundleModel.active.is_(True))
+ query = select(model).where(
+ model.bundle_name.in_(body.bundle_names),
+ model.bundle_name.in_(ready),
+ get_available_work_filter(model.processor_job_id),
+ )
+ if kind == "callbacks":
+ query = query.where(
+ or_(DagProcessorCallback.state.is_(None),
DagProcessorCallback.state == CallbackState.QUEUED)
+ )
+ if kind == "callbacks":
+ query = query.order_by(DagProcessorCallback.priority_weight.desc(),
DagProcessorCallback.id)
+ else:
+ query = query.order_by(DagPriorityParsingRequest.id)
+ if not replaying:
+ query = with_row_locks(query.limit(body.limit), session=session,
skip_locked=True)
+ replay = list(
+ cast("list[DagProcessorCallback | DagPriorityParsingRequest]",
session.scalars(query).all())
+ )
+ if not replaying:
+ for work in replay:
+ work.processor_job_id = job_id
+ work.processor_claim_id = body.claim_id
+ if isinstance(work, DagProcessorCallback):
+ work.state = CallbackState.QUEUED
Review Comment:
Setting `QUEUED` here takes these rows out of `airflow db clean`, since the
`callback` config only removes Dag processor callbacks whose state is NULL (or
terminal). A claim that never gets acked, say its owner retires and then the
bundle goes inactive so nothing reclaims it, now stays forever. Ownership
already lives in `processor_job_id`, so could we leave `state` alone? The
`QUEUED` branch in the filter above would go with it.
##########
airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_parsing.py:
##########
@@ -411,3 +421,132 @@ def is_removed(relative_fileloc: str | None) -> bool:
accepted_at=checkpoint.accepted_at,
revision=bundle.parse_revision,
)
+
+
[email protected](
+ "/{job_id}/requested-work/{kind}/claim",
+ responses=create_openapi_http_exception_doc(
+ [status.HTTP_403_FORBIDDEN, status.HTTP_404_NOT_FOUND,
status.HTTP_409_CONFLICT]
+ ),
+ dependencies=[Security(require_auth, scopes=["token:dag_processor"])],
+)
+def claim_work(
+ job_id: int,
+ kind: Literal["callbacks", "priority"],
+ body: ProcessorWorkClaimBody,
+ *,
+ session: SessionDep,
+ token=CurrentExecutionToken,
+) -> list[ProcessorWorkItem]:
+ """Replay outstanding work without letting an old or empty claim take new
requests."""
+ if not set(body.bundle_names) <= token.claims.dag_bundles:
+ raise HTTPException(status.HTTP_403_FORBIDDEN, detail={"reason":
"bundle_not_granted"})
+ _validate_job(job_id, token, session=session)
+ model = DagProcessorCallback if kind == "callbacks" else
DagPriorityParsingRequest
+ receipt = session.get(DagProcessorWorkClaim, (job_id, kind))
+ payload_hash = _compute_payload_hash(body)
+ replaying = receipt is not None and body.claim_sequence <=
receipt.claim_sequence
+ if receipt is not None and replaying:
+ if body.claim_sequence < receipt.claim_sequence:
+ return []
+ if receipt.claim_id != body.claim_id or receipt.payload_hash !=
payload_hash:
+ raise HTTPException(status.HTTP_409_CONFLICT, detail={"reason":
"claim_conflict"})
+ query = select(model).where(
+ model.processor_job_id == job_id, model.processor_claim_id ==
body.claim_id
+ )
+ else:
+ if receipt is None:
+ receipt = DagProcessorWorkClaim(job_id=job_id, kind=kind)
+ session.add(receipt)
+ receipt.claim_id = body.claim_id
+ receipt.claim_sequence = body.claim_sequence
+ receipt.payload_hash = payload_hash
+ ready =
select(DagBundleModel.name).where(DagBundleModel.active.is_(True))
+ query = select(model).where(
+ model.bundle_name.in_(body.bundle_names),
+ model.bundle_name.in_(ready),
+ get_available_work_filter(model.processor_job_id),
+ )
+ if kind == "callbacks":
+ query = query.where(
+ or_(DagProcessorCallback.state.is_(None),
DagProcessorCallback.state == CallbackState.QUEUED)
+ )
+ if kind == "callbacks":
+ query = query.order_by(DagProcessorCallback.priority_weight.desc(),
DagProcessorCallback.id)
+ else:
+ query = query.order_by(DagPriorityParsingRequest.id)
+ if not replaying:
+ query = with_row_locks(query.limit(body.limit), session=session,
skip_locked=True)
+ replay = list(
+ cast("list[DagProcessorCallback | DagPriorityParsingRequest]",
session.scalars(query).all())
Review Comment:
The `cast` here (and the one in `acknowledge_work`) comes from picking
`model` at runtime, and the two handlers branch on kind or type seven times
between them. A small typed helper per kind (query, item mapping, key parsing)
would drop both casts and keep each kind's rules in one place.
##########
airflow-core/src/airflow/dag_processing/maintenance.py:
##########
@@ -0,0 +1,94 @@
+# 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.
+
+"""Bounded metadata maintenance owned by the scheduler, including absent
bundles."""
+
+from __future__ import annotations
+
+import logging
+from datetime import timedelta
+
+from sqlalchemy import delete, select, tuple_
+
+from airflow._shared.timezones import timezone
+from airflow.configuration import conf
+from airflow.jobs.job import Job, JobState
+from airflow.models.dag import DagModel
+from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagwarning import DagWarning
+from airflow.utils.session import create_session
+from airflow.utils.sqlalchemy import with_row_locks
+
+log = logging.getLogger(__name__)
+
+
+def cleanup_processor_metadata(batch_size: int = 500) -> None:
+ """Commit each bounded batch; a later interval picks up remaining rows."""
+ with create_session() as session:
+ inactive =
select(DagBundleModel.name).where(DagBundleModel.active.is_(False))
+ query = (
+ select(DagModel)
+ .where(
+ DagModel.is_stale.is_(False),
+ DagModel.bundle_name.in_(inactive),
+ )
+ .order_by(DagModel.dag_id)
+ .limit(batch_size)
Review Comment:
This caps inactive-bundle deactivation at 500 Dags per
`parsing_cleanup_interval`, where `deactivate_stale_dags` used to mark all of
them in one pass. `dags_needing_dagruns` only filters on `is_paused` and
`is_stale`, so every Dag still waiting for its batch keeps getting runs
created. Removing a bundle with 10k Dags leaves the tail scheduling for about
20 minutes on a single scheduler. Could this loop over batches (committing
each) until one comes back short? A test with more Dags than `batch_size` would
pin it.
##########
airflow-core/src/airflow/config_templates/config.yml:
##########
@@ -3478,6 +3480,17 @@ dag_processor:
type: integer
example: ~
default: "30"
+ job_heartbeat_timeout:
Review Comment:
`register_job` already retires the same Job once it is past `[dag_processor]
health_check_threshold` (30s by default, via `is_alive()`), which releases its
claims with the same double-run risk described here. So 300s doesn't bound
duplicates on its own. Could expiry derive from `health_check_threshold`
instead, or at least cross-reference the two and reject a
`job_heartbeat_timeout` below it, where the scheduler would expire a Job that
still counts as alive?
##########
airflow-core/tests/unit/migrations/test_0151_add_processor_work_claims.py:
##########
@@ -0,0 +1,150 @@
+# 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 importlib import import_module
+from uuid import UUID
+
+import pytest
+import sqlalchemy as sa
+from alembic.migration import MigrationContext
+from alembic.operations import Operations
+from sqlalchemy.pool import NullPool
+
+pytestmark = pytest.mark.db_test
+
+migration =
import_module("airflow.migrations.versions.0151_3_5_0_add_processor_work_claims")
+CALLBACK_ID = UUID(int=1)
+CLAIM_ID = UUID(int=2)
+PRIORITY_ID = "priority-request"
+
+metadata = sa.MetaData()
+job = sa.Table("job", metadata, sa.Column("id", sa.Integer, primary_key=True))
+callback = sa.Table("callback", metadata, sa.Column("id", sa.Uuid,
primary_key=True))
+priority = sa.Table(
+ "dag_priority_parsing_request", metadata, sa.Column("id", sa.String(32),
primary_key=True)
+)
+deadline = sa.Table(
+ "deadline",
+ metadata,
+ sa.Column("id", sa.Integer, primary_key=True),
+ sa.Column("callback_id", sa.Uuid, sa.ForeignKey("callback.id",
ondelete="CASCADE"), nullable=False),
+)
+claim = sa.table(
+ "dag_processor_work_claim",
+ sa.column("job_id", sa.Integer),
+ sa.column("kind", sa.String(10)),
+ sa.column("claim_id", sa.Uuid),
+ sa.column("claim_sequence", sa.BigInteger),
+ sa.column("payload_hash", sa.String(64)),
+)
+
+
+def _run_migration(engine, migrate):
+ with engine.begin() as connection,
Operations.context(MigrationContext.configure(connection)):
+ migrate()
+
+
[email protected]
+def migrated_database(tmp_path):
+ engine = sa.create_engine(f"sqlite:///{tmp_path / 'work_claims.db'}",
poolclass=NullPool)
Review Comment:
This always builds its own SQLite engine, so the test runs the same SQLite
migration on every backend, and the `2**53 - 1` round trip would pass with a
plain `Integer` column too since SQLite integers are 64-bit. The 0142 and 0144
tests run against the configured backend. Could this one do the same, or else
be marked `backend("sqlite")`?
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_parsing.py:
##########
@@ -70,11 +77,17 @@
def clean_db(time_machine, session):
time_machine.move_to("2026-10-05T12:00:00Z", tick=False)
clear_db_jobs()
+ clear_db_callbacks()
+ session.execute(delete(DagPriorityParsingRequest))
Review Comment:
`clear_db_dag_parsing_requests()` in `tests_common.test_utils.db` already
does this (here and on line 89).
##########
airflow-core/src/airflow/dag_processing/maintenance.py:
##########
@@ -0,0 +1,94 @@
+# 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.
+
+"""Bounded metadata maintenance owned by the scheduler, including absent
bundles."""
+
+from __future__ import annotations
+
+import logging
+from datetime import timedelta
+
+from sqlalchemy import delete, select, tuple_
+
+from airflow._shared.timezones import timezone
+from airflow.configuration import conf
+from airflow.jobs.job import Job, JobState
+from airflow.models.dag import DagModel
+from airflow.models.dagbundle import DagBundleModel
+from airflow.models.dagwarning import DagWarning
+from airflow.utils.session import create_session
+from airflow.utils.sqlalchemy import with_row_locks
+
+log = logging.getLogger(__name__)
+
+
+def cleanup_processor_metadata(batch_size: int = 500) -> None:
+ """Commit each bounded batch; a later interval picks up remaining rows."""
+ with create_session() as session:
+ inactive =
select(DagBundleModel.name).where(DagBundleModel.active.is_(False))
+ query = (
+ select(DagModel)
+ .where(
+ DagModel.is_stale.is_(False),
+ DagModel.bundle_name.in_(inactive),
+ )
+ .order_by(DagModel.dag_id)
+ .limit(batch_size)
+ )
+ for dag in session.scalars(with_row_locks(query, session=session,
skip_locked=True)):
+ log.info(
+ "Deactivating Dag %s. Its bundle %s is no longer active.",
+ dag.dag_id,
+ dag.bundle_name,
+ )
+ dag.is_stale = True
+ with create_session() as session:
+ warning_query = (
+ select(DagWarning.dag_id, DagWarning.warning_type)
+ .join(DagModel)
+ .where(
+ DagModel.is_stale.is_(True),
+ )
+ .order_by(DagWarning.dag_id, DagWarning.warning_type)
+ # Each key needs two bind parameters; older SQLite versions allow
only 999.
+ .limit(min(batch_size, 499))
+ )
+ keys = session.execute(
+ with_row_locks(warning_query, of=DagWarning, session=session,
skip_locked=True)
+ ).all()
+ if keys:
+ session.execute(
+ delete(DagWarning).where(tuple_(DagWarning.dag_id,
DagWarning.warning_type).in_(keys))
+ )
+ with create_session() as session:
+ cutoff = timezone.utcnow() -
timedelta(seconds=conf.getint("dag_processor", "job_heartbeat_timeout"))
+ job_query = (
+ select(Job)
+ .where(
+ Job.job_type == "DagProcessorJob",
+ # API Jobs end with a terminal state; use its heartbeat index
to skip completed history.
+ Job.state == JobState.RUNNING,
+ Job.session_id.is_not(None),
+ Job.end_date.is_(None),
+ Job.latest_heartbeat < cutoff,
+ )
+ .order_by(Job.id)
+ .limit(batch_size)
+ )
+ for job in session.scalars(with_row_locks(job_query, session=session,
skip_locked=True)):
+ job.state = JobState.FAILED
Review Comment:
Can we log a warning here per Job, with its id, hostname, `latest_heartbeat`
and the `job_heartbeat_timeout` value? This is the path that releases claims
and can make a callback run twice, and right now it's silent, unlike the Dag
deactivation above.
It also leaves the Job's `DagParseCheckpoint` rows in place. `register_job`
and `complete_job` both delete them, and the model docstring says they go when
the Job completes.
##########
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_parsing.py:
##########
@@ -1114,3 +1163,288 @@ async def authenticate(request: Request):
json={"attempt_id": str(uuid7()), "dispatch_sequence": 1,
"expected_revision": None, "files": []},
)
assert response.status_code == 200, response.json()
+
+
[email protected]("kind", ["callbacks", "priority"])
+def test_requested_work_claim_survives_response_loss_and_ack_replay(client,
session, job, kind):
+ if kind == "callbacks":
+ work = DagProcessorCallback(
+ priority_weight=5,
+ callback=DagCallbackRequest(
+ filepath="a.py",
+ bundle_name="bundle",
+ bundle_version="v1",
+ dag_id="published",
+ run_id="run",
+ ),
+ )
+ else:
+ work = DagPriorityParsingRequest(bundle_name="bundle",
relative_fileloc="a.py")
+ session.add(work)
+ session.commit()
+ base = f"/execution/jobs/{job.id}/requested-work/{kind}"
+ claim = {"claim_id": str(uuid7()), "claim_sequence": 1, "bundle_names":
["bundle"], "limit": 1}
Review Comment:
None of the claim tests has more work items than `limit`, and none sets the
bundle inactive, so dropping `.limit(body.limit)` or the
`bundle_name.in_(ready)` filter in `claim_work` would still pass. Could one
case claim three items with `limit: 2`, and another claim against an inactive
bundle?
##########
airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dag_parsing.py:
##########
@@ -143,3 +143,32 @@ class DagBundleInventoryResponse(DagParseResultResponse):
"""The source revision to attach to imports dispatched from this
inventory."""
revision: UUID
+
+
+class ProcessorWorkClaimBody(StrictBaseModel):
+ """
+ A bounded claim whose sequence increases per Job and work kind.
+
+ Retry the same payload until its response arrives before sending a newer
sequence for that kind.
+ """
+
+ claim_id: UUID
+ claim_sequence: int = Field(ge=1, le=2**53 - 1)
+ bundle_names: list[str] = Field(min_length=1, max_length=1000)
+ limit: int = Field(default=20, ge=1, le=100)
+
+
+class ProcessorWorkItem(StrictBaseModel):
+ """A claimed callback or priority parse request."""
+
+ id: str
+ claim_id: UUID
+ bundle_name: str
+ relative_fileloc: str
+ callback: str | None = None
Review Comment:
This carries `CallbackRequest.to_json()` as an opaque string, so the
`TaskInstance` and `TIRunContext` inside it sit outside Cadwyn's versioning,
and a later change to those models can't be migrated for an older processor.
Since `v2027_02_28` hasn't shipped yet, would typing this as the callback
request (or a union discriminated on kind) be worth doing now?
--
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]
