sglang 0.1.14__py3-none-any.whl → 0.1.21__py3-none-any.whl

sglang/srt/managers/controller/manager_multi.py

@@ -0,0 +1,195 @@
+"""
+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.controller.dp_worker import (
+    DataParallelWorkerThread,
+    start_data_parallel_worker,
+)
+from sglang.srt.managers.io_struct import (
+    AbortReq,
+    FlushCacheReq,
+    TokenizedGenerateReqInput,
+)
+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:
+    """A controller that manages multiple data parallel workers."""
+
+    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.new_event_loop()
+    loop.set_default_executor(ThreadPoolExecutor(max_workers=256))
+
+    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,177 @@
+"""A controller that manages a group of tensor parallel workers."""
+
+import multiprocessing
+import logging
+import os
+import pickle
+
+import torch
+import torch.distributed as dist
+import zmq
+import zmq.asyncio
+
+from sglang.srt.managers.controller.tp_worker import ModelTpServer
+from sglang.srt.server_args import PortArgs, ServerArgs, ModelPortArgs
+from sglang.srt.utils import kill_parent_process
+from sglang.utils import get_exception_traceback
+
+
+logger = logging.getLogger("srt.controller")
+
+
+def run_tp_server(
+    gpu_id: int,
+    tp_rank: int,
+    server_args: ServerArgs,
+    model_port_args: ModelPortArgs,
+    model_overide_args: dict,
+):
+    """Run a tp server."""
+    try:
+        model_server = ModelTpServer(
+            gpu_id,
+            tp_rank,
+            server_args,
+            model_port_args,
+            model_overide_args,
+        )
+        tp_cpu_group = model_server.model_runner.tp_group.cpu_group
+
+        while True:
+            recv_reqs = broadcast_recv_input(None, tp_rank, tp_cpu_group)
+            model_server.exposed_step(recv_reqs)
+    except Exception:
+        logger.error("Exception in run_tp_server:\n" + get_exception_traceback())
+        raise
+
+
+def launch_tp_servers(gpu_ids, tp_rank_range, server_args,
+                      model_port_args, model_overide_args):
+    """Launch multiple tp servers."""
+    procs = []
+    for i in tp_rank_range:
+        proc = multiprocessing.Process(target=run_tp_server, args=(
+            gpu_ids[i], i, server_args, model_port_args, model_overide_args
+        ))
+        proc.start()
+        procs.append(proc)
+
+    return procs
+
+
+def broadcast_recv_input(data, rank, dist_group):
+    """Broadcast inputs from rank=0 to all other ranks with torch.dist backend."""
+
+    if rank == 0:
+        if len(data) == 0:
+            tensor_size = torch.tensor([0], dtype=torch.long)
+            dist.broadcast(tensor_size, src=0, group=dist_group)
+        else:
+            serialized_data = pickle.dumps(data)
+            size = len(serialized_data)
+            tensor_data = torch.ByteTensor(list(serialized_data))
+            tensor_size = torch.tensor([size], dtype=torch.long)
+
+            dist.broadcast(tensor_size, src=0, group=dist_group)
+            dist.broadcast(tensor_data, src=0, group=dist_group)
+    else:
+        tensor_size = torch.tensor([0], dtype=torch.long)
+        dist.broadcast(tensor_size, src=0, group=dist_group)
+        size = tensor_size.item()
+
+        if size == 0:
+            return []
+
+        tensor_data = torch.empty(size, dtype=torch.uint8)
+        dist.broadcast(tensor_data, src=0, group=dist_group)
+
+        serialized_data = bytes(tensor_data.tolist())
+        data = pickle.loads(serialized_data)
+        return data
+
+
+class ControllerSingle:
+    """A controller that manages a group of tensor parallel workers."""
+
+    def __init__(self, server_args: ServerArgs, port_args: PortArgs, model_overide_args: dict):
+        # Parse args
+        self.server_args = server_args
+
+        # Init communication
+        context = zmq.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 model server
+        tp_size_local = server_args.tp_size // server_args.nnodes
+        gpu_ids = [i for _ in range(server_args.nnodes) for i in range(tp_size_local)]
+
+        # Launch other tp ranks
+        if tp_size_local > 1:
+            tp_rank_range = range(1, tp_size_local)
+            self.tp_procs = launch_tp_servers(
+                gpu_ids, tp_rank_range, server_args,
+                port_args.model_port_args[0], model_overide_args)
+
+        # Launch tp rank 0
+        self.tp_server = ModelTpServer(
+            gpu_ids[0],
+            0,
+            server_args,
+            port_args.model_port_args[0],
+            model_overide_args,
+        )
+        self.tp_cpu_group = self.tp_server.model_runner.tp_group.cpu_group
+
+    def loop_for_forward(self):
+        while True:
+            recv_reqs = self.recv_requests()
+
+            if self.server_args.tp_size > 1:
+                broadcast_recv_input(recv_reqs, 0, self.tp_cpu_group)
+
+            out_pyobjs = self.tp_server.exposed_step(recv_reqs)
+
+            for obj in out_pyobjs:
+                self.send_to_detokenizer.send_pyobj(obj)
+
+    def recv_requests(self):
+        recv_reqs = []
+        while True:
+            try:
+                recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)
+                recv_reqs.append(recv_req)
+            except zmq.ZMQError:
+                break
+        return recv_reqs
+
+
+def start_controller_process(
+    server_args: ServerArgs, port_args: PortArgs, pipe_writer, model_overide_args: dict
+):
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    try:
+        controller = ControllerSingle(server_args, port_args, model_overide_args)
+    except Exception:
+        pipe_writer.send(get_exception_traceback())
+        raise
+
+    pipe_writer.send("init ok")
+
+    try:
+        controller.loop_for_forward()
+    except Exception:
+        logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
+    finally:
+        for t in controller.tp_procs:
+            os.kill(t.pid, 9)
+        kill_parent_process()