damccorm commented on code in PR #23506:
URL: https://github.com/apache/beam/pull/23506#discussion_r989168216
##########
sdks/python/apache_beam/ml/inference/base_test.py:
##########
@@ -251,6 +255,59 @@ def
test_run_inference_keyed_examples_with_unkeyed_model_handler(self):
| 'RunKeyed' >> base.RunInference(model_handler))
pipeline.run()
+ def test_model_handler_compatibility(self):
+ # ** IMPORTANT ** Do not change this test to make your PR pass without
+ # first reading below.
+ # Be certain that the modification will not break third party
+ # implementations of ModelHandler.
+ # See issue https://github.com/apache/beam/issues/23484
+ # If this test fails, likely third party implementations of
+ # ModelHandler will break.
+ class ThirdPartyHandler(base.ModelHandler[int, int, FakeModel]):
+ def __init__(self, custom_parameter=None):
+ pass
+
+ def load_model(self) -> FakeModel:
+ return FakeModel()
+
+ def run_inference(
+ self,
+ batch: Sequence[int],
+ model: FakeModel,
+ inference_args: Optional[Dict[str, Any]] = None) -> Iterable[int]:
+ yield 0
+
+ def get_num_bytes(self, batch: Sequence[int]) -> int:
+ return 1
+
+ def get_metrics_namespace(self) -> str:
+ return 'ThirdParty'
+
+ def get_resource_hints(self) -> dict:
+ return {}
+
+ def batch_elements_kwargs(self) -> Mapping[str, Any]:
+ return {}
+
+ def validate_inference_args(
+ self, inference_args: Optional[Dict[str, Any]]):
+ pass
+
+ # This test passes if calling these methods does not cause
+ # any runtime exceptions.
+ third_party_model_handler = ThirdPartyHandler(custom_parameter=0)
+ fake_model = third_party_model_handler.load_model()
+ third_party_model_handler.run_inference([], fake_model)
+ fake_inference_args = {some_arg: 1}
Review Comment:
```suggestion
fake_inference_args = {'some_arg': 1}
```
Looks like this was causing the GHA to fail as well.
--
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]
