(beam) branch master updated: Fix handling of cancelled tickets. (#40588)

[tv...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tvalentyn pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git


The following commit(s) were added to refs/heads/master by this push:
     new a3f3c8e01d4 Fix handling of cancelled tickets. (#40588)
a3f3c8e01d4 is described below

commit a3f3c8e01d4380ce2fd56e14adeacf7c09f51313
Author: tvalentyn <tv...@users.noreply.github.com>
AuthorDate: Fri Oct 16 10:32:07 2026 -0700

    Fix handling of cancelled tickets. (#40588)
---
 .../apache_beam/ml/inference/model_manager.py      | 16 ++++---
 .../apache_beam/ml/inference/model_manager_test.py | 54 ++++++++++++++++++++++
 2 files changed, 64 insertions(+), 6 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py b/sdks/python/apache_beam/ml/inference/model_manager.py
index 69db694cc97..bc16feb6f60 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager.py
@@ -470,12 +470,13 @@ class ModelManager:
     return False
 
   def _wake_next_in_queue(self):
+    # Clean up cancelled tickets at head of queue
+    while self._wait_queue and self._wait_queue[
+        0].ticket_num in self._cancelled_tickets:
+      self._cancelled_tickets.remove(self._wait_queue[0].ticket_num)
+      heapq.heappop(self._wait_queue)
+    # The queue may be empty if all remaining tickets were cancelled.
     if self._wait_queue:
-      # Clean up cancelled tickets at head of queue
-      while self._wait_queue and self._wait_queue[
-          0].ticket_num in self._cancelled_tickets:
-        self._cancelled_tickets.remove(self._wait_queue[0].ticket_num)
-        heapq.heappop(self._wait_queue)
       next_inline = self._wait_queue[0]
       next_inline.wake_event.set()
 
@@ -660,7 +661,10 @@ class ModelManager:
     # TODO: Also factor in the active counts to avoid thrashing
     demand_map = Counter()
     for item in self._wait_queue:
-      demand_map[item.tag] += 1
+      # Cancelled tickets linger in the queue until they reach the head;
+      # they no longer represent demand.
+      if item.ticket_num not in self._cancelled_tickets:
+        demand_map[item.tag] += 1
 
     my_demand = demand_map[requesting_tag]
     am_i_starving = len(self._models[requesting_tag]) == 0
diff --git a/sdks/python/apache_beam/ml/inference/model_manager_test.py b/sdks/python/apache_beam/ml/inference/model_manager_test.py
index fc84f4ba859..d1f2b35bc1d 100644
--- a/sdks/python/apache_beam/ml/inference/model_manager_test.py
+++ b/sdks/python/apache_beam/ml/inference/model_manager_test.py
@@ -30,6 +30,7 @@ try:
   from apache_beam.metrics.metricbase import MetricName
   from apache_beam.ml.inference.model_manager import GPUMonitor
   from apache_beam.ml.inference.model_manager import ModelManager
+  from apache_beam.ml.inference.model_manager import QueueTicket
   from apache_beam.ml.inference.model_manager import ResourceEstimator
 except ImportError as e:
   raise unittest.SkipTest("Model Manager dependencies are not installed")
@@ -436,6 +437,32 @@ class TestModelManager(unittest.TestCase):
     est_cost = self.manager._estimator.get_estimate(model_name)
     self.assertAlmostEqual(est_cost, model_cost, delta=100.0)
 
+  def test_wake_next_in_queue_all_cancelled(self):
+    """Draining a queue of only cancelled tickets must not raise."""
+    with self.manager._cv:
+      for ticket_num in (100, 101):
+        self.manager._wait_queue.append(QueueTicket(1, ticket_num, "m"))
+        self.manager._cancelled_tickets.add(ticket_num)
+
+      self.manager._wake_next_in_queue()
+
+      self.assertEqual(self.manager._wait_queue, [])
+      self.assertEqual(self.manager._cancelled_tickets, set())
+
+  def test_wake_next_in_queue_skips_cancelled_head(self):
+    """The first non-cancelled ticket is woken."""
+    cancelled = QueueTicket(1, 100, "m")
+    live = QueueTicket(1, 101, "m")
+    with self.manager._cv:
+      self.manager._wait_queue.extend([cancelled, live])
+      self.manager._cancelled_tickets.add(100)
+
+      self.manager._wake_next_in_queue()
+
+      self.assertEqual(self.manager._wait_queue, [live])
+      self.assertTrue(live.wake_event.is_set())
+      self.assertFalse(cancelled.wake_event.is_set())
+
 
 class TestModelManagerEviction(unittest.TestCase):
   def setUp(self):
@@ -583,6 +610,33 @@ class TestModelManagerEviction(unittest.TestCase):
     self.manager.release_model("B", model_b)
     executor.shutdown(wait=True)
 
+  def test_cancelled_tickets_do_not_count_as_demand(self):
+    self.manager._estimator.set_initial_estimate("A", 4000)
+    self.manager._estimator.set_initial_estimate("B", 4000)
+
+    model_a = self.manager.acquire_model("A", self.create_loader("A", 4000))
+    self.manager.release_model("A", model_a)
+    # B stays active so it is not starving and demand comparison applies.
+    model_b = self.manager.acquire_model("B", self.create_loader("B", 4000))
+
+    with self.manager._cv:
+      # One live request for B, two cancelled (timed out) requests for A.
+      self.manager._wait_queue.extend([
+          QueueTicket(1, 100, "B"),
+          QueueTicket(1, 101, "A"),
+          QueueTicket(1, 102, "A"),
+      ])
+      self.manager._cancelled_tickets.update({101, 102})
+
+      # 8000 used + 6000 requested > 12000, so A must be evicted. Without
+      # ignoring cancelled tickets, A's demand (2) >= B's (1) protects it.
+      made_space = self.manager._evict_to_make_space(
+          limit=12000, est_cost=6000, requesting_tag="B")
+
+    self.assertTrue(made_space)
+    self.assertTrue(model_a.deleted)
+    self.manager.release_model("B", model_b)
+
 
 class TestGPUMonitor(unittest.TestCase):
   def setUp(self):