damccorm commented on code in PR #28182:
URL: https://github.com/apache/beam/pull/28182#discussion_r1307691690
##########
sdks/python/apache_beam/ml/inference/base_test.py:
##########
@@ -319,6 +326,24 @@ def mult_two(example: str) -> int:
with self.assertRaises(ValueError):
base.KeyedModelHandler(mhs)
+ def test_keyed_model_handler_get_num_bytes(self):
+ mh = base.KeyedModelHandler(FakeModelHandler(num_bytes_per_element=10))
+ batch = [('key1', 1), ('key2', 2), ('key1', 3)]
+ expected = len(pickle.dumps(('key1', 'key2', 'key1'))) + 30
+ actual = mh.get_num_bytes(batch)
+ self.assertEqual(expected, actual)
+
+ def test_keyed_model_handler_multiple_models_get_num_bytes(self):
+ mhs = [
+ base.KeyMhMapping(['key1'],
FakeModelHandler(num_bytes_per_element=10)),
Review Comment:
Yes, I'm waiting until I don't have in flight PRs around this to change it
to avoid conflicts (right now https://github.com/apache/beam/pull/28026 uses
KeyMhMapping)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
