Prewarm GraphCast at FastAPI startup to avoid cold-start fallback
Browse filesMirrors 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>
- 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:
|