Repository: lm-sys/FastChat
File: fastchat/serve/huggingface_api_worker.py (Python)

Code before:
        default=5,
        help="Limit the model concurrency to prevent OOM.",
    )
    parser.add_argument("--stream-interval", type=int, default=2)
    parser.add_argument("--no-register", action="store_true")
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Overwrite the random seed for each generation.",
    )
    args = parser.parse_args()
    logger.info(f"args: {args}")

    with open(args.supported_models_file, "r") as f:
        supported_models = json.load(f)

    if args.model not in supported_models:
        raise ValueError(
            f"Model {args.model} not supported. Please add it to {args.supported_models_file}."
        )

    model_path = supported_models[args.model]["model_path"]
    api_base = supported_models[args.model]["api_base"]
    token = supported_models[args.model]["token"]
    context_length = supported_models[args.model]["context_length"]

    worker = HuggingfaceApiWorker(
        args.controller_address,
        args.worker_address,
        worker_id,
        model_path,
        api_base,
        token,
        context_length,
        args.model_names,
        args.limit_worker_concurrency,
        no_register=args.no_register,
        stream_interval=args.stream_interval,
        conv_template=args.conv_template,
        embed_in_truncate=args.embed_in_truncate,
        seed=args.seed,
    )
    return args, worker


if __name__ == "__main__":
    args, worker = create_model_worker()
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


Review comment (suggestion):
This single worker should be able to serve all models listed in `supported_models.json`.
It will register all workers to the controller and serve requests for all models.

Code after:
        model_path,
        api_base,
        token,
        context_length,
    ) in zip(
        model_names_list,
        conv_template_list,
        model_path_list,
        api_base_list,
        token_list,
        context_length_list,
    ):
        m = HuggingfaceApiWorker(
            args.controller_address,
            args.worker_address,
            worker_id,
            model_path,
            api_base,
            token,
            context_length,
            model_names,
            args.limit_worker_concurrency,
            no_register=args.no_register,
            conv_template=conv_template,
            seed=args.seed,
        )
        workers.append(m)
        for name in model_names:
            worker_map[name] = m

    # register all the models
    url = args.controller_address + "/register_worker"
    data = {
        "worker_name": workers[0].worker_addr,
        "check_heart_beat": not args.no_register,
        "worker_status": {
            "model_names": [m for w in workers for m in w.model_names],
            "speed": 1,
            "queue_length": sum([w.get_queue_length() for w in workers]),
        },
    }
    r = requests.post(url, json=data)
    assert r.status_code == 200

    return args, workers


if __name__ == "__main__":
    args, workers = create_huggingface_api_worker()
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")