amoghrajesh commented on code in PR #76021:
URL: https://github.com/apache/airflow/pull/76021#discussion_r4453042073
##########
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:
Moved to the tail of `finalize()`: `run()` now only sets
`ti._state_store_clear_owed`, and the delete happens after every callback and
listener.
Used `finalize()` rather than `main()` because the in-process path
(`supervisor.py:2668`, used by `dag.test()`) calls `run()` and `finalize()`
directly and never reaches `main()`.
The overtime budget is still shared, but the callbacks now go first.
##########
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:
Moved inside the `try`. Kept the `is None` check there rather than dropping
it, otherwise the
`log.info` above fires on every successful task when no custom backend is
configured.
After the change in the thread above this raises at the end of `finalize()`,
not from `run()`,
so `finalize()` is no longer skipped.
--
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]
