amoghrajesh commented on code in PR #76021:
URL: https://github.com/apache/airflow/pull/76021#discussion_r4492983845
##########
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:
Moved the scope onto the "Clearing task state" info line, taken from the
accessor's `TaskScope`, and logged before the delete so it survives a hang past
`task_success_overtime`. The exception line now carries only the message.
--
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]
