Add ExecuteCallback support to Edge Executor

[wjddn279]

related: #62887

NOTE

Make draft pr to discuss direction of implementation with code owners

---

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

• Yes (please specify the tool below)

---

• 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.

[wjddn279]

@jscheffl @dheerajturaga @ferruzzi

Hello! Following the multi-team work, I've been assigned to add functionality to the edge executor. As before, I've opened a draft PR first to discuss the direction before starting development.

This task involves changing the Deadline Alerts callback work so that it can be executed on the worker. The edge_executor is currently structured to directly update the task list in the edge_job table within the queue_workload method. It needs to be updated to handle not only the existing workloads.ExecuteTask but also the newly added workloads.ExecuteCallback — and I'd like us to align on how to store this class in the DB. I see three possible approaches:

• Fit it into the existing schema without changing the edge_job schema.
  This should work functionally, but it would require populating existing columns with meaningless values (dag_id, task_id, run_id).

• Modify the edge_job schema.
  By adding a workload_type field, we can implicitly distinguish the expected schema per value. This could also be an opportunity to make changes to the PK and nullable constraints.

• Add a separate edge_callback_job table.
  Fully separating the tables would guarantee clear and precise consistency.

All three approaches could work without any functional issues. I'd love to hear your thoughts.

[jscheffl]

@jscheffl @dheerajturaga @ferruzzi

Hello! Following the multi-team work, I've been assigned to add functionality to the edge executor. As before, I've opened a draft PR first to discuss the direction before starting development.

This task involves changing the Deadline Alerts callback work so that it can be executed on the worker. The edge_executor is currently structured to directly update the task list in the edge_job table within the queue_workload method. It needs to be updated to handle not only the existing workloads.ExecuteTask but also the newly added workloads.ExecuteCallback — and I'd like us to align on how to store this class in the DB. I see three possible approaches:

* Fit it into the existing schema without changing the edge_job schema.
  This should work functionally, but it would require populating existing columns with meaningless values (dag_id, task_id, run_id).

* Modify the edge_job schema.
  By adding a workload_type field, we can implicitly distinguish the expected schema per value. This could also be an opportunity to make changes to the PK and nullable constraints.

* Add a separate edge_callback_job table.
  Fully separating the tables would guarantee clear and precise consistency.

All three approaches could work without any functional issues. I'd love to hear your thoughts.

Thanks for raising a short check. Actually have not thought about this and have no strong opinion developed.

An callback also is belonging to a Dag, Task, correct? So these details still could be used to populate details?

Note that the table structure and PK is originating from the time of Airflow 2 before a Task had a UUID. Therefore Dag_id, Task_id and Run_id were primary keys as it was the job PK. I matter of change I'd be OK to change the DB layout and switch to Task instance UUID or some other artificial UUID would also be OK and keeping the existing fields informational. In this case if these are not PK anymore you could fill will empty or a placeholder if not relevant/existing. But this would also impact REST API services, maybe some adjustment there are also needed (e.g. /jobs/state/...) as they carry the existing UUID.

Note that Core and worker might a bit diverging so if REST API is changed at least the existing endpoints need to be kept for compatibility.

If technically not tooo much overhead I'd slightly prefer a single table for all, then only 1 query is needed to poll for jobs.

Other opinions?

[wjddn279]

@jscheffl

Yes, I have no objection to keeping it as a single table. I generally agree with the opinion you've provided, and I'll implement the details and discuss further from there. Thank you!

[jscheffl]

@jscheffl

Yes, I have no objection to keeping it as a single table. I generally agree with the opinion you've provided, and I'll implement the details and discuss further from there. Thank you!

Cool!

@dheerajturaga + @ferruzzi still other opinions welcome :-D

[ferruzzi]

Edge is more or less your baby, so I'm going to defer to whatever you think here.

An callback also is belonging to a Dag, Task, correct? So these details still could be used to populate details?

As of right now, all deadline callbacks are tied to a dagrun. They are created and cleaned up as part of the dagrun lifecycle, so that means they obviously also have access to dag, dagbundle, serialized_dag, etc but not to Task. That may (and likely will?) change in the future, but that's what's active right now.

One thing to note is that there is a lot of work going on to unify TaskInstance and ExedcutorCallback into a shared interface we're calling Workload. So pretty much anywhere that had TaskInstanceKey or something like that would ideally be migrated to use workload.key, for example. You can see (and review ;) ) what I think may be the last batch of those changes in #63491 and an example of one of the other executors getting this treatment over in #62984 if that helps you think through what @wjddn279 is in store for.

Edit: Jens - I know you are aware of the other one and have reviewed it, that was more just a general comment for anyone who comes across this and wants context.

[jscheffl]

Sorry, maybe coming a bit late during second pass: So far the API was always backwards compatible, means if you have a Edge worker with a previous version running it would happily continue to run and softly "drain" if there is a provider or Airflow core version mis-match.

Now here the return type for job fetching changes, means if the worker would attempt to fetch a job with an old version, the job would be assigned but the worker (probably?) fails in de-serializing the content and fails. Probably before draining.

Have you tested this and how does it behave?
(Compared to Task SDK there is currently no Cadwyn layer for versioning, this is in the backlog...)

If this is a problem can we make it somehow compatible or make a parallel endpoint such that it is not failing on old clients until they drain and jobs are not pulled and never executed?

[jscheffl]

Sorry in the second pass review came along one problem.

Additionally can I kindly request that you add one example /test callback into the integration_test.py Dag in providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py which serves as a (currently manual) system test?

[wjddn279]

@jscheffl

Yes — I found a few issues while testing locally. I'll apply the fixes and request another review.

[wjddn279]

pending until #66141