cn0303 commited on
Commit
2bb0dc4
·
verified ·
1 Parent(s): ebd4308

Fix spec-parser cold-start hang: pre-cache weights at boot

Browse files
Files changed (1) hide show
  1. 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=120) # cold path = 3.4GB download + load + generate
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: