[Python][RunInference] vLLM: prefix-aware batching and capacity-based autoscaling

CHANGES.md

@@ -99,6 +99,8 @@
 * (Go) Added the standard `millis_instant` and `decimal` logical types for cross-language rows ([#39684](https://github.com/apache/beam/issues/39684)).
 * [Beam SQL] Added `PRIMARY KEY (...) NOT ENFORCED` to `CREATE EXTERNAL TABLE`, and `ALTER TABLE ... ADD/DROP PRIMARY KEY` ([#40826](https://github.com/apache/beam/issues/40826)).
 * [Beam SQL] Added a `WITH (...)` clause to set per-statement connector options in `FROM` and `INSERT INTO` ([#40848](https://github.com/apache/beam/issues/40848)).
+* (Python) `BatchElements` can group elements by key across workers (`key_fn`, `num_shards`), and RunInference model handlers can use it through `batch_key_fn` / `batch_key_num_shards`. `apache_beam.ml.inference.prefix_batching` adds a `PrefixBatchingConfig` and prompt/chat prefix extractors so LLM handlers can batch requests that share a prompt prefix and hit the server's prefix cache ([#40905](https://github.com/apache/beam/pull/40905)).
+* (Python) vLLM model handlers accept a `prefix_batching_config` to batch requests that share a prompt prefix, and publish a capacity-based autoscaling signal (`num_loaded_models_<model_identifier>`, estimated from vLLM's completion rate and in-flight requests) so Dataflow's GPU autoscaling scales vLLM workers with load; tune or disable it with `autoscaling_config` / `enable_autoscaling_metric`. `apache_beam.ml.inference.autoscaling_metrics` provides the generic pieces for other model-server handlers ([#40908](https://github.com/apache/beam/pull/40908)).
 
 ### Performance
 

sdks/python/apache_beam/metrics/cells.py

@@ -181,6 +181,18 @@ def reset(self):
     # type: () -> None
     self.data = DistributionData.identity_element()
 
+  def reset_to(self, value):
+    """Atomically replaces the distribution with a single sample of value.
+
+    Like ``reset()`` followed by ``update(value)``, except that a concurrent
+    reader never sees the empty distribution in between. Useful for publishing
+    the latest value of a gauge-like quantity as a distribution whose
+    cumulative mean (sum / count) is that value.
+    """
+    ivalue = int(value)
+    with self._lock:
+      self.data = DistributionData(ivalue, 1, ivalue, ivalue)
+
   def combine(self, other):
     # type: (DistributionCell) -> DistributionCell
     result = DistributionCell()

sdks/python/apache_beam/metrics/cells_test.py

@@ -131,6 +131,37 @@ def test_basic_operations(self):
     d.update(900)
     self.assertEqual(d.get_cumulative(), DistributionData(912, 3, 2, 900))
 
+  def test_reset_to(self):
+    d = DistributionCell()
+    d.update(10)
+    d.update(2)
+    d.reset_to(7)
+    self.assertEqual(d.get_cumulative(), DistributionData(7, 1, 7, 7))
+    d.reset_to(3.6)
+    self.assertEqual(d.get_cumulative(), DistributionData(3, 1, 3, 3))
+    d.update(5)
+    self.assertEqual(d.get_cumulative(), DistributionData(8, 2, 3, 5))
+
+  def test_reset_to_never_exposes_an_empty_distribution(self):
+    d = DistributionCell()
+    d.reset_to(1)
+    stop = threading.Event()
+    counts = set()
+
+    def read():
+      while not stop.is_set():
+        counts.add(d.get_cumulative().count)
+
+    reader = threading.Thread(target=read)
+    reader.start()
+    try:
+      for i in range(20000):
+        d.reset_to(i % 7 + 1)
+    finally:
+      stop.set()
+      reader.join()
+    self.assertEqual(counts, {1})
+
   def test_integer_only(self):
     d = DistributionCell()
     d.update(3.1)