server_LLMgenerate.py

"""Qwen3.5 异构推理服务启动入口：模型加载与任务分配均在本文件完成。

设备按 --tp-size 分组为 Megatron 1D 张量并行组（TP 组，权重切分见
models/qwen35.py，基于 DTensor + ParallelStyle），模型层区间均分到
各 TP 组（流水线分片），经 DeviceConfig.compute 注入各 worker；ModelWorker
在基础 DeviceWorker 上扩展模型分片计算（prefill/decode_step/free）；
LLMDispatcher 在 Dispatcher_base 基础上实现 prefill/decode 的流水线调度
端点（/model、/llm_generate）。模型加载使用 models/qwen35.py。

用法：
    .venv/bin/python server_LLMgenerate.py [--model-dir ../Qwen3.5-0.8B] [--devices cpu,cuda] [--port 8000]
    # 2 路张量并行：cpu0,cpu1 组成一个 TP 组，共同承载全部 24 层
    .venv/bin/python server_LLMgenerate.py --devices cpu0,cpu1 --tp-size 2
    # 2 组 TP（每组 2 设备）× 2 段流水线：每组 12 层
    .venv/bin/python server_LLMgenerate.py --devices cpu0,cpu1,cpu2,cpu3 --tp-size 2

多节点（Ray 集群）部署：连接已有集群后自动把各设备轮转绑定到集群节点上
（节点亲和调度），同一 TP 组的各 rank 分散到不同节点，如
    --ray-address auto --devices cpu0,cpu1 --tp-size 2
各 TP 组的 rendezvous 地址自动取该组 rank 0 所在节点 IP
（可用 --tp-master-addr 覆盖）；
要求项目代码与模型目录在各节点上路径一致（如挂载到相同路径）。
"""

from __future__ import annotations

import argparse
import asyncio
import os
import socket
import time
import uuid
from typing import Any

import ray
import torch
from pydantic import BaseModel
from ray import serve

from heterogeneous_serve.config import DeviceConfig, FrameworkConfig, default_config
from heterogeneous_serve.dispatcher import Dispatcher_base, build_app, fastapi_app
from heterogeneous_serve.ops import get_op
from heterogeneous_serve.worker import DeviceWorker
from models.qwen35 import (
    Qwen35TextConfig,
    Qwen35TextModel,
    TensorParallelGroup,
    init_tensor_parallel,
    load_tokenizer,
)

DEFAULT_MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "Qwen3.5-0.8B")

"""
统计时间方法 
统计起点：start = time.perf_counter()
...
统计终点：end = time.perf_counter() 
花费时间：elapsed_ms = (end - start) * 1000 /单位ms
加入worker的统计池：self._record_task_time(elapsed_ms)
"""


class GenerateRequest(BaseModel):
    """Qwen3.5 文本生成：prompt 经分词后沿模型 worker 流水线逐层前向。"""

    prompt: str
    max_new_tokens: int = 20
    temperature: float = 0.0  # <=0 为贪心


def assign_model_shards(
    config: FrameworkConfig,
    model_dir: str,
    tp_size: int = 1,
    tp_master_addr: str | None = None,
    tp_port_base: int = 29500,
    nodes: list[tuple[str, str]] | None = None,
) -> list[list[str]]:
    """任务分配：设备按 tp_size 分组为 Megatron 1D 张量并行组（TP 组），
    模型层区间均分到各 TP 组（流水线分片），写入 DeviceConfig.compute。

    每个 TP 组作为一个流水线 stage：组内各 rank 持有同一层区间的 1D 权重分片
    （models/qwen35.py），组间按层区间流水线传递隐状态。

    nodes: [(node_id, node_ip), ...]，多节点部署时给出集群节点列表，
        设备按顺序轮流转绑到各节点（节点亲和调度），保证同一 TP 组的
        各 rank 分散到不同节点；此时各组 rendezvous 地址默认取该组
        rank 0 所在节点的 IP（可用 tp_master_addr 显式覆盖）。
    返回 TP 组列表，每组为该组内 rank 顺序的设备名列表。
    """
    num_layers = Qwen35TextConfig.from_pretrained(model_dir).num_hidden_layers
    devices = config.devices
    if len(devices) % tp_size != 0:
        raise ValueError(
            f"设备数 {len(devices)} 不能被 tp_size={tp_size} 整除，无法组成 TP 组"
        )
    groups = [devices[i : i + tp_size] for i in range(0, len(devices), tp_size)]
    bounds = [round(i * num_layers / len(groups)) for i in range(len(groups) + 1)]
    for g, group in enumerate(groups):
        master_addr = tp_master_addr
        for r, d in enumerate(group):
            if nodes:
                node_id, node_ip = nodes[(g * tp_size + r) % len(nodes)]
                d.node_id = node_id
                if master_addr is None and r == 0:
                    master_addr = node_ip  # rendezvous 指向本组 rank 0 所在节点
            d.compute = {
                **d.compute,
                "model": {
                    "model_dir": model_dir,
                    "layer_start": bounds[g],
                    "layer_end": bounds[g + 1],
                    "is_first": g == 0,
                    "is_last": g == len(groups) - 1,
                    "tp": {
                        "rank": r,
                        "world_size": tp_size,
                        "master_addr": master_addr or "127.0.0.1",
                        "master_port": tp_port_base + g,  # 每个 TP 组独立的进程组端口
                    },
                },
            }
    return [[d.device for d in group] for group in groups]


class ModelWorker(DeviceWorker):
    """在基础 DeviceWorker 上e扩展 Qwen3.5 模型分片计算（流水线分片的载体）。

    compute 配置在基础结构之上追加（由 assign_model_shards 注入）：
        {"model": {"model_dir": str, "layer_start": int, "layer_end": int,
                   "is_first": bool, "is_last": bool,
                   "tp": {"rank": int, "world_size": int,
                          "master_addr": str, "master_port": int}} | None}
    按 request_id 维护本层区间各层的 KV/递推 cache。
    tp.world_size > 1 时本 worker 作为 Megatron 1D 张量并行组的一个 rank，
    只持有本层区间的权重分片，前向中经集合通信得到完整结果。
    """

    def __init__(self, device: str, compute: dict[str, Any] | None = None):
        super().__init__(device, compute)
        self.model: Qwen35TextModel | None = None
        self.tp: TensorParallelGroup | None = None
        self._tp_cfg: dict[str, Any] | None = None
        self.layer_start = self.layer_end = 0
        self.is_first = self.is_last = False
        model_cfg = (compute or {}).get("model")
        if model_cfg:
            self.layer_start = model_cfg["layer_start"]
            self.layer_end = model_cfg["layer_end"]
            self.is_first = model_cfg["is_first"]
            self.is_last = model_cfg["is_last"]
            tp_cfg = model_cfg.get("tp") or {}
            if tp_cfg.get("world_size", 1) > 1:
                # 仅保存 TP 配置，进程组初始化推迟到 dispatcher 触发，
                # 避免 Ray Serve 滚动部署时单个 worker 在 __init__ 中阻塞等待其它 rank。
                self._tp_cfg = tp_cfg
            print(f"[{self.device_name}] 开始初始化模型分片: {model_cfg['model_dir']} "
                  f"layers=[{self.layer_start},{self.layer_end}), "
                  f"tp={tp_cfg.get('world_size', 1)}")
            self.model = Qwen35TextModel.from_pretrained(
                model_cfg["model_dir"], device=self.device,
                layer_start=self.layer_start, layer_end=self.layer_end,
                is_first=self.is_first, is_last=self.is_last,
            )
            print(f"[{self.device_name}] 模型分片加载完成，可接收请求")
        self._caches: dict[str, list[dict]] = {}
        self._past_len: dict[str, int] = {}

    def info(self) -> dict[str, Any]:
        info = super().info()
        info["compute"]["model"] = self.model is not None
        if self.model is not None:
            info["model"] = {
                "layers": [self.layer_start, self.layer_end],
                "is_first": self.is_first,
                "is_last": self.is_last,
            }
            if self.tp is not None:
                info["model"]["tp"] = {
                    "rank": self.tp.rank,
                    "world_size": self.tp.world_size,
                }
            elif self._tp_cfg is not None:
                info["model"]["tp"] = {
                    "rank": self._tp_cfg["rank"],
                    "world_size": self._tp_cfg["world_size"],
                    "initialized": False,
                }
        return info

    def init_tp(self) -> dict[str, Any]:
        """初始化张量并行进程组并在模型权重上应用切分。

        该操作从 __init__ 中拆出，避免 Ray Serve 滚动部署时 worker 在初始化阶段
        阻塞等待其它 rank 而导致死锁。
        """
        if self._tp_cfg is not None and self.tp is None:
            print(f"[{self.device_name}] 开始初始化 TP 进程组: "
                  f"rank={self._tp_cfg['rank']}, world_size={self._tp_cfg['world_size']}")
            self.tp = init_tensor_parallel(
                world_size=self._tp_cfg["world_size"],
                rank=self._tp_cfg["rank"],
                device=self.device,
                master_addr=self._tp_cfg["master_addr"],
                master_port=self._tp_cfg["master_port"],
            )
            if self.model is not None:
                self.model.apply_tensor_parallel(self.tp)
            print(f"[{self.device_name}] TP 进程组初始化完成")
        return self.info()

    # ------------------------------------------------------------------
    # 模型推理计算（Qwen3.5 流水线分片）
    # ------------------------------------------------------------------

    def _require_model(self) -> Qwen35TextModel:
        if self.model is None:
            raise RuntimeError(f"设备 {self.device_name} 未配置模型推理计算")
        return self.model

    def _forward(self, request_id: str, x):
        model = self._require_model()
        # TP 下每轮 forward 前清空通信统计，确保只记录本次 forward 内的通信耗时
        if self.tp is not None:
            self.tp.reset_comm_stats()
        cache = self._caches.setdefault(
            request_id, [{} for _ in range(self.layer_end - self.layer_start)]
        )
        pos_offset = self._past_len.get(request_id, 0)
        out, _ = model.forward_segment(
            x, cache, pos_offset,
            layer_start=self.layer_start, layer_end=self.layer_end,
            is_first=self.is_first, is_last=self.is_last,
        )
        return out

    def prefill(self, request_id: str, token_ids: list[int] | None = None, hidden=None):
        """首段收 token_ids，其余段收上游隐状态；末段返回 logits (list)。

        TP 组内各 rank 经集合通信后输出一致（logits 仅 rank 0 持有），
        只有 rank 0 需要回传结果，其余 rank 返回空 dict，
        避免整词表 logits / 隐状态的重复序列化与跨节点传输。
        """
        t1 = time.perf_counter()
        x = token_ids if self.is_first else hidden
        seq_len = len(token_ids) if self.is_first else hidden.shape[1]
        out = self._forward(request_id, x)
        self._past_len[request_id] = seq_len
        t2 = time.perf_counter()
        total_ms = (t2 - t1) * 1000
        comm_stats = self.tp.get_comm_stats() if self.tp is not None else {}
        comm_ms = comm_stats.get("total_ms", 0.0)
        local_ms = total_ms - comm_ms
        print(
            f"{self.device_name} prefill 总耗时: {total_ms:.3f}ms, "
            f"本地推理: {local_ms:.3f}ms, 通信: {comm_ms:.3f}ms, "
            f"seq_len: {seq_len}, is_first: {self.is_first}, is_last: {self.is_last}, "
            f"tp_rank: {self.tp.rank if self.tp else 0}, comm_stats: {comm_stats}"
        )

        if self.tp is not None and self.tp.rank != 0:
            return {}
        if self.is_last:
            return {"logits": out.detach().to("cpu", torch.float32).tolist()}
        return {"hidden": out}

    def decode_step(self, request_id: str, token_id: int | None = None, hidden=None):
        """单 token 解码。首段收 token_id，其余段收隐状态；末段返回 logits。"""
        t1 = time.perf_counter()
        x = [token_id] if self.is_first else hidden
        out = self._forward(request_id, x)
        self._past_len[request_id] += 1
        t2 = time.perf_counter()
        total_ms = (t2 - t1) * 1000
        comm_stats = self.tp.get_comm_stats() if self.tp is not None else {}
        comm_ms = comm_stats.get("total_ms", 0.0)
        local_ms = total_ms - comm_ms
        print(
            f"{self.device_name} decode 总耗时: {total_ms:.3f}ms, "
            f"本地推理: {local_ms:.3f}ms, 通信: {comm_ms:.3f}ms, "
            f"is_first: {self.is_first}, is_last: {self.is_last}, "
            f"tp_rank: {self.tp.rank if self.tp else 0}, comm_stats: {comm_stats}"
        )

        if self.tp is not None and self.tp.rank != 0:
            return {}
        if self.is_last:
            return {"logits": out.detach().to("cpu", torch.float32).tolist()}
        return {"hidden": out}

    def free(self, request_id: str) -> bool:
        self._caches.pop(request_id, None)
        self._past_len.pop(request_id, None)
        return True


@serve.deployment
@serve.ingress(fastapi_app)
class LLMDispatcher(Dispatcher_base):
    """异构推理调度：prefill 沿设备分片逐段前向，再逐 token 解码。

    流水线由若干 TP 组构成（tp_size=1 时每组一个设备）：同一 payload 扇出到
    组内所有 rank（张量并行，组内经集合通信得到一致结果，取 rank 0 的输出），
    组间按层区间依次传递隐状态（流水线并行）。
    """

    def __init__(
        self,
        workers: dict[str, Any],
        model_devices: list[list[str]],
        model_dir: str,
        endpoints: list[str] | None = None,
    ):
        super().__init__(workers, endpoints)
        self.model_devices = model_devices
        self.model_worker_groups = [
            [self.handles[d] for d in group] for group in model_devices
        ]
        self.model_dir = model_dir
        self._tokenizer = None
        self._eos_token_id: int | None = None
        # 张量并行进程组初始化任务：推迟到首次请求，避免 __init__ 中阻塞。
        self._tp_init_task: asyncio.Future | None = None

    def _ensure_tp_init(self):
        """启动所有模型 worker 的 TP 进程组初始化（幂等）。"""
        if self._tp_init_task is not None:
            return self._tp_init_task
        handles = [h for group in self.model_worker_groups for h in group]
        self._tp_init_task = asyncio.gather(*(h.init_tp.remote() for h in handles))
        return self._tp_init_task

    def _get_tokenizer(self):
        if self._tokenizer is None:
            self._tokenizer = load_tokenizer(self.model_dir)
            self._eos_token_id = Qwen35TextConfig.from_pretrained(self.model_dir).eos_token_id
        return self._tokenizer

    @fastapi_app.get("/model")
    async def model_info(self) -> list[dict[str, Any]]:
        """模型流水线各分片的设备与层区间信息（含 TP rank）。"""
        handles = [h for group in self.model_worker_groups for h in group]
        return await asyncio.gather(*(h.info.remote() for h in handles))

    @fastapi_app.post("/llm_generate")
    async def llm_generate(self, req: GenerateRequest) -> dict[str, Any]:
        """流水线推理：prefill 沿设备分片逐段前向，再逐 token 解码。"""
        await self._ensure_tp_init()
        tokenizer = self._get_tokenizer()
        request_id = uuid.uuid4().hex
        prompt_ids = tokenizer.encode(req.prompt).ids
        generated: list[int] = []
        try:
            # prefill：token ids 进入首段，隐状态依次流过各 TP 组（流水线 stage）
            payload: dict[str, Any] = {"token_ids": prompt_ids}
            for group in self.model_worker_groups:
                outs = await asyncio.gather(
                    *(h.prefill.remote(request_id, **payload) for h in group)
                )
                payload = outs[0]  # 组内各 rank 输出一致，取 rank 0
            logits = torch.tensor(payload["logits"], dtype=torch.float32)

            for _ in range(req.max_new_tokens):
                next_id = int(get_op("sample_token")(logits, req.temperature))
                generated.append(next_id)
                if next_id == self._eos_token_id:
                    break
                payload = {"token_id": next_id}
                for group in self.model_worker_groups:
                    outs = await asyncio.gather(
                        *(h.decode_step.remote(request_id, **payload) for h in group)
                    )
                    payload = outs[0]
                logits = torch.tensor(payload["logits"], dtype=torch.float32)

        finally:
            handles = [h for group in self.model_worker_groups for h in group]
            await asyncio.gather(*(h.free.remote(request_id) for h in handles))

        return {
            "prompt": req.prompt,
            "prompt_tokens": len(prompt_ids),
            "generated_token_ids": generated,
            "text": tokenizer.decode(generated),
            "decode_tokens": len(generated),
            "pipeline": self.model_devices,
        }


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Qwen3.5 异构流水线推理服务")
    p.add_argument("--model-dir", type=str, default=DEFAULT_MODEL_DIR, help="Qwen3.5 模型目录")
    p.add_argument(
        "--devices",
        type=str,
        default=None,
        help="逗号分隔的设备列表，如 cpu,cuda,npu；缺省自动探测",
    )
    p.add_argument("--replicas", type=int, default=1, help="每类设备的 worker 副本数")
    p.add_argument(
        "--dispatcher-cpus",
        type=int,
        default=2,
        help="为 dispatcher 预留的 CPU 核数：dispatcher 钉到首个节点，同节点 worker "
        "的 torch 线程数 = 容器可见核数 - 该值，避免 worker 计算线程占满节点时 "
        "dispatcher 的 RPC/序列化被饿死；设为 0 关闭",
    )
    p.add_argument(
        "--tp-size",
        type=int,
        default=1,
        help="Megatron 1D 张量并行组大小：每 tp_size 个设备组成一个 TP 组"
        "（组内按 models/qwen35.py 切分权重），各 TP 组构成流水线 stage",
    )
    p.add_argument(
        "--tp-master-addr",
        type=str,
        default=None,
        help="TP 进程组 rendezvous 地址；缺省时多节点自动取各组 rank 0 所在节点 IP，"
        "单机为 127.0.0.1",
    )
    p.add_argument("--host", type=str, default="127.0.0.1")
    p.add_argument("--port", type=int, default=8000)
    p.add_argument(
        "--ray-address",
        type=str,
        default=None,
        help="已有 Ray 集群地址（如 auto 或 <head-ip>:6379）；缺省在本机新起单机集群",
    )
    return p.parse_args()


def _free_port() -> int:
    """取一个本机空闲端口（作为 TP 进程组 rendezvous 端口基准）。"""
    with socket.socket() as s:
        s.bind(("", 0))
        return s.getsockname()[1]


def main() -> None:
    args = parse_args()

    if args.tp_size < 1:
        raise ValueError("--tp-size 必须 >= 1")
    if args.tp_size > 1 and args.replicas != 1:
        raise ValueError("启用张量并行（--tp-size > 1）时要求 --replicas 1，"
                         "以保证各 TP rank 与设备一一对应")

    if args.devices:
        config = FrameworkConfig(
            devices=[
                DeviceConfig(device=d.strip(), replicas=args.replicas)
                for d in args.devices.split(",")
                if d.strip()
            ]
        )
    else:
        config = default_config(replicas_per_device=args.replicas)

    # 指定 --ray-address 时连接已有（多节点）Ray 集群，否则在本机新起单机集群；
    # 先初始化再分配分片，以便多节点时按集群节点做 TP rank 放置
    if args.ray_address:
        ray.init(address=args.ray_address, ignore_reinit_error=True)
    else:
        ray.init(ignore_reinit_error=True)

    # 收集存活节点 [(node_id, ip), ...]（按 IP 排序保证各次启动一致）。
    # 即使未显式传入 --ray-address，如果环境已经存在 Ray 集群（ray.init 自动连接），
    # 只要存活节点数 > 1 就按多节点处理，保证 TP rendezvous 使用真实节点 IP 而非 127.0.0.1。
    nodes: list[tuple[str, str]] | None = None
    alive_nodes = sorted(
        (
            (n["NodeID"], n["NodeManagerAddress"])
            for n in ray.nodes()
            if n["Alive"]
        ),
        key=lambda t: t[1],
    )
    if len(alive_nodes) > 1 or args.ray_address:
        nodes = alive_nodes
        print(f"集群节点: {[ip for _, ip in nodes]}")

    # 模型加载与任务分配：设备按 tp_size 组成 TP 组，层区间均分到各组，
    # 注入各 worker 的 compute 配置（含 TP rank 与进程组 rendezvous 信息）
    model_devices = assign_model_shards(
        config,
        args.model_dir,
        tp_size=args.tp_size,
        tp_master_addr=args.tp_master_addr,
        tp_port_base=_free_port(),
        nodes=nodes,
    )

    # 为 dispatcher 预留 CPU：多节点时把 dispatcher 钉到首个节点（通常 head），
    # 同节点的 worker 按 --dispatcher-cpus 预留核数（torch 线程数相应减少），
    # 避免 worker 计算线程占满节点时 dispatcher 的 RPC/序列化被饿死；
    # 单机时所有 worker 与 dispatcher 同机，统一预留
    dispatcher_actor_options = None
    if args.dispatcher_cpus > 0:
        dispatcher_node = nodes[0][0] if nodes else None
        if dispatcher_node:
            dispatcher_actor_options = {
                "num_cpus": 1,
                "label_selector": {"ray.io/node-id": dispatcher_node},
            }
        for d in config.devices:
            if dispatcher_node is None or d.node_id == dispatcher_node:
                d.compute = {**d.compute, "reserve_cpus": args.dispatcher_cpus}

    serve.start(detached=True, http_options={"host": args.host, "port": args.port})
    app = build_app(
        config,
        dispatcher=LLMDispatcher,
        worker_cls=ModelWorker,
        model_devices=model_devices,
        model_dir=args.model_dir,
        dispatcher_ray_actor_options=dispatcher_actor_options,
    )
    serve.run(
        app, name=config.app_name, route_prefix=config.route_prefix, blocking=False
    )
    print(
        f"服务已启动: http://{args.host}:{args.port}  "
        f"模型流水线: {model_devices}  模型: {args.model_dir}  (Ctrl+C 退出)"
    )
    try:
        while True:
            time.sleep(10)
    except KeyboardInterrupt:
        pass
    finally:
        serve.shutdown()
        ray.shutdown()


if __name__ == "__main__":
    main()