Fix handling of cancelled tickets.

sdks/python/apache_beam/ml/inference/model_manager.py

@@ -470,12 +470,13 @@ def should_spawn_model(self, tag: str, ticket_num: int) -> bool:
     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 @@ def _evict_to_make_space(
     # 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

sdks/python/apache_beam/ml/inference/model_manager_test.py

@@ -30,6 +30,7 @@
   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 @@ def run_inference():
     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 @@ def acquire_b():
     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):