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

sglang/srt/managers/controller/manager_multi.py

@@ -0,0 +1,191 @@
+"""
+A controller that manages multiple data parallel workers.
+Each data parallel worker can manage multiple tensor parallel workers.
+"""
+
+import asyncio
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from enum import Enum, auto
+from typing import Dict
+
+import zmq
+import zmq.asyncio
+
+from sglang.global_config import global_config
+from sglang.srt.managers.io_struct import (
+    AbortReq,
+    FlushCacheReq,
+    TokenizedGenerateReqInput,
+)
+from sglang.srt.managers.controller.dp_worker import (
+    DataParallelWorkerThread,
+    start_data_parallel_worker,
+)
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.utils import get_exception_traceback
+
+logger = logging.getLogger("srt.controller")
+
+
+class LoadBalanceMethod(Enum):
+    ROUND_ROBIN = auto()
+    SHORTEST_QUEUE = auto()
+
+    @classmethod
+    def from_str(cls, method: str):
+        method = method.upper()
+        try:
+            return cls[method]
+        except KeyError as exc:
+            raise ValueError(f"Invalid load balance method: {method}") from exc
+
+
+class Controller:
+    def __init__(
+        self,
+        load_balance_method: str,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+        model_overide_args,
+    ):
+        self.load_balance_method = LoadBalanceMethod.from_str(load_balance_method)
+        self.server_args = server_args
+        self.port_args = port_args
+
+        if self.load_balance_method == LoadBalanceMethod.ROUND_ROBIN:
+            self.round_robin_counter = 0
+
+        self.dispatch_lookup = {
+            LoadBalanceMethod.ROUND_ROBIN: self.round_robin_scheduler,
+            LoadBalanceMethod.SHORTEST_QUEUE: self.shortest_queue_scheduler,
+        }
+        self.dispatching = self.dispatch_lookup[self.load_balance_method]
+
+        # Init communication
+        context = zmq.asyncio.Context()
+        self.recv_from_tokenizer = context.socket(zmq.PULL)
+        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.router_port}")
+
+        # Init status
+        self.recv_reqs = []
+
+        # Start data parallel workers
+        self.workers: Dict[int, DataParallelWorkerThread] = {}
+        tp_size = server_args.tp_size
+
+        def start_dp_worker(i):
+            try:
+                gpu_ids = list(range(i * tp_size, (i + 1) * tp_size))
+                worker_thread = start_data_parallel_worker(
+                    server_args, port_args, model_overide_args, gpu_ids, i
+                )
+                self.workers[i] = worker_thread
+            except Exception:
+                logger.error(
+                    f"Failed to start local worker {i}\n{get_exception_traceback()}"
+                )
+
+        for i in range(server_args.dp_size):
+            start_dp_worker(i)
+
+        # Parallel launch is slower, probably due to the disk bandwidth limitations.
+        # with ThreadPoolExecutor(server_args.dp_size) as executor:
+        #     executor.map(start_dp_worker, range(server_args.dp_size))
+
+    def have_any_live_worker(self):
+        return any(worker_thread.liveness for worker_thread in self.workers.values())
+
+    def put_req_to_worker(self, worker_id, req):
+        self.workers[worker_id].request_queue.put(req)
+
+    async def round_robin_scheduler(self, input_requests):
+        available_workers = list(self.workers.keys())
+        for r in input_requests:
+            self.put_req_to_worker(available_workers[self.round_robin_counter], r)
+            self.round_robin_counter = (self.round_robin_counter + 1) % len(
+                available_workers
+            )
+        return
+
+    async def shortest_queue_scheduler(self, input_requests):
+        for r in input_requests:
+            worker = min(
+                self.workers, key=lambda w: self.workers[w].request_queue.qsize()
+            )
+            self.put_req_to_worker(worker, r)
+        return
+
+    async def remove_dead_workers(self):
+        for i in list(self.workers.keys()):
+            worker_thread = self.workers[i]
+            if not worker_thread.liveness:
+                worker_thread.join()
+                # move unsuccessful requests back to the queue
+                while not worker_thread.request_queue.empty():
+                    self.recv_reqs.append(worker_thread.request_queue.get())
+                del self.workers[i]
+                logger.info(f"Stale worker {i} removed")
+
+    async def loop_for_forward(self):
+        while True:
+            await self.remove_dead_workers()
+
+            if self.have_any_live_worker():
+                next_step_input = list(self.recv_reqs)
+                self.recv_reqs = []
+                if next_step_input:
+                    await self.dispatching(next_step_input)
+            #else:
+            #    logger.error("There is no live worker.")
+
+            await asyncio.sleep(global_config.wait_for_new_request_delay)
+
+    async def loop_for_recv_requests(self):
+        while True:
+            recv_req = await self.recv_from_tokenizer.recv_pyobj()
+            if isinstance(recv_req, FlushCacheReq):
+                # TODO(lsyin): apply more specific flushCacheReq
+                for worker_thread in self.workers.values():
+                    worker_thread.request_queue.put(recv_req)
+            elif isinstance(recv_req, TokenizedGenerateReqInput):
+                self.recv_reqs.append(recv_req)
+            elif isinstance(recv_req, AbortReq):
+                in_queue = False
+                for i, req in enumerate(self.recv_reqs):
+                    if req.rid == recv_req.rid:
+                        self.recv_reqs[i] = recv_req
+                        in_queue = True
+                        break
+                if not in_queue:
+                    # Send abort req to all TP groups
+                    for worker in list(self.workers.keys()):
+                        self.put_req_to_worker(worker, recv_req)
+            else:
+                logger.error(f"Invalid object: {recv_req}")
+
+
+def start_controller_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    pipe_writer,
+    model_overide_args=None,
+):
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    try:
+        controller = Controller(
+            server_args.load_balance_method, server_args, port_args, model_overide_args
+        )
+    except Exception:
+        pipe_writer.send(get_exception_traceback())
+        raise
+
+    pipe_writer.send("init ok")
+    loop = asyncio.get_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.create_task(controller.loop_for_recv_requests())
+    loop.run_until_complete(controller.loop_for_forward())

sglang/srt/managers/controller/manager_single.py

@@ -0,0 +1,97 @@
+"""A controller that manages a group of tensor parallel workers."""
+import asyncio
+import logging
+import time
+
+import uvloop
+import zmq
+import zmq.asyncio
+
+from sglang.global_config import global_config
+from sglang.srt.managers.controller.tp_worker import ModelTpClient
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import kill_parent_process
+from sglang.utils import get_exception_traceback
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
+logger = logging.getLogger("srt.controller")
+
+
+class ControllerSingle:
+    def __init__(self, model_client: ModelTpClient, port_args: PortArgs):
+        # Init communication
+        context = zmq.asyncio.Context(2)
+        self.recv_from_tokenizer = context.socket(zmq.PULL)
+        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.router_port}")
+
+        self.send_to_detokenizer = context.socket(zmq.PUSH)
+        self.send_to_detokenizer.connect(
+            f"tcp://127.0.0.1:{port_args.detokenizer_port}"
+        )
+
+        # Init status
+        self.model_client = model_client
+        self.recv_reqs = []
+
+        # Init some configs
+        self.request_dependency_delay = global_config.request_dependency_delay
+
+    async def loop_for_forward(self):
+        while True:
+            next_step_input = list(self.recv_reqs)
+            self.recv_reqs = []
+            out_pyobjs = await self.model_client.step(next_step_input)
+
+            for obj in out_pyobjs:
+                self.send_to_detokenizer.send_pyobj(obj)
+
+            # async sleep for receiving the subsequent request and avoiding cache miss
+            slept = False
+            if len(out_pyobjs) != 0:
+                has_finished = any([obj.finished_reason is not None for obj in out_pyobjs])
+                if has_finished:
+                    if self.request_dependency_delay > 0:
+                        slept = True
+                        await asyncio.sleep(self.request_dependency_delay)
+
+            if not slept:
+                await asyncio.sleep(global_config.wait_for_new_request_delay)
+
+    async def loop_for_recv_requests(self):
+        while True:
+            recv_req = await self.recv_from_tokenizer.recv_pyobj()
+            self.recv_reqs.append(recv_req)
+
+
+def start_controller_process(
+    server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args
+):
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    try:
+        model_client = ModelTpClient(
+            list(range(server_args.tp_size)),
+            server_args,
+            port_args.model_port_args[0],
+            model_overide_args,
+        )
+        controller = ControllerSingle(model_client, port_args)
+    except Exception:
+        pipe_writer.send(get_exception_traceback())
+        raise
+
+    pipe_writer.send("init ok")
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.create_task(controller.loop_for_recv_requests())
+    try:
+        loop.run_until_complete(controller.loop_for_forward())
+    except Exception:
+        logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
+    finally:
+        kill_parent_process()