Move UPSTREAM_FAILED to Terminal State and Add State Validation Tests

[seanghaeli]

PR for #52999:

Move UPSTREAM_FAILED task instance state from intermediate to terminal state class and add tests to validate that state class relationships remain consistent (failed + success = terminal, terminal + intermediate = all task states).

Thoughts on this change? Motivated from #51719, @eladkal's comment:

there is no protection here against possible future addition of new state to task instance. For example we are discussing #12199

I suggest to add defensive test around adding more states so we'll know to modify code here or maybe we can consider adding more classes to categorized states similar to

airflow/task-sdk/src/airflow/sdk/api/datamodels/_generated.py

Lines 415 to 419 in 083e03a

	class TerminalTIState(str, Enum):
	SUCCESS = "success"
	FAILED = "failed"
	SKIPPED = "skipped"
	REMOVED = "removed"

Wanted to make this a separate PR for it since its scope it quite different.

Most likely this code will break some integration tests because there's various places that depend on the task instance states in interesting ways, but wanted to field feedback first before fixing them.

[seanghaeli]

@o-nikolas if you don't mind, let me know what you think

[o-nikolas]

@o-nikolas if you don't mind, let me know what you think

I like this change. All the sets that we have for specific terminal states (I think it's just success and failure states at the moment) should definitely sum up to the full set of terminal states. This makes the set math you're trying to do in your original PR much cleaner, future proof and not hardcoded to specific sets which is all great stuff.

Only concern I'd have is what you mention, the changes propagating in ways we don't expect, but I don't think this is insurmountable and I'd be interested to see hear how many tests failed because of this change.

[seanghaeli]

@o-nikolas Tests passing, ready for review.

[eladkal]

It feels strange that none of our current tests needed modification due to this PR.

Like Niko I do have concerns that this change might cause unexpected behaviour. I see that the TerminalStateNonSuccess class is referenced in:

airflow/task-sdk/src/airflow/sdk/api/client.py

Lines 151 to 159 in ac9e968

	def finish(self, id: uuid.UUID, state: TerminalStateNonSuccess, when: datetime, rendered_map_index):
	"""Tell the API server that this TI has reached a terminal state."""
	if state == TaskInstanceState.SUCCESS:
	raise ValueError("Logic error. SUCCESS state should call the `succeed` function instead")
	# TODO: handle the naming better. finish sounds wrong as "even" deferred is essentially finishing.
	body = TITerminalStatePayload(
	end_date=when, state=TerminalStateNonSuccess(state), rendered_map_index=rendered_map_index
	)
	self.client.patch(f"task-instances/{id}/state", content=body.model_dump_json())

This means that after this change UPSTREAM_FAILED will be processed differently. Do we know the implication of this?
If I read this right it may implicate on the end date of the task which can affect metrics collection. This may be a welcomed change but we need to be very sure that the all behavioral changes are expected.

This is one of the PRs that we should get several approval just to make sure we are making the right change. I would also ask to include newsfragment for this change.

[eladkal]

Setting temporary block to avoid accidental merge
Any maintainer may lift this block if

This is one of the PRs that we should get several approvals just to make sure we are making the right change.

If satisfied or no longer relevant.

[seanghaeli]

@kaxil Since you created the IntermediateTIState and TerminalTIState classes, what are your thoughts about moving the Task Instance State UPSTREAM_FAILED from IntermediateTIState to TerminalTIState?

[kaxil]

UPSTREAM_FAILED change looks good but adding of removed to success_states does not

[kaxil]

Looks good but need to address 2 comments

Stale

[seanghaeli]

UPSTREAM_FAILED change looks good but adding of removed to success_states does not

In my opinion, REMOVED semantically makes sense as a failure state. It represents work that we expect to be completed at the time the DAG is initially called, but then is no longer available when it comes time to execute the task. Would you agree @kaxil ?