add controller parameter. used to run model-worker

This commit is contained in:
liunux4odoo 2023-08-17 14:11:21 +08:00
parent f92b002342
commit 7dfc337bfa
1 changed files with 18 additions and 2 deletions

View File

@ -182,7 +182,12 @@ def run_controller(q: Queue, run_seq: int = 1):
uvicorn.run(app, host=host, port=port) 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 import uvicorn
kwargs = FSCHAT_MODEL_WORKERS[LLM_MODEL].copy() 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", "") model_path = llm_model_dict[model_name].get("local_model_path", "")
kwargs["model_path"] = model_path kwargs["model_path"] = model_path
kwargs["model_names"] = [model_name] 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() kwargs["worker_address"] = fschat_model_worker_address()
app = create_model_worker_app(**kwargs) app = create_model_worker_app(**kwargs)
@ -238,32 +243,43 @@ def run_webui():
def parse_args() -> argparse.ArgumentParser: def parse_args() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument( parser.add_argument(
"-a",
"--all", "--all",
action="store_true", action="store_true",
help="run fastchat's controller/model_worker/openai_api servers, run api.py and webui.py", help="run fastchat's controller/model_worker/openai_api servers, run api.py and webui.py",
) )
parser.add_argument( parser.add_argument(
"-o",
"--openai-api", "--openai-api",
action="store_true", action="store_true",
help="run fastchat controller/openai_api servers", help="run fastchat controller/openai_api servers",
) )
parser.add_argument( parser.add_argument(
"-m",
"--model-worker", "--model-worker",
action="store_true", action="store_true",
help="run fastchat model_worker server with specified model name. specify --model-name if not using default LLM_MODEL", help="run fastchat model_worker server with specified model name. specify --model-name if not using default LLM_MODEL",
) )
parser.add_argument( parser.add_argument(
"-n"
"--model-name", "--model-name",
type=str, type=str,
default=LLM_MODEL, default=LLM_MODEL,
help="specify model name for model worker.", 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( parser.add_argument(
"--api", "--api",
action="store_true", action="store_true",
help="run api.py server", help="run api.py server",
) )
parser.add_argument( parser.add_argument(
"-w",
"--webui", "--webui",
action="store_true", action="store_true",
help="run webui.py server", help="run webui.py server",