[Improvement]support base-port & flask-port (#7668)

* support base-port & flask-port

* update flask_port

Author: wj-Mcat

llm/flask_server.py

@@ -53,7 +53,8 @@ def __free_port(port):
 @dataclass
 class ServerArgument:
     port: int = field(default=8011, metadata={"help": "The port of ui service"})
-    base_port: int = field(default=8010, metadata={"help": "The port of flask 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"})
 
@@ -64,8 +65,8 @@ def __init__(self, args: ServerArgument, predictor: BasePredictor):
         self.predictor = predictor
         self.args = args
         scan_l, scan_u = (
-            self.args.base_port + port_interval * predictor.tensor_parallel_rank,
-            self.args.base_port + port_interval * (predictor.tensor_parallel_rank + 1),
+            self.args.flask_port + port_interval * predictor.tensor_parallel_rank,
+            self.args.flask_port + port_interval * (predictor.tensor_parallel_rank + 1),
         )
 
         if self.predictor.tensor_parallel_rank == 0:
@@ -174,6 +175,14 @@ def start_ui_service(self, args):
 
     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)

llm/gradio_ui.py

@@ -95,7 +95,7 @@ def infer(utterance, state, top_k, top_p, temperature, repetition_penalty, max_l
             "max_length": max_length,
             "min_length": 1,
         }
-        res = requests.post(f"http://0.0.0.0:{args.base_port}/api/chat", json=data, stream=True)
+        res = requests.post(f"http://0.0.0.0:{args.flask_port}/api/chat", json=data, stream=True)
         for line in res.iter_lines():
             result = json.loads(line)
             bot_response = result["result"]["response"]

tests/llm/test_gradio.py

@@ -39,11 +39,11 @@ def is_port_in_use(port):
 class UITest(unittest.TestCase):
     def setUp(self):
         # start web ui
-        self.base_port = self.avaliable_free_port()
-        self.port = self.avaliable_free_port([self.base_port])
+        self.flask_port = self.avaliable_free_port()
+        self.port = self.avaliable_free_port([self.flask_port])
         self.model_path = "__internal_testing__/tiny-random-llama"
-        command = 'cd llm && python flask_server.py --model_name_or_path {model_path} --port {port} --base_port {base_port} --src_length 1024 --dtype "float16"'.format(
-            base_port=self.base_port, port=self.port, model_path=self.model_path
+        command = 'cd llm && python flask_server.py --model_name_or_path {model_path} --port {port} --flask_port {flask_port} --src_length 1024 --dtype "float16"'.format(
+            flask_port=self.flask_port, port=self.port, model_path=self.model_path
         )
         self.ui_process = subprocess.Popen(command, shell=True, stdout=sys.stdout, stderr=sys.stderr)
         self.tokenizer = LlamaTokenizer.from_pretrained(self.model_path)
@@ -66,7 +66,7 @@ def avaliable_free_port(self, exclude=None):
 
     def wait_until_server_is_ready(self):
         while True:
-            if is_port_in_use(self.base_port) and is_port_in_use(self.port):
+            if is_port_in_use(self.flask_port) and is_port_in_use(self.port):
                 break
 
             print("waiting for server ...")
@@ -84,7 +84,7 @@ def test_argument(self):
         self.wait_until_server_is_ready()
 
         def get_response(data):
-            res = requests.post(f"http://localhost:{self.base_port}/api/chat", json=data, stream=True)
+            res = requests.post(f"http://localhost:{self.flask_port}/api/chat", json=data, stream=True)
             result_ = ""
             for line in res.iter_lines():
                 print(line)