flask_server.py

# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import json
import os
import socket
from contextlib import closing
from dataclasses import asdict, dataclass, field
from time import sleep

import requests
from filelock import FileLock
from predictor import BasePredictor, ModelArgument, PredictorArgument, create_predictor

from paddlenlp.trainer import PdArgumentParser
from paddlenlp.utils.log import logger

STOP_SIGNAL = "[END]"
port_interval = 200
PORT_FILE = "port-info"
FILE_LOCK = "port-lock"


def find_free_ports(port_l, port_u):
 def __free_port(port):
 with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
 try:
 s.bind(("", port))
 return port
 except:
 return -1

 for port in range(port_l, port_u):
 port = __free_port(port)
 if port != -1:
 return port

 return -1


@dataclass
class ServerArgument:
 port: int = field(default=8011, metadata={"help": "The port of ui service"})
 base_port: int = field(default=None, metadata={"help": "The port of flask service"})
 flask_port: int = field(default=None, metadata={"help": "The port of flask service"})
 title: str = field(default="LLM", metadata={"help": "The title of gradio"})
 sub_title: str = field(default="LLM-subtitle", metadata={"help": "The sub-title of gradio"})


class PredictorServer:
 def __init__(self, args: ServerArgument, predictor: BasePredictor):

 self.predictor = predictor
 self.args = args
 scan_l, scan_u = (
 self.args.flask_port + port_interval * predictor.tensor_parallel_rank,
 self.args.flask_port + port_interval * (predictor.tensor_parallel_rank + 1),
 )
 self.total_max_length = predictor.config.src_length + predictor.config.max_length

 if self.predictor.tensor_parallel_rank == 0:
 # fetch port info
 self.port = find_free_ports(scan_l, scan_u)
 self.peer_ports = {}
 while True and self.predictor.tensor_parallel_degree > 1:
 if os.path.exists(PORT_FILE):
 with FileLock(FILE_LOCK), open(PORT_FILE, "r") as f:
 cnt = 1
 for line in f:
 data = json.loads(line)
 self.peer_ports[data["rank"]] = data["port"]
 cnt += 1

 if cnt == predictor.tensor_parallel_degree:
 break
 else:
 print("waiting for port reach", cnt)
 sleep(1)
 else:
 # save port info
 self.port = find_free_ports(scan_l, scan_u)
 data = {"rank": predictor.tensor_parallel_rank, "port": self.port}
 with FileLock(FILE_LOCK), open(PORT_FILE, "a") as f:
 f.write(json.dumps(data) + "\n")
 print("rank: ", predictor.tensor_parallel_rank, " port info saving done.")

 def predict(self, input_texts: str | list[str]):
 return self.predictor.stream_predict(input_texts)

 def broadcast_msg(self, data):
 for _, peer_port in self.peer_ports.items():
 if peer_port != self.port:
 _ = requests.post(f"http://0.0.0.0:{peer_port}/api/chat", json=data)

 def start_flask_server(self):
 from flask import Flask, request, stream_with_context

 app = Flask(__name__)

 @app.post("/api/chat")
 def _server():
 data = request.get_json()
 logger.info(f"Request: {json.dumps(data, indent=2, ensure_ascii=False)}")

 if self.predictor.tensor_parallel_rank == 0:
 self.broadcast_msg(data)

 def streaming(data):
 query = data.pop("context", "")
 history = data.pop("history", "")
 data.pop("extra_info", None)

 # build chat template
 if self.predictor.tokenizer.chat_template is not None:
 if not history:
 history = []
 # also support history data
 elif isinstance(history, str):
 history = json.loads(history)

 assert len(history) % 2 == 0
 chat_query = []
 for idx in range(0, len(history), 2):
 if isinstance(history[idx], str):
 chat_query.append([history[idx], history[idx + 1]])
 elif isinstance(history[idx], dict):
 chat_query.append([history[idx]["utterance"], history[idx + 1]["utterance"]])
 else:
 raise ValueError(
 "history data should be list[str] or list[dict], eg: ['sentence-1', 'sentece-2', ...], or "
 "[{'utterance': 'sentence-1'}, {'utterance': 'sentence-2'}, ...]"
 )

 # the input of predictor should be batched.
 # batched query: [ [[user, bot], [user, bot], ..., [user]] ]
 query = [chat_query + [[query]]]

 generation_args = data
 self.predictor.config.max_length = generation_args["max_length"]
 if "src_length" in generation_args:
 self.predictor.config.src_length = generation_args["src_length"]

 if self.predictor.config.src_length + self.predictor.config.max_length > self.total_max_length:
 output = {
 "error_code": 1,
 "error_msg": f"The sum of src_length<{self.predictor.config.src_length}> and "
 f"max_length<{self.predictor.config.max_length}> should be smaller than or equal to "
 f"the max-total-length<{self.total_max_length}>",
 }
 yield json.dumps(output, ensure_ascii=False) + "\n"
 return

 self.predictor.config.top_p = generation_args["top_p"]
 self.predictor.config.temperature = generation_args["temperature"]
 self.predictor.config.top_k = generation_args["top_k"]
 self.predictor.config.repetition_penalty = generation_args["repetition_penalty"]

 for key, value in generation_args.items():
 setattr(self.args, key, value)

 streamer = self.predict(query)
 if self.predictor.tensor_parallel_rank == 0:
 for new_text in streamer:
 if not new_text:
 continue

 output = {
 "error_code": 0,
 "error_msg": "Success",
 "result": {"response": {"role": "bot", "utterance": new_text}},
 }
 yield json.dumps(output, ensure_ascii=False) + "\n"
 else:
 return "done"

 return app.response_class(stream_with_context(streaming(data)))

 # set single thread to do prediction
 # refer to: https://github.com/pallets/flask/blob/main/src/flask/app.py#L605
 app.run(host="0.0.0.0", port=self.port, threaded=False)

 def start_ui_service(self, args, predictor_args):
 # do not support start ui service in one command
 from multiprocessing import Process

 from gradio_ui import main

 p = Process(target=main, args=(args, predictor_args))
 p.daemon = True
 p.start()


if __name__ == "__main__":

 parser = PdArgumentParser((PredictorArgument, ModelArgument, ServerArgument))
 predictor_args, model_args, server_args = parser.parse_args_into_dataclasses()
 # check port
 if server_args.base_port is not None:
 logger.warning("`--base_port` is deprecated, please use `--flask_port` instead after 2023.12.30.")

 if server_args.flask_port is None:
 server_args.flask_port = server_args.base_port
 else:
 logger.warning("`--base_port` and `--flask_port` are both set, `--base_port` will be ignored.")

 log_dir = os.getenv("PADDLE_LOG_DIR", "./")
 PORT_FILE = os.path.join(log_dir, PORT_FILE)
 if os.path.exists(PORT_FILE):
 os.remove(PORT_FILE)

 predictor = create_predictor(predictor_args, model_args)

 server = PredictorServer(server_args, predictor)

 if server.predictor.tensor_parallel_rank == 0:
 server.start_ui_service(server_args, asdict(predictor.config))

 server.start_flask_server()