Fix spec-parser cold-start hang: pre-cache weights at boot
Browse files- spec_brick.py +16 -1
spec_brick.py
CHANGED
|
@@ -100,7 +100,7 @@ if _should_load():
|
|
| 100 |
_state["tok"] = tok
|
| 101 |
_state["model"] = model.to("cuda").eval()
|
| 102 |
|
| 103 |
-
@spaces.GPU(duration=
|
| 104 |
def _generate(text: str) -> str:
|
| 105 |
if _state["model"] is None:
|
| 106 |
_load()
|
|
@@ -121,6 +121,21 @@ if _should_load():
|
|
| 121 |
|
| 122 |
_GENERATE = _generate
|
| 123 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 124 |
# Eager boot path (flagged): resident-load now only if cached; never
|
| 125 |
# downloads at boot (local_files_only=True), so a cold cache is harmless.
|
| 126 |
if EAGER_LOAD:
|
|
|
|
| 100 |
_state["tok"] = tok
|
| 101 |
_state["model"] = model.to("cuda").eval()
|
| 102 |
|
| 103 |
+
@spaces.GPU(duration=180) # warm path: load from pre-cached weights + generate
|
| 104 |
def _generate(text: str) -> str:
|
| 105 |
if _state["model"] is None:
|
| 106 |
_load()
|
|
|
|
| 121 |
|
| 122 |
_GENERATE = _generate
|
| 123 |
|
| 124 |
+
# Pre-cache the weights at BOOT, on CPU, with no time limit, so the first
|
| 125 |
+
# /parse call loads from a warm cache and finishes inside its GPU window.
|
| 126 |
+
# Without this a cold cache (e.g. after a redeploy) makes the first call
|
| 127 |
+
# try to pull ~3.4GB during the GPU slot, which overruns it and hangs
|
| 128 |
+
# every request. snapshot_download is a no-op once the files are cached.
|
| 129 |
+
try:
|
| 130 |
+
from huggingface_hub import snapshot_download
|
| 131 |
+
snapshot_download(BASE_ID, revision=BASE_REVISION)
|
| 132 |
+
snapshot_download(ADAPTER_ID, revision=ADAPTER_REVISION)
|
| 133 |
+
print("[FitCheck] spec parser weights pre-cached at boot",
|
| 134 |
+
file=sys.stderr, flush=True)
|
| 135 |
+
except Exception as pe: # noqa: BLE001
|
| 136 |
+
print(f"[FitCheck] spec parser pre-download skipped "
|
| 137 |
+
f"({type(pe).__name__}: {pe})", file=sys.stderr, flush=True)
|
| 138 |
+
|
| 139 |
# Eager boot path (flagged): resident-load now only if cached; never
|
| 140 |
# downloads at boot (local_files_only=True), so a cold cache is harmless.
|
| 141 |
if EAGER_LOAD:
|