tomron commented on code in PR #74755:
URL: https://github.com/apache/airflow/pull/74755#discussion_r4279033107
##########
airflow-core/newsfragments/74755.bugfix.rst:
##########
@@ -0,0 +1 @@
+Enforce ``execution_timeout`` on stub tasks that run in a language SDK (Go,
Java, TypeScript). The Execution API now returns the timeout in the task run
context and the supervisor arms its deadline from it, so an overrunning stub
task is stopped like a Python task.
Review Comment:
Removed, thanks.
##########
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:
Renamed it to get_language_sdk_execution_timeout, and the docstring and
caller comment now cover stub tasks and native Go/TypeScript/Java Dag tasks.
##########
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:
Updated. Python tasks start the deadline just before execute(), language-SDK
tasks when the runtime connects, and the kill comes after the timeout plus the
cleanup grace period.
##########
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:
Rewrote it to assert the actual kill behavior (no kill before timeout plus
grace, then SIGTERM with escalation, only once) instead of private deadline
fields, following test_execution_timeout_kills_task_that_does_not_stop.
##########
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:
Good catch. The helper now uses `if not task.execution_timeout`, so zero
means no timeout like in Python. Added a zero-timeout regression case to the
Execution API test.
--
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]
