From 7dfc337bfad6915016be5030c11acc70888aa1e3 Mon Sep 17 00:00:00 2001
From: liunux4odoo <liunu@qq.com>
Date: Thu, 17 Aug 2023 14:11:21 +0800
Subject: [PATCH] add controller parameter. used to run model-worker

---
 startup.py | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/startup.py b/startup.py
index 975bc92..4d84bf8 100644
--- a/startup.py
+++ b/startup.py
@@ -182,7 +182,12 @@ def run_controller(q: Queue, run_seq: int = 1):
     uvicorn.run(app, host=host, port=port)
 
 
-def run_model_worker(model_name: str = LLM_MODEL, q: Queue = None, run_seq: int = 2):
+def run_model_worker(
+    model_name: str = LLM_MODEL,
+    controller_address: str = "",
+    q: Queue = None,
+    run_seq: int = 2,
+):
     import uvicorn
 
     kwargs = FSCHAT_MODEL_WORKERS[LLM_MODEL].copy()
@@ -191,7 +196,7 @@ def run_model_worker(model_name: str = LLM_MODEL, q: Queue = None, run_seq: int
     model_path = llm_model_dict[model_name].get("local_model_path", "")
     kwargs["model_path"] = model_path
     kwargs["model_names"] = [model_name]
-    kwargs["controller_address"] = fschat_controller_address()
+    kwargs["controller_address"] = controller_address or fschat_controller_address()
     kwargs["worker_address"] = fschat_model_worker_address()
 
     app = create_model_worker_app(**kwargs)
@@ -238,32 +243,43 @@ def run_webui():
 def parse_args() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser()
     parser.add_argument(
+        "-a",
         "--all",
         action="store_true",
         help="run fastchat's controller/model_worker/openai_api servers, run api.py and webui.py",
     )
     parser.add_argument(
+        "-o",
         "--openai-api",
         action="store_true",
         help="run fastchat controller/openai_api servers",
     )
     parser.add_argument(
+        "-m",
         "--model-worker",
         action="store_true",
         help="run fastchat model_worker server with specified model name. specify --model-name if not using default LLM_MODEL",
     )
     parser.add_argument(
+        "-n"
         "--model-name",
         type=str,
         default=LLM_MODEL,
         help="specify model name for model worker.",
     )
+    parser.add_argument(
+        "-c"
+        "--controller",
+        type=str,
+        help="specify controller address the worker is registered to. default is server_config.FSCHAT_CONTROLLER",
+    )
     parser.add_argument(
         "--api",
         action="store_true",
         help="run api.py server",
     )
     parser.add_argument(
+        "-w",
         "--webui",
         action="store_true",
         help="run webui.py server",