You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by da...@apache.org on 2023/05/16 19:59:31 UTC
[beam] 01/01: allow model handlers to consume multi_process_shared for large models
This is an automated email from the ASF dual-hosted git repository.
damccorm pushed a commit to branch users/damccorm/consumeMps
in repository https://gitbox.apache.org/repos/asf/beam.git
commit 016c967e14a1c011fa6d67e19ab96976c1f63e1e
Author: Danny McCormick <da...@google.com>
AuthorDate: Tue May 16 15:59:11 2023 -0400
allow model handlers to consume multi_process_shared for large models
---
.../apache_beam/ml/inference/onnx_inference.py | 9 +++
.../ml/inference/onnx_inference_test.py | 32 ++++++++
.../apache_beam/ml/inference/pytorch_inference.py | 19 +++++
.../ml/inference/pytorch_inference_test.py | 87 +++++++++++++++++++++
.../apache_beam/ml/inference/sklearn_inference.py | 18 +++++
.../ml/inference/sklearn_inference_test.py | 68 +++++++++++++++++
.../ml/inference/tensorflow_inference.py | 18 +++++
.../ml/inference/tensorflow_inference_test.py | 89 ++++++++++++++++++++++
.../apache_beam/ml/inference/tensorrt_inference.py | 9 +++
.../ml/inference/tensorrt_inference_test.py | 62 +++++++++++++++
10 files changed, 411 insertions(+)
diff --git a/sdks/python/apache_beam/ml/inference/onnx_inference.py b/sdks/python/apache_beam/ml/inference/onnx_inference.py
index baeb5f3a445..18e115a6188 100644
--- a/sdks/python/apache_beam/ml/inference/onnx_inference.py
+++ b/sdks/python/apache_beam/ml/inference/onnx_inference.py
@@ -62,6 +62,7 @@ class OnnxModelHandlerNumpy(ModelHandler[numpy.ndarray,
provider_options=None,
*,
inference_fn: NumpyInferenceFn = default_numpy_inference_fn,
+ large_model: bool = False,
**kwargs):
""" Implementation of the ModelHandler interface for onnx
using numpy arrays as input.
@@ -75,6 +76,10 @@ class OnnxModelHandlerNumpy(ModelHandler[numpy.ndarray,
model_uri: The URI to where the model is saved.
inference_fn: The inference function to use on RunInference calls.
default=default_numpy_inference_fn
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
"""
@@ -84,6 +89,7 @@ class OnnxModelHandlerNumpy(ModelHandler[numpy.ndarray,
self._provider_options = provider_options
self._model_inference_fn = inference_fn
self._env_vars = kwargs.get('env_vars', {})
+ self._large_model = large_model
def load_model(self) -> ort.InferenceSession:
"""Loads and initializes an onnx inference session for processing."""
@@ -134,3 +140,6 @@ class OnnxModelHandlerNumpy(ModelHandler[numpy.ndarray,
A namespace for metrics collected by the RunInference transform.
"""
return 'BeamML_Onnx'
+
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
diff --git a/sdks/python/apache_beam/ml/inference/onnx_inference_test.py b/sdks/python/apache_beam/ml/inference/onnx_inference_test.py
index 71ca307c42f..1c09a12462f 100644
--- a/sdks/python/apache_beam/ml/inference/onnx_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/onnx_inference_test.py
@@ -172,6 +172,7 @@ class TestOnnxModelHandler(OnnxModelHandlerNumpy):
provider_options=None,
*,
inference_fn=default_numpy_inference_fn,
+ large_model=False,
**kwargs):
self._model_uri = model_uri
self._session_options = session_options
@@ -179,6 +180,7 @@ class TestOnnxModelHandler(OnnxModelHandlerNumpy):
self._provider_options = provider_options
self._model_inference_fn = inference_fn
self._env_vars = kwargs.get('env_vars', {})
+ self._large_model = large_model
class OnnxTestBase(unittest.TestCase):
@@ -344,6 +346,36 @@ class OnnxPytorchRunInferencePipelineTest(OnnxTestBase):
self.assertTrue('FOO' in os.environ)
self.assertTrue('bar'.equals(os.environ['FOO']))
+ def test_model_handler_large_model(self):
+ with TestPipeline() as pipeline:
+
+ def onxx_numpy_inference_fn(
+ inference_session: ort.InferenceSession, batch, inference_args=None):
+ multi_process_shared_loaded = "multi_process_shared" in str(
+ type(inference_session))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded model of type {type(model)}, was ' +
+ 'expecting multi_process_shared_model')
+ return default_numpy_inference_fn(
+ inference_session, batch, inference_args)
+
+ path = os.path.join(self.tmpdir, 'my_onnx_pytorch_path')
+ model = self.test_data_and_model.get_torch_two_feature_model()
+ self.exportModelToOnnx(model, path)
+ model_handler = OnnxModelHandlerNumpy(
+ model_uri=path,
+ env_vars={'FOO': 'bar'},
+ inference_fn=onxx_numpy_inference_fn,
+ large_model=True)
+ self.assertFalse('FOO' in os.environ)
+ _ = (
+ pipeline
+ | 'start' >> beam.Create(
+ self.test_data_and_model.get_two_feature_examples())
+ | RunInference(model_handler))
+ pipeline.run()
+
@unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed')
def test_pipeline_gcs_model(self):
with TestPipeline() as pipeline:
diff --git a/sdks/python/apache_beam/ml/inference/pytorch_inference.py b/sdks/python/apache_beam/ml/inference/pytorch_inference.py
index 14abc202bdf..d8ab10725fe 100644
--- a/sdks/python/apache_beam/ml/inference/pytorch_inference.py
+++ b/sdks/python/apache_beam/ml/inference/pytorch_inference.py
@@ -190,6 +190,7 @@ class PytorchModelHandlerTensor(ModelHandler[torch.Tensor,
torch_script_model_path: Optional[str] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
+ large_model: bool = False,
**kwargs):
"""Implementation of the ModelHandler interface for PyTorch.
@@ -222,6 +223,10 @@ class PytorchModelHandlerTensor(ModelHandler[torch.Tensor,
batch will be fed into the inference_fn as a Sequence of Tensors.
max_batch_size: the maximum batch size to use when batching inputs. This
batch will be fed into the inference_fn as a Sequence of Tensors.
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -245,6 +250,7 @@ class PytorchModelHandlerTensor(ModelHandler[torch.Tensor,
self._batching_kwargs['max_batch_size'] = max_batch_size
self._torch_script_model_path = torch_script_model_path
self._env_vars = kwargs.get('env_vars', {})
+ self._large_model = large_model
_validate_constructor_args(
state_dict_path=self._state_dict_path,
@@ -324,6 +330,9 @@ class PytorchModelHandlerTensor(ModelHandler[torch.Tensor,
def batch_elements_kwargs(self):
return self._batching_kwargs
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
+
def default_keyed_tensor_inference_fn(
batch: Sequence[Dict[str, torch.Tensor]],
@@ -404,6 +413,7 @@ class PytorchModelHandlerKeyedTensor(ModelHandler[Dict[str, torch.Tensor],
torch_script_model_path: Optional[str] = None,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
+ large_model: bool = False,
**kwargs):
"""Implementation of the ModelHandler interface for PyTorch.
@@ -441,6 +451,10 @@ class PytorchModelHandlerKeyedTensor(ModelHandler[Dict[str, torch.Tensor],
batch will be fed into the inference_fn as a Sequence of Keyed Tensors.
max_batch_size: the maximum batch size to use when batching inputs. This
batch will be fed into the inference_fn as a Sequence of Keyed Tensors.
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -464,6 +478,8 @@ class PytorchModelHandlerKeyedTensor(ModelHandler[Dict[str, torch.Tensor],
self._batching_kwargs['max_batch_size'] = max_batch_size
self._torch_script_model_path = torch_script_model_path
self._env_vars = kwargs.get('env_vars', {})
+ self._large_model = large_model
+
_validate_constructor_args(
state_dict_path=self._state_dict_path,
model_class=self._model_class,
@@ -543,3 +559,6 @@ class PytorchModelHandlerKeyedTensor(ModelHandler[Dict[str, torch.Tensor],
def batch_elements_kwargs(self):
return self._batching_kwargs
+
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
diff --git a/sdks/python/apache_beam/ml/inference/pytorch_inference_test.py b/sdks/python/apache_beam/ml/inference/pytorch_inference_test.py
index 76eea44fe55..db4bf65e269 100644
--- a/sdks/python/apache_beam/ml/inference/pytorch_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/pytorch_inference_test.py
@@ -491,6 +491,45 @@ class PytorchRunInferencePipelineTest(unittest.TestCase):
equal_to(
TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result))
+ def test_pipeline_local_model_large(self):
+ with TestPipeline() as pipeline:
+
+ def batch_validator_tensor_inference_fn(
+ batch,
+ model,
+ device,
+ inference_args,
+ model_id,
+ ):
+ multi_process_shared_loaded = "multi_process_shared" in str(type(model))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded model of type {type(model)}, was ' +
+ 'expecting multi_process_shared_model')
+ return default_tensor_inference_fn(
+ batch, model, device, inference_args, model_id)
+
+ state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0, 3]])),
+ ('linear.bias', torch.Tensor([0.5]))])
+ path = os.path.join(self.tmpdir, 'my_state_dict_path')
+ torch.save(state_dict, path)
+
+ model_handler = PytorchModelHandlerTensor(
+ state_dict_path=path,
+ model_class=PytorchLinearRegression,
+ model_params={
+ 'input_dim': 2, 'output_dim': 1
+ },
+ inference_fn=batch_validator_tensor_inference_fn,
+ large_model=True)
+
+ pcoll = pipeline | 'start' >> beam.Create(TWO_FEATURES_EXAMPLES)
+ predictions = pcoll | RunInference(model_handler)
+ assert_that(
+ predictions,
+ equal_to(
+ TWO_FEATURES_PREDICTIONS, equals_fn=_compare_prediction_result))
+
def test_pipeline_local_model_extra_inference_args(self):
with TestPipeline() as pipeline:
inference_args = {
@@ -522,6 +561,54 @@ class PytorchRunInferencePipelineTest(unittest.TestCase):
equal_to(
KEYED_TORCH_PREDICTIONS, equals_fn=_compare_prediction_result))
+ def test_pipeline_local_model_extra_inference_args_large(self):
+ with TestPipeline() as pipeline:
+ inference_args = {
+ 'prediction_param_array': torch.from_numpy(
+ np.array([1, 2], dtype="float32")),
+ 'prediction_param_bool': True
+ }
+
+ state_dict = OrderedDict([('linear.weight', torch.Tensor([[2.0]])),
+ ('linear.bias', torch.Tensor([0.5]))])
+ path = os.path.join(self.tmpdir, 'my_state_dict_path')
+ torch.save(state_dict, path)
+
+ def batch_validator_keyed_tensor_inference_fn(
+ batch,
+ model,
+ device,
+ inference_args,
+ model_id,
+ ):
+ multi_process_shared_loaded = "multi_process_shared" in str(type(model))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded model of type {type(model)}, was ' +
+ 'expecting multi_process_shared_model')
+ return default_keyed_tensor_inference_fn(
+ batch, model, device, inference_args, model_id)
+
+ model_handler = PytorchModelHandlerKeyedTensor(
+ state_dict_path=path,
+ model_class=PytorchLinearRegressionKeyedBatchAndExtraInferenceArgs,
+ model_params={
+ 'input_dim': 1, 'output_dim': 1
+ },
+ inference_fn=batch_validator_keyed_tensor_inference_fn,
+ large_model=True)
+
+ pcoll = pipeline | 'start' >> beam.Create(KEYED_TORCH_EXAMPLES)
+ inference_args_side_input = (
+ pipeline | 'create side' >> beam.Create(inference_args))
+ predictions = pcoll | RunInference(
+ model_handler=model_handler,
+ inference_args=beam.pvalue.AsDict(inference_args_side_input))
+ assert_that(
+ predictions,
+ equal_to(
+ KEYED_TORCH_PREDICTIONS, equals_fn=_compare_prediction_result))
+
def test_pipeline_local_model_extra_inference_args_batching_args(self):
with TestPipeline() as pipeline:
inference_args = {
diff --git a/sdks/python/apache_beam/ml/inference/sklearn_inference.py b/sdks/python/apache_beam/ml/inference/sklearn_inference.py
index 2c17dcf4494..c2bd2cee66e 100644
--- a/sdks/python/apache_beam/ml/inference/sklearn_inference.py
+++ b/sdks/python/apache_beam/ml/inference/sklearn_inference.py
@@ -90,6 +90,7 @@ class SklearnModelHandlerNumpy(ModelHandler[numpy.ndarray,
inference_fn: NumpyInferenceFn = _default_numpy_inference_fn,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
+ large_model: bool = False,
**kwargs):
""" Implementation of the ModelHandler interface for scikit-learn
using numpy arrays as input.
@@ -110,6 +111,10 @@ class SklearnModelHandlerNumpy(ModelHandler[numpy.ndarray,
max_batch_size: the maximum batch size to use when batching inputs. This
batch will be fed into the inference_fn as a Sequence of Numpy
ndarrays.
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
"""
@@ -122,6 +127,7 @@ class SklearnModelHandlerNumpy(ModelHandler[numpy.ndarray,
if max_batch_size is not None:
self._batching_kwargs['max_batch_size'] = max_batch_size
self._env_vars = kwargs.get('env_vars', {})
+ self._large_model = large_model
def load_model(self) -> BaseEstimator:
"""Loads and initializes a model for processing."""
@@ -174,6 +180,9 @@ class SklearnModelHandlerNumpy(ModelHandler[numpy.ndarray,
def batch_elements_kwargs(self):
return self._batching_kwargs
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
+
PandasInferenceFn = Callable[
[BaseEstimator, Sequence[pandas.DataFrame], Optional[Dict[str, Any]]], Any]
@@ -203,6 +212,7 @@ class SklearnModelHandlerPandas(ModelHandler[pandas.DataFrame,
inference_fn: PandasInferenceFn = _default_pandas_inference_fn,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
+ large_model: bool = False,
**kwargs):
"""Implementation of the ModelHandler interface for scikit-learn that
supports pandas dataframes.
@@ -226,6 +236,10 @@ class SklearnModelHandlerPandas(ModelHandler[pandas.DataFrame,
max_batch_size: the maximum batch size to use when batching inputs. This
batch will be fed into the inference_fn as a Sequence of Pandas
Dataframes.
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
"""
@@ -238,6 +252,7 @@ class SklearnModelHandlerPandas(ModelHandler[pandas.DataFrame,
if max_batch_size is not None:
self._batching_kwargs['max_batch_size'] = max_batch_size
self._env_vars = kwargs.get('env_vars', {})
+ self._large_model = large_model
def load_model(self) -> BaseEstimator:
"""Loads and initializes a model for processing."""
@@ -291,3 +306,6 @@ class SklearnModelHandlerPandas(ModelHandler[pandas.DataFrame,
def batch_elements_kwargs(self):
return self._batching_kwargs
+
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
diff --git a/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py b/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py
index ba253d4a6cd..c2ea9fa1e95 100644
--- a/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/sklearn_inference_test.py
@@ -314,6 +314,38 @@ class SkLearnRunInferenceTest(unittest.TestCase):
assert_that(
actual, equal_to(expected, equals_fn=_compare_prediction_result))
+ def test_pipeline_pickled_large_model(self):
+ temp_file_name = self.tmpdir + os.sep + 'pickled_file'
+ with open(temp_file_name, 'wb') as file:
+ pickle.dump(build_model(), file)
+
+ def large_model_validator_numpy_inference_fn(
+ model: BaseEstimator,
+ batch: Sequence[numpy.ndarray],
+ inference_args: Optional[Dict[str, Any]] = None) -> Any:
+ multi_process_shared_loaded = "multi_process_shared" in str(type(model))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded model of type {type(model)}, was ' +
+ 'expecting multi_process_shared_model')
+ return _default_numpy_inference_fn(model, batch, inference_args)
+
+ with TestPipeline() as pipeline:
+ examples = [numpy.array([0, 0]), numpy.array([1, 1])]
+
+ pcoll = pipeline | 'start' >> beam.Create(examples)
+ actual = pcoll | RunInference(
+ SklearnModelHandlerNumpy(
+ model_uri=temp_file_name,
+ inference_fn=large_model_validator_numpy_inference_fn,
+ large_model=True))
+ expected = [
+ PredictionResult(numpy.array([0, 0]), 0),
+ PredictionResult(numpy.array([1, 1]), 1)
+ ]
+ assert_that(
+ actual, equal_to(expected, equals_fn=_compare_prediction_result))
+
def test_pipeline_joblib(self):
temp_file_name = self.tmpdir + os.sep + 'joblib_file'
with open(temp_file_name, 'wb') as file:
@@ -444,6 +476,42 @@ class SkLearnRunInferenceTest(unittest.TestCase):
assert_that(
actual, equal_to(expected, equals_fn=_compare_dataframe_predictions))
+ def test_pipeline_pandas_large_model(self):
+ temp_file_name = self.tmpdir + os.sep + 'pickled_file'
+ with open(temp_file_name, 'wb') as file:
+ pickle.dump(build_pandas_pipeline(), file)
+
+ def large_model_validator_pandas_inference_fn(
+ model: BaseEstimator,
+ batch: Sequence[numpy.ndarray],
+ inference_args: Optional[Dict[str, Any]] = None) -> Any:
+ multi_process_shared_loaded = "multi_process_shared" in str(type(model))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded model of type {type(model)}, was ' +
+ 'expecting multi_process_shared_model')
+ return _default_pandas_inference_fn(model, batch, inference_args)
+
+ with TestPipeline() as pipeline:
+ dataframe = pandas_dataframe()
+ splits = [dataframe.loc[[i]] for i in dataframe.index]
+ pcoll = pipeline | 'start' >> beam.Create(splits)
+ actual = pcoll | RunInference(
+ SklearnModelHandlerPandas(
+ model_uri=temp_file_name,
+ inference_fn=large_model_validator_pandas_inference_fn,
+ large_model=True))
+
+ expected = [
+ PredictionResult(splits[0], 5),
+ PredictionResult(splits[1], 8),
+ PredictionResult(splits[2], 1),
+ PredictionResult(splits[3], 1),
+ PredictionResult(splits[4], 2),
+ ]
+ assert_that(
+ actual, equal_to(expected, equals_fn=_compare_dataframe_predictions))
+
def test_pipeline_pandas_custom_inference(self):
temp_file_name = self.tmpdir + os.sep + 'pickled_file'
with open(temp_file_name, 'wb') as file:
diff --git a/sdks/python/apache_beam/ml/inference/tensorflow_inference.py b/sdks/python/apache_beam/ml/inference/tensorflow_inference.py
index cfa36f05f61..d1f236dc53b 100644
--- a/sdks/python/apache_beam/ml/inference/tensorflow_inference.py
+++ b/sdks/python/apache_beam/ml/inference/tensorflow_inference.py
@@ -102,6 +102,7 @@ class TFModelHandlerNumpy(ModelHandler[numpy.ndarray,
inference_fn: TensorInferenceFn = default_numpy_inference_fn,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
+ large_model: bool = False,
**kwargs):
"""Implementation of the ModelHandler interface for Tensorflow.
@@ -123,6 +124,10 @@ class TFModelHandlerNumpy(ModelHandler[numpy.ndarray,
once the model is loaded.
inference_fn: inference function to use during RunInference.
Defaults to default_numpy_inference_fn.
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -141,6 +146,7 @@ class TFModelHandlerNumpy(ModelHandler[numpy.ndarray,
self._batching_kwargs['min_batch_size'] = min_batch_size
if max_batch_size is not None:
self._batching_kwargs['max_batch_size'] = max_batch_size
+ self._large_model = large_model
def load_model(self) -> tf.Module:
"""Loads and initializes a Tensorflow model for processing."""
@@ -203,6 +209,9 @@ class TFModelHandlerNumpy(ModelHandler[numpy.ndarray,
def batch_elements_kwargs(self):
return self._batching_kwargs
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
+
class TFModelHandlerTensor(ModelHandler[tf.Tensor, PredictionResult,
tf.Module]):
@@ -217,6 +226,7 @@ class TFModelHandlerTensor(ModelHandler[tf.Tensor, PredictionResult,
inference_fn: TensorInferenceFn = default_tensor_inference_fn,
min_batch_size: Optional[int] = None,
max_batch_size: Optional[int] = None,
+ large_model: bool = False,
**kwargs):
"""Implementation of the ModelHandler interface for Tensorflow.
@@ -239,6 +249,10 @@ class TFModelHandlerTensor(ModelHandler[tf.Tensor, PredictionResult,
once the model is loaded.
inference_fn: inference function to use during RunInference.
Defaults to default_numpy_inference_fn.
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -257,6 +271,7 @@ class TFModelHandlerTensor(ModelHandler[tf.Tensor, PredictionResult,
self._batching_kwargs['min_batch_size'] = min_batch_size
if max_batch_size is not None:
self._batching_kwargs['max_batch_size'] = max_batch_size
+ self._large_model = large_model
def load_model(self) -> tf.Module:
"""Loads and initializes a tensorflow model for processing."""
@@ -319,3 +334,6 @@ class TFModelHandlerTensor(ModelHandler[tf.Tensor, PredictionResult,
def batch_elements_kwargs(self):
return self._batching_kwargs
+
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
diff --git a/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py b/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py
index dca40eb7010..31dde594010 100644
--- a/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/tensorflow_inference_test.py
@@ -170,6 +170,52 @@ class TFRunInferenceTest(unittest.TestCase):
expected_predictions,
equals_fn=_compare_tensor_prediction_result))
+ def test_predict_tensor_with_large_model(self):
+ model = _create_mult2_model()
+ model_path = os.path.join(self.tmpdir, 'mult2')
+ tf.keras.models.save_model(model, model_path)
+ with TestPipeline() as pipeline:
+
+ def fake_batching_inference_fn(
+ model: tf.Module,
+ batch: Union[Sequence[numpy.ndarray], Sequence[tf.Tensor]],
+ inference_args: Dict[str, Any],
+ model_id: Optional[str] = None) -> Iterable[PredictionResult]:
+ multi_process_shared_loaded = "multi_process_shared" in str(type(model))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded model of type {type(model)}, was ' +
+ 'expecting multi_process_shared_model')
+ batch = tf.stack(batch, axis=0)
+ predictions = model(batch)
+ return utils._convert_to_result(batch, predictions, model_id)
+
+ model_handler = TFModelHandlerTensor(
+ model_uri=model_path,
+ inference_fn=fake_batching_inference_fn,
+ large_model=True)
+ examples = [
+ tf.convert_to_tensor(numpy.array([1.1, 2.2, 3.3], dtype='float32')),
+ tf.convert_to_tensor(
+ numpy.array([10.1, 20.2, 30.3], dtype='float32')),
+ tf.convert_to_tensor(
+ numpy.array([100.1, 200.2, 300.3], dtype='float32')),
+ tf.convert_to_tensor(
+ numpy.array([200.1, 300.2, 400.3], dtype='float32')),
+ ]
+ expected_predictions = [
+ PredictionResult(ex, pred) for ex,
+ pred in zip(examples, [tf.math.multiply(n, 2) for n in examples])
+ ]
+
+ pcoll = pipeline | 'start' >> beam.Create(examples)
+ predictions = pcoll | RunInference(model_handler)
+ assert_that(
+ predictions,
+ equal_to(
+ expected_predictions,
+ equals_fn=_compare_tensor_prediction_result))
+
def test_predict_numpy_with_batch_size(self):
model = _create_mult2_model()
model_path = os.path.join(self.tmpdir, 'mult2_numpy')
@@ -212,6 +258,49 @@ class TFRunInferenceTest(unittest.TestCase):
expected_predictions,
equals_fn=_compare_tensor_prediction_result))
+ def test_predict_numpy_with_large_model(self):
+ model = _create_mult2_model()
+ model_path = os.path.join(self.tmpdir, 'mult2_numpy')
+ tf.keras.models.save_model(model, model_path)
+ with TestPipeline() as pipeline:
+
+ def fake_inference_fn(
+ model: tf.Module,
+ batch: Sequence[numpy.ndarray],
+ inference_args: Dict[str, Any],
+ model_id: Optional[str] = None) -> Iterable[PredictionResult]:
+ multi_process_shared_loaded = "multi_process_shared" in str(type(model))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded model of type {type(model)}, was ' +
+ 'expecting multi_process_shared_model')
+ vectorized_batch = numpy.stack(batch, axis=0)
+ predictions = model.predict(vectorized_batch, **inference_args)
+ return utils._convert_to_result(batch, predictions, model_id)
+
+ model_handler = TFModelHandlerNumpy(
+ model_uri=model_path,
+ inference_fn=fake_inference_fn,
+ large_model=True)
+ examples = [
+ numpy.array([1.1, 2.2, 3.3], dtype='float32'),
+ numpy.array([10.1, 20.2, 30.3], dtype='float32'),
+ numpy.array([100.1, 200.2, 300.3], dtype='float32'),
+ numpy.array([200.1, 300.2, 400.3], dtype='float32'),
+ ]
+ expected_predictions = [
+ PredictionResult(ex, pred) for ex,
+ pred in zip(examples, [numpy.multiply(n, 2) for n in examples])
+ ]
+
+ pcoll = pipeline | 'start' >> beam.Create(examples)
+ predictions = pcoll | RunInference(model_handler)
+ assert_that(
+ predictions,
+ equal_to(
+ expected_predictions,
+ equals_fn=_compare_tensor_prediction_result))
+
def test_predict_tensor_with_args(self):
fake_model = FakeTFTensorModel()
inference_runner = TFModelHandlerTensor(
diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py
index 7dab645f559..ff9bb78d579 100644
--- a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py
+++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py
@@ -229,6 +229,7 @@ class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
max_batch_size: int,
*,
inference_fn: TensorRTInferenceFn = _default_tensorRT_inference_fn,
+ large_model: bool = False,
**kwargs):
"""Implementation of the ModelHandler interface for TensorRT.
@@ -248,6 +249,10 @@ class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
max_batch_size: maximum accepted batch size.
inference_fn: the inference function to use on RunInference calls.
default: _default_tensorRT_inference_fn
+ large_model: set to true if your model is large enough to run into
+ memory pressure if you load multiple copies. Given a model that
+ consumes N memory and a machine with W cores and M memory, you should
+ set this to True if N*W > M.
kwargs: Additional arguments like 'engine_path' and 'onnx_path' are
currently supported. 'env_vars' can be used to set environment variables
before loading the model.
@@ -263,6 +268,7 @@ class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
elif 'onnx_path' in kwargs:
self.onnx_path = kwargs.get('onnx_path')
self._env_vars = kwargs.get('env_vars', {})
+ self._large_model = large_model
def batch_elements_kwargs(self):
"""Sets min_batch_size and max_batch_size of a TensorRT engine."""
@@ -321,3 +327,6 @@ class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
Returns a namespace for metrics collected by the RunInference transform.
"""
return 'BeamML_TensorRT'
+
+ def share_model_across_processes(self) -> bool:
+ return self._large_model
diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
index cf97af8c5b8..86bb7f695d3 100644
--- a/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py
@@ -32,6 +32,7 @@ from apache_beam.testing.util import equal_to
# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
try:
import tensorrt as trt
+ from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import PredictionResult, RunInference
from apache_beam.ml.inference.tensorrt_inference import \
TensorRTEngineHandlerNumPy
@@ -371,6 +372,67 @@ class TensorRTRunInferencePipelineTest(unittest.TestCase):
equal_to(
SINGLE_FEATURE_PREDICTIONS, equals_fn=_compare_prediction_result))
+ @unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed')
+ def test_pipeline_single_tensor_feature_built_engine_large_engine(self):
+ with TestPipeline() as pipeline:
+
+ def fake_inference_fn(batch, engine, inference_args=None):
+ multi_process_shared_loaded = "multi_process_shared" in str(
+ type(engine))
+ if not multi_process_shared_loaded:
+ raise Exception(
+ f'Loaded engine of type {type(engine)}, was ' +
+ 'expecting multi_process_shared engine')
+ from cuda import cuda
+ (
+ engine,
+ context,
+ context_lock,
+ inputs,
+ outputs,
+ gpu_allocations,
+ cpu_allocations,
+ stream) = engine.get_engine_attrs()
+
+ # Process I/O and execute the network
+ with context_lock:
+ _assign_or_fail(
+ cuda.cuMemcpyHtoDAsync(
+ inputs[0]['allocation'],
+ np.ascontiguousarray(batch),
+ inputs[0]['size'],
+ stream))
+ context.execute_async_v2(gpu_allocations, stream)
+ for output in range(len(cpu_allocations)):
+ _assign_or_fail(
+ cuda.cuMemcpyDtoHAsync(
+ cpu_allocations[output],
+ outputs[output]['allocation'],
+ outputs[output]['size'],
+ stream))
+ _assign_or_fail(cuda.cuStreamSynchronize(stream))
+
+ predictions = []
+ for idx in range(len(batch)):
+ predictions.append(
+ [prediction[idx] for prediction in cpu_allocations])
+
+ return utils._convert_to_result(batch, predictions)
+
+ engine_handler = TensorRTEngineHandlerNumPy(
+ min_batch_size=4,
+ max_batch_size=4,
+ engine_path=
+ 'gs://apache-beam-ml/models/single_tensor_features_engine.trt',
+ inference_fn=fake_inference_fn,
+ large_model=True)
+ pcoll = pipeline | 'start' >> beam.Create(SINGLE_FEATURE_EXAMPLES)
+ predictions = pcoll | RunInference(engine_handler)
+ assert_that(
+ predictions,
+ equal_to(
+ SINGLE_FEATURE_PREDICTIONS, equals_fn=_compare_prediction_result))
+
@unittest.skipIf(GCSFileSystem is None, 'GCP dependencies are not installed')
def test_pipeline_sets_env_vars_correctly(self):
with TestPipeline() as pipeline: