sglang 0.1.16__py3-none-any.whl → 0.1.17__py3-none-any.whl

sglang/srt/managers/controller/dp_worker.py

@@ -0,0 +1,110 @@
+"""A data parallel worker thread."""
+import asyncio
+import logging
+import queue
+import threading
+from typing import List, Callable
+
+import uvloop
+import zmq
+
+from sglang.global_config import global_config
+from sglang.srt.managers.controller.tp_worker import ModelTpClient
+from sglang.srt.managers.io_struct import BatchTokenIDOut
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import kill_parent_process
+from sglang.utils import get_exception_traceback
+
+logger = logging.getLogger("srt.controller")
+CHECKING_INTERVAL = 5
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
+
+class DataParallelWorkerThread(threading.Thread):
+    def __init__(
+        self,
+        worker_id: int,
+        request_queue: queue.Queue,
+        detokenizer_port: int,
+        step_func: Callable,
+    ):
+        super(DataParallelWorkerThread, self).__init__()
+        self.worker_id = worker_id
+        self.request_queue = request_queue
+        self.liveness = True
+        self.request_dependency_delay = global_config.request_dependency_delay
+
+        context = zmq.asyncio.Context()
+        self.send_to_detokenizer = context.socket(zmq.PUSH)
+        self.send_to_detokenizer.connect(f"tcp://127.0.0.1:{detokenizer_port}")
+
+        self.step = step_func
+
+    async def loop_for_forward(self):
+        while self.liveness:
+            requests = []
+            while not self.request_queue.empty():
+                requests.append(self.request_queue.get())
+
+            out_pyobjs: List[BatchTokenIDOut] = []
+            try:
+                out_pyobjs = await self.step(requests)
+            except Exception:
+                for r in requests:
+                    self.request_queue.put(r)
+                logger.error(
+                    f"Worker thread {self.worker_id}: "
+                    f"failed to get back from Model Server\n"
+                    f"{get_exception_traceback()}"
+                )
+                self.liveness = False
+                # Crash the whole server when there are any errors.
+                # TODO(lianmin): make this an option.
+                kill_parent_process()
+                return
+
+            for obj in out_pyobjs:
+                self.send_to_detokenizer.send_pyobj(obj)
+
+            # async sleep for receiving the subsequent request and avoiding cache miss
+            if len(out_pyobjs) != 0:
+                has_finished = any([obj.finished_reason is not None for obj in out_pyobjs])
+                if has_finished:
+                    await asyncio.sleep(self.request_dependency_delay)
+            await asyncio.sleep(global_config.wait_for_new_request_delay)
+
+    async def monitoring(self):
+        while True:
+            await asyncio.sleep(CHECKING_INTERVAL)
+            # can plug in monitoring logic here
+
+    def run(self):
+        logger.info(f"DataParallelWorkerThread {self.worker_id} start")
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.create_task(self.monitoring())
+        loop.run_until_complete(self.loop_for_forward())
+
+
+def start_data_parallel_worker(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    model_overide_args,
+    gpu_ids: List[int],
+    worker_id: int,
+):
+    model_tp_client = ModelTpClient(
+        gpu_ids,
+        server_args,
+        port_args.model_port_args[worker_id],
+        model_overide_args,
+    )
+    worker_thread = DataParallelWorkerThread(
+        worker_id=worker_id,
+        request_queue=queue.Queue(),
+        detokenizer_port=port_args.detokenizer_port,
+        step_func=model_tp_client.step,
+    )
+    worker_thread.start()
+    return worker_thread