kaxil commented on code in PR #76021:
URL: https://github.com/apache/airflow/pull/76021#discussion_r4467490253
##########
task-sdk/tests/task_sdk/execution_time/test_task_runner.py:
##########
@@ -7422,39 +7423,135 @@ def execute(self, context):
]
assert ClearTaskStateStore not in sent_types
+ @staticmethod
+ @contextlib.contextmanager
+ def _patch_state_backend(backend):
+ # Both modules resolve the backend independently, so both lookups need
patching.
+ with (
+ mock.patch(
+
"airflow.sdk.execution_time.context._get_worker_state_store_backend",
+ return_value=backend,
+ ),
+ mock.patch(
+
"airflow.sdk.execution_time.task_runner._get_worker_state_store_backend",
+ return_value=backend,
+ ),
+ ):
+ yield backend
+
+ @staticmethod
+ def _run_and_finalize(runtime_ti, log):
+ # The clear lives at the tail of finalize(), so run() alone would not
reach it.
+ context = runtime_ti.get_template_context()
+ state, msg, error = run(runtime_ti, context=context, log=log)
+ finalize(runtime_ti, state, context, log, error)
+ return state, msg, error
+
@conf_vars({("state_store", "clear_on_success"): "True"})
def test_clear_on_success_clears_backend_without_comms_roundtrip(
self, create_runtime_ti, mock_supervisor_comms
):
"""clear_on_success calls backend.clear() directly without sending
ClearTaskStateStore comms."""
- mock_backend = mock.MagicMock()
+ backend = mock.create_autospec(BaseStoreBackend, instance=True)
class MyOperator(BaseOperator):
def execute(self, context):
pass
- task = MyOperator(task_id="t")
- runtime_ti = create_runtime_ti(task=task)
+ 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,
- ),
- ):
- run(runtime_ti, context=runtime_ti.get_template_context(),
log=mock.MagicMock())
+ with self._patch_state_backend(backend):
+ self._run_and_finalize(runtime_ti, mock.MagicMock())
- mock_backend.clear.assert_called_once()
+ backend.clear.assert_called_once()
sent_types = [
type(call.kwargs.get("msg") or (call.args[0] if call.args else
None))
for call in mock_supervisor_comms.send.call_args_list
]
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."""
+ backend = mock.create_autospec(BaseStoreBackend, instance=True)
+ order = []
+ mock_supervisor_comms.send.side_effect = lambda msg, **kwargs:
order.append(type(msg).__name__)
+ 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 self._patch_state_backend(backend):
+ state, _, _ = self._run_and_finalize(runtime_ti, mock.MagicMock())
+
+ assert state == TaskInstanceState.SUCCESS
+ assert order.index("SucceedTask") < order.index("clear")
Review Comment:
This still passes with the clear back in `run()` right after the send, since
nothing runs between `SucceedTask` and `clear` here. An `on_success_callback`
that appends to `order`, with `SucceedTask < callback < clear`, would pin the
move to the end of `finalize()`.
##########
airflow-core/docs/core-concepts/task-state-store.rst:
##########
@@ -316,4 +316,18 @@ 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, once the task's callbacks and
listeners have run. If that
+ deletion fails the task stays successful and the failure is logged with the
Dag id, run id and
+ task id.
+
+ Those entries are then orphaned: the API server has already removed the
metadata-DB references,
+ ``airflow state-store clean`` covers the metastore backend only, and
+ :meth:`~airflow.sdk.state.BaseStoreBackend.cleanup` is a no-op unless your
backend implements
+ it. Because the storage scope includes the run id, the next Dag run writes
elsewhere rather than
+ overwriting them. Removing them is up to your storage's own lifecycle or
TTL policy, or a
+ ``cleanup()`` implementation on your backend.
Review Comment:
Nothing in Airflow calls `cleanup()` for a worker backend though, since
`airflow state-store clean` skips it. Could this say it has to be called from
the user's own maintenance process, like the state-store cleanup page does?
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1835,12 +1838,49 @@ 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:
+ ti._state_store_clear_owed = True
# 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.
+
+ Called at the end of finalize(), never before the terminal state is
reported: the
+ backend is remote storage, so a failure deleting from it must not turn a
task that
+ succeeded into a failed one, and a slow delete must not spend the
task_success_overtime
+ budget that the success callbacks and listeners need.
+ """
+ if not conf.getboolean("state_store", "clear_on_success"):
+ return
+ try:
+ # Inside the try: the backend is built on first use, so a bad [workers]
+ # state_store_backend class path raises here for a task that never
touched
+ # task_state_store, and @cache does not cache the raise.
+ if _get_worker_state_store_backend() is None:
+ return
+ log.info(
+ "Clearing task state from custom backend as clear_on_success is
enabled. The database references will be cleared by the API server."
+ )
+ context["task_state_store"]._clear_backend_only()
+ except Exception:
+ ti = context["ti"]
+ log.exception(
+ "Failed to clear task state from custom backend; the task stays
successful. Nothing in Airflow "
+ "reclaims these entries -- the database references are already
gone and the next run uses a "
+ "different run_id -- so only the storage's own lifecycle or TTL
policy will remove them.",
+ dag_id=ti.dag_id,
+ run_id=ti.run_id,
+ task_id=ti.task_id,
+ map_index=ti.map_index,
Review Comment:
The backend key comes from the `TaskScope` in `get_template_context`, which
uses `region_index` and `region_id` rather than `ti.map_index`, so for a task
inside a region this can point at the wrong entries. Could the scope go on the
`Clearing task state` info line instead? That line is also the only record when
the delete hangs past `task_success_overtime`, since `_on_term` only calls
`on_kill()` and this `log.exception` never runs.
##########
task-sdk/src/airflow/sdk/execution_time/task_runner.py:
##########
@@ -1835,12 +1838,49 @@ 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:
+ ti._state_store_clear_owed = True
# 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.
+
+ Called at the end of finalize(), never before the terminal state is
reported: the
+ backend is remote storage, so a failure deleting from it must not turn a
task that
+ succeeded into a failed one, and a slow delete must not spend the
task_success_overtime
+ budget that the success callbacks and listeners need.
+ """
+ if not conf.getboolean("state_store", "clear_on_success"):
+ return
+ try:
+ # Inside the try: the backend is built on first use, so a bad [workers]
+ # state_store_backend class path raises here for a task that never
touched
+ # task_state_store, and @cache does not cache the raise.
+ if _get_worker_state_store_backend() is None:
+ return
+ log.info(
+ "Clearing task state from custom backend as clear_on_success is
enabled. The database references will be cleared by the API server."
+ )
+ context["task_state_store"]._clear_backend_only()
Review Comment:
Now that this runs after `on_success_callback` and the success listeners,
anything they write through `task_state_store` gets a fresh DB ref (the set
route only checks `working_set`), and this then deletes the object behind it
and leaves the ref. Could this call `task_state_store.clear()` instead, so the
refs go first? It's one extra message, and the `_clear_backend_only` docstring
about the round-trip being redundant no longer holds anyway.
--
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]
