This commit is contained in:
parent
9e0c85b93a
commit
61962c074b
37
app/main.py
37
app/main.py
|
|
@ -179,26 +179,31 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
|
||||||
app = FastAPI()
|
app = FastAPI()
|
||||||
logger.info("Using SAFE_MIN_FREE_MB = %d MB", SAFE_MIN_FREE_MB)
|
logger.info("Using SAFE_MIN_FREE_MB = %d MB", SAFE_MIN_FREE_MB)
|
||||||
|
|
||||||
|
|
||||||
|
def _warm_worker(t, q):
|
||||||
|
try:
|
||||||
|
_ = model.encode(t, return_dense=True, num_processes=1)
|
||||||
|
q.put("ok")
|
||||||
|
except Exception as e:
|
||||||
|
q.put(str(e))
|
||||||
|
|
||||||
# ② -------- FastAPI 启动预热 --------
|
# ② -------- FastAPI 启动预热 --------
|
||||||
@app.on_event("startup")
|
@app.on_event("startup")
|
||||||
def warm_up():
|
def warm_up():
|
||||||
def _warm_worker(t, q):
|
logger.info("Warm-up on %s", DEVICE)
|
||||||
try:
|
try:
|
||||||
_ = model.encode(t, return_dense=True, num_processes=1)
|
texts = ["warmup"]
|
||||||
q.put("ok")
|
q = mp.Queue()
|
||||||
except Exception as e:
|
p = mp.Process(target=_warm_worker, args=(texts, q))
|
||||||
q.put(str(e))
|
p.start()
|
||||||
|
p.join(timeout=60)
|
||||||
|
|
||||||
texts = ["warmup"]
|
if not q.empty() and q.get() == "ok":
|
||||||
q = mp.Queue()
|
logger.info("Warm-up complete.")
|
||||||
p = mp.Process(target=_warm_worker, args=(texts, q))
|
else:
|
||||||
p.start()
|
logger.warning("Warm-up failed or timed out.")
|
||||||
p.join(timeout=60)
|
except Exception as e:
|
||||||
|
logger.warning("Warm-up exception: %s", e)
|
||||||
if not q.empty() and q.get() == "ok":
|
|
||||||
logger.info("Warm-up complete.")
|
|
||||||
else:
|
|
||||||
logger.warning("Warm-up failed or timed out.")
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue