From 0e0eae8efb3c3f3409f8315dc6480fffa718ce80 Mon Sep 17 00:00:00 2001
From: Alisha Gupta <alishagupta2004gupta@gmail.com>
Date: Thu, 22 Oct 2026 09:45:29 +0530
Subject: [PATCH 1/6] [RunInference] Validate model handler batching arguments

---
 sdks/python/apache_beam/ml/inference/base.py  | 13 +++++++
 .../apache_beam/ml/inference/base_test.py     | 34 +++++++++++++------
 2 files changed, 37 insertions(+), 10 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 3d2f2657c0a8..8875569184cb 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -227,6 +227,19 @@ def __init__(
         before loading the model.
     """
     self._env_vars = kwargs.get('env_vars', {})
+
+    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.')
+
     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..51818c14546b 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -2621,16 +2621,30 @@ 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

From 3c70baf9b8ee5374538e0dfda632e9a2e439271b Mon Sep 17 00:00:00 2001
From: Alisha Gupta <alishagupta2004gupta@gmail.com>
Date: Thu, 22 Oct 2026 12:52:10 +0530
Subject: [PATCH 2/6] Preserve batching weight argument test coverage

---
 .../apache_beam/ml/inference/base_test.py     | 23 ++++++++++++-------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py
index 51818c14546b..911b351be827 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -2524,24 +2524,31 @@ 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)
+        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):

From df7f7e94e081bf8fe0978e55b56f4ec543dcf879 Mon Sep 17 00:00:00 2001
From: Alisha Gupta <alishagupta2004gupta@gmail.com>
Date: Thu, 22 Oct 2026 13:13:59 +0530
Subject: [PATCH 3/6] [RunInference] Validate batching arguments across model
 handlers

---
 .../ml/inference/anthropic_inference.py       |  4 +++
 .../ml/inference/anthropic_inference_test.py  |  8 +++++
 sdks/python/apache_beam/ml/inference/base.py  | 36 +++++++++++++------
 .../ml/inference/gemini_inference.py          |  6 ++++
 .../ml/inference/gemini_inference_test.py     | 29 +++++++++++++++
 .../ml/inference/openai_inference.py          |  6 ++++
 .../ml/inference/openai_inference_test.py     | 29 ++++++++++++++-
 .../ml/inference/vertex_ai_inference.py       |  6 ++++
 .../ml/inference/vertex_ai_inference_test.py  | 29 +++++++++++++++
 9 files changed, 141 insertions(+), 12 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/anthropic_inference.py b/sdks/python/apache_beam/ml/inference/anthropic_inference.py
index 55bfff784ef8..284a8ceece28 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,9 @@ 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..b64f102bdd3b 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 8875569184cb..8034d07d358e 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -179,6 +179,26 @@ 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__(
@@ -228,17 +248,11 @@ def __init__(
     """
     self._env_vars = kwargs.get('env_vars', {})
 
-    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.')
+    _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:
diff --git a/sdks/python/apache_beam/ml/inference/gemini_inference.py b/sdks/python/apache_beam/ml/inference/gemini_inference.py
index 9d8cf685f79d..1e6ebc8742ad 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")
 
@@ -165,6 +166,11 @@ def __init__(
       batch_bucket_boundaries: optional. a sorted list of positive boundary
         values for length-aware batching buckets.
     """
+    _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..e77df54000f1 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/openai_inference.py b/sdks/python/apache_beam/ml/inference/openai_inference.py
index a32b57f56ea7..9d5ecdd69358 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',
@@ -309,6 +310,11 @@ def __init__(
       **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..f8c62f9bd524 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/vertex_ai_inference.py b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
index 7757146ab8dc..2b537de96f2c 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")
 
@@ -122,6 +123,11 @@ def __init__(
       batch_bucket_boundaries: optional. a sorted list of positive boundary
         values for length-aware batching buckets.
     """
+    _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..fa1cad5b5ede 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,

From 8d3616c6e0b861bf56f3075263bad58eb5cf4cc9 Mon Sep 17 00:00:00 2001
From: Alisha Gupta <alishagupta2004gupta@gmail.com>
Date: Thu, 22 Oct 2026 22:34:10 +0530
Subject: [PATCH 4/6] [RunInference] Format batching argument validation
 changes

---
 .../ml/inference/anthropic_inference.py           |  3 +--
 sdks/python/apache_beam/ml/inference/base.py      |  3 +--
 sdks/python/apache_beam/ml/inference/base_test.py | 15 +++++----------
 3 files changed, 7 insertions(+), 14 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/anthropic_inference.py b/sdks/python/apache_beam/ml/inference/anthropic_inference.py
index 284a8ceece28..e35ae52b9a22 100644
--- a/sdks/python/apache_beam/ml/inference/anthropic_inference.py
+++ b/sdks/python/apache_beam/ml/inference/anthropic_inference.py
@@ -229,8 +229,7 @@ def __init__(
         (weight) of an element.
     """
     _validate_batching_args(
-        max_batch_weight=max_batch_weight,
-        element_size_fn=element_size_fn)
+        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/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 8034d07d358e..c3a4eec241f7 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -187,8 +187,7 @@ def _validate_batching_args(
 ):
   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.')
+      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 '
diff --git a/sdks/python/apache_beam/ml/inference/base_test.py b/sdks/python/apache_beam/ml/inference/base_test.py
index 911b351be827..992f0899a19f 100644
--- a/sdks/python/apache_beam/ml/inference/base_test.py
+++ b/sdks/python/apache_beam/ml/inference/base_test.py
@@ -2526,9 +2526,7 @@ class ModelHandlerBatchingArgsTest(unittest.TestCase):
   def test_batch_elements_kwargs_all_args(self):
     """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)
+        min_batch_size=5, max_batch_size=20, max_batch_duration_secs=30)
 
     kwargs = handler.batch_elements_kwargs()
 
@@ -2542,8 +2540,7 @@ def size_fn(x):
       return 10
 
     handler = FakeModelHandlerForBatching(
-        max_batch_weight=100,
-        element_size_fn=size_fn)
+        max_batch_weight=100, element_size_fn=size_fn)
 
     kwargs = handler.batch_elements_kwargs()
 
@@ -2633,17 +2630,15 @@ def test_batch_bucket_boundaries_without_batch_length_fn(self):
       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'):
+    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)
+          max_batch_weight=10, element_size_fn=len, batch_length_fn=len)
 
   def test_max_batch_weight_rejects_bucket_boundaries(self):
     with self.assertRaisesRegex(

From b3598bdcb611ef7f56d028704093aea3172b5127 Mon Sep 17 00:00:00 2001
From: Alisha Gupta <alishagupta2004gupta@gmail.com>
Date: Fri, 23 Oct 2026 00:23:29 +0530
Subject: [PATCH 5/6] [RunInference] Format model handler validation tests

---
 .../apache_beam/ml/inference/anthropic_inference_test.py      | 4 ++--
 sdks/python/apache_beam/ml/inference/gemini_inference_test.py | 4 ++--
 sdks/python/apache_beam/ml/inference/openai_inference_test.py | 4 ++--
 .../apache_beam/ml/inference/vertex_ai_inference_test.py      | 4 ++--
 4 files changed, 8 insertions(+), 8 deletions(-)

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 b64f102bdd3b..eb5102edad3f 100644
--- a/sdks/python/apache_beam/ml/inference/anthropic_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/anthropic_inference_test.py
@@ -153,8 +153,8 @@ 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"):
+    with self.assertRaisesRegex(ValueError,
+                                "max_batch_weight requires element_size_fn"):
       AnthropicModelHandler(
           model_name=_TEST_MODEL,
           request_fn=message_from_string,
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 e77df54000f1..4f65fa52eed1 100644
--- a/sdks/python/apache_beam/ml/inference/gemini_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/gemini_inference_test.py
@@ -45,8 +45,8 @@ 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"):
+    with self.assertRaisesRegex(ValueError,
+                                "max_batch_weight requires element_size_fn"):
       GeminiModelHandler(
           model_name="gemini-model-123",
           request_fn=generate_from_string,
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 f8c62f9bd524..112fcda3d9a9 100644
--- a/sdks/python/apache_beam/ml/inference/openai_inference_test.py
+++ b/sdks/python/apache_beam/ml/inference/openai_inference_test.py
@@ -380,8 +380,8 @@ 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"):
+    with self.assertRaisesRegex(ValueError,
+                                "max_batch_weight requires element_size_fn"):
       OpenAIModelHandler(
           model_name=_TEST_MODEL,
           request_fn=chat_completion_from_string,
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 fa1cad5b5ede..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
@@ -39,8 +39,8 @@ 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"):
+    with self.assertRaisesRegex(ValueError,
+                                "max_batch_weight requires element_size_fn"):
       VertexAIModelHandlerJSON(
           endpoint_id="1",
           project="testproject",

From 92ad9f49c360d80fad08e586dc345c6e04a25da9 Mon Sep 17 00:00:00 2001
From: Alisha Gupta <alishagupta2004gupta@gmail.com>
Date: Fri, 23 Oct 2026 12:38:47 +0530
Subject: [PATCH 6/6] Clarify incompatible inference batching options

---
 sdks/python/apache_beam/ml/inference/base.py  |  4 +++-
 .../ml/inference/gemini_inference.py          |  9 ++++---
 .../ml/inference/huggingface_inference.py     | 24 ++++++++++++-------
 .../ml/inference/onnx_inference.py            |  8 ++++---
 .../ml/inference/openai_inference.py          | 15 ++++++------
 .../ml/inference/pytorch_inference.py         | 16 ++++++++-----
 .../ml/inference/sklearn_inference.py         | 16 ++++++++-----
 .../ml/inference/tensorflow_inference.py      | 16 ++++++++-----
 .../ml/inference/tensorrt_inference.py        |  8 ++++---
 .../ml/inference/vertex_ai_inference.py       |  9 ++++---
 .../ml/inference/vllm_inference.py            | 18 +++++++++-----
 .../ml/inference/xgboost_inference.py         |  9 ++++---
 12 files changed, 96 insertions(+), 56 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index c3a4eec241f7..586838e40a78 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -223,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.
@@ -233,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]).
diff --git a/sdks/python/apache_beam/ml/inference/gemini_inference.py b/sdks/python/apache_beam/ml/inference/gemini_inference.py
index 1e6ebc8742ad..4bbb2af1940b 100644
--- a/sdks/python/apache_beam/ml/inference/gemini_inference.py
+++ b/sdks/python/apache_beam/ml/inference/gemini_inference.py
@@ -158,13 +158,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.
     """
     _validate_batching_args(
         max_batch_weight=max_batch_weight,
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 9d5ecdd69358..bd3ec3b379dc 100644
--- a/sdks/python/apache_beam/ml/inference/openai_inference.py
+++ b/sdks/python/apache_beam/ml/inference/openai_inference.py
@@ -300,13 +300,14 @@ 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).
     """
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 2b537de96f2c..59106ac3580a 100644
--- a/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
+++ b/sdks/python/apache_beam/ml/inference/vertex_ai_inference.py
@@ -115,13 +115,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.
     """
     _validate_batching_args(
         max_batch_weight=max_batch_weight,
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.