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#37186
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? |
| # limitations under the License. | ||
|
|
||
| name: Inference Python Benchmarks Dataflow | ||
| name: Inference Python Benchmarks Dataflow (1 part) |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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?)
| # ============ DoFns ============ | ||
|
|
||
|
|
||
| class RateLimitDoFn(beam.DoFn): |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
| def run_inference( | ||
| self, batch: List[Dict[str, Any]], model, inference_args=None): | ||
|
|
||
| if model is not None: |
There was a problem hiding this comment.
When could model be none?
| if model is not None: | ||
| self._model = model | ||
| self._model.to(self.device) | ||
| self._model.eval() |
There was a problem hiding this comment.
Haven't we already called this in load_model?
| self._model.eval() | ||
| if self._processor is None: | ||
| from transformers import BlipProcessor | ||
| self._processor = BlipProcessor.from_pretrained(self.model_name) |
There was a problem hiding this comment.
A better pattern here might just be to return a class which contains both the processor and the model from load_model
| return "blip_captioning" | ||
|
|
||
|
|
||
| class ClipRankModelHandler(ModelHandler): |
There was a problem hiding this comment.
Same general comments as the Blip model handler apply here
| lines = ( | ||
| pipeline | ||
| | | ||
| 'ReadGCSFile' >> beam.Create(list(read_gcs_file_lines(known_args.input))) |
There was a problem hiding this comment.
Can we just use built in Beam transforms to read from gcs instead of doing it all locally?
There was a problem hiding this comment.
Or alternately, if we're doing it locally, can we just do all of our Pub/Sub hydration without a beam pipeline
| # ============ DoFns ============ | ||
|
|
||
|
|
||
| class RateLimitDoFn(beam.DoFn): |
There was a problem hiding this comment.
I have the same questions about this file as the previous one - also, can we split the shared functions out into a helper class?
| return torch.from_numpy(arr).float() # float32, shape (3,224,224) | ||
|
|
||
|
|
||
| class RateLimitDoFn(beam.DoFn): |
There was a problem hiding this comment.
Same general questions about this file
| # ============ Main pipeline ============ | ||
|
|
||
|
|
||
| def run( |
There was a problem hiding this comment.
I see this is a rightfitting pipeline, but does it actually use resource hints? https://docs.cloud.google.com/dataflow/docs/guides/right-fitting#python
|
/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.
| lines = ( | ||
| pipeline | ||
| | | ||
| 'ReadGCSFile' >> beam.Create(list(read_gcs_file_lines(known_args.input))) |
There was a problem hiding this comment.
Reading the entire input file into a list with list(read_gcs_file_lines(...)) and passing it to beam.Create can cause memory issues on the client submitting the job, especially with large input files. This approach doesn't scale. A better approach is to use beam.io.ReadFromText which creates a distributed PCollection from the file.
| 'ReadGCSFile' >> beam.Create(list(read_gcs_file_lines(known_args.input))) | |
| 'ReadGCSFile' >> beam.io.ReadFromText(known_args.input) |
| except Exception as e: | ||
| last_err = e | ||
| logging.warning("Batch size %s failed during warmup: %s", bs, e) | ||
| continue |
There was a problem hiding this comment.
Catching a broad Exception during the batch size warmup can mask critical errors like model loading failures, configuration issues, or network problems. It's safer to catch more specific exceptions, such as RuntimeError for CUDA out-of-memory errors, and allow other unexpected errors to fail fast.
| lines = ( | ||
| pipeline | ||
| | | ||
| 'ReadGCSFile' >> beam.Create(list(read_gcs_file_lines(known_args.input))) |
There was a problem hiding this comment.
Reading the entire input file into a list with list(read_gcs_file_lines(...)) and passing it to beam.Create can cause memory issues on the client submitting the job, especially with large input files. This approach doesn't scale. A better approach is to use beam.io.ReadFromText which creates a distributed PCollection from the file.
| 'ReadGCSFile' >> beam.Create(list(read_gcs_file_lines(known_args.input))) | |
| 'ReadGCSFile' >> beam.io.ReadFromText(known_args.input) |
| # element is bytes message, assume it includes | ||
| # {"image_id": "...", "bytes": base64?} or just raw bytes. | ||
| import hashlib | ||
| b = element if isinstance(element, (bytes, bytearray)) else bytes(element) |
There was a problem hiding this comment.
If element is a string, bytes(element) will raise a TypeError because it requires an encoding. You probably meant element.encode('utf-8'). The pipeline is constructed such that MakeKeyDoFn always receives a string (URI), so this bytes input mode appears to be unused and would fail if ever invoked with string inputs.
| b = element if isinstance(element, (bytes, bytearray)) else bytes(element) | |
| b = element if isinstance(element, (bytes, bytearray)) else element.encode('utf-8') |
| lines = ( | ||
| pipeline | ||
| | | ||
| 'ReadGCSFile' >> beam.Create(list(read_gcs_file_lines(known_args.input))) |
There was a problem hiding this comment.
Reading the entire input file into a list with list(read_gcs_file_lines(...)) and passing it to beam.Create can cause memory issues on the client submitting the job, especially with large input files. This approach doesn't scale. A better approach is to use beam.io.ReadFromText which creates a distributed PCollection from the file.
| 'ReadGCSFile' >> beam.Create(list(read_gcs_file_lines(known_args.input))) | |
| 'ReadGCSFile' >> beam.io.ReadFromText(known_args.input) |
| This streaming pipeline performs image classification using an open-source PyTorch EfficientNet-B0 model optimized for T4 GPUs. | ||
| It reads image URIs from Pub/Sub, decodes and preprocesses them in parallel, and runs inference with adaptive batch sizing for optimal GPU utilization. | ||
|
|
||
| The following graphs show various metrics when running PyTorch Sentiment Analysis Streaming using Hugging Face DistilBERT base uncased model pipeline. |
There was a problem hiding this comment.
This description seems to be a copy-paste error from another benchmark page. It should describe the Image Classification pipeline, not Sentiment Analysis.
| The following graphs show various metrics when running PyTorch Sentiment Analysis Streaming using Hugging Face DistilBERT base uncased model pipeline. | |
| The following graphs show various metrics when running the PyTorch Image Classification EfficientNet-B0 Streaming (Right-fitting) GPU pipeline. |
| @@ -0,0 +1,44 @@ | |||
| --- | |||
| title: "PyTorch Language Modeling BERT base Performance" | |||
There was a problem hiding this comment.
The title of this page appears to be incorrect. It refers to 'PyTorch Language Modeling BERT base Performance', but the content is about the Image Classification EfficientNet-B0 pipeline.
| title: "PyTorch Language Modeling BERT base Performance" | |
| title: "PyTorch Image Classification EfficientNet-B0 Streaming (Right-fitting, Exactly-once) CPU Performance" |
| It reads image URIs from Pub/Sub, decodes and preprocesses them in parallel, and runs inference with adaptive batch sizing. | ||
| The pipeline ensures exactly-once semantics via stateful deduplication and idempotent BigQuery writes, allowing stable and reproducible performance measurements under continuous load. | ||
|
|
||
| The following graphs show various metrics when running PyTorch Sentiment Analysis Streaming using Hugging Face DistilBERT base uncased model pipeline. |
There was a problem hiding this comment.
This description seems to be a copy-paste error from another benchmark page. It should describe the Image Classification pipeline, not Sentiment Analysis.
| The following graphs show various metrics when running PyTorch Sentiment Analysis Streaming using Hugging Face DistilBERT base uncased model pipeline. | |
| The following graphs show various metrics when running the PyTorch Image Classification EfficientNet-B0 Streaming (Right-fitting, Exactly-once) CPU pipeline. |
| @@ -0,0 +1,44 @@ | |||
| --- | |||
| title: "PyTorch Language Modeling BERT base Performance" | |||
There was a problem hiding this comment.
The title of this page appears to be incorrect. It refers to 'PyTorch Language Modeling BERT base Performance', but the content is about the Image Classification EfficientNet-B0 pipeline.
| title: "PyTorch Language Modeling BERT base Performance" | |
| title: "PyTorch Image Classification EfficientNet-B0 Streaming (Right-fitting, Exactly-once) GPU Performance" |
| It reads image URIs from Pub/Sub, decodes and preprocesses them in parallel, and runs inference with adaptive batch sizing for optimal GPU utilization. | ||
| The pipeline ensures exactly-once semantics via stateful deduplication and idempotent BigQuery writes, allowing stable and reproducible performance measurements under continuous load. | ||
|
|
||
| The following graphs show various metrics when running PyTorch Sentiment Analysis Streaming using Hugging Face DistilBERT base uncased model pipeline. |
There was a problem hiding this comment.
This description seems to be a copy-paste error from another benchmark page. It should describe the Image Classification pipeline, not Sentiment Analysis.
| The following graphs show various metrics when running PyTorch Sentiment Analysis Streaming using Hugging Face DistilBERT base uncased model pipeline. | |
| The following graphs show various metrics when running the PyTorch Image Classification EfficientNet-B0 Streaming (Right-fitting, Exactly-once) GPU pipeline. |
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.