This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 1d8e8373848 [RunInference] Validate model handler batching arguments
(#40718)
1d8e8373848 is described below
commit 1d8e837384809dd41529b8d007a08905a69c70cc
Author: alisha-1000 <[email protected]>
AuthorDate: Fri Oct 23 21:09:52 2026 +0530
[RunInference] Validate model handler batching arguments (#40718)
* [RunInference] Validate model handler batching arguments
* Preserve batching weight argument test coverage
* [RunInference] Validate batching arguments across model handlers
* [RunInference] Format batching argument validation changes
* [RunInference] Format model handler validation tests
* Clarify incompatible inference batching options
---------
Co-authored-by: Alisha Gupta <[email protected]>
---
.../ml/inference/anthropic_inference.py | 3 ++
.../ml/inference/anthropic_inference_test.py | 8 ++++
sdks/python/apache_beam/ml/inference/base.py | 30 +++++++++++-
sdks/python/apache_beam/ml/inference/base_test.py | 56 ++++++++++++++--------
.../apache_beam/ml/inference/gemini_inference.py | 15 ++++--
.../ml/inference/gemini_inference_test.py | 29 +++++++++++
.../ml/inference/huggingface_inference.py | 24 ++++++----
.../apache_beam/ml/inference/onnx_inference.py | 8 ++--
.../apache_beam/ml/inference/openai_inference.py | 21 +++++---
.../ml/inference/openai_inference_test.py | 29 ++++++++++-
.../apache_beam/ml/inference/pytorch_inference.py | 16 ++++---
.../apache_beam/ml/inference/sklearn_inference.py | 16 ++++---
.../ml/inference/tensorflow_inference.py | 16 ++++---
.../apache_beam/ml/inference/tensorrt_inference.py | 8 ++--
.../ml/inference/vertex_ai_inference.py | 15 ++++--
.../ml/inference/vertex_ai_inference_test.py | 29 +++++++++++
.../apache_beam/ml/inference/vllm_inference.py | 18 ++++---
.../apache_beam/ml/inference/xgboost_inference.py | 9 ++--
18 files changed, 273 insertions(+), 77 deletions(-)
diff --git a/sdks/python/apache_beam/ml/inference/anthropic_inference.py
b/sdks/python/apache_beam/ml/inference/anthropic_inference.py
index 55bfff784ef..e35ae52b9a2 100644
--- a/sdks/python/apache_beam/ml/inference/anthropic_inference.py
+++ b/sdks/python/apache_beam/ml/inference/anthropic_inference.py
@@ -80,6 +80,7 @@ from anthropic import APIStatusError
from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RemoteModelHandler
+from apache_beam.ml.inference.base import _validate_batching_args
LOGGER = logging.getLogger("AnthropicModelHandler")
@@ -227,6 +228,8 @@ class AnthropicModelHandler(RemoteModelHandler[Any,
PredictionResult,
element_size_fn: optional. a function that returns the size
(weight) of an element.
"""
+ _validate_batching_args(
+ max_batch_weight=max_batch_weight, element_size_fn=element_size_fn)
self._batching_kwargs = {}
self._env_vars = kwargs.get('env_vars', {})
if min_batch_size is not None:
diff --git a/sdks/python/apache_beam/ml/inference/anthropic_inference_test.py
b/sdks/python/apache_beam/ml/inference/anthropic_inference_test.py
index 86e527d515f..eb5102edad3 100644
--- a/sdks/python/apache_beam/ml/inference/anthropic_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/anthropic_inference_test.py
@@ -152,6 +152,14 @@ class MessageFromConversationTest(unittest.TestCase):
class AnthropicModelHandlerTest(unittest.TestCase):
+ def test_max_batch_weight_requires_element_size_fn(self):
+ with self.assertRaisesRegex(ValueError,
+ "max_batch_weight requires element_size_fn"):
+ AnthropicModelHandler(
+ model_name=_TEST_MODEL,
+ request_fn=message_from_string,
+ max_batch_weight=100)
+
@mock.patch('apache_beam.ml.inference.anthropic_inference.Anthropic')
def test_create_client_with_api_key(self, mock_anthropic):
handler = AnthropicModelHandler(
diff --git a/sdks/python/apache_beam/ml/inference/base.py
b/sdks/python/apache_beam/ml/inference/base.py
index a417bb9560f..c2e084d0e03 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -179,6 +179,25 @@ class KeyModelPathMapping(Generic[KeyT]):
model_id: str = ''
+def _validate_batching_args(
+ max_batch_weight=None,
+ element_size_fn=None,
+ batch_length_fn=None,
+ batch_bucket_boundaries=None,
+):
+ if max_batch_weight is not None:
+ if element_size_fn is None:
+ raise ValueError('max_batch_weight requires element_size_fn to be set.')
+ if batch_length_fn is not None or batch_bucket_boundaries is not None:
+ raise ValueError(
+ 'max_batch_weight cannot be combined with length-based batching '
+ '(batch_length_fn or batch_bucket_boundaries).')
+
+ if batch_bucket_boundaries is not None and batch_length_fn is None:
+ raise ValueError(
+ 'batch_bucket_boundaries requires batch_length_fn to be set.')
+
+
class ModelHandler(Generic[ExampleT, PredictionT, ModelT]):
"""Has the ability to load and apply an ML model."""
def __init__(
@@ -204,7 +223,8 @@ class ModelHandler(Generic[ExampleT, PredictionT, ModelT]):
max_batch_weight: the maximum total weight of a batch, where the weight
of an element is given by element_size_fn, which must be set. The
batch size limits (min_batch_size and max_batch_size) then count
- elements.
+ elements. Cannot be combined with batch_length_fn or
+ batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
Without max_batch_weight, the batch size limits are expressed in this
size instead of in number of elements.
@@ -214,6 +234,7 @@ class ModelHandler(Generic[ExampleT, PredictionT, ModelT]):
padding waste for variable-length inputs. Bucket assignment uses
bisect_right so boundaries are lower-inclusive: e.g., for boundaries
[10, 50], buckets are (-inf, 10), [10, 50), [50, inf).
+ Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
length bucketing. Boundaries are lower-inclusive (bisect_right
semantics): bucket i covers lengths in [boundaries[i-1],
boundaries[i]).
@@ -227,6 +248,13 @@ class ModelHandler(Generic[ExampleT, PredictionT, ModelT]):
before loading the model.
"""
self._env_vars = kwargs.get('env_vars', {})
+
+ _validate_batching_args(
+ max_batch_weight=max_batch_weight,
+ element_size_fn=element_size_fn,
+ batch_length_fn=batch_length_fn,
+ batch_bucket_boundaries=batch_bucket_boundaries)
+
self._batching_kwargs: dict[str, Any] = {}
if min_batch_size is not None:
self._batching_kwargs['min_batch_size'] = min_batch_size
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py
b/sdks/python/apache_beam/ml/inference/base_test.py
index dbf81035cf7..f1bb5d1245e 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -2665,24 +2665,28 @@ class
FakeModelHandlerForBatching(base.ModelHandler[int, int, FakeModel]):
class ModelHandlerBatchingArgsTest(unittest.TestCase):
"""Tests for ModelHandler.__init__ batching parameters."""
def test_batch_elements_kwargs_all_args(self):
- """All batching args passed to __init__ are in batch_elements_kwargs."""
- def size_fn(x):
- return 10
-
+ """All non-weight batching args are passed to batch_elements_kwargs."""
handler = FakeModelHandlerForBatching(
- min_batch_size=5,
- max_batch_size=20,
- max_batch_duration_secs=30,
- max_batch_weight=100,
- element_size_fn=size_fn)
+ min_batch_size=5, max_batch_size=20, max_batch_duration_secs=30)
kwargs = handler.batch_elements_kwargs()
self.assertEqual(kwargs['min_batch_size'], 5)
self.assertEqual(kwargs['max_batch_size'], 20)
self.assertEqual(kwargs['max_batch_duration_secs'], 30)
+
+ def test_batch_elements_kwargs_weight_args(self):
+ """Weight-based batching args are passed through to kwargs."""
+ def size_fn(x):
+ return 10
+
+ handler = FakeModelHandlerForBatching(
+ max_batch_weight=100, element_size_fn=size_fn)
+
+ kwargs = handler.batch_elements_kwargs()
+
self.assertEqual(kwargs['max_batch_weight'], 100)
- self.assertIn('element_size_fn', kwargs)
+ self.assertIs(kwargs['element_size_fn'], size_fn)
self.assertEqual(kwargs['element_size_fn'](1), 10)
def test_batch_elements_kwargs_partial_args(self):
@@ -2762,16 +2766,28 @@ class ModelHandlerBatchingArgsTest(unittest.TestCase):
self.assertNotIn('bucket_boundaries', kwargs)
def test_batch_bucket_boundaries_without_batch_length_fn(self):
- """Passing batch_bucket_boundaries without batch_length_fn should fail in
- BatchElements.
-
- Note: ModelHandler.__init__ doesn't validate this; the error is raised
- by BatchElements when batch_elements_kwargs are used."""
- handler = FakeModelHandlerForBatching(batch_bucket_boundaries=[10, 20])
- kwargs = handler.batch_elements_kwargs()
- # The kwargs are stored, but BatchElements will reject them
- self.assertEqual(kwargs['bucket_boundaries'], [10, 20])
- self.assertNotIn('length_fn', kwargs)
+ with self.assertRaisesRegex(
+ ValueError, 'batch_bucket_boundaries requires batch_length_fn'):
+ FakeModelHandlerForBatching(batch_bucket_boundaries=[10, 20])
+
+ def test_max_batch_weight_requires_element_size_fn(self):
+ with self.assertRaisesRegex(ValueError,
+ 'max_batch_weight requires element_size_fn'):
+ FakeModelHandlerForBatching(max_batch_weight=10)
+
+ def test_max_batch_weight_rejects_length_based_batching(self):
+ with self.assertRaisesRegex(
+ ValueError, 'max_batch_weight cannot be combined with length-based'):
+ FakeModelHandlerForBatching(
+ max_batch_weight=10, element_size_fn=len, batch_length_fn=len)
+
+ def test_max_batch_weight_rejects_bucket_boundaries(self):
+ with self.assertRaisesRegex(
+ ValueError, 'max_batch_weight cannot be combined with length-based'):
+ FakeModelHandlerForBatching(
+ max_batch_weight=10,
+ element_size_fn=len,
+ batch_bucket_boundaries=[10, 20])
def test_batching_kwargs_none_values_omitted(self):
"""None values for batch_length_fn and batch_bucket_boundaries are not in
diff --git a/sdks/python/apache_beam/ml/inference/gemini_inference.py
b/sdks/python/apache_beam/ml/inference/gemini_inference.py
index 9d8cf685f79..4bbb2af1940 100644
--- a/sdks/python/apache_beam/ml/inference/gemini_inference.py
+++ b/sdks/python/apache_beam/ml/inference/gemini_inference.py
@@ -33,6 +33,7 @@ from PIL.Image import Image
from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RemoteModelHandler
+from apache_beam.ml.inference.base import _validate_batching_args
LOGGER = logging.getLogger("GeminiModelHandler")
@@ -157,14 +158,22 @@ class GeminiModelHandler(RemoteModelHandler[Any,
PredictionResult,
inputs.
max_batch_duration_secs: optional. the maximum amount of time to buffer
a batch before emitting; used in streaming contexts.
- max_batch_weight: optional. the maximum total weight of a batch.
+ max_batch_weight: optional. the maximum total weight of a batch. Cannot
+ be combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: optional. a function that returns the size (weight)
of an element.
batch_length_fn: optional. a callable that returns the length of an
- element for length-aware batching.
+ element for length-aware batching. Cannot be used together with
+ max_batch_weight.
batch_bucket_boundaries: optional. a sorted list of positive boundary
- values for length-aware batching buckets.
+ values for length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
"""
+ _validate_batching_args(
+ max_batch_weight=max_batch_weight,
+ element_size_fn=element_size_fn,
+ batch_length_fn=batch_length_fn,
+ batch_bucket_boundaries=batch_bucket_boundaries)
self._batching_kwargs = {}
self._env_vars = kwargs.get('env_vars', {})
if min_batch_size is not None:
diff --git a/sdks/python/apache_beam/ml/inference/gemini_inference_test.py
b/sdks/python/apache_beam/ml/inference/gemini_inference_test.py
index 012287e98f3..4f65fa52eed 100644
--- a/sdks/python/apache_beam/ml/inference/gemini_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/gemini_inference_test.py
@@ -44,6 +44,35 @@ class RetryOnClientErrorTest(unittest.TestCase):
class ModelHandlerArgConditions(unittest.TestCase):
+ def test_max_batch_weight_requires_element_size_fn(self):
+ with self.assertRaisesRegex(ValueError,
+ "max_batch_weight requires element_size_fn"):
+ GeminiModelHandler(
+ model_name="gemini-model-123",
+ request_fn=generate_from_string,
+ api_key="123456789",
+ max_batch_weight=100)
+
+ def test_max_batch_weight_cannot_combine_with_length_batching(self):
+ with self.assertRaisesRegex(
+ ValueError, "cannot be combined with length-based batching"):
+ GeminiModelHandler(
+ model_name="gemini-model-123",
+ request_fn=generate_from_string,
+ api_key="123456789",
+ max_batch_weight=100,
+ element_size_fn=len,
+ batch_length_fn=len)
+
+ def test_batch_bucket_boundaries_require_batch_length_fn(self):
+ with self.assertRaisesRegex(
+ ValueError, "batch_bucket_boundaries requires batch_length_fn"):
+ GeminiModelHandler(
+ model_name="gemini-model-123",
+ request_fn=generate_from_string,
+ api_key="123456789",
+ batch_bucket_boundaries=[16, 32])
+
def test_all_params_set(self):
self.assertRaises(
ValueError,
diff --git a/sdks/python/apache_beam/ml/inference/huggingface_inference.py
b/sdks/python/apache_beam/ml/inference/huggingface_inference.py
index 93ffc3fc1f5..87c9f203ba6 100644
--- a/sdks/python/apache_beam/ml/inference/huggingface_inference.py
+++ b/sdks/python/apache_beam/ml/inference/huggingface_inference.py
@@ -306,12 +306,14 @@ class HuggingFaceModelHandlerKeyedTensor(
model_copies: The exact number of models that you would like loaded
onto your machine. This can be useful if you exactly know your CPU or
GPU capacity and want to maximize resource utilization.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -502,12 +504,14 @@ class
HuggingFaceModelHandlerTensor(ModelHandler[Union['tf.Tensor',
model_copies: The exact number of models that you would like loaded
onto your machine. This can be useful if you exactly know your CPU or
GPU capacity and want to maximize resource utilization.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -686,12 +690,14 @@ class HuggingFacePipelineModelHandler(ModelHandler[str,
model_copies: The exact number of models that you would like loaded
onto your machine. This can be useful if you exactly know your CPU or
GPU capacity and want to maximize resource utilization.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
diff --git a/sdks/python/apache_beam/ml/inference/onnx_inference.py
b/sdks/python/apache_beam/ml/inference/onnx_inference.py
index c525abbff45..df281b998a8 100644
--- a/sdks/python/apache_beam/ml/inference/onnx_inference.py
+++ b/sdks/python/apache_beam/ml/inference/onnx_inference.py
@@ -94,12 +94,14 @@ class OnnxModelHandlerNumpy(ModelHandler[numpy.ndarray,
max_batch_size: the maximum batch size to use when batching inputs.
max_batch_duration_secs: the maximum amount of time to buffer a batch
before emitting; used in streaming contexts.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
"""
diff --git a/sdks/python/apache_beam/ml/inference/openai_inference.py
b/sdks/python/apache_beam/ml/inference/openai_inference.py
index a32b57f56ea..bd3ec3b379d 100644
--- a/sdks/python/apache_beam/ml/inference/openai_inference.py
+++ b/sdks/python/apache_beam/ml/inference/openai_inference.py
@@ -82,6 +82,7 @@ from openai import OpenAI
from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RemoteModelHandler
+from apache_beam.ml.inference.base import _validate_batching_args
__all__ = [
'OpenAIModelHandler',
@@ -299,16 +300,22 @@ class OpenAIModelHandler(RemoteModelHandler[Any,
PredictionResult, OpenAI]):
batching inputs.
max_batch_duration_secs: optional. the maximum amount of time to
buffer a batch before emitting; used in streaming contexts.
- max_batch_weight: optional. the maximum total weight of a batch.
- element_size_fn: optional. a function that returns the size
- (weight) of an element.
- batch_length_fn: optional. a callable that returns the length of an
- element for length-aware batching.
- batch_bucket_boundaries: optional. a sorted list of positive boundary
- values for length-aware batching buckets.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
+ element_size_fn: a function that returns the size (weight) of an element.
+ batch_length_fn: a callable that returns the length of an element for
+ length-aware batching. Cannot be used together with max_batch_weight.
+ batch_bucket_boundaries: a sorted list of positive boundary values for
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
**kwargs: optional arguments passed to RemoteModelHandler (e.g.
num_retries, throttle_delay_secs, rate_limiter).
"""
+ _validate_batching_args(
+ max_batch_weight=max_batch_weight,
+ element_size_fn=element_size_fn,
+ batch_length_fn=batch_length_fn,
+ batch_bucket_boundaries=batch_bucket_boundaries)
self._batching_kwargs = {}
self._env_vars = kwargs.get('env_vars', {})
if min_batch_size is not None:
diff --git a/sdks/python/apache_beam/ml/inference/openai_inference_test.py
b/sdks/python/apache_beam/ml/inference/openai_inference_test.py
index 23a13980ea8..112fcda3d9a 100644
--- a/sdks/python/apache_beam/ml/inference/openai_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/openai_inference_test.py
@@ -379,6 +379,32 @@ class EmbeddingFromStringTest(unittest.TestCase):
class OpenAIModelHandlerTest(unittest.TestCase):
+ def test_max_batch_weight_requires_element_size_fn(self):
+ with self.assertRaisesRegex(ValueError,
+ "max_batch_weight requires element_size_fn"):
+ OpenAIModelHandler(
+ model_name=_TEST_MODEL,
+ request_fn=chat_completion_from_string,
+ max_batch_weight=100)
+
+ def test_max_batch_weight_cannot_combine_with_length_batching(self):
+ with self.assertRaisesRegex(
+ ValueError, "cannot be combined with length-based batching"):
+ OpenAIModelHandler(
+ model_name=_TEST_MODEL,
+ request_fn=chat_completion_from_string,
+ max_batch_weight=100,
+ element_size_fn=len,
+ batch_length_fn=len)
+
+ def test_batch_bucket_boundaries_require_batch_length_fn(self):
+ with self.assertRaisesRegex(
+ ValueError, "batch_bucket_boundaries requires batch_length_fn"):
+ OpenAIModelHandler(
+ model_name=_TEST_MODEL,
+ request_fn=chat_completion_from_string,
+ batch_bucket_boundaries=[16, 32])
+
@mock.patch('apache_beam.ml.inference.openai_inference.OpenAI')
def test_create_client_with_api_key(self, mock_openai):
handler = OpenAIModelHandler(
@@ -445,7 +471,8 @@ class OpenAIModelHandlerTest(unittest.TestCase):
min_batch_size=2,
max_batch_size=10,
max_batch_duration_secs=5,
- max_batch_weight=100)
+ max_batch_weight=100,
+ element_size_fn=len)
kwargs = handler.batch_elements_kwargs()
self.assertEqual(kwargs['min_batch_size'], 2)
self.assertEqual(kwargs['max_batch_size'], 10)
diff --git a/sdks/python/apache_beam/ml/inference/pytorch_inference.py
b/sdks/python/apache_beam/ml/inference/pytorch_inference.py
index 5c6b41568c5..81234cb5c4f 100644
--- a/sdks/python/apache_beam/ml/inference/pytorch_inference.py
+++ b/sdks/python/apache_beam/ml/inference/pytorch_inference.py
@@ -258,12 +258,14 @@ class PytorchModelHandlerTensor(ModelHandler[torch.Tensor,
GPU capacity and want to maximize resource utilization.
load_model_args: a dictionary of parameters passed to the torch.load
function to specify custom config for loading models.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -503,12 +505,14 @@ class
PytorchModelHandlerKeyedTensor(ModelHandler[dict[str, torch.Tensor],
GPU capacity and want to maximize resource utilization.
load_model_args: a dictionary of parameters passed to the torch.load
function to specify custom config for loading models.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
diff --git a/sdks/python/apache_beam/ml/inference/sklearn_inference.py
b/sdks/python/apache_beam/ml/inference/sklearn_inference.py
index 8d02cacaac3..65220689283 100644
--- a/sdks/python/apache_beam/ml/inference/sklearn_inference.py
+++ b/sdks/python/apache_beam/ml/inference/sklearn_inference.py
@@ -126,12 +126,14 @@ class SklearnModelHandlerNumpy(ModelHandler[numpy.ndarray,
model_copies: The exact number of models that you would like loaded
onto your machine. This can be useful if you exactly know your CPU or
GPU capacity and want to maximize resource utilization.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
"""
@@ -266,12 +268,14 @@ class
SklearnModelHandlerPandas(ModelHandler[pandas.DataFrame,
model_copies: The exact number of models that you would like loaded
onto your machine. This can be useful if you exactly know your CPU or
GPU capacity and want to maximize resource utilization.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
"""
diff --git a/sdks/python/apache_beam/ml/inference/tensorflow_inference.py
b/sdks/python/apache_beam/ml/inference/tensorflow_inference.py
index 3563b1113ba..d290ecb2bce 100644
--- a/sdks/python/apache_beam/ml/inference/tensorflow_inference.py
+++ b/sdks/python/apache_beam/ml/inference/tensorflow_inference.py
@@ -144,13 +144,15 @@ class TFModelHandlerNumpy(ModelHandler[numpy.ndarray,
model_copies: The exact number of models that you would like loaded
onto your machine. This can be useful if you exactly know your CPU or
GPU capacity and want to maximize resource utilization.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an
element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
@@ -285,13 +287,15 @@ class TFModelHandlerTensor(ModelHandler[tf.Tensor,
PredictionResult,
model_copies: The exact number of models that you would like loaded
onto your machine. This can be useful if you exactly know your CPU or
GPU capacity and want to maximize resource utilization.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an
element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
diff --git a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py
b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py
index b6089788eae..5eaa6b029af 100644
--- a/sdks/python/apache_beam/ml/inference/tensorrt_inference.py
+++ b/sdks/python/apache_beam/ml/inference/tensorrt_inference.py
@@ -318,12 +318,14 @@ class TensorRTEngineHandlerNumPy(ModelHandler[np.ndarray,
GPU capacity and want to maximize resource utilization.
max_batch_duration_secs: the maximum amount of time to buffer
a batch before emitting; used in streaming contexts.
- max_batch_weight: the maximum total weight of a batch.
+ max_batch_weight: the maximum total weight of a batch. Cannot be
+ combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: a function that returns the size (weight) of an element.
batch_length_fn: a callable that returns the length of an element for
- length-aware batching.
+ length-aware batching. Cannot be used together with max_batch_weight.
batch_bucket_boundaries: a sorted list of positive boundary values for
- length-aware batching buckets.
+ length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
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.
diff --git a/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
index 7757146ab8d..59106ac3580 100644
--- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
@@ -31,6 +31,7 @@ from google.cloud import aiplatform
from apache_beam.ml.inference import utils
from apache_beam.ml.inference.base import PredictionResult
from apache_beam.ml.inference.base import RemoteModelHandler
+from apache_beam.ml.inference.base import _validate_batching_args
LOGGER = logging.getLogger("VertexAIModelHandlerJSON")
@@ -114,14 +115,22 @@ class VertexAIModelHandlerJSON(RemoteModelHandler[Any,
inputs.
max_batch_duration_secs: optional. the maximum amount of time to buffer
a batch before emitting; used in streaming contexts.
- max_batch_weight: optional. the maximum total weight of a batch.
+ max_batch_weight: optional. the maximum total weight of a batch. Cannot
+ be combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: optional. a function that returns the size (weight)
of an element.
batch_length_fn: optional. a callable that returns the length of an
- element for length-aware batching.
+ element for length-aware batching. Cannot be used together with
+ max_batch_weight.
batch_bucket_boundaries: optional. a sorted list of positive boundary
- values for length-aware batching buckets.
+ values for length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
"""
+ _validate_batching_args(
+ max_batch_weight=max_batch_weight,
+ element_size_fn=element_size_fn,
+ batch_length_fn=batch_length_fn,
+ batch_bucket_boundaries=batch_bucket_boundaries)
self._batching_kwargs = {}
self._env_vars = kwargs.get('env_vars', {})
self._invoke_route = invoke_route
diff --git a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_test.py
b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_test.py
index 8aa638ebe7c..4d2bf540840 100644
--- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference_test.py
@@ -38,6 +38,35 @@ class RetryOnClientErrorTest(unittest.TestCase):
class ModelHandlerArgConditions(unittest.TestCase):
+ def test_max_batch_weight_requires_element_size_fn(self):
+ with self.assertRaisesRegex(ValueError,
+ "max_batch_weight requires element_size_fn"):
+ VertexAIModelHandlerJSON(
+ endpoint_id="1",
+ project="testproject",
+ location="us-central1",
+ max_batch_weight=100)
+
+ def test_max_batch_weight_cannot_combine_with_length_batching(self):
+ with self.assertRaisesRegex(
+ ValueError, "cannot be combined with length-based batching"):
+ VertexAIModelHandlerJSON(
+ endpoint_id="1",
+ project="testproject",
+ location="us-central1",
+ max_batch_weight=100,
+ element_size_fn=len,
+ batch_length_fn=len)
+
+ def test_batch_bucket_boundaries_require_batch_length_fn(self):
+ with self.assertRaisesRegex(
+ ValueError, "batch_bucket_boundaries requires batch_length_fn"):
+ VertexAIModelHandlerJSON(
+ endpoint_id="1",
+ project="testproject",
+ location="us-central1",
+ batch_bucket_boundaries=[16, 32])
+
def test_exception_on_private_without_network(self):
self.assertRaises(
ValueError,
diff --git a/sdks/python/apache_beam/ml/inference/vllm_inference.py
b/sdks/python/apache_beam/ml/inference/vllm_inference.py
index e5d8918d5ec..e3bc203494c 100644
--- a/sdks/python/apache_beam/ml/inference/vllm_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vllm_inference.py
@@ -415,13 +415,16 @@ class VLLMCompletionsModelHandler(ModelHandler[str,
inputs.
max_batch_duration_secs: optional. the maximum amount of time to buffer
a batch before emitting; used in streaming contexts.
- max_batch_weight: optional. the maximum total weight of a batch.
+ max_batch_weight: optional. the maximum total weight of a batch. Cannot
+ be combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: optional. a function that returns the size (weight) of
an element.
batch_length_fn: optional. a callable that returns the length of an
- element for length-aware batching.
+ element for length-aware batching. Cannot be used together with
+ max_batch_weight.
batch_bucket_boundaries: optional. a sorted list of positive boundary
- values for length-aware batching buckets.
+ values for length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
"""
super().__init__(
min_batch_size=min_batch_size,
@@ -558,13 +561,16 @@ class
VLLMChatModelHandler(ModelHandler[Sequence[OpenAIChatMessage],
inputs.
max_batch_duration_secs: optional. the maximum amount of time to buffer
a batch before emitting; used in streaming contexts.
- max_batch_weight: optional. the maximum total weight of a batch.
+ max_batch_weight: optional. the maximum total weight of a batch. Cannot
+ be combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: optional. a function that returns the size (weight) of
an element.
batch_length_fn: optional. a callable that returns the length of an
- element for length-aware batching.
+ element for length-aware batching. Cannot be used together with
+ max_batch_weight.
batch_bucket_boundaries: optional. a sorted list of positive boundary
- values for length-aware batching buckets.
+ values for length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
"""
super().__init__(
min_batch_size=min_batch_size,
diff --git a/sdks/python/apache_beam/ml/inference/xgboost_inference.py
b/sdks/python/apache_beam/ml/inference/xgboost_inference.py
index b27db8b2046..dd6b15f1b1e 100644
--- a/sdks/python/apache_beam/ml/inference/xgboost_inference.py
+++ b/sdks/python/apache_beam/ml/inference/xgboost_inference.py
@@ -109,13 +109,16 @@ class XGBoostModelHandler(ModelHandler[ExampleT,
PredictionT, ModelT], ABC):
inputs.
max_batch_duration_secs: optional. the maximum amount of time to buffer
a batch before emitting; used in streaming contexts.
- max_batch_weight: optional. the maximum total weight of a batch.
+ max_batch_weight: optional. the maximum total weight of a batch. Cannot
+ be combined with batch_length_fn or batch_bucket_boundaries.
element_size_fn: optional. a function that returns the size (weight)
of an element.
batch_length_fn: optional. a callable that returns the length of an
- element for length-aware batching.
+ element for length-aware batching. Cannot be used together with
+ max_batch_weight.
batch_bucket_boundaries: optional. a sorted list of positive boundary
- values for length-aware batching buckets.
+ values for length-aware batching buckets. Cannot be used together with
+ max_batch_weight.
kwargs: 'env_vars' can be used to set environment variables
before loading the model.
