Fix server launch for jupyter notebook (#186)

This commit is contained in:
Lianmin Zheng
2024-02-12 04:43:14 -08:00
committed by GitHub
parent bb824da41a
commit 06008bc295

View File

@@ -464,7 +464,6 @@ def launch_server(server_args, pipe_finish_writer):
assert proc_router.is_alive() and proc_detoken.is_alive()
def _launch_server():
# Launch api server
uvicorn.run(
app,
host=server_args.host,
@@ -474,49 +473,54 @@ def launch_server(server_args, pipe_finish_writer):
loop="uvloop",
)
t = threading.Thread(target=_launch_server)
t.start()
def _wait_and_warmup():
url = server_args.url()
for _ in range(60):
time.sleep(1)
try:
requests.get(url + "/get_model_info", timeout=5)
break
except requests.exceptions.RequestException as e:
pass
else:
if pipe_finish_writer is not None:
pipe_finish_writer.send(str(e))
else:
print(e, flush=True)
return
url = server_args.url()
for _ in range(60):
time.sleep(1)
# Warmup
try:
requests.get(url + "/get_model_info", timeout=5)
break
except requests.exceptions.RequestException as e:
pass
else:
if pipe_finish_writer is not None:
pipe_finish_writer.send(str(e))
else:
print(e, flush=True)
return
# Warmup
try:
# print("Warmup...", flush=True)
res = requests.post(
url + "/generate",
json={
"text": "Say this is a warmup request.",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 16,
# print("Warmup...", flush=True)
res = requests.post(
url + "/generate",
json={
"text": "Say this is a warmup request.",
"sampling_params": {
"temperature": 0,
"max_new_tokens": 16,
},
},
},
timeout=60,
)
# print(f"Warmup done. model response: {res.json()['text']}")
# print("=" * 20, "Server is ready", "=" * 20, flush=True)
except requests.exceptions.RequestException as e:
if pipe_finish_writer is not None:
pipe_finish_writer.send(str(e))
else:
print(e, flush=True)
return
timeout=60,
)
# print(f"Warmup done. model response: {res.json()['text']}")
# print("=" * 20, "Server is ready", "=" * 20, flush=True)
except requests.exceptions.RequestException as e:
if pipe_finish_writer is not None:
pipe_finish_writer.send(str(e))
else:
print(e, flush=True)
return
if pipe_finish_writer is not None:
pipe_finish_writer.send("init ok")
if pipe_finish_writer is not None:
pipe_finish_writer.send("init ok")
t = threading.Thread(target=_wait_and_warmup)
t.start()
try:
_launch_server()
finally:
t.join()
class Runtime: