diff --git a/sdks/python/apache_beam/ml/inference/model_manager.py b/sdks/python/apache_beam/ml/inference/model_manager.py index 9c0d90d7a02f..0c65701d185b 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 @@ 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 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 370961eb61dd..4f59905c9881 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 @@ 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):