diff --git a/sdks/python/apache_beam/ml/inference/anthropic_inference.py b/sdks/python/apache_beam/ml/inference/anthropic_inference.py index 55bfff784ef8..e35ae52b9a22 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 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 @@ def __init__( 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 86e527d515f1..eb5102edad3f 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 @@ def test_sends_conversation(self): 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 3d2f2657c0a8..586838e40a78 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 d1c37811d9d3..992f0899a19f 100644 --- a/sdks/python/apache_beam/ml/inference/base_test.py +++ b/sdks/python/apache_beam/ml/inference/base_test.py @@ -2524,24 +2524,28 @@ def run_inference(self, batch, model, inference_args=None): 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): @@ -2621,16 +2625,28 @@ def test_batch_length_fn_only(self): 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 9d8cf685f79d..4bbb2af1940b 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 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 @@ def __init__( 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 012287e98f3e..4f65fa52eed1 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 @@ def test_retry_on_server_error(self): 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 93ffc3fc1f52..87c9f203ba69 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 c525abbff45e..df281b998a8c 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 @@ def __init__( #pylint: disable=dangerous-default-value 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 a32b57f56ea7..bd3ec3b379dc 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 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 @@ def __init__( 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 23a13980ea8b..112fcda3d9a9 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 @@ def test_passes_inference_args(self): 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 @@ def test_batch_elements_kwargs(self): 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 5c6b41568c54..81234cb5c4f8 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 @@ def __init__( 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 @@ def __init__( 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 8d02cacaac34..65220689283e 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 @@ def __init__( 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 @@ def __init__( 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 3563b1113ba4..d290ecb2bce7 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 @@ def __init__( 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 @@ def __init__( 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 b6089788eae9..5eaa6b029af0 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 @@ def __init__( 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 7757146ab8dc..59106ac3580a 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 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 @@ def __init__( 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 8aa638ebe7c2..4d2bf540840b 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 @@ def test_retry_on_client_error_negative(self): 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 e5d8918d5ec5..e3bc203494c7 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 @@ def __init__( 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 @@ def __init__( 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 f11b892062e8..a17f33719a23 100644 --- a/sdks/python/apache_beam/ml/inference/xgboost_inference.py +++ b/sdks/python/apache_beam/ml/inference/xgboost_inference.py @@ -108,13 +108,16 @@ def __init__( 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.