dheerajturaga commented on code in PR #72057:
URL: https://github.com/apache/airflow/pull/72057#discussion_r4074224925
##########
providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py:
##########
@@ -44,11 +45,19 @@
jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs")
+# ``ExecuteCallback`` is also a valid Dag id, so checking ``dag_id`` alone
would mistake that Dag's
+# tasks for callbacks. ``EdgeExecutor.queue_workload()`` gives a callback this
whole identity.
+_IS_CALLBACK_JOB = and_(
Review Comment:
`models/types.py` already has `CALLBACK_JOB_MAP_INDEX`,
`CALLBACK_JOB_TRY_NUMBER` and `build_callback_run_id()`, which
`queue_workload()` uses when writing a callback row. This repeats them as `-1`,
`0` and `f"{EXECUTE_CALLBACK_TAG}-"`. Could we use the constants here? Even
better, move this SQL expression into `types.py` next to `is_callback_job()`,
so the identity is defined once and `fetch()` stays in step if
`queue_workload()` changes.
---
Drafted-by: Claude Code (Opus 5); reviewed by @dheerajturaga before posting
##########
providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py:
##########
@@ -83,21 +92,32 @@ def fetch(
if not worker:
raise HTTPException(status.HTTP_404_NOT_FOUND, "Worker not found")
- query = (
- select(EdgeJobModel)
- .where(
- EdgeJobModel.state == TaskInstanceState.QUEUED,
- EdgeJobModel.concurrency_slots <= body.free_concurrency,
- )
- .order_by(EdgeJobModel.queued_dttm)
+ query = select(EdgeJobModel).where(
+ EdgeJobModel.state == TaskInstanceState.QUEUED,
+ EdgeJobModel.concurrency_slots <= body.free_concurrency,
)
if body.queues:
query = query.where(EdgeJobModel.queue.in_(body.queues))
if worker.team_name is not None:
query = query.where(EdgeJobModel.team_name == worker.team_name)
query = query.limit(1)
query = query.with_for_update(skip_locked=True)
- job: EdgeJobModel | None = session.scalar(query)
+
+ # Callbacks finish work that is already running, so a worker takes them
before any task, and
Review Comment:
The comment says each query can use an existing index for its ORDER BY. The
callback query orders by `queued_dttm` alone, though, and `rj_order` is now
`(state, priority_weight DESC, queued_dttm, queue)`, so that index can't serve
it. It's still cheap, because `dag_id` is the first column of the primary key
and only callback rows get read. Could the comment say that, and be shorter?
One or two lines would do.
---
Drafted-by: Claude Code (Opus 5); reviewed by @dheerajturaga before posting
##########
providers/edge3/src/airflow/providers/edge3/worker_api/routes/jobs.py:
##########
@@ -44,11 +45,19 @@
jobs_router = AirflowRouter(tags=["Jobs"], prefix="/jobs")
+# ``ExecuteCallback`` is also a valid Dag id, so checking ``dag_id`` alone
would mistake that Dag's
+# tasks for callbacks. ``EdgeExecutor.queue_workload()`` gives a callback this
whole identity.
+_IS_CALLBACK_JOB = and_(
+ EdgeJobModel.dag_id == EXECUTE_CALLBACK_TAG,
+ EdgeJobModel.run_id == literal(f"{EXECUTE_CALLBACK_TAG}-") +
EdgeJobModel.task_id,
+ EdgeJobModel.map_index == -1,
+ EdgeJobModel.try_number == 0,
+)
+
def parse_command(command: str, dag_id: str, run_id: str) -> ExecuteTypeBody:
if AIRFLOW_V_3_3_PLUS:
from airflow.executors.workloads import ExecuteCallback
- from airflow.providers.edge3.models.types import EXECUTE_CALLBACK_TAG
if dag_id == EXECUTE_CALLBACK_TAG and
run_id.startswith(EXECUTE_CALLBACK_TAG):
Review Comment:
Follow-up, non-blocking: `parse_command()` still uses `dag_id ==
EXECUTE_CALLBACK_TAG and run_id.startswith(...)`. Now that `fetch()` checks the
full identity, a task from a Dag named `ExecuteCallback` whose `run_id` starts
with that tag would still be parsed as a callback. This predates the PR, so
it's fine as a follow-up, but switching to `is_callback_job()` there would make
the two consistent.
---
Drafted-by: Claude Code (Opus 5); reviewed by @dheerajturaga before posting
##########
providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py:
##########
@@ -177,6 +178,163 @@ def
test_state_finish_metric_omits_team_name_for_global_job(self, mock_stats_inc
},
)
+ @patch(f"{Stats.__module__}.Stats.incr")
+ def test_fetch_returns_highest_priority_job_first(self, mock_stats_incr,
session: Session):
+ with create_session() as session:
+ session.add(
+ EdgeWorkerModel(
+ worker_name="worker1", state=EdgeWorkerState.IDLE,
queues=[QUEUE], team_name=None
+ )
+ )
+ # The oldest job carries the lowest priority, so a FIFO-only query
returns "low" first.
+ queued_dttm = timezone.utcnow()
+ for offset, (task_id, priority_weight) in enumerate([("low", 1),
("high", 100), ("medium", 10)]):
+ session.add(
+ EdgeJobModel(
+ dag_id=DAG_ID,
+ task_id=task_id,
+ run_id=RUN_ID,
+ try_number=1,
+ map_index=-1,
+ state=TaskInstanceState.QUEUED,
+ queue=QUEUE,
+ concurrency_slots=1,
+ command=MOCK_COMMAND_STR,
+ priority_weight=priority_weight,
+ queued_dttm=queued_dttm + timedelta(seconds=offset),
+ )
+ )
+ session.commit()
+
+ body = WorkerQueuesBody(free_concurrency=1, queues=[QUEUE],
team_name=None)
+ fetched = [fetch("worker1", body, session) for _ in range(3)]
+
+ assert [job.task_id for job in fetched if job] == ["high",
"medium", "low"]
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, reason="ExecuteCallback
requires Airflow 3.3+")
+ @patch(f"{Stats.__module__}.Stats.incr")
+ def test_fetch_returns_callbacks_before_higher_priority_tasks(self,
mock_stats_incr, session: Session):
Review Comment:
When I swapped `main`'s `jobs.py` back in, this test still passed. Both
callbacks are queued before the heavy task, so plain oldest-first ordering
already gives the expected result. If the heavy task is queued *before* the
callbacks, the test would actually show that callbacks jump ahead of tasks. The
other two new fetch tests fail on `main` as expected.
---
Drafted-by: Claude Code (Opus 5); reviewed by @dheerajturaga before posting
##########
providers/edge3/tests/unit/edge3/models/test_db.py:
##########
@@ -358,6 +358,40 @@ def test_migration_adds_concurrency_column(self, session):
assert "concurrency" in columns, "Migration 0002 should have added the
concurrency column"
assert "team_name" in columns, "Migration 0003 should have added the
team_name column"
+ def test_migration_adds_priority_weight_column(self, session):
Review Comment:
This drops the `rj_order` index and the `priority_weight` column on the
shared test database before upgrading. If the test fails partway, later tests
run against a broken schema and fail in confusing ways. Could the teardown
restore the schema (`try/finally`, or a fixture that runs `upgradedb()` at the
end)?
---
Drafted-by: Claude Code (Opus 5); reviewed by @dheerajturaga before posting
##########
providers/edge3/tests/unit/edge3/worker_api/routes/test_jobs.py:
##########
@@ -177,6 +178,163 @@ def
test_state_finish_metric_omits_team_name_for_global_job(self, mock_stats_inc
},
)
+ @patch(f"{Stats.__module__}.Stats.incr")
+ def test_fetch_returns_highest_priority_job_first(self, mock_stats_incr,
session: Session):
Review Comment:
Nit: the test names already say what each test checks, so most of the inline
comments ("The oldest job carries the lowest priority…", "Queued last and far
heavier…", etc.) could go, or shrink to one short line per test.
---
Drafted-by: Claude Code (Opus 5); reviewed by @dheerajturaga before posting
##########
providers/edge3/src/airflow/providers/edge3/migrations/versions/0006_4_4_0_add_priority_weight_to_edge_job.py:
##########
@@ -0,0 +1,68 @@
+#
+# 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.
+
+"""
+Add priority_weight column to edge_job table.
+
+Revision ID: d5a2f8b41c07
+Revises: c6b3c3d093fd
+Create Date: 2026-08-25 00:00:00.000000
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "d5a2f8b41c07"
+down_revision = "c6b3c3d093fd"
+branch_labels = None
+depends_on = None
+edge3_version = "4.4.0"
+
+NEW_INDEX_COLUMNS: list[str | sa.TextClause] = [
+ "state",
+ sa.text("priority_weight DESC"),
+ "queued_dttm",
+ "queue",
+]
+OLD_INDEX_COLUMNS: list[str | sa.TextClause] = ["state", "queued_dttm",
"queue"]
+
+
+def _recreate_rj_order(columns: list[str | sa.TextClause]) -> None:
+ inspector = sa.inspect(op.get_bind())
+ if "rj_order" in {idx["name"] for idx in
inspector.get_indexes("edge_job")}:
+ op.drop_index("rj_order", table_name="edge_job")
+ op.create_index("rj_order", "edge_job", columns)
+
+
+def upgrade() -> None:
Review Comment:
Nit: the earlier edge3 migrations don't check whether the column or index
already exists before changing them. Is there an install path that needs these
checks here, e.g. tables created with `create_all` before the migration runs?
If so, a one-line comment would help. If not, dropping them would match the
other migrations.
---
Drafted-by: Claude Code (Opus 5); reviewed by @dheerajturaga before posting
--
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]
