Add deferrable support to AzureBatchOperator by Ankurdeewan · Pull Request #59798 · apache/airflow

Ankurdeewan · 2025-12-25T07:21:47Z

This PR adds an opt-in deferrable mode to AzureBatchOperator.

What this changes

Adds a deferrable flag (default: False).
When deferrable=True, the operator:
- submits the Azure Batch pool/job/task
- immediately defers
- moves all waiting and polling to a trigger
The worker no longer blocks on sleep-based polling in deferrable mode.

What stays the same

Default behavior is unchanged.
With deferrable=False, the operator still runs exactly as before:
- same polling logic
- same exceptions
- same cleanup behavior on both success and failure
Existing DAGs are not affected unless deferrable=True is explicitly enabled.

Cleanup behavior

Cleanup semantics are preserved:

Cleanup runs on both success and failure, matching historical behavior.
Cleanup does not run inside execute_complete().
Cleanup runs exactly once in all paths (sync, deferrable, failure, timeout).

Why this is needed

Previously, the operator blocked a worker while polling Azure Batch using sleep-based loops.
This change aligns the Azure Batch operator with existing deferrable patterns and reduces unnecessary worker usage.

Related issue: #59779

Copilot

Pull request overview

This PR adds opt-in deferrable support to AzureBatchOperator, allowing the operator to release worker resources while waiting for Azure Batch jobs to complete instead of blocking with sleep-based polling.

Introduces a deferrable flag (default False) to maintain backward compatibility
Adds AzureBatchJobTrigger for asynchronous polling of Azure Batch job completion
Implements cleanup logic in post_execute() to handle resource deletion in both sync and async modes

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
`providers/microsoft/azure/src/airflow/providers/microsoft/azure/triggers/batch.py`	New trigger implementation for async polling of Azure Batch job completion status
`providers/microsoft/azure/src/airflow/providers/microsoft/azure/operators/batch.py`	Updated operator to support deferrable mode with conditional execution paths and centralized cleanup

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-25T07:29:54Z

+
+ async def run(self) -> AsyncIterator[TriggerEvent]:
+ hook = AzureBatchHook(self.azure_batch_conn_id)
+ timeout_time = timezone.utcnow() + timedelta(minutes=self.timeout)


Using timezone.utcnow() is deprecated. Use timezone.utcnow() consistently if it's the project standard, but consider that datetime.now(timezone.utc) is the recommended approach in modern Python for timezone-aware UTC timestamps.

Copilot · 2025-12-25T07:29:55Z

+ self.clean_up(self.batch_pool_id)
+ self._cleanup_done = True
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str:


The execute_complete method lacks a docstring. Add documentation explaining the trigger callback, expected event structure (status, fail_task_ids, message), return value, and possible exceptions.
Suggested change
defexecute_complete(self, context: Context, event: dict[str, Any] |None=None) ->str:
defexecute_complete(self, context: Context, event: dict[str, Any] |None=None) ->str:
"""
CallbackexecutedwhenthedeferrableAzureBatchjobtriggerfires.
Thismethodisinvokedafter :meth:`defer`whenrunningindeferrablemode.
Itprocessestheeventemittedby :class:`~airflow.providers.microsoft.azure.triggers.batch.AzureBatchJobTrigger`
andconvertsitintothefinaltaskoutcome.
Theeventisexpectedtobeamappingwith (atleast) thefollowingkeys:
*``status``: Astringindicatingthejoboutcome. Recognizedvaluesare:
*``"success"``: alltaskscompletedsuccessfully.
*``"failure"``: oneormoretasksfailed.
*``"timeout"``: thejobdidnotcompletewithintheconfiguredtimeout.
*``"error"``: aninternalerroroccurredwhilewaitingfortaskcompletion.
*``fail_task_ids``: OptionaliterableoftaskIDsthatfailed. Defaultstoanemptylist
whennofailuresarereported.
*``message``: Optionalhuman-readabledescriptionprovidingadditionalcontextfor
``"timeout"``or``"error"``statuses.
:paramcontext: Airflowtaskcontextprovidedbythescheduleronresumption.
:paramevent: Eventdictionarysentfrom :class:`AzureBatchJobTrigger`describingthe
terminalstateofthemonitoredAzureBatchjob.
:returns: TheAzureBatchjobIDassociatedwiththisoperator, onsuccessfulcompletion.
:raisesAirflowException: Ifnoeventisreceived, ifthestatusis``"error"``or
``"failure"``, ifanyfailedtaskIDsarereported, orifanunexpectedstatusis
encountered.
:raisesTimeoutError: Iftheeventindicatesthattheoperationtimedout.
"""

Copilot · 2025-12-25T07:29:55Z

+ fail_task_ids = event.get("fail_task_ids", [])
+
+ if status == "timeout":
+ raise TimeoutError(event.get("message", "Timed out waiting for tasks to complete"))


Using built-in TimeoutError is inconsistent with Airflow patterns. Consider using AirflowException with a timeout-specific message to maintain consistency with other exception handling in this operator.
Suggested change
raiseTimeoutError(event.get("message", "Timed out waiting for tasks to complete"))
raiseAirflowException(event.get("message", "Timed out waiting for tasks to complete"))

Copilot · 2025-12-25T07:29:55Z

+ if status == "failure" or fail_task_ids:
+ raise AirflowException(f"Job failed. Failed tasks: {fail_task_ids}")
+ if status != "success":
+ raise AirflowException(f"Unexpected event status: {event}")


The error message could be more actionable. Consider clarifying what statuses are expected and provide the actual status received, e.g., f\"Unexpected event status '{status}'. Expected 'success', 'failure', 'timeout', or 'error'.\"
Suggested change
raiseAirflowException(f"Unexpected event status: {event}")
raiseAirflowException(
f"Unexpected event status '{status}'. Expected 'success', 'failure', 'timeout', or 'error'. "
f"Full event: {event}"
)

Ankurdeewan · 2025-12-26T06:55:31Z

I’ve pushed a set of updates to the PR to align the implementation closely with the patterns you mentioned on the issue.

In particular:

-deferrable default is now config-driven

-polling is fully trigger-based with an operator-level poll_interval

-cleanup semantics are centralized and preserved

-added comprehensive unit tests for deferrable and non-deferrable paths

This should now be consistent with AWSBatchOperator and existing Azure deferrable operators.

@potiuk tagging for visibility after the updates.

Ankurdeewan · 2026-01-04T07:00:25Z

@potiuk would you be able to take a look at this PR when you have bandwidth?

This is a clean PR with only AzureBatchOperator deferrable changes (previous PR #59787 was closed due to an unrelated diff).

dabla · 2026-01-06T12:19:08Z

+from airflow.triggers.base import TriggerEvent
+
+
+def test_azure_batch_job_trigger_serialize():


Maybe put all tests under a class named TestAzureBatchJobTrigger.

dabla · 2026-01-06T12:21:18Z

+ # Verify pool and nodes are in terminal state before deferral
+ pool = self.hook.connection.pool.get(self.batch_pool_id)
+ nodes = list(self.hook.connection.compute_node.list(self.batch_pool_id))
+ if pool.resize_errors:


This check could be done before computing the node list no?

Yess, that makes sense! I’ve reordered it so we check resize_errors first and only list compute nodes if the pool looks healthy.

dabla · 2026-01-06T12:22:47Z


+ def post_execute(self, context: Context, result: Any | None = None) -> None: # type: ignore[override]
+ """Perform cleanup after task completion in both deferrable and non-deferrable modes."""
+ if getattr(self, "_cleanup_done", False):


How come we have to use getattr here for _cleanup_done?

Yeah, this is mainly because _cleanup_done might not exist in some paths; especially after deferral or when the task gets deserialized on resume. Using getattr(..., False) just keeps it safe and still guarantees cleanup only runs once.

Ankurdeewan · 2026-01-17T07:10:03Z

@dabla I’ve pushed updates that address all the points ; tests are grouped, the pool check now short-circuits before listing nodes, and I added a comment explaining the getattr guard for cleanup. Would appreciate another look when you get a chance.

dabla · 2026-01-22T18:34:52Z

You'll need to add trigger entry in provider.yaml as well.

Ankurdeewan · 2026-02-01T07:02:14Z

@dabla please take a look when you can! thanks :)

…e/triggers/batch.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

potiuk · 2026-03-12T01:32:42Z

@Ankurdeewan This PR has been converted to draft because it does not yet meet our Pull Request quality criteria.

Issues found:

❌ Merge conflicts: This PR has merge conflicts with the main branch. Your branch is 591 commits behind main. Please rebase your branch (git fetch origin && git rebase origin/main), resolve the conflicts, and push again. See contributing quick start.
❌ Pre-commit / static checks: Failing: CI image checks / Static checks. Run prek run --from-ref main locally to find and fix issues. See Pre-commit / static checks docs.
❌ mypy (type checking): Failing: CI image checks / MyPy checks (mypy-providers). Run prek --stage manual mypy-providers --all-files locally to reproduce. You need breeze ci-image build --python 3.10 for Docker-based mypy. See mypy (type checking) docs.
❌ Build docs: Failing: CI image checks / Build documentation (--docs-only). Run breeze build-docs locally to reproduce. See Build docs docs.
❌ Provider tests: Failing: provider distributions tests / Providers wheel tests, provider distributions tests / Providers sdist tests, provider distributions tests / Compat 2.11.0:P3.10:, provider distributions tests / Compat 3.0.6:P3.10:, provider distributions tests / Compat 3.1.7:P3.10: (+4 more). Run provider tests with breeze run pytest <provider-test-path> -xvs. See Provider tests docs.
❌ Other failing CI checks: Failing: Airflow CTL tests / airflow-ctl:P3.10 tests, Generate constraints / Generate constraints for 3.10 on linux/amd64, CI image checks / Build documentation (--spellcheck-only), CI image checks / Test Python API client. Run prek run --from-ref main locally to reproduce. See static checks docs.
⚠️Unresolved review comments: This PR has 3 unresolved review threads from maintainers. Please review and resolve all inline review comments before requesting another review. You can resolve a conversation by clicking 'Resolve conversation' on each thread after addressing the feedback. See pull request guidelines.

Note: Your branch is 591 commits behind main. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.

What to do next:

The comment informs you what you need to do.
Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
Maintainers will then proceed with a normal review.

Converting a PR to draft is not a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. If you have questions, feel free to ask on the Airflow Slack.

Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this two-stage triage process so that our maintainers' limited time is spent where it matters most: the conversation with you.

potiuk · 2026-04-02T17:00:16Z

This pull request has been converted to draft due to quality issues more than a week ago and there has been no response from the author.

@Ankurdeewan, you are welcome to reopen this PR when you are ready to continue working on it. Thank you for your contribution!

CopilotAI review requested due to automatic review settings December 25, 2025 07:21

boring-cyborgBot added area:providers provider:microsoft-azure Azure-related issues labels Dec 25, 2025

Copilotstarted reviewing on behalf of Ankurdeewan December 25, 2025 07:22 View session

CopilotAI reviewed Dec 25, 2025
View reviewed changes

eladkal requested a review from dabla January 6, 2026 09:46

dabla reviewed Jan 6, 2026
View reviewed changes

Ankurdeewanand others added 6 commits February 21, 2026 11:15

Fix deferrable execution and cleanup for AzureBatchOperator
3c4014b

Update providers/microsoft/azure/src/airflow/providers/microsoft/azur…
1a6b0da
…e/triggers/batch.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Align AzureBatchOperator with deferrable trigger-based execution pattern
4e46eb8

fix(azure-batch): improve deferrable execution and test structure
38f60f0

providers/microsoft/azure/provider.yaml
379213a

Rebase on latest main to re-trigger CI
32424b6

Ankurdeewan force-pushed the azure-batch-deferrable-clean branch from 4454d6f to 32424b6Compare February 21, 2026 05:45

tsanton mentioned this pull request Feb 24, 2026
Add parameter deferrable to AzureContainerInstanceOperator #62433
Closed
2 tasks

Ankurdeewan mentioned this pull request Mar 6, 2026
Add parameter deferrable to AzureBatchOperator#59779
Closed
2 tasks

potiuk marked this pull request as draft March 12, 2026 01:32

potiuk closed this Apr 2, 2026

potiuk added the closed because of multiple quality violations Label used to close the PRs when there are multiple quality violations label Apr 2, 2026

Ankurdeewan deleted the azure-batch-deferrable-clean branch April 21, 2026 06:38

SameerMesiah97 mentioned this pull request May 12, 2026
Add deferrable mode support to AzureBatchOperator #66815
Merged

potiuk mentioned this pull request Jun 3, 2026
Status of testing Providers that were prepared on June 02, 2026 #67957
Closed

-defexecute_complete(self, context: Context, event: dict[str, Any] |None=None) ->str:
+defexecute_complete(self, context: Context, event: dict[str, Any] |None=None) ->str:
+ """
+CallbackexecutedwhenthedeferrableAzureBatchjobtriggerfires.
+Thismethodisinvokedafter :meth:`defer`whenrunningindeferrablemode.
+Itprocessestheeventemittedby :class:`~airflow.providers.microsoft.azure.triggers.batch.AzureBatchJobTrigger`
+andconvertsitintothefinaltaskoutcome.
+Theeventisexpectedtobeamappingwith (atleast) thefollowingkeys:
+*``status``: Astringindicatingthejoboutcome. Recognizedvaluesare:
+*``"success"``: alltaskscompletedsuccessfully.
+*``"failure"``: oneormoretasksfailed.
+*``"timeout"``: thejobdidnotcompletewithintheconfiguredtimeout.
+*``"error"``: aninternalerroroccurredwhilewaitingfortaskcompletion.
+*``fail_task_ids``: OptionaliterableoftaskIDsthatfailed. Defaultstoanemptylist
+whennofailuresarereported.
+*``message``: Optionalhuman-readabledescriptionprovidingadditionalcontextfor
+``"timeout"``or``"error"``statuses.
+ :paramcontext: Airflowtaskcontextprovidedbythescheduleronresumption.
+ :paramevent: Eventdictionarysentfrom :class:`AzureBatchJobTrigger`describingthe
+terminalstateofthemonitoredAzureBatchjob.
+ :returns: TheAzureBatchjobIDassociatedwiththisoperator, onsuccessfulcompletion.
+ :raisesAirflowException: Ifnoeventisreceived, ifthestatusis``"error"``or
+``"failure"``, ifanyfailedtaskIDsarereported, orifanunexpectedstatusis
+encountered.
+ :raisesTimeoutError: Iftheeventindicatesthattheoperationtimedout.
+ """

	raiseTimeoutError(event.get("message", "Timed out waiting for tasks to complete"))
	raiseAirflowException(event.get("message", "Timed out waiting for tasks to complete"))

-raiseAirflowException(f"Unexpected event status: {event}")
+raiseAirflowException(
+f"Unexpected event status '{status}'. Expected 'success', 'failure', 'timeout', or 'error'. "
+f"Full event: {event}"
+ )

		from airflow.triggers.base import TriggerEvent


		def test_azure_batch_job_trigger_serialize():

Uh oh!

Conversation

Ankurdeewan commented Dec 25, 2025

Uh oh!

CopilotAI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

CopilotAIDec 25, 2025

Choose a reason for hiding this comment

Uh oh!

CopilotAIDec 25, 2025

Choose a reason for hiding this comment

Uh oh!

CopilotAIDec 25, 2025

Choose a reason for hiding this comment

Uh oh!

CopilotAIDec 25, 2025

Choose a reason for hiding this comment

Uh oh!

Ankurdeewan commented Dec 26, 2025

Uh oh!

Ankurdeewan commented Jan 4, 2026

Uh oh!

dablaJan 6, 2026

Choose a reason for hiding this comment

Uh oh!

dablaJan 6, 2026

Choose a reason for hiding this comment

Uh oh!

AnkurdeewanJan 17, 2026

Choose a reason for hiding this comment

Uh oh!

dablaJan 6, 2026

Choose a reason for hiding this comment

Uh oh!

AnkurdeewanJan 17, 2026

Choose a reason for hiding this comment

Uh oh!

Ankurdeewan commented Jan 17, 2026

Uh oh!

dabla commented Jan 22, 2026

Uh oh!

Ankurdeewan commented Feb 1, 2026

Uh oh!

potiuk commented Mar 12, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

potiuk commented Apr 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

potiuk commented Mar 12, 2026•
edited
Loading