[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.
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.
You are going to need to import Mapping
zwestrick
changed the title
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 |
|
R: @yeandy |
| return FakeInferenceRunnerNeedsBigBatch() | ||
|
|
||
| def batch_elements_kwargs(self): | ||
| return {'min_batch_size': 9999} |
There was a problem hiding this comment.
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.
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.
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.