amoghrajesh commented on code in PR #76021:
URL: https://github.com/apache/airflow/pull/76021#discussion_r4492965865
##########
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:
Switched to `task_state_store.clear()`, so the DB refs are cleared before
the backend objects. Removed `_clear_backend_only`, which had no other callers.
The test now asserts ClearTaskStateStore is sent before backend.clear.
--
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]
