egotist0 commented on issue #60858: URL: https://github.com/apache/airflow/issues/60858#issuecomment-6054204991
> [@potiuk](https://github.com/potiuk), I'd like to propose a narrower approach than a new callback-runner process, and get your thoughts before writing a PR. > > **OOM often does not kill the whole pod.** On cgroup v2 with `memory.oom.group=0` (the Kubernetes default), the kernel kills only the largest process in the container. That is the task process, not the supervisor. We see this in production on 3.3 with the KubernetesExecutor: > > * The supervisor logs `Process terminated by signal. Likely out of memory error (OOM). signal=-9`. > * It reports the TI as `failed` about 30 ms later. > * The pod exits `Succeeded`. > > So the case where the supervisor survives is the common one. It is also the only case that is broken. > > **The whole-pod case already works.** If both processes die, the TI stays `running`. The scheduler then catches the mismatch, either from the executor event or the heartbeat timeout, and sends a `TaskCallbackRequest` to the DAG processor. The DAG processor already parses the DAG and runs `on_failure_callback`. The callback goes missing only because the surviving supervisor reports `failed` itself. After that the scheduler sees no mismatch, so no callback request is ever created. In the same case no `on_task_instance_failed` listener fires either, so OpenLineage gets no task-level FAIL event for that try. > > **Proposal:** route the case where the supervisor survives into the callback path that already exists, instead of having the supervisor run the callback. > > 1. When the child exits without reporting a terminal state and the supervisor sends the final state itself (`update_task_state_if_needed` in `supervisor.py`), it marks the request as a supervisor-originated failure. For example, an optional field on the terminal-state payload, set only on that path. > 2. When `ti_update_state` receives that marker, it does two things in the same transaction: > > * Enqueues a `TaskCallbackRequest` (`task_callback_type=FAILED`) via `DbCallbackRequest`. That is the same sink the scheduler uses, and the DAG processor runs the callback with the full TI context. > * Fires `on_task_instance_failed`, the same way `_emit_state_listener_hooks` already does for manual state changes in `core_api/services/public/task_instances.py`. > > The worker stays DB-less. The supervisor only adds one field to the HTTP call it already makes. The DB writes and the listener call both happen in the API server, which already writes the TI state in that same handler. > > This needs no new process, no DAG parsing in the supervisor, and nothing specific to an executor or deployment. Whether the task was killed by OOM, SIGKILL or SEGV, it takes the same path. The "OOM during DAG parsing" edge case is handled the way it already is for the whole-pod case, because the callback runs in the DAG processor rather than the worker. > > **A related gap with the same root cause:** when retries remain, the supervisor classifies the exit as `up_for_retry`. That state is in `STATES_SENT_DIRECTLY`, so it sends nothing, and the TI sits in `running` until the 300s heartbeat timeout. The same marker could let the supervisor send `up_for_retry` and enqueue an `UP_FOR_RETRY` callback right away. That would remove the 5-minute delay. I'd keep it as a separate follow-up unless you'd rather address both together. > > Costs I'm aware of: > > * The new payload field is an execution API change, so it needs a new versioned API change. Older Task SDK clients simply don't send the field and keep today's behaviour. > * Listeners would fire in the API server process for this case, the same as for manual state changes today. OpenLineage would emit the task FAIL event from the API server, not the worker. > > Open questions: > > * Is it acceptable for the execution API to enqueue callback requests, or should this go through the scheduler instead? One alternative: write a marker on the TI that the scheduler picks up in `_process_executor_events`. > * Should the failure type be visible to callbacks, as you suggested earlier ("extend on_failure_callback and add type of failure")? The marker would make that easy to expose in the context. > > If this direction looks reasonable, I'm happy to open a PR. @hkc-8010 thanks for writing this up. The production case above is ours (Airflow 3.3.2, Python 3.12, `KubernetesExecutor`), and +1 on this direction over a separate callback-runner process, it can fix both the missing callback and the missing listener @potiuk on the concern that this has to be executor or deployment specific: the broken case isn't tied to OOM or to `memory.oom.group`. It's any try where the task process exits non-zero without sending a terminal state while the supervisor is still alive in the worker pod. The supervisor then infers the state from the exit code in [update_task_state_if_needed](https://github.com/apache/airflow/blob/08205150190b3819a411b24f564328e590ab18ac/task-sdk/src/airflow/sdk/execution_time/supervisor.py#L1764) and calls `finish(FAILED)` itself when no retries remain. Besides OOM, the same path is hit by a SIGSEGV in native code, a SIGKILL from any other source, and (on main) execution_timeout when the task is stuck in native code and the supervisor has to SIGKILL it. In all of these the supervisor sees the task process exit non-zero without a terminal state, and this handling is the same on every executor. If the whole pod is killed, the TI stays `running`. The scheduler notices this from the executor event or the heartbeat timeout and asks the dag processor to run `on_failure_callback` via a `TaskCallbackRequest`, , and `handle_failure` fires `on_task_instance_failed`. The callback and listener are only lost when the supervisor survives and writes failed itself, so the scheduler never notices anything wrong. That write is where the marker goes. A few things I'd want in the PR: - [_fail_task_before_runtime](https://github.com/apache/airflow/blob/08205150190b3819a411b24f564328e590ab18ac/task-sdk/src/airflow/sdk/execution_time/supervisor.py#L2967) also writes failed when the task runtime fails to start (`TaskLaunchError`), and no callback runs there either. It could set the same marker. - Fire `on_task_instance_failed` after the commit, so OpenLineage never gets a FAIL event for a state update that rolled back. On the retry gap, I'd prefer fixing it in the same PR. We hit it too: the TI stayed `running` until the heartbeat timeout. [_fail_task_before_runtime](https://github.com/apache/airflow/blob/08205150190b3819a411b24f564328e590ab18ac/task-sdk/src/airflow/sdk/execution_time/supervisor.py#L2967) already calls `retry()` when retries remain, so when the supervisor infers `up_for_retry` it could do the same with the same marker. Otherwise fixing the delay would also drop the `on_retry_callback` and listener that the heartbeat-timeout path runs today. -- 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]
