From cce0a1244b5a3ee60411fb403f65f856a1d012f0 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Sun, 23 Aug 2026 01:24:58 +0800 Subject: [PATCH] refactor(disagg): hoist staging helper imports out of the bootstrap loops (#35980) --- python/sglang/srt/disaggregation/mooncake/conn.py | 10 ++-------- python/sglang/srt/disaggregation/nixl/conn.py | 10 ++-------- 2 files changed, 4 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 4b34b3f43..6f4a2425f 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -29,6 +29,8 @@ from sglang.srt.disaggregation.common.staging_handler import ( PrefillStagingContext, StagingManagerMixin, StagingTransferInfo, + handle_staging_rsp, + handle_watermark_msg, ) from sglang.srt.disaggregation.common.utils import ( AuxDataCodec, @@ -1944,18 +1946,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): room = waiting_req_bytes[0].decode("ascii") # Staging: decode reports consumption watermark back to prefill if room == "WATERMARK": - from sglang.srt.disaggregation.common.staging_handler import ( - handle_watermark_msg, - ) - handle_watermark_msg(self._staging_ctx, waiting_req_bytes) continue # Staging: decode replies with allocated staging offset if room == "STAGING_RSP": - from sglang.srt.disaggregation.common.staging_handler import ( - handle_staging_rsp, - ) - handle_staging_rsp(waiting_req_bytes, self.transfer_infos) continue # Decode-side abort notification: mark room as failed and ACK diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index eae55f6cd..2f1b43f40 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -28,6 +28,8 @@ from sglang.srt.disaggregation.common.conn import ( from sglang.srt.disaggregation.common.staging_handler import ( STAGING_WATERMARK_WAIT_S, StagingManagerMixin, + handle_staging_rsp, + handle_watermark_msg, ) from sglang.srt.disaggregation.common.utils import ( FastQueue, @@ -2674,20 +2676,12 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): # Staging: decode reports consumption watermark back to prefill if waiting_req_bytes[0] == b"WATERMARK": if self.enable_staging: - from sglang.srt.disaggregation.common.staging_handler import ( - handle_watermark_msg, - ) - handle_watermark_msg(self._staging_ctx, waiting_req_bytes) continue # Staging: decode replies with allocated staging offset if waiting_req_bytes[0] == b"STAGING_RSP": if self.enable_staging: - from sglang.srt.disaggregation.common.staging_handler import ( - handle_staging_rsp, - ) - handle_staging_rsp(waiting_req_bytes, self.transfer_infos) continue