Enforce execution_timeout on stub tasks run by a language SDK by tomron · Pull Request #74755 · apache/airflow

tomron · 2026-10-14T17:57:36Z

A language-SDK task's execution_timeout is serialized with the Dag, but the runtime never got it. This covers both @task.stub tasks and native language-SDK Dag tasks, since both serialize with is_stub=True. TIRunContext had no field for the timeout, and only the Python task runner sends SetExecutionTimeout, which language-SDK runtimes skip. So a Go, Java or TypeScript task ran past its timeout.

What changed:

TIRunContext gets execution_timeout (seconds). ti_run sets it only for language-SDK tasks (stub and native Dag tasks) that have a timeout.
ActivitySubprocess arms its existing deadline from that value when it starts the runtime, the same way SetExecutionTimeout does. A task still running at execution_timeout + killed_task_cleanup_time gets SIGTERM, then SIGKILL after another killed_task_cleanup_time.
The field is added to the Execution API and the supervisor schema at the in-progress 2026-10-30 version, with downgrade for older clients. Go and TypeScript generated models and vendored schemas are updated.
Python tasks are unchanged.

Not covered: the runtimes do not enforce the timeout themselves (for example Go context cancel), see #74717. I did not run the Go repro end to end.

Tests: Execution API head and backward-compat tests, supervisor startup tests (with and without a server-provided timeout), supervisor schema migrator test.

Was generative AI tooling used to co-author this PR?

Yes (please specify the tool below)

Generated-by: an LLM-based AI coding assistant, reviewed and tested by me, following the guidelines

A stub task's execution_timeout was serialized with the Dag but never reached the runtime: TIRunContext had no field for it, and only the Python task runner sends SetExecutionTimeout, which stub tasks skip. A Go, Java or TypeScript task therefore ran past its timeout. The Execution API now returns the stub task's execution_timeout, in seconds, in TIRunContext, and the supervisor arms its existing execution_timeout deadline from it when it starts the runtime. An overrunning task gets SIGTERM, then SIGKILL, after the cleanup grace period. Python tasks are unchanged: the field is only set for stubs. The new field is added to both the Execution API and the supervisor schema at the in-progress 2026-10-30 version, and the Go and TypeScript generated models and vendored schemas are updated. Signed-off-by: Tom Ron <rontom@gmail.com>

Signed-off-by: Tom Ron <rontom@gmail.com>

Andrushika · 2026-10-15T02:04:31Z

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


No newsfragment needed

Removed, thanks.

Andrushika

Thanks for picking this up. Left some comments.

Andrushika · 2026-10-15T03:18:58Z

    return task.arg_bindings
+
+
+def get_stub_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float | None:


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.

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.

Andrushika · 2026-10-15T03:19:52Z

            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:


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.

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.

Andrushika · 2026-10-15T03:20:28Z

+            pytest.param(None, None, id="no_timeout"),
+        ],
+    )
+    def test_execution_timeout_from_server_arms_deadline_at_startup(


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.

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.

Andrushika · 2026-10-15T03:24:04Z

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


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

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.

Signed-off-by: Tom Ron <rontom@gmail.com>

Andrushika

Thanks for the quick update, the earlier points look good to me. One thing for the PR body: SIGTERM is sent at execution_timeout + killed_task_cleanup_time, not at execution_timeout, and native Dag tasks are covered too. Could we say both, since the body becomes the commit message?

Drafted-by: Claude Code (Opus 5.5); reviewed by @Andrushika before posting

Andrushika · 2026-10-16T05:54:28Z

        AddArgBindingsToTIRunContext,
        AddCallbackRunEndpoint,
        AddDagRunNoteUpdateEndpoint,
+        AddExecutionTimeoutToTIRunContext,


We might need to decide which release this goes into before merging. 2026-10-30 already shipped in 3.4.0b1 and its published supervisor schema has no execution_timeout, while main now adds new changes at 2027-02-28 (#74370). So either this is backported to v3-4-test before 3.4.0, or it moves to 2027-02-28 here and in the supervisor schema.

Thanks. Moving to 2027-02-28 is my preference, but on the supervisor-schema side it has a knock-on effect I'd like to check first.

A new 2027-02-28 supervisor-schema version moves schema.json's api_version, and the sync-go-sdk-schemas / sync-ts-sdk-schemas hooks keep the Go and TS vendored copies identical to it. So both SDKs would also have to declare supervisor schema 2027-02-28 (SupervisorSchemaVersion, capabilities.yaml, the generated TS constant), which means SDK builds from main would need an Airflow release newer than 3.4.

Is that acceptable, or would you rather keep 2026-10-30 and backport to v3-4-test before 3.4.0? I have the 2027-02-28 move ready locally and can push either way. The other three comments are addressed in 5454b4f, and the PR body now spells out the SIGTERM timing and the native Dag task coverage.

Drafted-by: Claude Code (Opus 5.5) (no human review before posting)

Andrushika · 2026-10-16T05:55:19Z

+            pytest.param(None, id="no_timeout"),
+        ],
+    )
+    def test_execution_timeout_from_server_kills_task_that_does_not_stop(self, mocker, server_timeout):


Could we split the no_timeout case into its own test? With the if server_timeout is None branch it is hard to see what each case expects.

Andrushika · 2026-10-16T05:55:53Z


+    execution_timeout: float | None = None
+    """
+    The task's ``execution_timeout`` in seconds, for stub (foreign-runtime) tasks.


Nit: this docstring and the two AddExecutionTimeout... descriptions still say stub tasks only. Maybe mention native Dag tasks too, like the helper docstring now does.

Andrushika · 2026-10-16T05:56:18Z

    return task.arg_bindings
+
+
+def get_language_sdk_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float | None:


Nit: the first lines here are the same lookup as get_arg_bindings. Maybe a small helper that returns the stub task, so both can use it?

The stub-task lookup was duplicated between the arg-bindings and timeout helpers, the docstrings still described the timeout as stub-only although native language-SDK Dag tasks get it too, and the parametrized supervisor test hid what the no-timeout case expects behind a branch.

Andrushika · 2026-10-16T11:41:03Z

Drafted-by: Claude Code (Opus 5.5) (no human review before posting)

WOW that's not good... we're fixing task-sdk -- one of the most sensitive parts in Airflow, please make sure you reviewed every line you submitted!

tomron requested review from amoghrajesh, ashb, jason810496 and kaxil as code owners October 14, 2026 17:57

boring-cyborg Bot added area:API Airflow's REST/HTTP API area:go-sdk area:task-sdk area:ts-sdk labels Oct 14, 2026

tomron added 2 commits October 14, 2026 20:57

Add newsfragment for execution_timeout on stub tasks

46574f1

Signed-off-by: Tom Ron <rontom@gmail.com>

Regenerate Go SDK models for the execution_timeout schema change

d2ab85e

Signed-off-by: Tom Ron <rontom@gmail.com>

Andrushika reviewed Oct 15, 2026

View reviewed changes

Andrushika mentioned this pull request Oct 15, 2026

Language SDK runtimes do not handle execution_timeout in the task process #74717

Open

1 of 5 tasks

Andrushika suggested changes Oct 15, 2026

View reviewed changes

Address execution timeout review feedback for language SDK tasks

32e2ea0

Signed-off-by: Tom Ron <rontom@gmail.com>

Andrushika mentioned this pull request Oct 15, 2026

Fix flaky in-process Execution API transport lifecycle test on MySQL #74785

Merged

1 task done

Andrushika suggested changes Oct 16, 2026

View reviewed changes

tomron closed this Oct 16, 2026

ChickenBenny mentioned this pull request Oct 16, 2026

Enforce execution_timeout for tasks run by a language SDK #74869

Merged

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

		return task.arg_bindings


		def get_stub_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float \| None:

		return task.arg_bindings


		def get_language_sdk_execution_timeout(dag_bag: DBDagBag, ti: Any, *, session: Session) -> float \| None:

Conversation

tomron commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Was generative AI tooling used to co-author this PR?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Andrushika left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Andrushika left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Andrushika commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tomron commented Oct 14, 2026 •

edited

Loading

Andrushika commented Oct 16, 2026 •

edited

Loading