server.py

"""启动入口：python serverRun.py [--devices cpu,cuda,npu] [--replicas N]

仅启动最基础的异构计算服务（/run、/split_matmul 等）；
Qwen3.5 异构推理由 serverRun_LLMgenerate.py 启动。
"""

from __future__ import annotations

import argparse

import ray
from ray import serve

from heterogeneous_serve.dispatcher import build_app
from heterogeneous_serve.config import DeviceConfig, FrameworkConfig, default_config
from heterogeneous_serve.dispatcher import Dispatcher_base,RunRequest,fastapi_app

import asyncio
from typing import Any
from fastapi import HTTPException


@serve.deployment
@serve.ingress(fastapi_app)
class Dispatcher(Dispatcher_base):
    """自定义分布式调度的算法，继承于Dispatcher_base，添加新的分布调度函数;
    建议不要重写Dispatcher_base的任何函数
    """

    # 测试算法
    @fastapi_app.post("/tests")
    async def test(self, req: RunRequest) -> dict[str, Any]:
        req.op = "add"
        payload = req.model_dump(exclude={"op"})
        responses = [
            h.run.remote(req.op, **payload) for h in self.handles.values()
        ]
        try:
            results = await asyncio.gather(*responses)
        except Exception as e:
            raise HTTPException(status_code=400, detail=str(e)[:500]) from e
        
        return {"mode": "replicate", "op": req.op, "results": results}



def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="异构并行矩阵运算服务")
    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("--host", type=str, default="127.0.0.1")
    p.add_argument("--port", type=int, default=8000)
    return p.parse_args()


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

    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.init(ignore_reinit_error=True)
    serve.start(detached=True, http_options={"host": args.host, "port": args.port})
    app = build_app(config, dispatcher=Dispatcher)
    serve.run(
        app, name=config.app_name, route_prefix=config.route_prefix, blocking=False
    )
    print(
        f"服务已启动: http://{args.host}:{args.port}  "
        f"设备: {config.device_names()}  (Ctrl+C 退出)"
    )
    try:
        import time

        while True:
            time.sleep(10)
    except KeyboardInterrupt:
        pass
    finally:
        serve.shutdown()
        ray.shutdown()


if __name__ == "__main__":
    main()