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 32 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
0148ee2 to
7dfdc4a
Compare
|
Reminder, please take a look at this pr: @shunping @liferoad @kennknowles |
|
waiting on author |
github-actions
Bot
added
Next Action: Author
and removed
Next Action: Reviewers
slow-review
labels
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces several new PyTorch-based inference example pipelines, including image captioning, object detection, and image classification, along with associated benchmarks and performance documentation. Feedback from the reviewer highlighted a critical architectural issue where adaptive batch sizing logic was incorrectly placed on the launcher instead of the workers. Additionally, the reviewer suggested optimizing model loading by moving device placement to the load_model method, improving error handling by avoiding broad exceptions, and replacing print statements with proper logging to ensure compatibility with Beam runners.
| with FileSystems.open(uri) as f: | ||
| image_bytes = f.read() | ||
| yield uri, {"image_bytes": image_bytes} | ||
| except Exception as e: |
There was a problem hiding this comment.
Catching a broad Exception is discouraged as it can mask unexpected errors and make debugging difficult. Consider catching more specific exceptions (e.g., IOError, GoogleCloudError) or logging the exception details with the stack trace.
| def load_model(self): | ||
| from transformers import BlipForConditionalGeneration, BlipProcessor | ||
| processor = BlipProcessor.from_pretrained(self.model_name) | ||
| model = BlipForConditionalGeneration.from_pretrained(self.model_name) | ||
| return (model, processor) |
There was a problem hiding this comment.
Moving the model to the device and setting it to evaluation mode should be done once during model loading rather than in every inference call to improve performance.
| def load_model(self): | |
| from transformers import BlipForConditionalGeneration, BlipProcessor | |
| processor = BlipProcessor.from_pretrained(self.model_name) | |
| model = BlipForConditionalGeneration.from_pretrained(self.model_name) | |
| return (model, processor) | |
| def load_model(self): | |
| from transformers import BlipForConditionalGeneration, BlipProcessor | |
| processor = BlipProcessor.from_pretrained(self.model_name) | |
| model = BlipForConditionalGeneration.from_pretrained(self.model_name) | |
| model.to(self.device) | |
| model.eval() | |
| return (model, processor) |
| model, processor = model_bundle | ||
| model.to(self.device) | ||
| model.eval() |
There was a problem hiding this comment.
Redundant model device placement and evaluation mode setting. These operations should be handled in load_model.
| model, processor = model_bundle | |
| model.to(self.device) | |
| model.eval() | |
| model, processor = model_bundle |
| def load_model(self): | ||
| from transformers import CLIPModel, CLIPProcessor | ||
| processor = CLIPProcessor.from_pretrained(self.model_name) | ||
| model = CLIPModel.from_pretrained(self.model_name) | ||
| return (model, processor) |
There was a problem hiding this comment.
Moving the model to the device and setting it to evaluation mode should be done once during model loading to avoid redundant operations during inference.
| def load_model(self): | |
| from transformers import CLIPModel, CLIPProcessor | |
| processor = CLIPProcessor.from_pretrained(self.model_name) | |
| model = CLIPModel.from_pretrained(self.model_name) | |
| return (model, processor) | |
| def load_model(self): | |
| from transformers import CLIPModel, CLIPProcessor | |
| processor = CLIPProcessor.from_pretrained(self.model_name) | |
| model = CLIPModel.from_pretrained(self.model_name) | |
| model.to(self.device) | |
| model.eval() | |
| return (model, processor) |
| model, processor = model_bundle | ||
| model.to(self.device) | ||
| model.eval() |
There was a problem hiding this comment.
Redundant model device placement and evaluation mode setting. These operations should be handled in load_model.
| model, processor = model_bundle | |
| model.to(self.device) | |
| model.eval() | |
| model, processor = model_bundle |
| try: | ||
| subscriber.delete_subscription( | ||
| request={"subscription": full_subscription_path}) | ||
| print(f"Deleted subscription: {subscription_name}") |
There was a problem hiding this comment.
Using print for logging in a Beam pipeline is discouraged as it may not be captured correctly by the runner's logging system. Use the logging module instead.
| print(f"Deleted subscription: {subscription_name}") | |
| logging.info(f"Deleted subscription: {subscription_name}") |
| if logits is None: | ||
| # fallback: try first value if dict shape differs | ||
| try: | ||
| logits = next(iter(inference_obj.values())) |
There was a problem hiding this comment.
Falling back to the first value in the model output dictionary is fragile. If the model returns multiple outputs (e.g., auxiliary logits or metadata), this might select the wrong tensor. It is safer to require an explicit key or handle known output structures for the supported models.
| | 'RunInference' >> RunInference(KeyedModelHandler(model_handler))) | ||
| | 'RunInference' >> RunInference( | ||
| KeyedModelHandler(model_handler)).with_resource_hints( | ||
| accelerator="type:nvidia-tesla-t4;count:1;install-nvidia-driver")) |
There was a problem hiding this comment.
Thanks for adding the resource hint. However, because of how resource hints work, this will currently add a GPU for the whole pipeline. The way resource hints work is they ask for the requested resource for the whole stage containing this step. Because this pipeline doesn't contain any shuffles, it will all end up being in the same stage, so it will all get the T4.
To fix this, could you please add a Reshuffle step directly before RunInference?
I'm not sure if this will make the pipeline run more efficiently or not; because it is primarily bound on IO, not compute, the benefits might not outweigh the shuffle costs. That is good for us to understand regardless, though.
|
Thanks - could you please address my comment and gemini's? Then I think this should be good to go |
|
waiting on author |
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.