[BEAM-13970] RunInference V1

• Type: New Feature
• Status: Open
• Priority: P2
• Resolution: Unresolved
• Component/s: sdk-py-core
• Labels: run-inference
• Assignee: Unassigned
• Reporter: Andy Ye
• Created: 17/Feb/22 18:51
• Updated: 04/Jun/22 23:31

Description

Users of machine learning frameworks must currently implement their own transforms for running ML inferences. The exception is the TensorFlow RunInference transform. However, this is hosted in its own repo, and has an API that is exclusively geared towards the TensorFlow TFX library. Our goal is to add new implementations of RunInference for the two other popular machine learning frameworks: scikit-learn and Pytorch.

Please see main design document here.

Sub-Tasks

• Provide GPU support: Open
• Add Metrics to base class: Open
• Hook In Batching DoFn Apis to RunInference: Open
• Add documentation page to Beam website: Open
• RunInference Benchmarking tests: Open
• Implement E2E tests for RunInference classes: Open
• Investigate releasing models in inference base class: Open
• Add batch_size back off in RunInferenceBase: Open
• Figure out how type hints should work: Open
• Add resource location hints to base RunInference Implementation: Open
• Correctly comment or remove Metrics Cache: Open
• Drop the clock abastraction and just use time.time for time measurements: Open
• Remove base_run_inference in TFX-BSL and use inference.base in Beam: Open
• Support **kwargs for PyTorch models.: In Progress
• Explore versions of pytorch to test in Tox: Open
• Investigate windows Permissions Error on Jenkins Windows Unit tests: Open
• Add flag to drop example from PredicitonResult: Open
• Investigate load state_dict vs loading whole model: Open
• Add support for Pandas Dataframes to sklearn RunInference Implementation: Open