-
Notifications
You must be signed in to change notification settings - Fork 4.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[BEAM-14044] Allow ModelLoader to forward BatchElements args #17527
Conversation
Can one of the admins verify this patch? |
2 similar comments
Can one of the admins verify this patch? |
Can one of the admins verify this patch? |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think it should be possible to add a unit test. LMK if you need some help with that.
Probably a good consistent unit test to add would be to set min and max size to 2, and then expect sizes of 2 in a mock/fake. I suggest 2 because I think default might be 1.
@@ -80,6 +80,10 @@ def load_model(self) -> T: | |||
def get_inference_runner(self) -> InferenceRunner: | |||
"""Returns an implementation of InferenceRunner for this model.""" | |||
raise NotImplementedError(type(self)) | |||
|
|||
def batch_elements_kwargs(self) -> Mapping[str, Any]: |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You are going to need to import Mapping
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Added
Codecov Report
@@ Coverage Diff @@
## master #17527 +/- ##
==========================================
+ Coverage 73.81% 73.94% +0.13%
==========================================
Files 690 694 +4
Lines 90944 91527 +583
==========================================
+ Hits 67129 67683 +554
- Misses 22595 22624 +29
Partials 1220 1220
Flags with carried forward coverage won't be shown. Click here to find out more.
Continue to review full report at Codecov.
|
Done, added a test that batches 100 elements so that we're very far outside the range of random initialization |
You will need to look through the tests and linter that are failing and fix those issues. If you look at the bottom of the comments you section can see all the tests that are failing. |
Done. It looks like the remaining test failure is likely flaky |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah, those tests that are failing are just flaky.
R: @yeandy |
return FakeInferenceRunnerNeedsBigBatch() | ||
|
||
def batch_elements_kwargs(self): | ||
return {'min_batch_size': 9999} |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If the goal is to be able to set a max batch size of 1, I think it might also be useful to explicitly test the case in which batch_elements_kwargs
returns {max_batch_size=1}
.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It makes the test much more difficult to write, since the actual batch size chosen in that case would be nondeterministic, so it's difficult to distinguish between size 1 batches due to randomness and size 1 batches due to forwarding the correct args.
Run Python PreCommit |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It seems like with the API, selecting batch_elements_kwargs is up to implementers of ModelLoader
/InferenceRunner
.
What if the implementer wants to enable users to set the values as appropriate for their model? Then each implementation would need to decide on a way to expose it, right?
Yes, although the motivating context for this is that we have a particular ModelLoader (in TFX-BSL) for handling pre-batched inputs, and would like a way to limit subsequent batching for all users of that ModelLoader. The goal isn't really to expose knobs on a per-model basis, although you might want to (e.g., if a model is very big). |
Got it, thanks. I suppose we can re-visit some shared way to expose this to end users if the implementations start doing it independently. |
Adds a method to ModelLoader to allow overriding of arguments to beam.BatchElements. The main goal here is to set a max batch size (e.g., 1) when the caller has particular requirements.