Make EdgeExecutor respect [core] parallelism

[FrankYang0529]

Why

• Other distributed executors follow the BaseExecutor workload pipeline and record each workload in self.running. KubernetesExecutor does it in _process_workloads(), so slots_available reflects what they have in flight.
• EdgeExecutor writes queued tasks straight into the edge_job table and never records them in self.running. Nothing has written to that set since Remove Airflow 2 code path in executors #51009 removed the Airflow 2 dispatch path. slots_available is therefore always parallelism and [core] parallelism has no effect on EdgeExecutor.
• The empty set also makes the if job.key in self.running: block in _purge_jobs() unreachable, so the executor never reports task state to the scheduler.

How

• queue_workload() records the key of an ExecuteTask workload in self.running.
• _purge_jobs() reconciles self.running against every job row of the team, queued ones included. The previous reconciliation covered only the six non-queued states, so a key added at queue time was dropped again on the next sync(), before a worker could pick the job up. The _get_tracked_job_keys read takes no row locks: an edge worker fetches its next job with FOR UPDATE SKIP LOCKED, so locking queued rows here would make it come back empty.
• _purge_jobs() no longer reports RESTARTING as a failure. The fetch endpoint parks a claimed job in that state until the worker reports RUNNING, so the executor now records it as a transition and the slot stays taken.
• EdgeJobModel.key returns the airflow.models TaskInstanceKey instead of the airflow.sdk one. state_class_for_key() dispatches with isinstance and the two are unrelated classes, so running_state() would raise TypeError as soon as running became non-empty.
• try_adopt_task_instances() restores running from the team's edge_job rows after a scheduler restart and returns task instances without a job row as not adopted.

Verification

• uv run --project providers/edge3 pytest providers/edge3/tests/unit

---

Was generative AI tooling used to co-author this PR?

• Yes - Claude Code

---

• Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
• For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
• When adding dependency, check compliance with the ASF 3rd Party License Policy.
• For significant user-facing changes create newsfragment: {pr_number}.significant.rst, in airflow-core/newsfragments. You can add this file in a follow-up commit after the PR is created so you know the PR number.

[dheerajturaga]

Requesting changes for two correctness issues:

1. The worker claim transition can incorrectly fail a task. queue_workload() now adds the task key to self.running, while the worker's fetch endpoint commits the job as RESTARTING before it reports RUNNING. _purge_jobs() treats every RESTARTING job in self.running as failed. A scheduler heartbeat in that window therefore emits a FAILED executor event and removes the slot even though the worker is starting the task. This needs a non-terminal claim state or equivalent handling that distinguishes a worker claim from an actual failed/retry transition.

2. Scheduler failover loses the new parallelism accounting. try_adopt_task_instances() returns an empty list, declaring all Edge tasks adopted, but it never restores their keys to the new executor's self.running set. _get_tracked_job_keys() only intersects the existing set, so it cannot restore those entries. After a scheduler restart, existing Edge tasks consume no executor slots and the scheduler can queue up to parallelism additional tasks. Adoption should populate self.running for the Edge jobs actually present for the executor's team and return any missing jobs as not adopted.

I reproduced both cases with focused Edge provider tests: the claim-state test received a FAILED event, and the adoption test left running empty with the slot still available.

---

Drafted-by: Codex (GPT-5); reviewed by @dheerajturaga before posting

[FrankYang0529]

@dheerajturaga Good catch! Thanks for the review. I did following changes.

1. _purge_jobs() no longer treats RESTARTING as a failure. A RESTARTING job now only records the transition in last_reported_state, so the slot stays taken and the RUNNING that follows is reported as before. A worker that dies right after claiming still falls back to [scheduler] task_queued_timeout and then revoke_task().
2. try_adopt_task_instances() now reads the edge_job rows of the executor's team, puts the matching keys back into running, and returns task instances without a job row as not adopted, so the scheduler clears and re-schedules them.

[dheerajturaga]

Thanks for the quick turnaround on the previous round — both fixes look correct: the RESTARTING-vs-failure change in _purge_jobs() matches the fetch endpoint's claim semantics, and try_adopt_task_instances() now restores running from the team's edge_job rows.

Going through the rest of the diff turned up a few more correctness issues, most to least severe:

1. Callback workloads bypass the new parallelism enforcement (edge_executor.py:106) — the callback-execute branch of queue_workload() (lines 106-134) never calls self.running.add(...), unlike the ExecuteTask branch (line 171). So DAG/task callback workloads still bypass the parallelism enforcement this PR introduces — an unbounded number can run concurrently regardless of [core] parallelism, even though slots_available's docstring says it should count "tasks, callbacks, and connection tests." This is the one I'd want addressed before merge, since it undermines the PR's own goal for that workload type.

2. try_adopt_task_instances() can adopt already-finished jobs (edge_executor.py:419) — _get_tracked_job_keys() (lines 264-282) has no state filter, so it returns keys for SUCCESS/FAILED/REMOVED rows too, as long as they haven't yet been purged by _purge_jobs()'s success/fail purge timers. After a scheduler restart, this can re-adopt an already-finished job and hold a parallelism slot for it until the next purge cycle.

3. self.running updated before the caller's session commit (edge_executor.py:171) — self.running.add(key) runs before the caller (e.g. the scheduler's _critical_section_enqueue_task_instances, which batches multiple TIs per commit) commits the session. If that transaction rolls back, the DB insert is undone but the in-memory key isn't — self.running over-reports occupied slots until the next _purge_jobs() sync intersects it with the DB. Self-healing within one heartbeat interval, so lower severity than the other two.

Apologies for the back-and-forth here — appreciate you working through these with me.

---

Drafted-by: Claude Code (Sonnet 5); reviewed by @dheerajturaga before posting

[FrankYang0529]

Thanks for the second review. Learn a lot from you comment.

1. Callbacks: queue_workload() now records the key of both workload types, so slots_available and slots_occupied count callbacks too. For that to work, a callback row has to map back to its CallbackKey. Otherwise the reconciliation in _purge_jobs() drops the key on the next sync and the state block never matches the row.

2. Adoption: try_adopt_task_instances() only considers rows in QUEUED, RESTARTING or RUNNING and doesn't return finished rows (SUCCESS, FAILED, or REMOVED).

3. running.add() before the commit: keep the behavior with a comment at the call site. BaseExecutor.queue_workload() has the same pattern: it fills queued_tasks before the caller commits, and nothing removes that entry on rollback. Here the DB reconciliation in _purge_jobs() corrects an over-report within one heartbeat.

[dheerajturaga]

Thanks for addressing the earlier findings. One blocking regression remains:

build_job_key() treats every row with dag_id == "ExecuteCallback" as a callback. Since ExecuteCallback is a valid Dag ID, normal tasks in such a Dag are reconstructed as CallbackKeys, removed from self.running during synchronization, and no longer report terminal states.

Please identify callbacks using the complete synthetic identity—such as the expected run_id, task_id, try_number, and map_index combination—or persist an explicit workload type. Please also add a regression test for a normal task in a Dag named ExecuteCallback.

Once the key collision is fixed, this looks ready to merge.

---

Drafted-by: Codex (GPT-5); reviewed by @dheerajturaga before posting

[FrankYang0529]

Thanks for the review. build_job_key() now recognises a callback row only by the full identity queue_workload() writes for it: dag_id == "ExecuteCallback", run_id == "ExecuteCallback-<task_id>", try_number == 0 and map_index == -1.

Besides the key collision, this push fixes one more problem and add a changelog about this PR:

1. A job can end up in a state _purge_jobs() never handles. _update_orphaned_jobs() copies whatever state the task instance has (scheduled, deferred, up_for_reschedule, ...) into the job row. The task instance row only holds the latest try, so that state can come from a later try than the job's. Such a job row is never reported or deleted, so with the new accounting it would hold a slot until the scheduler restarts. The reconciliation now keeps a slot only for a row that is queued or in a state _purge_jobs() handles.
2. The changelog warning about [core] parallelism taking effect on Edge was missing from the branch.

[dheerajturaga]

Great! thanks for implementing this and going back and forth with reviews.