feat: add sglang worker task monkey patch

README.md

@@ -122,8 +122,8 @@ import sglang as sgl
 
 sgl_engine = sgl.Engine(model_path="xxx", tp_size=2, random_seed=42)
 awex_config = InferenceConfig.from_sgl_engine(sgl_engine, comm_backend="nccl")
-# for sglang support, you must ensure https://github.com/sgl-project/sglang/pull/13595
-# is included in your sglang version
+# Awex installs a runtime SGLang patch when awex.engine.sglang is imported, so
+# SGLang does not need sgl-project/sglang#13595.
 inference_engine = SGLangEngine(awex_config, sgl_engine)
 reader = WeightsReader(inference_engine)
 reader.initialize()

awex/__init__.py

@@ -16,13 +16,26 @@
 # under the License.
 
 from awex.config import InferenceConfig as InferenceConfig
-from awex.reader.nccl_reader import NCCLWorkerWeightsReader as NCCLWorkerWeightsReader
-from awex.reader.weights_reader import WeightsReader as WeightsReader
-from awex.writer.nccl_writer import NCCLWeightsWriter as NCCLWeightsWriter
 
 __all__ = [
     "InferenceConfig",
     "NCCLWeightsWriter",
     "WeightsReader",
     "NCCLWorkerWeightsReader",
 ]
+
+
+def __getattr__(name):
+    if name == "NCCLWeightsWriter":
+        from awex.writer.nccl_writer import NCCLWeightsWriter
+
+        return NCCLWeightsWriter
+    if name == "WeightsReader":
+        from awex.reader.weights_reader import WeightsReader
+
+        return WeightsReader
+    if name == "NCCLWorkerWeightsReader":
+        from awex.reader.nccl_reader import NCCLWorkerWeightsReader
+
+        return NCCLWorkerWeightsReader
+    raise AttributeError(f"module 'awex' has no attribute {name!r}")

awex/engine/sglang.py

@@ -22,13 +22,17 @@
 from awex.config import InferenceConfig
 from awex.engine.core import InferenceEngine
 from awex.reader.weights_reader import get_weights_exchange_reader
+from awex.sglang_patch import ensure_sglang_patched
 from awex.util.gpu import get_gpu_status
 
 logger = logging.getLogger(__name__)
 
+ensure_sglang_patched()
+
 
 class SGLangEngine(InferenceEngine):
     def __init__(self, config: Union[Dict[str, Any], InferenceConfig], sgl_engine):
+        ensure_sglang_patched()
         super().__init__(sgl_engine.tokenizer_manager.model_config)
         if isinstance(config, dict):
             config = InferenceConfig.from_dict(config)
@@ -140,6 +144,7 @@ def resume_memory_occupation(self, tags: Optional[List[str]] = None) -> None:
             logger.info(f"GPU status after resume:\n{get_gpu_status()}")
 
     def execute_task_in_model_worker(self, fn, **kwargs):
+        ensure_sglang_patched()
         if not self._initialized:
             raise RuntimeError("Engine not initialized. Call `initialize` first.")
         if self.node_rank != 0:
@@ -159,6 +164,7 @@ def engine_rank(self):
 
 
 def extract_sgl_config(config: Dict[str, Any]) -> Dict[str, Any]:
+    ensure_sglang_patched()
     from sglang.srt.server_args import ServerArgs
 
     engine_kwargs = {