Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

AIP 72: Handling "deferrable" tasks in execution_api and task SDK #44241

Merged
merged 2 commits into from
Nov 27, 2024

Conversation

amoghrajesh
Copy link
Contributor

@amoghrajesh amoghrajesh commented Nov 21, 2024

closes: #44137

This PR is trying to port the "deferral" logic from airflow 2 to the airflow 3 (execution api + task sdk)

Summary of changes:

Server side changes (execution api):

  1. Handling TIDeferredStatePayload in ti_update_state -> covered by unit test: test_ti_update_state_to_deferred
    a. Didn't piggy back on ti.defer_task() as it extracts the trigger out of TaskDeferred exception. It is much more expensive to send across multiple models like TaskInstance, TaskDeferred, Trigger instead of just the required minimal properties
    b. returning and not proceeding with query execution as we already do it above https://github.com/apache/airflow/pull/44241/files#diff-d44a72566870079ee943e24bac2af74fb84c426c54d210561a251549a7078ed7L129
  2. Defining a datamodel for TIDeferredStatePayload and adding it to the discriminator: ti_state_discriminator

Client side changes (task sdk):

HTTP client:

Added a new function defer that sends a patch request to the task-instances/{id}/state execution api with payload: PatchTIToDeferred

Comms:

Defining a new data model to send a request to patch ti as "deferred" from task runner to supervisor: PatchTIToDeferred (Added to ToSupervisor)

Supervisor:

  1. Adding a new class property as _final_state to support @property final_state which is final state of a TI
    a. Added a setter to set values for this final state for cases like deferred so that the finish is not called for tasks those aren't in terminal stage: https://github.com/apache/airflow/pull/44241/files#diff-c2651fdee1a25e091e2a9d4f937f8032ca3d289d0de76f38ed88aee5df0f880dL392-L394
  2. Deferred tasks first enter the "wait" in supervisor and then mark themselves as deferred by setting up a trigger. So skipping finish() when final_state is not TerminalTIState
  3. Extended handle_requests to receive requests from task runner and forwarding the message to http client to call defer

TaskRunner:

Task runner executes: ti.task.execute and raises TaskDeferred for deferral. This sends a request to supervisor using SUPERVISOR_COMMS

How was this tested?

  1. The end to end flow isn't available as of now, but I have tried to cover up the ground using unit tests
  2. New case under test_handle_requests covers the supervisor + client side of things along with a mock of the message from task runner
  3. test_ti_update_state_to_deferred covers the scenario for execution API
  4. Added an integration test for task runner component:
  • It creates a DAG that raises the deferral exception:
    image
  • test_run_deferred_basic tests if the DAG raised a task exception and sent the right message across the SUPERVISOR_COMMS or not.

Additional sanity tests

Testing the "task runner" entries and exits (client side)

  1. Wrote a async operator DAG:
    image

  2. Breakpoint inside task_runner.py#run

  3. Validated the exception entrypoint with all the variables to be set as required by running the test: test_run_basic
    image

Testing the DB state in the server side (execution API)

  1. Used the test: test_ti_update_state_to_deferred

  2. Debugger inside ti_update_state method

  3. Validated the TaskDeferred and related logic
    image

  4. Checked state of the DB on execution
    image

image


^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

Copy link
Member

@ashb ashb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A good start!

@amoghrajesh amoghrajesh added the area:task-execution-interface-aip72 AIP-72: Task Execution Interface (TEI) aka Task SDK label Nov 22, 2024
@amoghrajesh amoghrajesh changed the title AIP 72: Handling deferrable tasks in execution API as well as TASK SDK AIP 72: Handling "deferrable" tasks in execution_api and task SDK Nov 22, 2024
@amoghrajesh amoghrajesh marked this pull request as ready for review November 22, 2024 11:39
@amoghrajesh amoghrajesh self-assigned this Nov 25, 2024
@amoghrajesh amoghrajesh requested review from kaxil and ashb November 25, 2024 07:33
airflow/api_fastapi/execution_api/routes/task_instances.py Outdated Show resolved Hide resolved
airflow/api_fastapi/execution_api/routes/task_instances.py Outdated Show resolved Hide resolved
airflow/api_fastapi/execution_api/routes/task_instances.py Outdated Show resolved Hide resolved
airflow/api_fastapi/execution_api/routes/task_instances.py Outdated Show resolved Hide resolved
task_sdk/src/airflow/sdk/execution_time/comms.py Outdated Show resolved Hide resolved
task_sdk/src/airflow/sdk/execution_time/supervisor.py Outdated Show resolved Hide resolved
task_sdk/src/airflow/sdk/execution_time/supervisor.py Outdated Show resolved Hide resolved
task_sdk/src/airflow/sdk/execution_time/task_runner.py Outdated Show resolved Hide resolved
Copy link
Member

@kaxil kaxil left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Approving pre-emptively as the code looks good but the tests needs a bit more work.

@amoghrajesh
Copy link
Contributor Author

@kaxil I wrote a new integration test that tests out the deferred exception functionality in task runner as well as tests the message sent across to the SUPERVISOR_COMMS through the send_request function

In this commit: dc39226

@kaxil
Copy link
Member

kaxil commented Nov 27, 2024

Squashed commits into 1 and resolved conflicts.

Doing one final pass now

@kaxil kaxil force-pushed the AIP72-task-deferred branch from 1051fc5 to c2b6003 Compare November 27, 2024 13:01
@kaxil kaxil force-pushed the AIP72-task-deferred branch from 187c6ee to 5bdb4e1 Compare November 27, 2024 14:39
@kaxil kaxil force-pushed the AIP72-task-deferred branch from 5bdb4e1 to 4bd9d05 Compare November 27, 2024 14:39
@kaxil kaxil merged commit 761cedd into apache:main Nov 27, 2024
50 checks passed
@kaxil kaxil deleted the AIP72-task-deferred branch November 27, 2024 15:07
prabhusneha pushed a commit to astronomer/airflow that referenced this pull request Nov 28, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
area:task-execution-interface-aip72 AIP-72: Task Execution Interface (TEI) aka Task SDK
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Handle "deferral" in ti_update_state endpoint in Execution API
3 participants