This is an automated email from the ASF dual-hosted git repository.
ferruzzi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 27e807b5f1e Add Execution API endpoint for callback terminal state
(#75253)
27e807b5f1e is described below
commit 27e807b5f1ee3471017954a41b349384763de42a
Author: D. Ferruzzi <[email protected]>
AuthorDate: Tue Nov 10 14:43:13 2026 -0700
Add Execution API endpoint for callback terminal state (#75253)
* Add Execution API endpoint for callback terminal state
Callback state was written only by the scheduler, reading each executor's
in-process event buffer. Executors whose workers report back over HTTP never
reach that channel, so their callbacks sat in QUEUED forever.
The executor-event path stays as a fallback, now narrowed to rows still in
flight so a late event cannot overwrite an outcome the supervisor already
reported, or the output captured with it.
Replaces #66141, whose /run half landed separately in #71192.
---------
Co-authored-by: wjddn279 <[email protected]>
---
.../execution_api/datamodels/callback.py | 35 ++++
.../api_fastapi/execution_api/routes/callbacks.py | 58 ++++++-
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2027_02_28.py | 10 ++
.../src/airflow/jobs/scheduler_job_runner.py | 53 +++++-
.../execution_api/versions/head/test_callbacks.py | 188 +++++++++++++++++++++
.../versions/v2027_02_28/test_callbacks.py | 43 +++++
airflow-core/tests/unit/jobs/test_scheduler_job.py | 137 ++++++++++++++-
task-sdk/src/airflow/sdk/api/client.py | 7 +
.../src/airflow/sdk/api/datamodels/_generated.py | 21 +++
.../sdk/execution_time/callback_supervisor.py | 26 ++-
task-sdk/tests/task_sdk/api/test_client.py | 34 ++++
.../execution_time/test_callback_supervisor.py | 88 +++++++++-
13 files changed, 686 insertions(+), 16 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/callback.py
b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/callback.py
new file mode 100644
index 00000000000..bf0154be92f
--- /dev/null
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/callback.py
@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from enum import Enum
+
+from airflow.api_fastapi.core_api.base import StrictBaseModel
+
+
+class CallbackTerminalState(str, Enum):
+ """Terminal states a callback can transition to from RUNNING."""
+
+ SUCCESS = "success"
+ FAILED = "failed"
+
+
+class CallbackTerminalStatePayload(StrictBaseModel):
+ """Payload for transitioning a callback from RUNNING to a terminal
state."""
+
+ state: CallbackTerminalState
+ output: str | None = None
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py
index 2f6e386b0b3..655967b665a 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/callbacks.py
@@ -16,13 +16,15 @@
# under the License.
from __future__ import annotations
+from typing import Annotated
from uuid import UUID
from cadwyn import VersionedAPIRouter
-from fastapi import HTTPException, Response, Security, status
+from fastapi import Body, HTTPException, Response, Security, status
from airflow.api_fastapi.common.db.common import SessionDep
from airflow.api_fastapi.core_api.openapi.exceptions import
create_openapi_http_exception_doc
+from airflow.api_fastapi.execution_api.datamodels.callback import
CallbackTerminalStatePayload
from airflow.api_fastapi.execution_api.deps import DepContainer
from airflow.api_fastapi.execution_api.security import (
ExecutionAPIRoute,
@@ -35,7 +37,7 @@ from airflow.utils.state import CallbackState
router = VersionedAPIRouter(
route_class=ExecutionAPIRoute,
dependencies=[
- Security(require_auth, scopes=["cb:self", "token:callback"]),
+ Security(require_auth, scopes=["cb:self"]),
],
)
@@ -43,6 +45,7 @@ router = VersionedAPIRouter(
@router.patch(
"/{callback_id}/run",
status_code=status.HTTP_204_NO_CONTENT,
+ dependencies=[Security(require_auth, scopes=["token:callback"])],
responses=create_openapi_http_exception_doc(
[
(status.HTTP_404_NOT_FOUND, "Callback not found"),
@@ -83,3 +86,54 @@ def run_callback(
callback.state = CallbackState.RUNNING
issue_execution_token(services, response, sub=str(callback_id))
+
+
[email protected](
+ "/{callback_id}/state",
+ status_code=status.HTTP_204_NO_CONTENT,
+ responses=create_openapi_http_exception_doc(
+ [
+ (status.HTTP_404_NOT_FOUND, "Callback not found"),
+ (status.HTTP_409_CONFLICT, "The callback is not running and is not
already in that state"),
+ ]
+ ),
+)
+def update_callback_state(
+ callback_id: UUID,
+ payload: Annotated[CallbackTerminalStatePayload, Body()],
+ session: SessionDep,
+) -> None:
+ """Report the final state of a callback that has finished running."""
+ callback = session.get(Callback, callback_id, with_for_update=True)
+ if callback is None:
+ raise HTTPException(
+ status_code=status.HTTP_404_NOT_FOUND,
+ detail={
+ "reason": "not_found",
+ "message": f"Callback {callback_id} not found",
+ },
+ )
+
+ # Accepting the state the row already holds keeps the report idempotent:
the SDK client retries
+ # a lost response (Client.request defaults to retry=True), and a 409 there
makes the supervisor
+ # log a failure for an outcome that is correct. Safe to replay because
this route issues no
+ # token, unlike /run above.
+ reported_state = CallbackState(payload.state)
+ if callback.state not in (CallbackState.RUNNING, reported_state):
+ raise HTTPException(
+ status_code=status.HTTP_409_CONFLICT,
+ detail={
+ "reason": "invalid_state",
+ "message": (
+ f"Callback {callback_id} is in state {callback.state}; a
terminal state can only "
+ f"be reported for a RUNNING callback, or for one already
in {reported_state}."
+ ),
+ "previous_state": callback.state,
+ },
+ )
+
+ callback.state = reported_state
+ # Written on the idempotent path too: the scheduler may have reached this
same terminal state
+ # first and stored its own placeholder output.
+ if payload.output is not None:
+ callback.output = payload.output
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index c0dd008cecf..65c0be914f1 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -69,6 +69,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_10_30
import (
ReturnNotFoundForMissingDagRun,
)
from airflow.api_fastapi.execution_api.versions.v2027_02_28 import (
+ AddCallbackStateEndpoint,
AddDagProcessorJobEndpoints,
AddExecutionTimeoutToTIRunContext,
MakeVariableValueNonNullable,
@@ -81,6 +82,7 @@ bundle = VersionBundle(
AddDagProcessorJobEndpoints,
AddExecutionTimeoutToTIRunContext,
MakeVariableValueNonNullable,
+ AddCallbackStateEndpoint,
),
Version(
"2026-10-30",
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py
index 8447ed2756f..96cbfbab984 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2027_02_28.py
@@ -83,3 +83,13 @@ class MakeVariableValueNonNullable(VersionChange):
"""Keep the old behavior of storing ``str(None)`` when an older client
sets null."""
if "val" in request.body and request.body["val"] is None:
request.body["val"] = "None"
+
+
+class AddCallbackStateEndpoint(VersionChange):
+ """Add the callback terminal-state endpoint."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ endpoint("/callbacks/{callback_id}/state", ["PATCH"]).didnt_exist,
+ )
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 8566da6cb83..7cc8f06dd0e 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -1578,7 +1578,8 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
# Handle callback state events
for callback_id in callback_keys_with_events:
state, info = event_buffer.pop(callback_id)
- callback = session.get(Callback, UUID(str(callback_id)))
+ callback_uuid = UUID(str(callback_id))
+ callback = session.get(Callback, callback_uuid)
if not callback:
# This should not normally happen - we just received an event
for this callback.
# Only possible if callback was deleted mid-execution (e.g.,
cascade delete from DagRun deletion).
@@ -1587,14 +1588,52 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
continue
+ # The supervisor reports the outcome through PATCH
/callbacks/{id}/state, this is
+ # in case a supervisor dies before reporting.
+ if callback.state not in (CallbackState.QUEUED,
CallbackState.RUNNING):
+ cls.logger().debug(
+ "Callback %s is already %s; ignoring executor event with
state %s",
+ callback_id,
+ callback.state,
+ state,
+ )
+ continue
+
+ new_values: dict[str, Any] = {"state": state}
+ if state == CallbackState.FAILED:
+ new_values["output"] = str(info) if info else "Execution
failed"
+
+ # The QUEUED/RUNNING check above reads an ORM snapshot, and the
supervisor can report a
+ # terminal state through the API before this write lands.
Repeating it here stops a
+ # terminal row moving.
+ updated = cast(
+ "CursorResult",
+ session.execute(
+ update(Callback)
+ .where(
+ Callback.id == callback_uuid,
+ Callback.state.in_((CallbackState.QUEUED,
CallbackState.RUNNING)),
+ )
+ .values(**new_values)
+ .execution_options(synchronize_session=False)
+ ),
+ )
+ # The row moved under us, so the ORM copy is stale either way.
+ session.expire(callback)
+
+ if not updated.rowcount:
+ cls.logger().info(
+ "Callback %s reached a terminal state before its executor
event (%s) was handled; "
+ "keeping the reported result",
+ callback_id,
+ state,
+ )
+ continue
+
if state == CallbackState.SUCCESS:
- callback.state = CallbackState.SUCCESS
cls.logger().info("Callback %s completed successfully",
callback_id)
- elif state == CallbackState.FAILED:
- callback.state = CallbackState.FAILED
- callback.output = str(info) if info else "Execution failed"
- cls.logger().error("Callback %s failed: %s", callback_id,
callback.output)
- session.add(callback)
+ else:
+ cls.logger().error("Callback %s failed: %s", callback_id,
new_values["output"])
# Return if no finished tasks
if not tis_with_right_state:
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py
index da24c5c5f9f..d7a6cbafd4b 100644
---
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_callbacks.py
@@ -94,6 +94,148 @@ class TestRunCallback:
assert response.json()["detail"]["reason"] == "not_found"
[email protected]
+def make_running_callback(session):
+ """Build a callback already in RUNNING, the only state the state endpoint
accepts."""
+
+ def _make():
+ callback = ExecutorCallback(
+ callback_def=SyncCallback(sync_callback, kwargs={}),
+ fetch_method=CallbackFetchMethod.IMPORT_PATH,
+ )
+ callback.data["dag_id"] = "test_dag"
+ callback.state = CallbackState.RUNNING
+ session.add(callback)
+ session.commit()
+ return callback
+
+ return _make
+
+
[email protected]
+def running_callback(make_running_callback):
+ return make_running_callback()
+
+
[email protected]
+def other_callback(make_running_callback):
+ """A second, unrelated callback, so a token can name a real id that is not
this one."""
+ return make_running_callback()
+
+
+class TestUpdateCallbackState:
+ @pytest.mark.parametrize(
+ ("state", "output"),
+ [
+ pytest.param("success", None, id="success-without-output"),
+ pytest.param("success", "all good", id="success-with-output"),
+ pytest.param("failed", "exited with code 3",
id="failed-with-output"),
+ ],
+ )
+ def test_running_callback_reaches_its_terminal_state(
+ self, client, session, running_callback, state, output
+ ):
+ body = {"state": state}
+ if output is not None:
+ body["output"] = output
+
+ response =
client.patch(f"/execution/callbacks/{running_callback.id}/state", json=body)
+
+ assert response.status_code == 204
+ session.expire_all()
+ session.refresh(running_callback)
+ assert running_callback.state == state
+ assert running_callback.output == output
+
+ def test_output_is_left_alone_when_omitted(self, client, session,
running_callback):
+ running_callback.output = "captured earlier"
+ session.commit()
+
+ response =
client.patch(f"/execution/callbacks/{running_callback.id}/state",
json={"state": "failed"})
+
+ assert response.status_code == 204
+ session.expire_all()
+ session.refresh(running_callback)
+ assert running_callback.output == "captured earlier"
+
+ @pytest.mark.parametrize(
+ "state",
+ [CallbackState.QUEUED, CallbackState.PENDING, CallbackState.FAILED],
+ )
+ def test_only_a_running_callback_can_be_reported(self, client, session,
running_callback, state):
+ """A row that is neither RUNNING nor already in the reported state is
refused."""
+ running_callback.state = state
+ session.commit()
+
+ response = client.patch(
+ f"/execution/callbacks/{running_callback.id}/state",
json={"state": "success"}
+ )
+
+ assert response.status_code == 409
+ detail = response.json()["detail"]
+ assert detail["reason"] == "invalid_state"
+ assert detail["previous_state"] == state
+
+ session.expire_all()
+ session.refresh(running_callback)
+ assert running_callback.state == state
+
+ @pytest.mark.parametrize("state", [CallbackState.SUCCESS,
CallbackState.FAILED])
+ def test_repeat_report_of_the_same_state_is_accepted(self, client,
session, running_callback, state):
+ """A retry after a lost response must not look like a failure to the
supervisor."""
+ running_callback.state = state
+ running_callback.output = "reported on the first attempt"
+ session.commit()
+
+ response = client.patch(
+ f"/execution/callbacks/{running_callback.id}/state",
+ json={"state": str(state), "output": "reported on the first
attempt"},
+ )
+
+ assert response.status_code == 204
+ session.expire_all()
+ session.refresh(running_callback)
+ assert running_callback.state == state
+ assert running_callback.output == "reported on the first attempt"
+
+ def test_repeat_report_replaces_output_written_by_the_scheduler(self,
client, session, running_callback):
+ """The scheduler can reach FAILED first with a placeholder; the real
output still wins."""
+ running_callback.state = CallbackState.FAILED
+ running_callback.output = "Execution failed"
+ session.commit()
+
+ response = client.patch(
+ f"/execution/callbacks/{running_callback.id}/state",
+ json={"state": "failed", "output": "exited with code 3"},
+ )
+
+ assert response.status_code == 204
+ session.expire_all()
+ session.refresh(running_callback)
+ assert running_callback.state == CallbackState.FAILED
+ assert running_callback.output == "exited with code 3"
+
+ def test_returns_404_for_nonexistent(self, client):
+ response = client.patch(
+ "/execution/callbacks/00000000-0000-0000-0000-000000000000/state",
+ json={"state": "success"},
+ )
+ assert response.status_code == 404
+ assert response.json()["detail"]["reason"] == "not_found"
+
+ @pytest.mark.parametrize(
+ "body",
+ [
+ pytest.param({"state": "running"}, id="non-terminal-state"),
+ pytest.param({"state": "success", "extra": "nope"},
id="unknown-field"),
+ pytest.param({}, id="missing-state"),
+ ],
+ )
+ def test_rejects_a_malformed_payload(self, client, running_callback, body):
+ response =
client.patch(f"/execution/callbacks/{running_callback.id}/state", json=body)
+ assert response.status_code == 422
+
+
@pytest.fixture
def _use_real_jwt_bearer(exec_app):
"""Remove the mock require_auth override so the real JWT validation runs
end-to-end."""
@@ -129,3 +271,49 @@ def test_run_validates_token_scope_and_sub(client,
queued_callback, token_sub, t
resp = client.patch(f"/execution/callbacks/{queued_callback.id}/run")
assert resp.status_code == expected_status, resp.json()
+
+
[email protected]("_use_real_jwt_bearer")
[email protected](
+ ("token_sub", "token_scope", "expected_status"),
+ [
+ pytest.param("self", "execution", 204,
id="execution-token-matching-sub-accepted"),
+ pytest.param("other", "execution", 403,
id="execution-token-for-other-callback-rejected"),
+ pytest.param(
+ "00000000-0000-0000-0000-000000000000",
+ "execution",
+ 404,
+ id="execution-token-for-unknown-subject-rejected",
+ ),
+ pytest.param("self", "callback", 403,
id="callback-scope-token-rejected"),
+ ],
+)
+def test_state_validates_token_scope_and_sub(
+ client, running_callback, other_callback, token_sub, token_scope,
expected_status
+):
+ """
+ Reporting state needs the execution token minted by /run, not the
admission ticket.
+
+ A callback-scope token is refused here, which is the point of the swap:
the single-use
+ ticket buys one transition into RUNNING and nothing else.
+
+ The two rejected execution-token cases differ by design. A sub naming
another *existing*
+ callback clears the live-attempt admission check and is then refused 403
by ``cb:self``,
+ while a sub naming nothing at all is refused 404 by the admission check
itself, which runs
+ first because the parent ``authenticated_router`` declares
``require_auth`` with no scopes
+ and FastAPI resolves the outermost dependency before the ``cb:self`` one.
+ """
+ subs = {"self": str(running_callback.id), "other": str(other_callback.id)}
+ sub = subs.get(token_sub, token_sub)
+ validator = mock.AsyncMock(spec=JWTValidator)
+ validator.avalidated_claims.return_value = {
+ "sub": sub,
+ "scope": token_scope,
+ "exp": 9999999999,
+ "iat": 1000000000,
+ "nbf": 1000000000,
+ }
+ lifespan.registry.register_value(JWTValidator, validator)
+
+ resp = client.patch(f"/execution/callbacks/{running_callback.id}/state",
json={"state": "success"})
+ assert resp.status_code == expected_status, resp.json()
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_callbacks.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_callbacks.py
new file mode 100644
index 00000000000..4e85d24b3b4
--- /dev/null
+++
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2027_02_28/test_callbacks.py
@@ -0,0 +1,43 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import pytest
+
+pytestmark = pytest.mark.db_test
+
+MISSING_CALLBACK_URL =
"/execution/callbacks/00000000-0000-0000-0000-000000000000/state"
+
+
+class TestUpdateCallbackStateEndpointVersioning:
+ """The callbacks/{callback_id}/state endpoint didn't exist before the
2027-02-28 API version."""
+
+ def test_old_version_returns_404(self, client):
+ """Before 2027-02-28 the route is absent, so routing itself 404s (no
endpoint-shaped detail)."""
+ client.headers["Airflow-API-Version"] = "2026-10-30"
+
+ response = client.patch(MISSING_CALLBACK_URL, json={"state":
"success"})
+
+ assert response.status_code == 404
+ assert response.json() == {"detail": "Not Found"}
+
+ def test_head_version_routes_to_endpoint(self, client):
+ """At head the route exists: the same request reaches the endpoint's
own 404 handling."""
+ response = client.patch(MISSING_CALLBACK_URL, json={"state":
"success"})
+
+ assert response.status_code == 404
+ assert response.json()["detail"]["reason"] == "not_found"
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index b15a82ace30..4104c8740f1 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -86,7 +86,7 @@ from airflow.models.asset import (
PartitionedAssetKeyLog,
)
from airflow.models.backfill import Backfill, BackfillDagRun,
ReprocessBehavior, _create_backfill
-from airflow.models.callback import Callback, ExecutorCallback
+from airflow.models.callback import Callback, CallbackKey, ExecutorCallback
from airflow.models.connection_test import (
ConnectionTestKey,
ConnectionTestRequest,
@@ -223,6 +223,9 @@ DEFAULT_DATE = timezone.datetime(2016, 1, 1)
DEFAULT_LOGICAL_DATE = timezone.coerce_datetime(DEFAULT_DATE)
TRY_NUMBER = 1
+# Output already on a callback row before an executor event arrives for it.
+PRIOR_OUTPUT = "reported by the supervisor"
+
def get_callback_ti(ti):
return TIDataModel.model_validate(ti, from_attributes=True).model_copy(
@@ -1207,6 +1210,138 @@ class TestSchedulerJob:
assert session.get(ExecutorCallback, queued_callback.id).state ==
CallbackState.QUEUED
assert session.get(ExecutorCallback, running_callback.id).state ==
CallbackState.RUNNING
+ @pytest.mark.parametrize(
+ ("db_state", "event_state", "event_info", "expected_state",
"expected_output"),
+ [
+ pytest.param(
+ CallbackState.RUNNING,
+ CallbackState.SUCCESS,
+ None,
+ CallbackState.SUCCESS,
+ PRIOR_OUTPUT,
+ id="running-advances-on-success",
+ ),
+ pytest.param(
+ CallbackState.QUEUED,
+ CallbackState.FAILED,
+ "worker died",
+ CallbackState.FAILED,
+ "worker died",
+ id="queued-advances-when-the-supervisor-never-reported",
+ ),
+ pytest.param(
+ CallbackState.SUCCESS,
+ CallbackState.FAILED,
+ "worker died",
+ CallbackState.SUCCESS,
+ PRIOR_OUTPUT,
+ id="terminal-row-keeps-its-state-and-output",
+ ),
+ pytest.param(
+ CallbackState.SCHEDULED,
+ CallbackState.FAILED,
+ "worker died",
+ CallbackState.SCHEDULED,
+ PRIOR_OUTPUT,
+ id="never-enqueued-row-is-left-alone",
+ ),
+ ],
+ )
+ def test_process_executor_events_only_advances_callbacks_still_in_flight(
+ self, dag_maker, session, db_state, event_state, event_info,
expected_state, expected_output
+ ):
+ """
+ The executor event is a fallback behind PATCH /callbacks/{id}/state.
+
+ Both write terminal callback state, so the first writer wins: once a
row is terminal, a
+ later executor event must not replace the outcome the supervisor
reported, nor its output.
+ """
+
+ def test_callback():
+ pass
+
+ with dag_maker(dag_id="test_callback_forward_only"):
+ pass
+ dag_run = dag_maker.create_dagrun()
+
+ callback = Deadline(
+ deadline_time=timezone.utcnow(),
+ callback=SyncCallback(test_callback),
+ dagrun_id=dag_run.id,
+ deadline_alert_id=None,
+ ).callback
+ callback.state = db_state
+ callback.output = PRIOR_OUTPUT
+ callback.data["dag_run_id"] = dag_run.id
+ callback.data["dag_id"] = dag_run.dag_id
+ session.add(callback)
+ session.flush()
+
+ executor = MockExecutor(do_update=False)
+ self.job_runner = SchedulerJobRunner(Job(), executors=[executor])
+ executor.event_buffer[CallbackKey(str(callback.id))] = event_state,
event_info
+
+ self.job_runner._process_executor_events(executor=executor,
session=session)
+
+ # Flush before expiring: the scheduler must not commit, so an
unflushed change would be
+ # discarded by expire_all() and every case would read as "not
advanced".
+ session.flush()
+ session.expire_all()
+ stored = session.get(ExecutorCallback, callback.id)
+ assert stored.state == expected_state
+ assert stored.output == expected_output
+
+ def
test_process_executor_events_does_not_overwrite_a_concurrent_terminal_report(
+ self, dag_maker, session
+ ):
+ """
+ The supervisor can report through the API between the scheduler's read
and its write.
+
+ The in-flight check reads an ORM snapshot, so it cannot see that. Here
the row is moved to
+ SUCCESS behind the session's back, leaving the snapshot saying RUNNING
while the database
+ says otherwise -- the same divergence a committed `PATCH
/callbacks/{id}/state` produces,
+ and the reason the write itself has to carry the in-flight predicate.
+ """
+
+ def test_callback():
+ pass
+
+ with dag_maker(dag_id="test_callback_concurrent_report"):
+ pass
+ dag_run = dag_maker.create_dagrun()
+
+ callback = Deadline(
+ deadline_time=timezone.utcnow(),
+ callback=SyncCallback(test_callback),
+ dagrun_id=dag_run.id,
+ deadline_alert_id=None,
+ ).callback
+ callback.state = CallbackState.RUNNING
+ callback.data["dag_run_id"] = dag_run.id
+ callback.data["dag_id"] = dag_run.dag_id
+ session.add(callback)
+ session.flush()
+
+ session.execute(
+ update(Callback)
+ .where(Callback.id == callback.id)
+ .values(state=CallbackState.SUCCESS, output=PRIOR_OUTPUT)
+ .execution_options(synchronize_session=False)
+ )
+ assert callback.state == CallbackState.RUNNING, "snapshot must still
be stale"
+
+ executor = MockExecutor(do_update=False)
+ self.job_runner = SchedulerJobRunner(Job(), executors=[executor])
+ executor.event_buffer[CallbackKey(str(callback.id))] =
CallbackState.FAILED, "worker died"
+
+ self.job_runner._process_executor_events(executor=executor,
session=session)
+
+ session.flush()
+ session.expire_all()
+ stored = session.get(ExecutorCallback, callback.id)
+ assert stored.state == CallbackState.SUCCESS
+ assert stored.output == PRIOR_OUTPUT
+
@mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest")
@mock.patch("airflow._shared.observability.metrics.stats._get_backend")
def test_process_executor_events_with_callback(
diff --git a/task-sdk/src/airflow/sdk/api/client.py
b/task-sdk/src/airflow/sdk/api/client.py
index 0eb04302736..008ae64f8ee 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -50,6 +50,8 @@ from airflow.sdk.api.datamodels._generated import (
AssetResponse,
AssetStateStorePutBody,
AssetStateStoreResponse,
+ CallbackTerminalState,
+ CallbackTerminalStatePayload,
ConnectionResponse,
ConnectionTestConnectionResponse,
ConnectionTestResultBody,
@@ -1201,6 +1203,11 @@ class CallbackOperations:
"""Exchange the single-use callback token for an execution token."""
self.client.patch(f"callbacks/{callback_id}/run")
+ def finish(self, callback_id: uuid.UUID, state: CallbackTerminalState,
output: str | None = None) -> None:
+ """Report the terminal state of a callback to the API server."""
+ body = CallbackTerminalStatePayload(state=state, output=output)
+ self.client.patch(f"callbacks/{callback_id}/state",
content=body.model_dump_json())
+
class BearerAuth(httpx2.Auth):
def __init__(self, token: str):
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index 5351dd8ed37..a2d53f76159 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -64,6 +64,27 @@ class AssetProfile(BaseModel):
type: Annotated[str, Field(title="Type")]
+class CallbackTerminalState(str, Enum):
+ """
+ Terminal states a callback can transition to from RUNNING.
+ """
+
+ SUCCESS = "success"
+ FAILED = "failed"
+
+
+class CallbackTerminalStatePayload(BaseModel):
+ """
+ Payload for transitioning a callback from RUNNING to a terminal state.
+ """
+
+ model_config = ConfigDict(
+ extra="forbid",
+ )
+ state: CallbackTerminalState
+ output: Annotated[str | None, Field(title="Output")] = None
+
+
class ConnectionResponse(BaseModel):
"""
Connection schema for responses with fields that are needed for Runtime.
diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
index a2861de413c..842aea9bfd2 100644
--- a/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/callback_supervisor.py
@@ -33,6 +33,7 @@ import structlog
from pydantic import Field, TypeAdapter
from airflow.sdk._shared.module_loading import UNUSUAL_MODULE_PREFIX,
accepts_context, accepts_keyword_args
+from airflow.sdk.api.datamodels._generated import CallbackTerminalState
from airflow.sdk.exceptions import ErrorType
from airflow.sdk.execution_time.comms import (
ErrorResponse,
@@ -503,8 +504,14 @@ def supervise_callback(
logger = structlog.get_logger(logger_name="callback").bind()
# Swap the single-use callback token for an execution token before any
context read.
+ # A failure here leaves the row QUEUED; the scheduler's executor-event
fallback in
+ # _process_executor_events is what marks it failed, so finish() is
deliberately not
+ # reached on this path.
client.callbacks.run(UUID(id))
+ terminal_state = CallbackTerminalState.FAILED
+ terminal_output: str | None = None
+
try:
process = CallbackSubprocess.start(
id=id,
@@ -527,9 +534,22 @@ def supervise_callback(
exit_code=exit_code,
duration=end - start,
)
- if exit_code != 0:
- raise RuntimeError(f"Callback subprocess exited with code
{exit_code}")
- return exit_code
+ if exit_code == 0:
+ terminal_state = CallbackTerminalState.SUCCESS
+ else:
+ terminal_output = f"Callback subprocess exited with code
{exit_code}"
+ except Exception as e:
+ terminal_output = f"Callback supervisor error: {type(e).__name__}:
{e}"
+ raise
finally:
+ try:
+ client.callbacks.finish(UUID(id), state=terminal_state,
output=terminal_output)
+ except Exception:
+ log.exception("Failed to report final callback state",
workload_id=id)
+
if log_path and log_file_descriptor:
log_file_descriptor.close()
+
+ if exit_code != 0:
+ raise RuntimeError(f"Callback subprocess exited with code {exit_code}")
+ return exit_code
diff --git a/task-sdk/tests/task_sdk/api/test_client.py
b/task-sdk/tests/task_sdk/api/test_client.py
index ca0c2793a52..ab1441b0c7d 100644
--- a/task-sdk/tests/task_sdk/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -38,6 +38,7 @@ from airflow.sdk.api.datamodels._generated import (
AssetEventsResponse,
AssetResponse,
AssetStateStoreResponse,
+ CallbackTerminalState,
ConnectionResponse,
DagResponse,
DagRunState,
@@ -2637,3 +2638,36 @@ class TestCallbackOperations:
with pytest.raises(ServerResponseError) as err:
client.callbacks.run(callback_id)
assert err.value.response.status_code == 409
+
+ @pytest.mark.parametrize(
+ ("state", "output"),
+ [
+ pytest.param(CallbackTerminalState.SUCCESS, None,
id="success-without-output"),
+ pytest.param(CallbackTerminalState.FAILED, "exited with code 3",
id="failed-with-output"),
+ ],
+ )
+ def test_finish_reports_the_terminal_state(self, state, output):
+ callback_id = uuid7()
+
+ def handle_request(request: httpx2.Request) -> httpx2.Response:
+ assert request.method == "PATCH"
+ assert request.url.path == f"/callbacks/{callback_id}/state"
+ assert json.loads(request.content) == {"state": state.value,
"output": output}
+ return httpx2.Response(status_code=204)
+
+ client = make_client(transport=httpx2.MockTransport(handle_request))
+ client.callbacks.finish(callback_id, state=state, output=output)
+
+ def test_finish_raises_when_the_callback_is_not_running(self):
+ callback_id = uuid7()
+
+ def handle_request(request: httpx2.Request) -> httpx2.Response:
+ return httpx2.Response(
+ status_code=409,
+ json={"detail": {"reason": "invalid_state", "previous_state":
"queued"}},
+ )
+
+ client = make_client(transport=httpx2.MockTransport(handle_request))
+ with pytest.raises(ServerResponseError) as err:
+ client.callbacks.finish(callback_id,
state=CallbackTerminalState.FAILED)
+ assert err.value.response.status_code == 409
diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
index 36c0cf52ee8..564d9cd1b8d 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_callback_supervisor.py
@@ -31,7 +31,8 @@ from unittest.mock import ANY, Mock, patch
import pytest
import structlog
-from airflow.sdk.api.datamodels._generated import XComResponse
+from airflow.sdk.api.client import CallbackOperations, Client
+from airflow.sdk.api.datamodels._generated import CallbackTerminalState,
XComResponse
from airflow.sdk.execution_time.callback_supervisor import (
CallbackSubprocess,
Path,
@@ -731,6 +732,13 @@ class TestCallbackSubprocessStart:
assert calls == ["import task runner", "fork"]
[email protected]
+def callback_client():
+ client = Mock(spec=Client)
+ client.callbacks = Mock(spec=CallbackOperations)
+ return client
+
+
class TestSuperviseCallbackExchangesTokenFirst:
"""No callback code may run until the single-use token has been
exchanged."""
@@ -764,14 +772,19 @@ class TestSuperviseCallbackExchangesTokenFirst:
@patch("airflow.sdk.execution_time.callback_supervisor._make_process_nondumpable")
@patch.object(CallbackSubprocess, "start")
- def test_rejected_exchange_prevents_the_subprocess_from_starting(self,
mock_start, _nondumpable):
- client = Mock()
+ def test_rejected_exchange_prevents_the_subprocess_from_starting(
+ self, mock_start, _nondumpable, callback_client
+ ):
+ client = callback_client
client.callbacks.run.side_effect = RuntimeError("409 Conflict: token
already exchanged")
with pytest.raises(RuntimeError, match="already exchanged"):
self._supervise(client)
mock_start.assert_not_called()
+ # The row is still QUEUED, which the state endpoint rejects; the
scheduler's
+ # executor-event fallback is what fails it.
+ client.callbacks.finish.assert_not_called()
@pytest.mark.parametrize("callback_kwargs_encoding", [None, "serde"])
@patch("airflow.sdk.execution_time.callback_supervisor._make_process_nondumpable")
@@ -791,3 +804,72 @@ class TestSuperviseCallbackExchangesTokenFirst:
)
assert mock_start.call_args.kwargs["callback_kwargs_encoding"] ==
callback_kwargs_encoding
+
+
+class TestSuperviseCallbackReportsTerminalState:
+ """Every path that starts the subprocess must report a terminal state to
the API."""
+
+ CALLBACK_ID = "01890a5d-ac70-7a5b-b7d5-0dd5b1c7be47"
+
+ def _supervise(self, client):
+ return supervise_callback(
+ id=self.CALLBACK_ID,
+ callback_path="does.not.matter",
+ callback_kwargs={},
+ dag_rel_path=Path("dag.py"),
+ client=client,
+ )
+
+
@patch("airflow.sdk.execution_time.callback_supervisor._make_process_nondumpable")
+ @patch.object(CallbackSubprocess, "start")
+ def test_clean_exit_reports_success(self, mock_start, _nondumpable,
callback_client):
+ client = callback_client
+ mock_start.return_value = Mock(wait=Mock(return_value=0))
+
+ assert self._supervise(client) == 0
+
+ client.callbacks.finish.assert_called_once_with(
+ uuid.UUID(self.CALLBACK_ID),
+ state=CallbackTerminalState.SUCCESS,
+ output=None,
+ )
+
+
@patch("airflow.sdk.execution_time.callback_supervisor._make_process_nondumpable")
+ @patch.object(CallbackSubprocess, "start")
+ def test_nonzero_exit_reports_failed_before_raising(self, mock_start,
_nondumpable, callback_client):
+ client = callback_client
+ mock_start.return_value = Mock(wait=Mock(return_value=3))
+
+ with pytest.raises(RuntimeError, match="exited with code 3"):
+ self._supervise(client)
+
+ client.callbacks.finish.assert_called_once_with(
+ uuid.UUID(self.CALLBACK_ID),
+ state=CallbackTerminalState.FAILED,
+ output="Callback subprocess exited with code 3",
+ )
+
+
@patch("airflow.sdk.execution_time.callback_supervisor._make_process_nondumpable")
+ @patch.object(CallbackSubprocess, "start")
+ def test_supervisor_error_reports_failed_with_the_cause(self, mock_start,
_nondumpable, callback_client):
+ client = callback_client
+ mock_start.side_effect = ValueError("boom")
+
+ with pytest.raises(ValueError, match="boom"):
+ self._supervise(client)
+
+ client.callbacks.finish.assert_called_once_with(
+ uuid.UUID(self.CALLBACK_ID),
+ state=CallbackTerminalState.FAILED,
+ output="Callback supervisor error: ValueError: boom",
+ )
+
+
@patch("airflow.sdk.execution_time.callback_supervisor._make_process_nondumpable")
+ @patch.object(CallbackSubprocess, "start")
+ def test_a_failure_to_report_does_not_mask_the_outcome(self, mock_start,
_nondumpable, callback_client):
+ """The scheduler fallback still has the executor event, so a reporting
blip is not fatal."""
+ client = callback_client
+ client.callbacks.finish.side_effect = RuntimeError("API unreachable")
+ mock_start.return_value = Mock(wait=Mock(return_value=0))
+
+ assert self._supervise(client) == 0
