[Python][RunInference] Add a capacity-based autoscaling signal for model-server handlers

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)