Repository: hpcaitech/ColossalAI
File: colossalai/inference/server/api_server.py (Python)

Code before:
    parser.add_argument(
        "--response-role",
        type=str,
        default="assistant",
        help="The role name to return if " "`request.add_generation_prompt=true`.",
    )
    parser = add_engine_config(parser)

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    inference_config = InferenceConfig.from_dict(vars(args))
    model = AutoModelForCausalLM.from_pretrained(args.model)
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    async_engine = AsyncInferenceEngine(
        start_engine_loop=True, model=model, tokenizer=tokenizer, inference_config=inference_config
    )
    engine = async_engine.engine
    completion_serving = CompletionServing(async_engine, served_model=model.__class__.__name__)
    chat_serving = ChatServing(
        async_engine,
        served_model=model.__class__.__name__,
        tokenizer=tokenizer,
        response_role=args.response_role,
        chat_template=args.chat_template,
    )
    app.root_path = args.root_path
    uvicorn.run(
        app=app,
        host=args.host,
        port=args.port,
        log_level="debug",
        timeout_keep_alive=TIMEOUT_KEEP_ALIVE,
        ssl_keyfile=args.ssl_keyfile,
        ssl_certfile=args.ssl_certfile,
    )


Review comment (suggestion):
![image](https://github.com/hpcaitech/ColossalAI/assets/26091437/9767755f-124c-4fae-aed8-b2ea5253e5ba)
Consider adding the above logic inside the engine to avoid initiating TP logic even when using a single GPU.

Code after:
        "specified, the model name will be the same as "
        "the huggingface name.",
    )

    parser.add_argument(
        "--chat-template",
        type=str,
        default=None,
        help="The file path to the chat template, " "or the template in single-line form " "for the specified model",
    )
    parser.add_argument(
        "--response-role",
        type=str,
        default="assistant",
        help="The role name to return if " "`request.add_generation_prompt=true`.",
    )
    parser = add_engine_config(parser)

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    inference_config = InferenceConfig.from_dict(vars(args))
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    colossalai_backend_port = find_available_ports(1)[0]
    colossalai.launch(
        rank=0,
        world_size=1,
        host=args.host,
        port=colossalai_backend_port,
        backend="nccl",
    )
    model = AutoModelForCausalLM.from_pretrained(args.model)
    async_engine = AsyncInferenceEngine(
        start_engine_loop=True, model_or_path=model, tokenizer=tokenizer, inference_config=inference_config
    )
    engine = async_engine.engine
    completion_serving = CompletionServing(async_engine, model.__class__.__name__)
    chat_serving = ChatServing(
        async_engine,
        served_model=model.__class__.__name__,
        tokenizer=tokenizer,
        response_role=args.response_role,
        chat_template=args.chat_template,
    )
    app.root_path = args.root_path
    uvicorn.run(
        app=app,
        host=args.host,
        port=args.port,