Handle AzureBatch task result correctly

airflow/providers/microsoft/azure/hooks/batch.py

@@ -366,17 +366,46 @@ def wait_for_job_tasks_to_complete(self, job_id: str, timeout: int) -> list[batc
             incomplete_tasks = [task for task in tasks if task.state != batch_models.TaskState.completed]
             if not incomplete_tasks:
                 # detect if any task in job has failed
-                fail_tasks = [
+                failed_tasks = [
                     task
                     for task in tasks
-                    if task.executionInfo.result == batch_models.TaskExecutionResult.failure
+                    if (task.execution_info.result == batch_models.TaskExecutionResult.failure) or (
+                        task.execution_info.exit_code != 0)
                 ]
-                return fail_tasks
+                return failed_tasks
             for task in incomplete_tasks:
                 self.log.info("Waiting for %s to complete, currently on %s state", task.id, task.state)
             time.sleep(15)
         raise TimeoutError("Timed out waiting for tasks to complete")
 
+    def wait_for_single_job_task_to_complete(self, job_id: str, task_id: str,
+                                             timeout: int) -> bool:
+        """
+        Wait for a single task in a particular job to complete, return False if it ultimately fails or True
+        if it succeeds.
+
+        :param job_id: A string that identifies the job
+        :param task_id: A string that identifies the task
+        :param timeout: The amount of time to wait before timing out in minutes
+        :return: A bool that represents whether the task failed (false) or succeeded (true)
+        """
+        timeout_time = timezone.utcnow() + timedelta(minutes=timeout)
+        while timezone.utcnow() < timeout_time:
+            task = self.connection.task.get(job_id, task_id)
+
+            if task.state == batch_models.TaskState.completed:
+                # detect if task in job has failed
+                self.log.info("Task %s completed with result %s and exit code %s", task.id,
+                              task.execution_info.result,
+                              str(task.execution_info.exit_code))
+                if (task.execution_info.result == batch_models.TaskExecutionResult.failure) or (
+                    task.execution_info.exit_code != 0):
+                    return False
+                return True
+            self.log.info("Waiting for %s to complete, currently on %s state", task.id, task.state)
+            time.sleep(15)
+        raise TimeoutError(f"Timed out waiting for task {task_id} to complete")
+
     def test_connection(self):
         """Test a configured Azure Batch connection."""
         try:

airflow/providers/microsoft/azure/operators/batch.py

@@ -293,16 +293,18 @@ def execute(self, context: Context) -> None:
         # Add task to job
         self.hook.add_single_task_to_job(job_id=self.batch_job_id, task=task)
         # Wait for tasks to complete
-        fail_tasks = self.hook.wait_for_job_tasks_to_complete(job_id=self.batch_job_id, timeout=self.timeout)
+        is_task_successful = self.hook.wait_for_single_job_task_to_complete(job_id=self.batch_job_id,
+                                                                            task_id=task.id,
+                                                                            timeout=self.timeout)
         # Clean up
         if self.should_delete_job:
             # delete job first
             self.clean_up(job_id=self.batch_job_id)
         if self.should_delete_pool:
             self.clean_up(self.batch_pool_id)
-        # raise exception if any task fail
-        if fail_tasks:
-            raise AirflowException(f"Job fail. The failed task are: {fail_tasks}")
+        # raise exception if the task failed
+        if not is_task_successful:
+            raise AirflowException(f"Task {task.id} has failed.")
 
     def on_kill(self) -> None:
         response = self.hook.connection.job.terminate(