[Feature Request]: Support loading a different model/key for RunInference

[damccorm]

What would you like to happen?

Today, many users have pipelines that choose a single model for inference from 100s or 1000s of models based on properties of the data. Unfortunately, RunInference does not support this use case. We should support a new use case for RunInference that allows a single keyed RunInference transform to serve a different model for each key.

See design doc here - https://docs.google.com/document/d/1kj3FyWRbJu1KhViX07Z0Gk0MU0842jhYRhI-DMhhcv4/edit?usp=sharing

Issue Priority

Priority: 2 (default / most feature requests should be filed as P2)

Issue Components

• Component: Python SDK
• Component: Java SDK
• Component: Go SDK
• Component: Typescript SDK
• Component: IO connector
• Component: Beam examples
• Component: Beam playground
• Component: Beam katas
• Component: Website
• Component: Spark Runner
• Component: Flink Runner
• Component: Samza Runner
• Component: Twister2 Runner
• Component: Hazelcast Jet Runner
• Component: Google Cloud Dataflow Runner

[damccorm]

Remaining work:

• Add support in the model handler - Add ability to run per key inference #27857
• Add an example and E2E test - Add per key inference integration test #28026
• Model updates - Per key model updates #28161
• Metrics - Implemented in Shell out to model handlers to collect byte sizes #28182 and Keep track of metrics at the KeyedModelHandler level #28228
• Add mechanism for limiting memory consumption - Add support for limiting number of models in memory #28263
• Docs - Add docs for per key inference #28243
• Runnable notebook - Add notebook for per key models #28327