Priyaj11 opened a new pull request, #72500:
URL: https://github.com/apache/airflow/pull/72500
What happened
EmrContainerSensor.poke() treats "not a failure state and not an
intermediate state" as success:
python
INTERMEDIATE_STATES = ("PENDING", "SUBMITTED", "RUNNING")
FAILURE_STATES = ("FAILED", "CANCELLED", "CANCEL_PENDING")
SUCCESS_STATES = ("COMPLETED",) # declared, never used
def poke(self, context: Context) -> bool:
state = self.hook.poll_query_status(...)
if state in self.FAILURE_STATES:
raise AirflowException(...)
if state in self.INTERMEDIATE_STATES:
return False
return True # anything else is "success"
SUCCESS_STATES is defined on the class and never referenced anywhere in the
file.
Why it matters
The fall-through is reachable through the hook's own documented behaviour:
EmrContainerHook.check_query_status catches a generic ClientError, logs it
and returns None. Its docstring says "Returns None or one of valid query
states."
EmrContainerHook.poll_query_status handles that case explicitly, logging
"Try %s: Invalid query state. Retrying again", and once max_polling_attempts is
reached it does return query_state, which is None.
poke() finds None in neither tuple and returns True.
So a `ThrottlingException`, or any transient `DescribeJobRun` failure that
outlasts `max_retries`, marks the Airflow task successful and releases every
downstream task while the EMR on EKS job is still running. The same applies to
any job-run state AWS adds that Airflow does not yet know about.
Worth noting the precondition: reaching the `None` return requires
`max_retries` to be set on the sensor, since that is what bounds the hook's
polling loop. Left unset, the loop spins indefinitely instead of returning.
The change
diff
if state in self.INTERMEDIATE_STATES:
return False
- return True
+ return state in self.SUCCESS_STATES
Why the existing tests did not catch this
test_emr_containers.py covers all seven documented EMR on EKS states:
PENDING, SUBMITTED, RUNNING, COMPLETED, FAILED, CANCELLED, CANCEL_PENDING.
Every one behaves correctly today. Nothing covered the None that the hook is
documented to return, or an unrecognised state.
This PR adds test_poke_unknown_state and test_poke_unrecognised_state to
close that gap.
Test results
Before the fix:
test_poke_unknown_state FAILED
test_poke_unrecognised_state FAILED
2 failed, 9 passed
with the hook logging Try 1: Invalid query state. Retrying again and Try 1:
Query is still in non-terminal state - SOME_FUTURE_STATE respectively,
immediately before poke() returned True.
After the fix:
11 passed
Behaviour change worth calling out
An unknown state now makes the sensor keep poking until its own timeout
rather than succeed immediately. That is the minimal, fail-safe change, but the
alternative would be to raise on a state that is in none of the three tuples so
the failure is loud instead of a timeout. Happy to change it to whichever the
maintainers prefer.
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes
Generated-by: Claude following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
### Gen-AI disclosure
I used Claude to assist with this contribution. The defect was found by
reviewing provider modules that carry no unit tests; the fix and the two added
tests were drafted with its help. I verified the call chain through
`EmrContainerHook.check_query_status` and `poll_query_status` myself, and
ran the test file locally both with and without the fix (2 failed / 9 passed
before, 11 passed after). I understand the change and take responsibility for
it.
--
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]
