Fix task retry logic to respect retries for all exit codes (#58384)

Author: amoghrajesh

task-sdk/src/airflow/sdk/execution_time/supervisor.py

@@ -1171,10 +1171,12 @@ def final_state(self):
         if self._exit_code != 0 and self._terminal_state == SERVER_TERMINATED:
             return SERVER_TERMINATED
 
-        # Any negative exit code indicates a signal kill
-        # We consider all signal kills as potentially retryable
-        # since they're often transient issues that could succeed on retry
-        if self._exit_code < 0 and self._should_retry:
+        # Any non zero exit code indicates a failure
+        # If retries are configured, mark as UP_FOR_RETRY
+        # Negative exit codes indicate signal kills (often transient)
+        # Positive exit codes can also be transient failures like network issues in a task communicating to
+        # external services
+        if self._exit_code != 0 and self._should_retry:
             return TaskInstanceState.UP_FOR_RETRY
 
         return TaskInstanceState.FAILED

task-sdk/tests/task_sdk/execution_time/test_supervisor.py

@@ -2695,7 +2695,7 @@ def mock_upload_to_remote(process_log, ti):
 
 
 class TestSignalRetryLogic:
-    """Test signal based retry logic in ActivitySubprocess."""
+    """Test retry logic for exit codes (signals and non-signal failures) in ActivitySubprocess."""
 
     @pytest.mark.parametrize(
         "signal",
@@ -2748,8 +2748,8 @@ def test_signals_without_retry_always_fail(self, mocker, signal):
         result = mock_watched_subprocess.final_state
         assert result == TaskInstanceState.FAILED
 
-    def test_non_signal_exit_code_goes_to_failed(self, mocker):
-        """Test that non signal exit codes go to failed regardless of task retries."""
+    def test_non_signal_exit_code_with_retry_goes_to_up_for_retry(self, mocker):
+        """Test that non-signal exit codes with retries enabled go to UP_FOR_RETRY."""
         mock_watched_subprocess = ActivitySubprocess(
             process_log=mocker.MagicMock(),
             id=TI_ID,
@@ -2761,6 +2761,21 @@ def test_non_signal_exit_code_goes_to_failed(self, mocker):
         mock_watched_subprocess._exit_code = 1
         mock_watched_subprocess._should_retry = True
 
+        assert mock_watched_subprocess.final_state == TaskInstanceState.UP_FOR_RETRY
+
+    def test_non_signal_exit_code_without_retry_goes_to_failed(self, mocker):
+        """Test that non-signal exit codes without retries enabled go to FAILED."""
+        mock_watched_subprocess = ActivitySubprocess(
+            process_log=mocker.MagicMock(),
+            id=TI_ID,
+            pid=12345,
+            stdin=mocker.Mock(),
+            process=mocker.Mock(),
+            client=mocker.Mock(),
+        )
+        mock_watched_subprocess._exit_code = 1
+        mock_watched_subprocess._should_retry = False
+
         assert mock_watched_subprocess.final_state == TaskInstanceState.FAILED