Add OOM handling for RunInference with model manager #37557
Conversation
Summary of ChangesHello @AMOOOMA, 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 the robustness of the RunInference transform within Apache Beam's ML inference capabilities. By introducing a dedicated Out-Of-Memory (OOM) protection mechanism, it ensures that model loading and inference execution can gracefully handle CUDA OOM errors, attempting memory cleanup and guaranteeing resource release to prevent system instability and improve overall reliability. 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. Changelog
Activity
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
|
|
R: @damccorm |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
|
/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 valuable improvements for handling out-of-memory errors in RunInference when using the model manager. The OOMProtectedFn wrapper is a neat way to encapsulate the cleanup logic without adding hard dependencies. Making the logging in the model manager optional is also a great change for reducing operational costs and log noise. The added exception handling for model deletion will certainly help with robustness.
My main feedback is on the new test case for OOM handling, which appears to be non-deterministic and doesn't fully exercise the new CUDA OOM cleanup path. I've left a specific comment with a suggestion to make the test more robust and comprehensive.
| class OOMFakeModelHandler(SimpleFakeModelHandler): | ||
| def run_inference( | ||
| self, | ||
| batch: Sequence[int], | ||
| model: FakeModel, | ||
| inference_args=None) -> Iterable[int]: | ||
| if random.random() < 0.8: | ||
| raise MemoryError("Simulated OOM") | ||
| for example in batch: | ||
| yield model.predict(example) | ||
|
|
||
| with self.assertRaises(Exception): | ||
| with TestPipeline() as pipeline: | ||
| examples = [1, 5, 3, 10] | ||
| pcoll = pipeline | 'start' >> beam.Create(examples) | ||
| actual = pcoll | base.RunInference( | ||
| OOMFakeModelHandler(), use_model_manager=True) | ||
| assert_that(actual, equal_to([2, 6, 4, 11]), label='assert:inferences') |
There was a problem hiding this comment.
This new test has a few issues that could be improved:
-
Non-determinism: The use of
random.random()makes this test non-deterministic, which can lead to flaky builds. It's better to make tests deterministic. -
Confusing structure: The
assert_thatcall is inside thewith self.assertRaises(Exception):block. This means the assertion is only checked if the test is about to fail anyway because no exception was raised. It's clearer to have separate tests for success and failure cases. -
Incomplete test coverage: The
OOMProtectedFnspecifically looks for'out of memory'and'CUDA'in the exception string to trigger the memory cleanup logic. TheMemoryError("Simulated OOM")raised here does not contain'out of memory', so the cleanup path is not actually being tested.
I'm suggesting a change to make this test deterministic and to correctly test the OOM cleanup path by raising a more specific error. This will make the test more reliable and ensure the new functionality is properly verified.
class OOMFakeModelHandler(SimpleFakeModelHandler):
def run_inference(
self,
batch: Sequence[int],
model: FakeModel,
inference_args=None) -> Iterable[int]:
# This will always raise to test the OOM path.
raise MemoryError("CUDA out of memory. Simulated OOM.")
with self.assertRaises(MemoryError):
with TestPipeline() as pipeline:
examples = [1, 5, 3, 10]
pcoll = pipeline | 'start' >> beam.Create(examples)
# The pipeline will fail, so we don't need to check the output.
_ = pcoll | base.RunInference(
OOMFakeModelHandler(), use_model_manager=True)
There was a problem hiding this comment.
I don't think the suggestion is good, but there is a 20% chance this test succeeds. Could we drop the batch size to 1? That would help since we'd get 4 run_inference calls instead of just one
There was a problem hiding this comment.
Ah yes good catch! I didn't notice the batch size, so was assuming 0.2^4, will update!
|
Outside of the gemini comment thread, this LGTM |
2 participants
Make sure the code is resilient against failures.
Made logging info optional so that we don't incur cost with the model manager by default.
Also tries to reduce testing flakes by adding a try catch statement on model deletion.
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.