Pack scattered scheduler IPC channel state into a dedicated container (#25714)

test/registered/unit/managers/test_priority_scheduling_disaggregation.py

@@ -30,7 +30,7 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase):
         scheduler.model_config = SimpleNamespace(num_key_value_heads=8)
         scheduler.disagg_prefill_bootstrap_queue = MagicMock()
         scheduler.disagg_decode_prealloc_queue = MagicMock()
-        scheduler.send_to_tokenizer = MagicMock()
+        scheduler.ipc_channels = MagicMock()
         return scheduler
 
     def _new_req(self, priority=None):
@@ -72,7 +72,7 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase):
         scheduler._add_request_to_queue(req)
 
         scheduler.disagg_decode_prealloc_queue.add.assert_not_called()
-        scheduler.send_to_tokenizer.send_output.assert_called_once()
+        scheduler.ipc_channels.send_to_tokenizer.send_output.assert_called_once()
         req.time_stats.trace_ctx.abort.assert_called_once()
 
 

test/registered/unit/managers/test_scheduler_flush_cache.py

@@ -16,7 +16,7 @@ class TestSchedulerFlushCache(unittest.TestCase):
     def _new_scheduler(self) -> Scheduler:
         scheduler = Scheduler.__new__(Scheduler)
         scheduler._pending_flush = None
-        scheduler.send_to_tokenizer = MagicMock()
+        scheduler.ipc_channels = MagicMock()
         scheduler.flush_cache = MagicMock(return_value=True)
         scheduler.is_fully_idle = MagicMock(return_value=False)
         return scheduler
@@ -82,7 +82,7 @@ class TestSchedulerFlushCache(unittest.TestCase):
 
         self.assertIsNone(scheduler._pending_flush)
         scheduler.flush_cache.assert_called_once()
-        out = scheduler.send_to_tokenizer.send_output.call_args.args[0]
+        out = scheduler.ipc_channels.send_to_tokenizer.send_output.call_args.args[0]
         self.assertTrue(out.success)
 
     def test_pending_flush_expires_on_timeout(self):
@@ -95,7 +95,7 @@ class TestSchedulerFlushCache(unittest.TestCase):
 
         self.assertIsNone(scheduler._pending_flush)
         scheduler.flush_cache.assert_not_called()
-        out = scheduler.send_to_tokenizer.send_output.call_args.args[0]
+        out = scheduler.ipc_channels.send_to_tokenizer.send_output.call_args.args[0]
         self.assertFalse(out.success)
 
     def test_pending_flush_survives_before_deadline(self):
@@ -107,7 +107,7 @@ class TestSchedulerFlushCache(unittest.TestCase):
             Scheduler._check_pending_flush(scheduler)
 
         self.assertIsNotNone(scheduler._pending_flush)
-        scheduler.send_to_tokenizer.send_output.assert_not_called()
+        scheduler.ipc_channels.send_to_tokenizer.send_output.assert_not_called()
 
 
 if __name__ == "__main__":