ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning, OSS Image Classification by Amar3tto · Pull Request #37186 · apache/beam

Amar3tto · 2025-12-25T08:15:28Z

Please add a meaningful description for your change here

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 · 2025-12-25T08:15:47Z

Summary of Changes

Hello @Amar3tto, 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 significantly enhances Apache Beam's machine learning capabilities by integrating a new PyTorch-based image object detection pipeline. The pipeline leverages the RunInference transform for efficient batched GPU inference with open-source TorchVision models, processing images from cloud storage and outputting structured detection results to BigQuery. This addition is complemented by a new performance benchmark and corresponding documentation, ensuring that the pipeline's efficiency and resource usage can be consistently monitored and evaluated.

Highlights

New PyTorch Object Detection Example: Introduced a new example pipeline for PyTorch image object detection using Apache Beam's RunInference, capable of processing image URIs from GCS, performing batched GPU inference with TorchVision models, and writing results to BigQuery.
Dedicated Performance Benchmark: Added a new benchmark test (PytorchImageObjectDetectionBenchmarkTest) to measure and track the performance of the PyTorch image object detection pipeline, focusing on stable GPU inference workloads.
Updated Documentation and Dependencies: Included new Python dependencies for PyTorch object detection and updated the project's website with a dedicated performance page for the new benchmark, including placeholders for Looker Studio metrics.

🧠 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.

Ignored Files

Ignored by pattern: .github/workflows/** (2)
- .github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml
- .github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Image_Detection.txt

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. ↩

codecov · 2025-12-31T12:30:55Z

Codecov Report

❌ Patch coverage is 0% with 28 lines in your changes missing coverage. Please review.
✅ Project coverage is 36.33%. Comparing base (358e007) to head (64187d9).
⚠️ Report is 13 commits behind head on master.

Files with missing lines	Patch %	Lines
...s/inference/pytorch_image_captioning_benchmarks.py	0.00%	14 Missing ⚠️
...rence/pytorch_image_object_detection_benchmarks.py	0.00%	14 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (358e007) and HEAD (64187d9). Click for more details.

HEAD has 3 uploads less than BASE

Flag BASE (358e007) HEAD (64187d9)

python 4 1

Additional details and impacted files

@@              Coverage Diff              @@
##             master   #37186       +/-   ##
=============================================
- Coverage     55.28%   36.33%   -18.96%     
  Complexity     1676     1676               
=============================================
  Files          1067     1069        +2     
  Lines        167148   167178       +30     
  Branches       1208     1208               
=============================================
- Hits          92415    60737    -31678     
- Misses        72551   104259    +31708     
  Partials       2182     2182

Flag	Coverage Δ
python	`40.60% <0.00%> (-40.46%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

…ting) ML pipeline

github-actions · 2026-02-07T06:48:50Z

Assigning reviewers:

R: @claudevdm for label python.
R: @liferoad for label build.
R: @shunping for label website.

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).

Amar3tto · 2026-02-17T13:07:07Z

@Abacn Could you please help with review?

damccorm · 2026-02-19T15:16:47Z

 # limitations under the License.

-name: Inference Python Benchmarks Dataflow
+name: Inference Python Benchmarks Dataflow (1 part)


Is there a reason to split these into different workflows? If it is just about minimizing the time it takes to run, could we do one workflow with 2 jobs?

Or if there is a reason, could you add a comment explaining it? (maybe what I'm suggesting would exhaust resources and we need different cron schedules?)

Yes, initially my idea was to have two parts with different cron schedules, because for example there is a quota for GPU machines.

damccorm · 2026-02-19T15:20:52Z

+# ============ DoFns ============
+
+
+class RateLimitDoFn(beam.DoFn):


Why do we have this function? It will not effectively provide a global rate limit since multiple instances of this will be running in parallel

Creating a rate limit is doable, but we would need to use a stateful DoFn to effectively do this. Basically the idea would be:

Key all the incoming data with a single (non-unique) key

For each incoming piece of data:

check stored state to see if it is ready to be released, and if not sleep until it is

Yield the element

Store the next release time (current time + delay) in state

Because this functionally single-threads the output, it may be too slow to achieve the target rate; if that's the case, in step (1) you can partition to N keys, and do the same thing for each of them, yielding at a rate of rate_per_sec/N

damccorm · 2026-02-19T15:22:25Z

+  def run_inference(
+      self, batch: List[Dict[str, Any]], model, inference_args=None):
+
+    if model is not None:


When could model be none?

damccorm · 2026-02-19T15:22:39Z

+    if model is not None:
+      self._model = model
+      self._model.to(self.device)
+      self._model.eval()


Haven't we already called this in load_model?

damccorm · 2026-02-19T15:24:17Z

+      self._model.eval()
+    if self._processor is None:
+      from transformers import BlipProcessor
+      self._processor = BlipProcessor.from_pretrained(self.model_name)


A better pattern here might just be to return a class which contains both the processor and the model from load_model

damccorm · 2026-02-19T15:25:48Z

+    return "blip_captioning"
+
+
+class ClipRankModelHandler(ModelHandler):


Same general comments as the Blip model handler apply here

damccorm · 2026-02-19T15:38:14Z

+# ============ DoFns ============
+
+
+class RateLimitDoFn(beam.DoFn):


I have the same questions about this file as the previous one - also, can we split the shared functions out into a helper class?

damccorm · 2026-02-19T15:38:48Z

+    return torch.from_numpy(arr).float()  # float32, shape (3,224,224)
+
+
+class RateLimitDoFn(beam.DoFn):


Same general questions about this file

damccorm · 2026-02-19T15:42:09Z

/gemini review

gemini-code-assist

Code Review

This pull request introduces three new ML inference pipelines for image classification, object detection, and image captioning using PyTorch, along with their corresponding benchmarks and documentation. The pipelines are well-structured and showcase advanced Beam features like RunInference with custom model handlers and stateful DoFns. My review focuses on improving scalability, robustness, and maintainability. I've identified a few key areas for improvement, including a scalability bottleneck in the data loading pipelines, several instances of broad exception handling that could mask errors, some potentially buggy logic, and a few copy-paste errors in the new documentation pages. Overall, this is a valuable contribution, and the suggested changes aim to make these examples more robust and easier to understand.

github-actions · 2026-02-27T12:23:36Z

Reminder, please take a look at this pr: @claudevdm @liferoad @shunping

github-actions · 2026-03-03T12:25:33Z

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @jrmccluskey for label python.
R: @damccorm for label build.
R: @damccorm for label website.

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)

github-actions · 2026-03-10T12:25:49Z

Reminder, please take a look at this pr: @jrmccluskey @damccorm

github-actions · 2026-03-13T12:25:26Z

Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment assign to next reviewer:

R: @shunping for label python.
R: @liferoad for label build.
R: @kennknowles for label website.

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)

damccorm · 2026-03-19T20:14:14Z

waiting on author

ML pipelines: RunInference - OSS Image Object detection

6a09efb

github-actions Bot added python build website infra examples labels Dec 25, 2025

Amar3tto added 2 commits December 31, 2025 15:30

Refactoring

281509a

Add PyTorch Image Captioning BLIP + CLIP Batch

a22d694

Amar3tto changed the title ~~ML pipelines: RunInference - OSS Image Object detection~~ ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning Dec 31, 2025

Amar3tto added 3 commits January 2, 2026 13:22

Refactoring

64187d9

Add PyTorch Image Classification EfficientNet-B0 Streaming (Right-fit…

b7edf62

…ting) ML pipeline

Fix lint

006f650

Amar3tto changed the title ~~ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning~~ ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning, OSS Image Classification Jan 21, 2026

Amar3tto mentioned this pull request Jan 21, 2026

ML pipelines: RunInference - OSS Image Classification #37082

Closed

3 tasks

Amar3tto added 14 commits January 21, 2026 21:58

Rename files

a3dafe4

Add project argument

0a80470

Fix inference_fn

e398220

Fix inference result

169b022

Update torch for image captioning

7dd82b3

Fix inference

1d654cd

Add Looker ids

9dca475

Fix float for imagenet rightfit

126d30b

Add all pipelines to separate workflow

8126ddb

Fix streaming mode

6c01f24

Fix clip blip

1bb5c36

Fix quotes

c76849d

Fix formatting

c5b356d

Fix lint

aa8ac62

Amar3tto requested a review from damccorm February 7, 2026 05:40

Amar3tto marked this pull request as ready for review February 7, 2026 05:41

github-actions Bot added the Next Action: Reviewers label Feb 7, 2026

damccorm reviewed Feb 11, 2026

View reviewed changes

Comment thread website/www/site/data/performance.yaml Outdated

Add Looker IDs

d86d3df

Amar3tto requested a review from damccorm February 12, 2026 14:40

damccorm reviewed Feb 19, 2026

View reviewed changes

gemini-code-assist Bot reviewed Feb 19, 2026

View reviewed changes

github-actions Bot added the slow-review label Feb 27, 2026

github-actions Bot added reassigned-reviewers and removed slow-review labels Mar 3, 2026

github-actions Bot added the slow-review label Mar 10, 2026

Amar3tto added 4 commits March 10, 2026 16:48

Resolve comments [1]

c965879

Resolve comments [2]

00ddd57

Resolve comments [3]

66b6452

Resolve comments [4]

0148ee2

github-actions Bot removed the slow-review label Mar 13, 2026

github-actions Bot added Next Action: Author and removed Next Action: Reviewers labels Mar 19, 2026

		# ============ DoFns ============


		class RateLimitDoFn(beam.DoFn):

		return "blip_captioning"


		class ClipRankModelHandler(ModelHandler):

		return torch.from_numpy(arr).float() # float32, shape (3,224,224)


		class RateLimitDoFn(beam.DoFn):

Conversation

Amar3tto commented Dec 25, 2025

GitHub Actions Tests Status (on master branch)

Uh oh!

gemini-code-assist Bot commented Dec 25, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

codecov Bot commented Dec 31, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

github-actions Bot commented Feb 7, 2026

Uh oh!

Uh oh!

Amar3tto commented Feb 17, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

damccorm commented Feb 19, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Feb 27, 2026

Uh oh!

github-actions Bot commented Mar 3, 2026

Uh oh!

github-actions Bot commented Mar 10, 2026

Uh oh!

github-actions Bot commented Mar 13, 2026

Uh oh!

damccorm commented Mar 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Dec 31, 2025 •

edited

Loading