Andrushika commented on code in PR #74755:
URL: https://github.com/apache/airflow/pull/74755#discussion_r4278658854
##########
airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py:
##########
@@ -59,3 +59,22 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session:
Session) -> list |
if (task := dag.task_dict.get(ti.task_id)) is None or not task.is_stub:
return None
return task.arg_bindings
+
+
+def get_stub_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session:
Session) -> float | None:
Review Comment:
We might need to mention native Dag tasks here too. The Go, TypeScript and
Java serializers also write `is_stub: true` on their tasks, so this already
returns their timeout, but the name and docstring read as stub only.
##########
task-sdk/src/airflow/sdk/execution_time/supervisor.py:
##########
@@ -1720,6 +1720,9 @@ def _on_child_started(
ti_context = self.client.task_instances.start(ti.id, self.pid,
datetime.now(tz=UTC))
self._should_retry = ti_context.should_retry
self._last_successful_heartbeat = time.monotonic()
+ if ti_context.execution_timeout is not None:
Review Comment:
We might need to update the `_handle_execution_timeout_if_needed` docstring
too. It still says the deadline is armed right before `execute()`, but for a
language SDK task it now starts when the runtime connects, so the task is
stopped `killed_task_cleanup_time` after the timeout, counted from there.
##########
task-sdk/tests/task_sdk/execution_time/test_supervisor.py:
##########
@@ -1664,6 +1664,54 @@ def
test_execution_timeout_kills_task_that_does_not_stop(self, mocker):
proc._handle_execution_timeout_if_needed()
mock_kill.assert_called_once()
+ @pytest.mark.parametrize(
+ ("server_timeout", "expected_timeout"),
+ [
+ pytest.param(30.0, 30.0, id="stub_task_with_timeout"),
+ pytest.param(None, None, id="no_timeout"),
+ ],
+ )
+ def test_execution_timeout_from_server_arms_deadline_at_startup(
Review Comment:
Could we assert the kill here instead of the private deadline fields? The
same shape as `test_execution_timeout_kills_task_that_does_not_stop` would work.
##########
airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py:
##########
@@ -59,3 +59,22 @@ def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session:
Session) -> list |
if (task := dag.task_dict.get(ti.task_id)) is None or not task.is_stub:
return None
return task.arg_bindings
+
+
+def get_stub_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session:
Session) -> float | None:
+ """
+ Return the stub task's ``execution_timeout`` in seconds.
+
+ Only stub (foreign-runtime) tasks need it from the server: a Python task
parses the Dag file and
+ reports its own timeout, but a lang-SDK runtime never sees the Dag.
``None`` for any other task
+ and for stubs that set no timeout.
+ """
+ if ti.dag_version_id is None:
+ return None
+ if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None:
+ return None
+ if (task := dag.task_dict.get(ti.task_id)) is None or not task.is_stub:
+ return None
+ if task.execution_timeout is None:
Review Comment:
Nit: the Python task runner treats a zero `execution_timeout` as no timeout
(`if task.execution_timeout:`), but here `timedelta(0)` returns `0.0` and the
supervisor kills the task after the grace period. Maybe `if not
task.execution_timeout:`?
--
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]
