To implement a readiness endpoint that becomes healthy only after the model file is loaded and a warmup inference returns within 200 ms, we can lazily load the model once and reuse it for subsequent calls. Use a module-level lock and a double-checked pattern to avoid reinitialization and to prevent race conditions when multiple requests call readiness concurrently.python
import os
import time
import threading
import torch
_MODEL_PATH = "/models/model.pt"
_model = None
_model_lock = threading.Lock()
_warmup_input = None # prepare once if needed
def _load_model():
global _model, _warmup_input
# load model (device choice depends on env)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
m = torch.jit.load(_MODEL_PATH, map_location=device) if _MODEL_PATH.endswith(".pt") else torch.load(_MODEL_PATH, map_location=device)
m.eval()
_model = m
# create a sample warmup input depending on model expected shape
# adjust shape/dtype to match your model
_warmup_input = torch.zeros((1, 3, 224, 224), device=device)
def readiness(timeout_ms=200):
"""
Returns dict: {"ready": bool, "reason": str}
Healthy only if model file exists, model loaded, and a warmup inference completes within timeout_ms.
Thread-safe: uses double-checked locking to ensure single initialization.
"""
global _model
# quick check: file must exist
if not os.path.exists(_MODEL_PATH):
return {"ready": False, "reason": "model file missing"}
# Double-checked locking: avoid taking lock on subsequent calls once loaded
if _model is None:
with _model_lock:
if _model is None: # check again inside lock
try:
_load_model()
except Exception as e:
return {"ready": False, "reason": f"model load failed: {e}"}
# perform warmup inference and time it
try:
start = time.perf_counter()
with torch.no_grad():
# ensure input on same device
inp = _warmup_input
out = _model(inp)
# optionally touch output to ensure computation finishes
if isinstance(out, (list, tuple)):
_ = out[0]
duration_ms = (time.perf_counter() - start) * 1000
if duration_ms <= timeout_ms:
return {"ready": True, "reason": f"warmup {duration_ms:.1f}ms"}
else:
return {"ready": False, "reason": f"warmup too slow: {duration_ms:.1f}ms"}
except Exception as e:
return {"ready": False, "reason": f"inference failed: {e}"}
