jtlevine Claude Opus 4.7 (1M context) commited on
Commit
5f87e82
·
1 Parent(s): 964e1bf

Prewarm GraphCast at FastAPI startup to avoid cold-start fallback

Browse files

Mirrors the MI Chronos prewarm pattern. load_model() was called lazily
on the first _step_predict, which on a freshly-rebuilt HF Space took
~30-120s (downloading the 1.4GB checkpoint from GCS plus normalization
stats). That puts the pipeline's first weekly trigger at risk of the
same silent-fallback path that bit MI before PR #5.

Fix: spawn a daemon thread in the lifespan context that calls
load_model() at startup. JAX compile + GCS download happen in the
background while FastAPI comes up; by the time someone hits
/api/pipeline/trigger the model is resident. Failures log and move on.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Files changed (1) hide show
  1. src/api.py +21 -0
src/api.py CHANGED
@@ -38,6 +38,21 @@ logger = logging.getLogger(__name__)
38
  _db_conn = None
39
 
40
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
  @asynccontextmanager
42
  async def lifespan(app: FastAPI):
43
  global _db_conn
@@ -66,6 +81,12 @@ async def lifespan(app: FastAPI):
66
  except Exception as e:
67
  logger.warning("DB init failed (non-fatal): %s", e)
68
  _db_conn = None
 
 
 
 
 
 
69
  scheduler = _start_scheduler()
70
  yield
71
  if scheduler:
 
38
  _db_conn = None
39
 
40
 
41
+ def _prewarm_graphcast() -> None:
42
+ """Load GraphCast model into memory at startup so the first pipeline
43
+ trigger doesn't pay the ~30-120s download/init cost. Runs in a
44
+ background thread — failures are logged, not fatal.
45
+ """
46
+ try:
47
+ from src.prediction.graphcast_inference import load_model
48
+ import time as _time
49
+ t0 = _time.time()
50
+ load_model()
51
+ logger.info("[PREWARM] GraphCast loaded at startup (%.1fs)", _time.time() - t0)
52
+ except Exception as exc:
53
+ logger.warning("[PREWARM] GraphCast prewarm threw: %s", exc)
54
+
55
+
56
  @asynccontextmanager
57
  async def lifespan(app: FastAPI):
58
  global _db_conn
 
81
  except Exception as e:
82
  logger.warning("DB init failed (non-fatal): %s", e)
83
  _db_conn = None
84
+
85
+ prewarm_thread = threading.Thread(
86
+ target=_prewarm_graphcast, daemon=True, name="graphcast-prewarm",
87
+ )
88
+ prewarm_thread.start()
89
+
90
  scheduler = _start_scheduler()
91
  yield
92
  if scheduler: