refactor(disagg): move _is_watermark_ready into StagingManagerMixin (#36030)

python/sglang/srt/disaggregation/common/staging_handler.py

@@ -833,6 +833,11 @@ class StagingManagerMixin:
     optionally ``kv_buffer_tensors``.
     """
 
+    def _is_watermark_ready(
+        self, session_id: str, alloc_round: int, alloc_end: int
+    ) -> bool:
+        return is_watermark_ready(self._staging_ctx, session_id, alloc_round, alloc_end)
+
     def _handle_staging_req(self, msg):
         room = int(msg[1].decode("ascii"))
         session_id = msg[4].decode("ascii")

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -364,15 +364,6 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
         )
         self.kv_buffer_tensors = None
 
-    def _is_watermark_ready(
-        self, session_id: str, alloc_round: int, alloc_end: int
-    ) -> bool:
-        from sglang.srt.disaggregation.common.staging_handler import (
-            is_watermark_ready,
-        )
-
-        return is_watermark_ready(self._staging_ctx, session_id, alloc_round, alloc_end)
-
     def _try_create_staging_strategy(self, staging_buffer):
         if not self.enable_staging or self.kv_buffer_tensors is None:
             return None

python/sglang/srt/disaggregation/nixl/conn.py

@@ -589,15 +589,6 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
         self._staging_ctx.room_bootstrap[room] = bootstrap_infos
         self._staging_ctx.room_receivers[room] = receiver
 
-    def _is_watermark_ready(
-        self, agent_name: str, alloc_round: int, alloc_end: int
-    ) -> bool:
-        from sglang.srt.disaggregation.common.staging_handler import (
-            is_watermark_ready,
-        )
-
-        return is_watermark_ready(self._staging_ctx, agent_name, alloc_round, alloc_end)
-
     def _start_decode_listener_thread(self):
         """Decode-side ZMQ listener for STAGING_REQ and ABORT_ACK. A thread, not
         NIXL notifs: the decode agent has no progress thread, so notifs only drain