amoghrajesh commented on code in PR #57778:
URL: https://github.com/apache/airflow/pull/57778#discussion_r2488954966
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
Review Comment:
This one too.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -498,8 +499,13 @@ def test_run_task_timeout(time_machine, create_runtime_ti,
mock_supervisor_comms
assert ti.state == TaskInstanceState.FAILED
+ # Verify that SetTaskExecutionTimeout was sent to supervisor
+
mock_supervisor_comms.send.assert_any_call(SetTaskExecutionTimeout(execution_timeout_seconds=0.01))
+
# this state can only be reached if the try block passed down the
exception to handler of AirflowTaskTimeout
-
mock_supervisor_comms.send.assert_called_with(TaskState(state=TaskInstanceState.FAILED,
end_date=instant))
+ mock_supervisor_comms.send.assert_called_with(
+ msg=TaskState(state=TaskInstanceState.FAILED, end_date=instant)
+ )
Review Comment:
This test is wrong now, this checs if task runner sent FAILED state to
supervisor which is not true anymore because the supervisor handles the
timeouts now. Test needs updates, you can move it to test_supervisor if you
like too
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -801,6 +802,12 @@ def _prepare(ti: RuntimeTaskInstance, log: Logger,
context: Context) -> ToSuperv
ti.hostname = get_hostname()
ti.task = ti.task.prepare_for_execution()
+ # Send execution timeout to supervisor so it can handle timeout in parent
process
+ if ti.task.execution_timeout:
+ timeout_seconds = ti.task.execution_timeout.total_seconds()
+ log.debug("Sending execution_timeout to supervisor",
timeout_seconds=timeout_seconds)
+
SUPERVISOR_COMMS.send(msg=SetTaskExecutionTimeout(execution_timeout_seconds=timeout_seconds))
Review Comment:
This is cool, but we will have to remove the handling in task runner under
`_execute_task`?
--
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]
