Add High Obtainability (HO) mode support to Dataproc operators by jigneshjain25 · Pull Request #74753 · apache/airflow

jigneshjain25 · 2026-10-14T17:04:52Z

Adds support for Dataproc High Obtainability (HO) mode declarative instance group configurations (master_flexible_config, worker_flexible_config, secondary_worker_flexible_config):

Bumps google-cloud-dataproc>=5.32.0 (and uv.lock) for native FlexibleInstanceGroupConfig support.
Adds system test / example DAG example_dataproc_cluster_high_obtainability.py demonstrating creating and scaling an HO cluster.
Updates dataproc.rst operator documentation with High Obtainability mode usage examples and links to the official Google Cloud documentation.

…ataproc operators Adds support for Dataproc High Obtainability (HO) mode declarative instance group configurations: - Adds FlexibleInstanceGroupConfig dataclass and HardwareClass enum (STANDARD, PERFORMANCE). - Adds master_flexible_config, worker_flexible_config, and secondary_worker_flexible_config parameters to ClusterGenerator. - Validates mutual exclusivity against standard machine types, disk configs, accelerators, instance flexibility policies, and zone. - Supports dual-path scaling in DataprocUpdateClusterOperator (flexible update masks config.worker_flexible_config.num_instances and config.secondary_worker_flexible_config.num_instances, with conflict checks against legacy paths). - Adds helper flexible_scale_cluster_update for scaling HO clusters. - Adds comprehensive unit tests, system test example DAG, and documentation.

boring-cyborg · 2026-10-14T17:04:57Z

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

…her than being regional

…terOperator cluster_config Address review feedback from Colin-Cheng: ClusterGenerator is deprecated and scheduled for removal after September 2027. Revert HO additions to ClusterGenerator and move flexible config validation and normalization directly into DataprocCreateClusterOperator(cluster_config=...). Update documentation, example DAG, and unit tests accordingly.

…dation - Safely access software_config and gce_cluster_config when None - Reject concrete worker_config on single-node HO clusters - Disallow boolean values for num_instances, core_count, boot_disk_size_gb, and worker scaling counts - Validate hardware_class type and values strictly - Add unit tests covering these edge cases

Colin-Cheng

Thanks for removing the ClusterGenerator changes and updating the tests! Left a few inline comments on the updated implementation (mainly a mypy [union-attr] issue in _validate_and_normalize_flexible_cluster_config, a leftover num_workers=0 reference in the docs, and a couple of minor validation edge cases).

Also, could you confirm whether the minimum google-cloud-dataproc version in providers/google/pyproject.toml (currently >=5.23.0) needs to be bumped so that ClusterConfig proto-plus messages recognize master_flexible_config / worker_flexible_config / secondary_worker_flexible_config?

…pport - Tighten _to_flexible_config return type annotation to avoid mypy union-attr errors. - Disallow scaling primary workers to 0 in flexible_scale_cluster_update while allowing 0 for secondary workers. - Avoid reading proto3 default 0 for unset submessages in proto-plus Cluster messages in _get_num_instances. - Explicitly exclude boolean values from being coerced as integer counts in _get_num_instances. - Skip parse-time conflicting scaling path validation in DataprocUpdateClusterOperator when cluster is templated or None. - Update docs to describe single-node HO cluster configuration via allow.zero.workers property. - Handle AirflowProviderDeprecationWarning in unit tests for ClusterGenerator.

Colin-Cheng

Thanks for addressing all the inline comments! One last thing regarding the google-cloud-dataproc dependency version:

… 5.32.0 - Raise minimum google-cloud-dataproc from >=5.27.0 to >=5.32.0 in pyproject.toml, README.rst, and docs/index.rst. - Support for FlexibleInstanceGroupConfig (master_flexible_config, worker_flexible_config, secondary_worker_flexible_config) in ClusterConfig was added in google-cloud-dataproc 5.32.0. - Reformat example_dataproc_cluster_high_obtainability.py to satisfy ruff format check. - Add unit test verifying proto Cluster accepts all flexible config fields.

olegkachur-e

Please take a look at the comments.

Overall, I think the cleanest way to introduce this feature is to provide docs, an example DAG and the dependency upgrade.

I don't think we should duplicate the API validation here.

…pport - Remove redundant client-side HO validation and custom dataclass FlexibleInstanceGroupConfig from dataproc.py, relying on native google-cloud-dataproc types and backend API validation. - Remove flexible_scale_cluster_update helper function and update_mask validation in DataprocUpdateClusterOperator. - Update uv.lock for google-cloud-dataproc>=5.32.0. - Update example_dataproc_cluster_high_obtainability.py with minimized resources (2 cores, 32GB boot disk, plain dicts) and single DataprocUpdateClusterOperator call scaling worker group to 3 and secondary worker to 0. - Update documentation in dataproc.rst linking to official Google Cloud documentation and API reference. - Refactor unit tests in test_dataproc_high_obtainability.py to verify operator pass-through and proto compatibility.

jigneshjain25 · 2026-10-20T16:23:55Z

Thanks for the review @olegkachur-e!

All feedback has been addressed in commit 4bb6288df1 on dataproc-ho-support:

Removed custom classes & redundant validation: Completely removed FlexibleInstanceGroupConfig, HardwareClass enum, flexible_scale_cluster_update, and client-side validation methods from dataproc.py. Airflow Dataproc operators now act as pass-through orchestrators passing cluster_config, cluster, and update_mask directly to the backend API via DataprocHook. Users can provide plain dictionaries or native google.cloud.dataproc_v1 proto types (FlexibleInstanceGroupConfig).
Cleaned up operator docstring: Removed the inline HO example block from DataprocCreateClusterOperator's cluster_config docstring, referencing the documentation guide and example DAG instead.
Official documentation links: Added direct links in dataproc.rst to the official Google Cloud documentation for Dataproc High Obtainability mode and the FlexibleInstanceGroupConfig API reference.
Updated uv.lock: Refreshed uv.lock with google-cloud-dataproc>=5.32.0 and exact PyPI release metadata and hashes.
Minimized example DAG resources: Updated example_dataproc_cluster_high_obtainability.py to use minimal resources (2 cores, 32GB boot disk per node across master, primary workers, and secondary worker), plain dictionaries, and consolidated scaling into a single DataprocUpdateClusterOperator call scaling worker_flexible_config to 3 and secondary_worker_flexible_config to 0.
Verified DAG execution & test suite:
- The example DAG was executed and passed end-to-end on GCP test project dataproc-test-jnj (create -> scale -> spark_pi -> delete all SUCCESS).
- Ran ruff check providers/google, ruff format --check providers/google, and pytest on both test_dataproc_high_obtainability.py and test_dataproc.py (all tests passing).

# Conflicts: # uv.lock

…ion link

olegkachur-e

Thanks for the updates, the PR looks much cleaner now!

A few thoughts on test_dataproc_high_obtainability.py:

TestProtoClusterFlexibleConfig tests the google-cloud-dataproc client rather than the Airflow operator, so I don't think it should be unit tested here.
DataprocCreateClusterOperator accepts any valid ClusterConfig and passes it to the API as is, and the validation happens on the API side.
We usually don't write unit tests for system tests (e.g. TestExampleDag).
These tests depend on the google-cloud-dataproc version rather than on the operator code, so they could fail with older client versions even though the operator code is unchanged.

I'd suggest removing this file completely. The example_dataproc_cluster_high_obtainability.py system test already covers this feature end to end.

Thanks again for your work on this!

…lity.py unit test file

jigneshjain25 · 2026-10-21T15:45:20Z

Thanks for the feedback @olegkachur-e! Agreed on all points — I have removed providers/google/tests/unit/google/cloud/operators/test_dataproc_high_obtainability.py completely (and updated worker_flexible_config.num_instances to 2 in example_dataproc_cluster_high_obtainability.py to match your resource snippet), leaving only the documentation update, google-cloud-dataproc>=5.32.0 dependency bump, and the example_dataproc_cluster_high_obtainability.py system test. Please take another look when you get a chance!

olegkachur-e

LGTM, thank you!

MaksYermak

Hello @jigneshjain25, thank you for the PR!
LGTM, could I ask you to provide a screenshots of successful run for example_dataproc_cluster_high_obtainability.py‎ system tests?
@shahar1 what do you think about this PR?

jigneshjain25 · 2026-10-23T21:20:39Z

Thank you for the review and approval @MaksYermak!

Here are the screenshots from the successful execution of the example_dataproc_cluster_high_obtainability.py (dataproc_cluster_high_obtainability) system test:

1. Airflow UI — DAG Graph (`Success`)

2. Airflow UI — DAG Run & Task Instances (`Success`)

3. Airflow UI — DAG Overview (`Success`)

4. System Test Execution Output (`pytest --system`)

shahar1

Overall LGTM, pleas fix failing CI job

…ordlist

jigneshjain25 · 2026-10-25T19:53:32Z

Hi @shahar1,

I have fixed the failing CI job (CI image checks / Build documentation):

Added Obtainability, obtainability, and stockouts to docs/spelling_wordlist.txt (sorted via scripts/ci/prek/sort_spelling_wordlist.py) to satisfy the Sphinx documentation spell check.
Pushed commit 9bf67f5f7a to the PR branch.

Could you please take another look and merge when you get a chance? Thank you!

shahar1 · 2026-10-26T03:52:42Z

Hi @shahar1,

I have fixed the failing CI job (CI image checks / Build documentation):

Added Obtainability, obtainability, and stockouts to docs/spelling_wordlist.txt (sorted via scripts/ci/prek/sort_spelling_wordlist.py) to satisfy the Sphinx documentation spell check.

Pushed commit 9bf67f5f7a to the PR branch.

Could you please take another look and merge when you get a chance? Thank you!

Approved, please do not merge from main and let the CI run until it finishes :)

jigneshjain25 · 2026-10-26T06:25:46Z

Thanks @shahar1! All CI checks have now completed and passed 100% green. Ready to be merged whenever you get a chance!

boring-cyborg · 2026-10-26T10:50:21Z

Awesome work, congrats on your first merged pull request! You are invited to check our Issue Tracker for additional contributions.

jigneshjain25 requested a review from shahar1 as a code owner October 14, 2026 17:04

boring-cyborg Bot added area:providers kind:documentation provider:google Google (including GCP) related issues labels Oct 14, 2026

jigneshjain25 added 3 commits October 14, 2026 10:05

Merge branch 'main' into dataproc-ho-support

27ad86d

Merge branch 'main' into dataproc-ho-support

87bd885

docs(providers/google): clarify that HO clusters require AutoZone rat…

cd012e2

…her than being regional

Colin-Cheng reviewed Oct 15, 2026

View reviewed changes

Comment thread providers/google/src/airflow/providers/google/cloud/operators/dataproc.py Outdated

jigneshjain25 added 2 commits October 15, 2026 19:38

jigneshjain25 requested a review from Colin-Cheng October 15, 2026 20:53

Merge branch 'main' into dataproc-ho-support

bd4aea4

Colin-Cheng reviewed Oct 15, 2026

View reviewed changes

jigneshjain25 added 2 commits October 15, 2026 22:52

Merge branch 'main' into dataproc-ho-support

f887819

jigneshjain25 requested a review from Colin-Cheng October 15, 2026 22:56

Colin-Cheng reviewed Oct 15, 2026

View reviewed changes

Comment thread providers/google/tests/unit/google/cloud/operators/test_dataproc_high_obtainability.py Outdated

Colin-Cheng approved these changes Oct 15, 2026

View reviewed changes

jigneshjain25 requested a review from Colin-Cheng October 16, 2026 18:26

Colin-Cheng approved these changes Oct 16, 2026

View reviewed changes

jigneshjain25 added 2 commits October 16, 2026 11:30

Merge branch 'main' into dataproc-ho-support

c61d1d7

Merge branch 'main' into dataproc-ho-support

e50fc65

shahar1 changed the title ~~feat(providers/google): add High Obtainability (HO) mode support to Dataproc operators~~ Add High Obtainability (HO) mode support to Dataproc operators Oct 17, 2026

olegkachur-e suggested changes Oct 20, 2026

View reviewed changes

Merge remote-tracking branch 'origin/main' into dataproc-ho-support

ef8849f

# Conflicts: # uv.lock

jigneshjain25 requested a review from olegkachur-e October 20, 2026 20:49

jigneshjain25 added 2 commits October 20, 2026 14:02

Merge branch 'main' into dataproc-ho-support

aec4551

docs(providers/google): update Dataproc High Obtainability documentat…

dc23070

…ion link

olegkachur-e suggested changes Oct 21, 2026

View reviewed changes

test(providers/google): remove redundant test_dataproc_high_obtainabi…

df7dead

…lity.py unit test file

Merge branch 'main' into dataproc-ho-support

9d7f464

jigneshjain25 requested a review from olegkachur-e October 21, 2026 18:56

jigneshjain25 added 2 commits October 21, 2026 19:01

Fix Dataproc High Obtainability system test schedule and docs

1ba51af

Merge branch 'main' into dataproc-ho-support

1ade66a

olegkachur-e approved these changes Oct 22, 2026

View reviewed changes

jigneshjain25 added 2 commits October 22, 2026 14:27

Merge branch 'main' into dataproc-ho-support

089aeed

Merge branch 'main' into dataproc-ho-support

ab058ab

MaksYermak approved these changes Oct 23, 2026

View reviewed changes

jigneshjain25 added a commit to jigneshjain25/airflow that referenced this pull request Oct 23, 2026

Add DAG Graph view screenshot for PR apache#74753

a22551d

shahar1 reviewed Oct 25, 2026

View reviewed changes

jigneshjain25 added 2 commits October 25, 2026 07:43

docs(providers/google): add Obtainability and stockouts to spelling w…

9bf67f5

…ordlist

Merge branch 'main' into dataproc-ho-support

2039cbd

github-actions Bot added the need workflow approval label Oct 25, 2026

Merge branch 'main' into dataproc-ho-support

0f5a8a9

shahar1 approved these changes Oct 26, 2026

View reviewed changes

shahar1 merged commit 6b26473 into apache:main Oct 26, 2026
91 checks passed

Conversation

jigneshjain25 commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

boring-cyborg Bot commented Oct 14, 2026

Uh oh!

Uh oh!

Colin-Cheng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Colin-Cheng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

olegkachur-e left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jigneshjain25 commented Oct 20, 2026

Uh oh!

olegkachur-e left a comment

Choose a reason for hiding this comment

Uh oh!

jigneshjain25 commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

olegkachur-e left a comment

Choose a reason for hiding this comment

Uh oh!

MaksYermak left a comment

Choose a reason for hiding this comment

Uh oh!

jigneshjain25 commented Oct 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

1. Airflow UI — DAG Graph (Success)

2. Airflow UI — DAG Run & Task Instances (Success)

3. Airflow UI — DAG Overview (Success)

4. System Test Execution Output (pytest --system)

Uh oh!

shahar1 left a comment

Choose a reason for hiding this comment

Uh oh!

jigneshjain25 commented Oct 25, 2026

Uh oh!

shahar1 commented Oct 26, 2026

Uh oh!

jigneshjain25 commented Oct 26, 2026

Uh oh!

Uh oh!

boring-cyborg Bot commented Oct 26, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

jigneshjain25 commented Oct 14, 2026 •

edited

Loading

olegkachur-e left a comment •

edited

Loading

jigneshjain25 commented Oct 21, 2026 •

edited

Loading

jigneshjain25 commented Oct 23, 2026 •

edited

Loading

1. Airflow UI — DAG Graph (`Success`)

2. Airflow UI — DAG Run & Task Instances (`Success`)

3. Airflow UI — DAG Overview (`Success`)

4. System Test Execution Output (`pytest --system`)