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

amoghrajesh · 2024-11-21T09:59:32Z

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):

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
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:

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
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
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?

The end to end flow isn't available as of now, but I have tried to cover up the ground using unit tests
New case under test_handle_requests covers the supervisor + client side of things along with a mock of the message from task runner
test_ti_update_state_to_deferred covers the scenario for execution API

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

ashb

A good start!

airflow/api_fastapi/execution_api/datamodels/taskinstance.py

airflow/api_fastapi/execution_api/routes/task_instances.py

ashb · 2024-11-21T10:08:31Z

tests/api_fastapi/execution_api/routes/test_task_instances.py

+            "state": "deferred",
+            "classpath": "my-class-path",
+            "kwargs": {},
+            "created_date": "2024-10-31T12:00:00Z",


Hmm, why is this being sent in the payload? Should it be the time the server received the request instead?

It is handled this way:

self.created_date = created_date or timezone.utcnow()

In case we do not pass it, it will take the UTC now. Should we allow the option to override it or remove it entirely?

tests/api_fastapi/execution_api/routes/test_task_instances.py

Co-authored-by: Ash Berlin-Taylor <[email protected]>

airflow/api_fastapi/execution_api/datamodels/taskinstance.py

airflow/api_fastapi/execution_api/routes/task_instances.py

task_sdk/src/airflow/sdk/execution_time/comms.py

amoghrajesh · 2024-11-22T10:52:06Z

task_sdk/src/airflow/sdk/execution_time/supervisor.py

+                # hmmm, can we do better here
+                # setting to "\n" as we do not have a response to return


Needs some thought here @ashb @kaxil

Two options:

Move self.stdin.write(resp + b"\n") to all the if's where we are sending a msg

Set resp to None. and change L529 to:

if resp: self.stdin.write(resp + b"\n")

task_sdk/tests/execution_time/test_supervisor.py

amoghrajesh · 2024-11-22T10:53:32Z

task_sdk/tests/execution_time/test_supervisor.py

+        if mock_response == "":
+            # for task instance endpoints, there won't be any response


We should handle this based on how we handle the client response for task instance endpoints

airflow/api_fastapi/execution_api/datamodels/taskinstance.py

airflow/api_fastapi/execution_api/routes/task_instances.py

airflow/api_fastapi/execution_api/datamodels/taskinstance.py

airflow/api_fastapi/execution_api/routes/task_instances.py

task_sdk/src/airflow/sdk/execution_time/comms.py

kaxil · 2024-11-25T12:20:38Z

task_sdk/src/airflow/sdk/execution_time/supervisor.py

-        self.client.task_instances.finish(
-            id=self.ti_id, state=self.final_state, when=datetime.now(tz=timezone.utc)
-        )
+        if self.final_state in TerminalTIState:


Not 100% certain if we need this check -- I am thinking by defn self.final_state has to be in TerminalState -- if not it is an error!

kaxil · 2024-11-25T12:23:11Z

task_sdk/src/airflow/sdk/execution_time/supervisor.py

+    @final_state.setter
+    def final_state(self, value):
+        """Setter for final_state for certain task instance stated present in IntermediateTIState."""
+        if value not in TerminalTIState:
+            self._final_state = value


Why do we need a setter? Shouldn't we set TI state to falil if exit_code != 0 like we already in

airflow/task_sdk/src/airflow/sdk/execution_time/supervisor.py

Lines 472 to 474 in 4d3140d

if self._exit_code == 0:

return self._terminal_state or TerminalTIState.SUCCESS

return TerminalTIState.FAILED

kaxil · 2024-11-25T12:25:29Z

task_sdk/src/airflow/sdk/execution_time/supervisor.py

+                # hmmm, can we do better here
+                # setting to "\n" as we do not have a response to return


Two options:

Move self.stdin.write(resp + b"\n") to all the if's where we are sending a msg

Set resp to None. and change L529 to:

if resp: self.stdin.write(resp + b"\n")

kaxil · 2024-11-25T12:27:03Z

task_sdk/src/airflow/sdk/execution_time/task_runner.py

-    except TaskDeferred:
-        ...
+    except TaskDeferred as defer:
+        trigger = Trigger.from_object(defer.trigger)


Why Trigger model here?

i.e. why do we need L164 to L169

It is needed because the trigger object provides the classpath, timeout and the kwargs from the exception. Done similarly in

airflow/airflow/models/taskinstance.py

Lines 1631 to 1634 in acf106b

ti.state = TaskInstanceState.DEFERRED

ti.trigger_id = trigger_row.id

ti.next_method = next_method

ti.next_kwargs = next_kwargs or {}

execution_api side changes (server)

d2ec946

ashb reviewed Nov 21, 2024

View reviewed changes

Update tests/api_fastapi/execution_api/routes/test_task_instances.py

6d54213

Co-authored-by: Ash Berlin-Taylor <[email protected]>

kaxil reviewed Nov 21, 2024

View reviewed changes

airflow/api_fastapi/execution_api/datamodels/taskinstance.py Outdated Show resolved Hide resolved

amoghrajesh added 3 commits November 22, 2024 12:13

initial review comments

9e669c8

Merge branch 'main' into AIP72-task-deferred

9ea3549

task_sdk side changes (client)

4caeee1

amoghrajesh added the area:task-execution-interface-aip72 AIP-72: Task Execution Interface (TEI) aka Task SDK label Nov 22, 2024

amoghrajesh added 2 commits November 22, 2024 15:30

_generated.py changes

da73a74

removing dirty code and fixing timeout

fb640a7

amoghrajesh commented Nov 22, 2024

View reviewed changes

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 marked this pull request as ready for review November 22, 2024 11:39

amoghrajesh requested review from ephraimbuddy and pierrejeambrun as code owners November 22, 2024 11:39

fixing static checks

6dd8f14

kaxil reviewed Nov 22, 2024

View reviewed changes

amoghrajesh self-assigned this Nov 25, 2024

amoghrajesh added 2 commits November 25, 2024 12:25

review comments from kaxil - part 1

38f483f

review comments from kaxil - part 2

4d3140d

amoghrajesh requested review from kaxil and ashb November 25, 2024 07:33

kaxil reviewed Nov 25, 2024

View reviewed changes

amoghrajesh added 3 commits November 25, 2024 18:16

removing the created_at from payload

03e83e2

adding generated.py

d5da767

removing execution_timeout and session.flush

28d0b9f

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

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

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

amoghrajesh commented Nov 21, 2024 •

edited

Loading

ashb left a comment

ashb Nov 21, 2024

amoghrajesh Nov 22, 2024

amoghrajesh Nov 22, 2024

kaxil Nov 25, 2024

amoghrajesh Nov 22, 2024

kaxil Nov 25, 2024

kaxil Nov 25, 2024

kaxil Nov 25, 2024

kaxil Nov 25, 2024

amoghrajesh Nov 25, 2024

		# hmmm, can we do better here
		# setting to "\n" as we do not have a response to return

		if mock_response == "":
		# for task instance endpoints, there won't be any response

	if self._exit_code == 0:
	return self._terminal_state or TerminalTIState.SUCCESS
	return TerminalTIState.FAILED

	ti.state = TaskInstanceState.DEFERRED
	ti.trigger_id = trigger_row.id
	ti.next_method = next_method
	ti.next_kwargs = next_kwargs or {}

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

Are you sure you want to change the base?

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

Conversation

amoghrajesh commented Nov 21, 2024 • edited Loading

Summary of changes:

Server side changes (execution api):

Client side changes (task sdk):

HTTP client:

Comms:

Supervisor:

TaskRunner:

How was this tested?

ashb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

amoghrajesh commented Nov 21, 2024 •

edited

Loading