diff --git a/app/main.py b/app/main.py index 66bea2e..6ba53c1 100644 --- a/app/main.py +++ b/app/main.py @@ -179,26 +179,31 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) app = FastAPI() 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 启动预热 -------- @app.on_event("startup") def warm_up(): - 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)) + logger.info("Warm-up on %s", DEVICE) + try: + texts = ["warmup"] + q = mp.Queue() + p = mp.Process(target=_warm_worker, args=(texts, q)) + p.start() + p.join(timeout=60) - texts = ["warmup"] - q = mp.Queue() - p = mp.Process(target=_warm_worker, args=(texts, q)) - p.start() - p.join(timeout=60) - - if not q.empty() and q.get() == "ok": - logger.info("Warm-up complete.") - else: - logger.warning("Warm-up failed or timed out.") + if not q.empty() and q.get() == "ok": + logger.info("Warm-up complete.") + else: + logger.warning("Warm-up failed or timed out.") + except Exception as e: + logger.warning("Warm-up exception: %s", e)