diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index 74316ca05..3132fdcff 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/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") diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 503818092..47bec0d76 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 2f1b43f40..49314582d 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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