This commit is contained in:
hailin 2025-08-05 17:01:05 +08:00
parent 9e0c85b93a
commit 61962c074b
1 changed files with 21 additions and 16 deletions

View File

@ -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.")