kaxil commented on code in PR #76021:
URL: https://github.com/apache/airflow/pull/76021#discussion_r4437471883
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1835,12 +1835,38 @@ def run(
# a retry policy turned the exception into a reschedule.
if error is not None or state in (TaskInstanceState.FAILED,
TaskInstanceState.UP_FOR_RETRY):
ti._terminal_state_send_failed = True
+ else:
+ if state == TaskInstanceState.SUCCESS:
+ _clear_task_state_store_on_success(context, log)
# Return the message to make unit tests easier too
ti.state = state
return state, msg, error
+def _clear_task_state_store_on_success(context: Context, log: Logger) -> None:
+ """
+ Clear the custom state store backend once the task is recorded as
successful.
+
+ Runs after the terminal state is reported, never before: the backend is
remote storage
+ and a failure deleting from it must not turn a task that succeeded into a
failed one.
+ Leftover entries are harmless -- retention and the next ``set()`` clean
them up.
+ """
+ if not conf.getboolean("state_store", "clear_on_success"):
+ return
+ if _get_worker_state_store_backend() is None:
Review Comment:
This lookup sits outside the `try`. `_get_worker_state_store_backend()`
builds the backend on first use and raises `ValueError` on a bad class path
(constructor errors pass through, and `@cache` doesn't cache a raise). For a
task that never touched `task_state_store`, this is the first call, so the
exception leaves `run()` from inside its `finally` after success was reported,
and `main()` exits 1 through "Top level error" without calling `finalize()`.
`_clear_backend_only()` already returns when the backend is `None`, so dropping
this check and keeping only the `clear_on_success` early return would put the
whole thing under the `try`.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1835,12 +1835,38 @@ def run(
# a retry policy turned the exception into a reschedule.
if error is not None or state in (TaskInstanceState.FAILED,
TaskInstanceState.UP_FOR_RETRY):
ti._terminal_state_send_failed = True
+ else:
+ if state == TaskInstanceState.SUCCESS:
+ _clear_task_state_store_on_success(context, log)
Review Comment:
This now runs after `SucceedTask` reaches the supervisor, and
`_handle_finished_task` starts the `task_success_overtime` clock (default 20s)
at that moment. So a remote delete that hangs rather than failing fast, or a
clear over many keys (the common.io backend globs then unlinks one by one),
gets the process SIGTERMed before `main()` reaches `finalize()`. Then
`on_success_callback` and the success listeners never run, and `on_kill()`
fires on a task that succeeded. Could the clear move into `main()` after
`finalize()`, only when the success was delivered? Callbacks see the same state
either way, since the server drops the DB refs on `SucceedTask`.
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -7455,6 +7455,66 @@ def execute(self, context):
]
assert ClearTaskStateStore not in sent_types
+ @conf_vars({("state_store", "clear_on_success"): "True"})
+ def test_clear_on_success_runs_after_success_is_reported(self,
create_runtime_ti, mock_supervisor_comms):
+ """The backend is only cleared once SucceedTask has been reported to
the supervisor."""
+ mock_backend = mock.MagicMock()
+ order = []
+ mock_supervisor_comms.send.side_effect = lambda msg, **kwargs:
order.append(type(msg).__name__)
+ mock_backend.clear.side_effect = lambda *a, **kw: order.append("clear")
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ pass
+
+ runtime_ti = create_runtime_ti(task=MyOperator(task_id="t"))
+
+ with (
+ mock.patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ mock.patch(
+
"airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ ):
+ state, _, _ = run(runtime_ti,
context=runtime_ti.get_template_context(), log=mock.MagicMock())
+
+ assert state == TaskInstanceState.SUCCESS
+ assert order.index("SucceedTask") < order.index("clear")
+
+ @conf_vars({("state_store", "clear_on_success"): "True"})
+ def test_clear_on_success_failure_does_not_fail_the_task(self,
create_runtime_ti, mock_supervisor_comms):
+ """A backend that cannot be cleared leaves the task successful instead
of failing it."""
+ mock_backend = mock.MagicMock()
+ mock_backend.clear.side_effect = ConnectionError("object storage is
down")
+
+ class MyOperator(BaseOperator):
+ def execute(self, context):
+ pass
+
+ runtime_ti = create_runtime_ti(task=MyOperator(task_id="t"))
+
+ with (
+ mock.patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ mock.patch(
+
"airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
+ return_value=mock_backend,
+ ),
+ ):
+ state, msg, error = run(
+ runtime_ti, context=runtime_ti.get_template_context(),
log=mock.MagicMock()
+ )
+
+ assert state == TaskInstanceState.SUCCESS
+ assert error is None
+ assert isinstance(msg, SucceedTask)
+ assert runtime_ti._terminal_state_send_failed is False
Review Comment:
This test would also pass if the clear was never attempted. Nothing asserts
`clear` was called, and `_terminal_state_send_failed is False` holds on any
code since send never raises here. `mock_backend.clear.assert_called_once()`
plus a check on `log.exception` would pin it. Nothing covers the new `state ==
SUCCESS` guard either: a case where `execute` raises, and one where the
`SucceedTask` send raises, each asserting `clear` isn't called, would catch a
FAILED/UP_FOR_RETRY task wiping the checkpoint its retry needs.
`mock.create_autospec(BaseStoreBackend, instance=True)` would also be better
than the bare `MagicMock()`.
##########
airflow-core/docs/core-concepts/task-state-store.rst:
##########
@@ -316,4 +316,10 @@ When ``[state_store] clear_on_success = True``, all task
state store keys for a
``clear_on_success`` clears **task state store only**. Asset store is
scoped to the asset, not the task instance, and is never affected by this
setting. Asset store persists across runs and must be cleared explicitly.
+.. note::
+
+ With a custom ``[workers] state_store_backend``, the worker deletes the
backend's entries only
+ after the success has been reported. If that deletion fails the task stays
successful and the
+ failure is logged; the leftover entries are removed by retention or
overwritten by the next run.
Review Comment:
I don't think anything removes these. `airflow state-store clean` skips any
backend that isn't `MetastoreBackend`, `BaseStoreBackend.cleanup()` is a no-op
that the common.io backend doesn't override, and the server already deleted the
DB refs on `SucceedTask`, so nothing in Airflow points at the objects any more.
The scope includes `run_id`, so the next Dag run writes elsewhere too. Could
this say the entries are only removed by the storage's own lifecycle/TTL
policy? The docstring at task_runner.py:1853 and the `log.exception` text at
1866 make the same promise. Logging the dag_id/run_id/task_id there would help
people find the orphans.
--
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]
