ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning, OSS Image Classification#37186
ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning, OSS Image Classification#37186Amar3tto wants to merge 39 commits into
Conversation
Summary of ChangesHello @Amar3tto, I'm Gemini Code Assist1! 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 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
Using Gemini Code AssistThe 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
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 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. Footnotes
|
Amar3tto
changed the title
Codecov Report❌ Patch coverage is
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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Amar3tto
changed the title
|
Assigning reviewers: R: @claudevdm for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
|
@Abacn Could you please help with review? |
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
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.
|
Reminder, please take a look at this pr: @claudevdm @liferoad @shunping |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @jrmccluskey for label python. Available commands:
|
|
Reminder, please take a look at this pr: @jrmccluskey @damccorm |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @shunping for label python. Available commands:
|
|
waiting on author |
Amar3tto
force-pushed
the
oss-image-detection
branch
from
a7ba9db to
439b5aa
Compare
|
Could you please fix the formatting failures? Also, please avoid rebasing when possible to avoid breaking GitHub's review features |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces several new PyTorch-based inference pipelines (image captioning, object detection, and image classification with adaptive batch sizing) along with their corresponding benchmarks, requirements, and performance documentation. The review feedback highlights critical bugs regarding shape mismatches during batching in the classification and object detection pipelines, a performance bottleneck from redundant image encoding in the CLIP model, an architectural flaw where model warmup runs on the submission client instead of the workers, and a robustness issue concerning guaranteed cleanup of Pub/Sub resources.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
/gemini review |
| self._batch_size = bs | ||
| self._inference_batch_size = bs |
There was a problem hiding this comment.
Why is this nondeterministic? Can't we set this up front based on what parameters a user passes in?
There was a problem hiding this comment.
When the user provides an explicit batch size, it is set up front through inference_batch_size. The runtime selection is only used for auto, because the submission client cannot know which batch size fits on the actual Dataflow worker GPU. The selection itself is deterministic: it tries the configured sizes in order and picks the first one that fits.
There was a problem hiding this comment.
Ideally we just wouldn't need this for any of the examples. The downside here is that it adds complexity and it doesn't feed into the batching done by BatchElements (which seems like a better place for this)
There was a problem hiding this comment.
Code Review
This pull request adds several new PyTorch-based inference pipelines and benchmarks to Apache Beam's Python SDK, covering image captioning (BLIP + CLIP), image object detection (Faster R-CNN ResNet-50), and image classification (EfficientNet-B0 with right-fitting), along with corresponding documentation and performance tracking configurations. The review feedback highlights several critical issues and improvement opportunities: a potential GPU OOM in the right-fitting handler due to runner-side batching which requires sub-batching in run_inference; hardcoded GPU resource hints that should be conditionally applied; potential mismatches when parsing Pub/Sub paths; race conditions from modifying shared pipeline_args in background threads; and opportunities to simplify device transfers using Hugging Face's built-in .to(device) method.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
waiting on author |
github-actions
Bot
added
Next Action: Author
and removed
Next Action: Reviewers
slow-review
labels
|
I will resolve the conflicts later to avoid breaking the reviewed code |
|
This mostly LGTM, my only lingering concerns are around the batching
Thanks! I will note that you can do this without rebasing if you create a merge commit instead: That will allow you to resolve the conflicts without breaking the review feature |
…age-detection # Conflicts: # .github/workflows/beam_Inference_Python_Benchmarks_Dataflow.yml # website/www/site/data/performance.yaml
damccorm
left a comment
damccorm
left a comment
There was a problem hiding this comment.
Thanks! Could you fix the formatter precommit? Then I can merge
2 participants
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:
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, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.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.