Add DiskProvisionedIops and DiskProvisionedThroughputMibps pipeline options by bambadiouf1 · Pull Request #37377 · apache/beam

bambadiouf1 · 2026-01-21T20:56:32Z

This pull request introduces two new pipeline options for the Google Cloud Dataflow runner for the Java, Python and Go SDKs. These options allow users to specify provisioned performance for worker VM root disks:

disk_provisioned_iops: Sets the provisioned IOPS for the root disk. If unspecified, the service chooses a default
disk_provisioned_throughput_mibps: Sets the provisioned throughput in MiB/s for the root disk.

Added getters and setters to DataflowPipelineWorkerPoolOptions and updated DataflowPipelineTranslator to pass these values to the worker pool configuration.
Will follow up with a PR for the Go and Python SDKs

Tests have been added/updated to verify that these options are correctly parsed and translated.

More context:
we need to add these pipeline options before submitting this cl: https://critique.corp.google.com/cl/858930428
Issue: #37374

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

gemini-code-assist · 2026-01-21T20:56:55Z

Summary of Changes

Hello @bambadiouf1, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request enhances Google Cloud Dataflow's capabilities by allowing users to specify detailed performance characteristics for worker VM root disks. By introducing options for provisioned IOPS and throughput, it provides greater control over the underlying infrastructure, which can be crucial for optimizing the performance and cost-efficiency of data processing pipelines. This change is integrated across all supported SDKs, ensuring a consistent experience for developers.

Highlights

New Pipeline Options: Introduced DiskProvisionedIops and DiskProvisionedThroughputMibps as new pipeline options for Google Cloud Dataflow.
Cross-SDK Implementation: These new options are implemented across the Python, Java, and Go SDKs, allowing users to specify provisioned performance for worker VM root disks regardless of the SDK used.
Worker Disk Performance Control: Users can now explicitly set the provisioned IOPS (Input/Output Operations Per Second) and throughput (in MiB/s) for Dataflow worker root disks, offering finer control over job performance.
Comprehensive Testing: New tests have been added for each SDK to ensure that these options are correctly parsed, translated, and applied to the Dataflow worker pool configuration.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

github-actions · 2026-01-21T21:43:02Z

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

…r python sdk

…rker pools for java sdk

…wPipelineTranslator

bambadiouf1 · 2026-04-14T03:13:43Z

assign set of reviewers

github-actions · 2026-04-14T03:14:58Z

Assigning reviewers:

R: @damccorm for label python.
R: @jrmccluskey for label go.
R: @Abacn for label build.

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

stop reviewer notifications - opt out of the automated review tooling
remind me after tests pass - tag the comment author after tests pass
waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).

damccorm · 2026-04-14T12:31:37Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces support for configuring provisioned IOPS and throughput for worker disks in the Dataflow runner across the Java, Go, and Python SDKs. The changes include updating the Dataflow API version, adding new pipeline options, and implementing the necessary translation logic to pass these parameters to the service. Review feedback highlights that the Go implementation is currently incomplete as it lacks the translation logic to the API request. Additionally, suggestions were made for the Java SDK to use Long instead of Integer for these fields for better consistency and to include basic validation to ensure values are greater than zero before being set.

gemini-code-assist · 2026-04-14T12:33:29Z

sdks/go/pkg/beam/runners/dataflow/dataflow.go

+	diskProvisionedIops            = flag.Int64("disk_provisioned_iops", 0, "Provisioned IOPS for the worker disk (optional).")
+	diskProvisionedThroughputMibps = flag.Int64("disk_provisioned_throughput_mibps", 0, "Provisioned throughput in MiB/s for the worker disk (optional).")


The Go implementation appears to be incomplete. While flags and JobOptions fields have been added, there is no corresponding logic to translate these values into the WorkerPool configuration within the Dataflow API request (typically handled in environment.go). Without this translation, these flags will have no effect on the provisioned workers.

I agree that I would've expected this to show up in

beam/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go

Line 188 in 3f5ce80

WorkerPools: []*df.WorkerPool{{

gemini-code-assist · 2026-04-14T12:33:29Z

...dataflow-java/src/main/java/org/apache/beam/runners/dataflow/DataflowPipelineTranslator.java

+      if (options.getDiskProvisionedIops() != null) {
+        workerPool.setDiskProvisionedIops(Long.valueOf(options.getDiskProvisionedIops()));
+      }
+      if (options.getDiskProvisionedThroughputMibps() != null) {
+        workerPool.setDiskProvisionedThroughputMibps(
+            Long.valueOf(options.getDiskProvisionedThroughputMibps()));
+      }


For consistency with how diskSizeGb is handled (lines 489-491), it is recommended to also check that these values are greater than zero. This provides basic client-side validation and ensures that a value of 0 (if provided) doesn't inadvertently override service defaults in an unexpected way. Additionally, the conversion to Long can be simplified.

Suggested change

if (options.getDiskProvisionedIops() != null) {

workerPool.setDiskProvisionedIops(Long.valueOf(options.getDiskProvisionedIops()));

}

if (options.getDiskProvisionedThroughputMibps() != null) {

workerPool.setDiskProvisionedThroughputMibps(

Long.valueOf(options.getDiskProvisionedThroughputMibps()));

}

if (options.getDiskProvisionedIops() != null && options.getDiskProvisionedIops() > 0) {

workerPool.setDiskProvisionedIops((long) options.getDiskProvisionedIops());

}

if (options.getDiskProvisionedThroughputMibps() != null && options.getDiskProvisionedThroughputMibps() > 0) {

workerPool.setDiskProvisionedThroughputMibps(

(long) options.getDiskProvisionedThroughputMibps());

}

gemini-code-assist · 2026-04-14T12:33:29Z

...rc/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineWorkerPoolOptions.java

+  Integer getDiskProvisionedIops();
+
+  void setDiskProvisionedIops(Integer value);


Consider using Long instead of Integer for these performance-related fields. While Integer is likely sufficient for current IOPS limits, using Long is more consistent with the underlying Dataflow API (which the translator casts to Long) and the Go implementation in this PR. It also future-proofs the options against very high performance disk specifications.

Suggested change

Integer getDiskProvisionedIops();

void setDiskProvisionedIops(Integer value);

Long getDiskProvisionedIops();

void setDiskProvisionedIops(Long value);

gemini-code-assist · 2026-04-14T12:33:29Z

...rc/main/java/org/apache/beam/runners/dataflow/options/DataflowPipelineWorkerPoolOptions.java

+  Integer getDiskProvisionedThroughputMibps();
+
+  void setDiskProvisionedThroughputMibps(Integer value);


Consider using Long for throughput as well, for consistency and future-proofing.

Suggested change

Integer getDiskProvisionedThroughputMibps();

void setDiskProvisionedThroughputMibps(Integer value);

Long getDiskProvisionedThroughputMibps();

void setDiskProvisionedThroughputMibps(Long value);

damccorm · 2026-04-14T12:34:47Z

sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go

+	DiskType                       string
+	DiskProvisionedIops            int64
+	DiskProvisionedThroughputMibps int64
+	MachineType                    string


These don't look correctly formatted - could you please run go fmt to fix?

damccorm · 2026-04-14T12:39:25Z

sdks/go/pkg/beam/runners/dataflow/dataflow.go

+	diskProvisionedIops            = flag.Int64("disk_provisioned_iops", 0, "Provisioned IOPS for the worker disk (optional).")
+	diskProvisionedThroughputMibps = flag.Int64("disk_provisioned_throughput_mibps", 0, "Provisioned throughput in MiB/s for the worker disk (optional).")


I agree that I would've expected this to show up in

beam/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go

Line 188 in 3f5ce80

WorkerPools: []*df.WorkerPool{{

github-actions bot added python go runners dataflow labels Jan 21, 2026

bambadiouf1 marked this pull request as draft January 21, 2026 22:12

damccorm mentioned this pull request Feb 2, 2026

[Feature Request]: Add pipeline options for diskProvisionedIops and diskProvisionedThroughputMipbs #37374

Open

17 tasks

bambadiouf1 marked this pull request as ready for review March 17, 2026 22:28

github-actions bot added go build and removed go python labels Mar 26, 2026

bambadiouf1 marked this pull request as draft April 8, 2026 16:44

bambadiouf1 marked this pull request as ready for review April 8, 2026 16:44

feat: add support for disk provisioned IOPS and throughput options fo…

6cb13d9

…r python sdk

bambadiouf1 force-pushed the iops_throughput_pipeline_options branch from 309f5fd to 6cb13d9 Compare April 8, 2026 22:12

github-actions bot added python and removed build dataflow labels Apr 8, 2026

bambadiouf1 force-pushed the iops_throughput_pipeline_options branch 2 times, most recently from efa4d20 to 6cb13d9 Compare April 13, 2026 20:10

feat: add disk provisioned IOPS and throughput options to Dataflow wo…

a1aa2dc

…rker pools for java sdk

github-actions bot added the dataflow label Apr 13, 2026

bambadiouf1 added 2 commits April 13, 2026 13:19

Add iops and throughput pipeline options for Go SDK

9e6e47b

Update dataflow api service versions

deefe29

github-actions bot added go build labels Apr 13, 2026

Update dataflow services api version to latest

d93ac19

bambadiouf1 added 4 commits April 13, 2026 13:41

cast disk provisioned values to Long and reformat comments in Dataflo…

55ab768

…wPipelineTranslator

remove unrelated lint changes

809c722

Fix python format issues

1f9d868

fix format

e2836ab

github-actions bot added the Next Action: Reviewers label Apr 14, 2026

gemini-code-assist bot reviewed Apr 14, 2026

View reviewed changes

damccorm reviewed Apr 14, 2026

View reviewed changes

		diskProvisionedIops = flag.Int64("disk_provisioned_iops", 0, "Provisioned IOPS for the worker disk (optional).")
		diskProvisionedThroughputMibps = flag.Int64("disk_provisioned_throughput_mibps", 0, "Provisioned throughput in MiB/s for the worker disk (optional).")

		Integer getDiskProvisionedIops();

		void setDiskProvisionedIops(Integer value);

		Integer getDiskProvisionedThroughputMibps();

		void setDiskProvisionedThroughputMibps(Integer value);

Conversation

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

Uh oh!

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist bot commented Jan 21, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

github-actions bot commented Jan 21, 2026

Uh oh!

bambadiouf1 commented Apr 14, 2026

Uh oh!

github-actions bot commented Apr 14, 2026

Uh oh!

damccorm commented Apr 14, 2026

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

damccorm Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

damccorm Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

damccorm Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

bambadiouf1 commented Jan 21, 2026 •

edited

Loading