[RunInference] Validate model handler batching arguments

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:

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(

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

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

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:

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,

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.
 

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.
     """

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:

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)