Skip to content

Migrate AzureBatchHook/Operator to azure-batch 15.x (track 2 SDK)#67600

Open
arieleli01212 wants to merge 2 commits into
apache:mainfrom
arieleli01212:fix/azure-batch-15x-sdk
Open

Migrate AzureBatchHook/Operator to azure-batch 15.x (track 2 SDK)#67600
arieleli01212 wants to merge 2 commits into
apache:mainfrom
arieleli01212:fix/azure-batch-15x-sdk

Conversation

@arieleli01212
Copy link
Copy Markdown

@arieleli01212 arieleli01212 commented May 27, 2026

Description

The azure-batch package underwent a complete API overhaul in version 15.x, moving from the old Azure SDK track 1 (BatchServiceClient, batch_auth) to track 2 (BatchClient, azure-identity). This PR updates AzureBatchHook and AzureBatchOperator to work with the new SDK and removes the <15.0.0 upper bound cap that was explicitly blocking users from upgrading.

Changes

hooks/batch.py

  • Replace BatchServiceClient with BatchClient(endpoint, credential)
  • Replace batch_auth.SharedKeyCredentials with AzureNamedKeyCredential
  • Catch ResourceExistsError (from azure.core.exceptions) instead of BatchErrorException
  • Rename model types: PoolAddParameterBatchPoolCreateOptions, JobAddParameterBatchJobCreateOptions, TaskAddParameterBatchTaskCreateOptions, ImageReferenceBatchVmImageReference, PoolInformationBatchPoolInfo
  • Update enum references to use SCREAMING_SNAKE_CASE values: BatchNodeState.IDLE, BatchTaskState.COMPLETED, BatchTaskExecutionResult.FAILURE
  • Replace the nested sub-client calls (client.pool.add, client.compute_node.list, etc.) with the new flat API (client.create_pool, client.list_nodes, etc.)
  • Use LRO methods begin_delete_pool, begin_delete_job, begin_terminate_job
  • Remove CloudServiceConfiguration support — Azure has retired cloud-service VM pools

operators/batch.py

  • Update type hints for batch_job_manager_task, batch_job_preparation_task, etc. to use the new Batch-prefixed names
  • Remove os_family and os_version parameters (no longer applicable without CloudServiceConfiguration)
  • Update node-state set to use BatchNodeState.START_TASK_FAILED / UNUSABLE / IDLE
  • Update on_kill() and clean_up() to use the LRO methods

pyproject.toml

  • Lift the azure-batch constraint from >=8.0.0,<15.0.0 to >=15.0.0

Tests

  • Rewrite hook and operator tests to use the new BatchClient, model types, and enum values

Relates

Fixes the blocking comment left in pyproject.toml about needing a hook rewrite before lifting the <15.0.0 cap.

closes: #66466

Testing

  • I ran pytest providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_batch.py and all tests pass
  • I ran pytest providers/microsoft/azure/tests/unit/microsoft/azure/operators/test_batch.py and all tests pass

…ack 2 SDK)

The azure-batch package 15.x is a ground-up rewrite of the Azure SDK for
Batch. The old BatchServiceClient and batch_auth module are gone; the new
SDK exposes a flat BatchClient that uses azure-identity / AzureNamedKeyCredential.

Changes:
- Replace BatchServiceClient with BatchClient(endpoint, credential)
- Replace batch_auth.SharedKeyCredentials with AzureNamedKeyCredential
- Catch ResourceExistsError (azure.core) instead of BatchErrorException
- Rename model types: PoolAddParameter → BatchPoolCreateOptions,
  JobAddParameter → BatchJobCreateOptions, TaskAddParameter → BatchTaskCreateOptions,
  ImageReference → BatchVmImageReference, PoolInformation → BatchPoolInfo,
  and Batch-prefixed task/job sub-types
- Rename enums: ComputeNodeState → BatchNodeState, TaskState → BatchTaskState,
  TaskExecutionResult → BatchTaskExecutionResult (SCREAMING_SNAKE_CASE values)
- Replace pool.list / compute_node.list / job.list / task.list hierarchy with
  flat client.list_nodes / list_tasks / list_jobs / list_supported_images calls
- Use LRO methods begin_delete_pool / begin_delete_job / begin_terminate_job
- Remove CloudServiceConfiguration support (Azure removed cloud-service VMs)
- Drop os_family / os_version parameters from AzureBatchOperator (no longer
  applicable without CloudServiceConfiguration)
- Lift azure-batch version cap from <15.0.0 to >=15.0.0 in pyproject.toml
- Update all unit tests to use the new SDK classes and method signatures

Closes apache#66466
@boring-cyborg
Copy link
Copy Markdown

boring-cyborg Bot commented May 27, 2026

Congratulations on your first Pull Request and welcome to the Apache Airflow community! If you have any issues or are unsure about any anything please check our Contributors' Guide
Here are some useful points:

  • Pay attention to the quality of your code (ruff, mypy and type annotations). Our prek-hooks will help you with that.
  • In case of a new feature add useful documentation (in docstrings or in docs/ directory). Adding a new operator? Check this short guide Consider adding an example Dag that shows how users should use it.
  • Consider using Breeze environment for testing locally, it's a heavy docker but it ships with a working Airflow and a lot of integrations.
  • Be patient and persistent. It might take some time to get a review or get the final approval from Committers.
  • Please follow ASF Code of Conduct for all communication including (but not limited to) comments on Pull Requests, Mailing list and Slack.
  • Be sure to read the Airflow Coding style.
  • Always keep your Pull Requests rebased, otherwise your build might fail due to changes not related to your commits.
    Apache Airflow is a community-driven project and together we are making it better 🚀.
    In case of doubts contact the developers at:
    Mailing List: dev@airflow.apache.org
    Slack: https://s.apache.org/airflow-slack

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Migrate microsoft.azure provider to azure-batch 15.x (BatchClient) — unblock <15 cap from #66452

1 participant