Preloading

app.py

@@ -41,6 +41,11 @@ REQUIRED_COLS = [
     "l2_assoc",
 ]
 
+# -------------------------------------------------
+# Global model cache
+# -------------------------------------------------
+MODEL_CACHE = {}
+
 # -------------------------------------------------
 # Model Download
 # -------------------------------------------------
@@ -48,11 +53,9 @@ def ensure_models():
     if not os.path.exists(MODEL_DIR):
         snapshot_download(
             repo_id=HF_REPO_ID,
-            local_dir='.',
-            allow_patterns="*.pkl")
-    
-        print(f"Model dir: {MODEL_DIR}")
-        print(os.listdir(MODEL_DIR))
+            local_dir=".",
+            allow_patterns="*.pkl",
+        )
 
 # -------------------------------------------------
 # Utilities
@@ -61,19 +64,39 @@ def resolve_workload(workload: str) -> str:
     return WORKLOAD_ALIAS.get(workload, workload)
 
 def load_model(workload: str, target: str):
-    model_path = os.path.join(MODEL_DIR, f"model_{workload}_{target}.pkl")
-    if not os.path.exists(model_path):
-        raise FileNotFoundError(f"Model not found: {model_path}")
-    payload = joblib.load(model_path)
-    return payload["model"], payload["log_target"]
+    try:
+        return MODEL_CACHE[(workload, target)]
+    except KeyError:
+        raise RuntimeError(f"Model not preloaded: {workload}, {target}")
 
 def physical_sanity_check(ipc, miss_rate):
-    warnings = []
+    warnings_out = []
     if ipc < 0 or ipc > 3.5:
-        warnings.append(f"IPC={ipc:.3f} out of physical range")
+        warnings_out.append(f"IPC={ipc:.3f} out of physical range")
     if miss_rate < 0 or miss_rate > 1:
-        warnings.append(f"L2 miss rate={miss_rate:.3f} out of [0,1]")
-    return warnings
+        warnings_out.append(f"L2 miss rate={miss_rate:.3f} out of [0,1]")
+    return warnings_out
+
+# -------------------------------------------------
+# Preload all models at startup
+# -------------------------------------------------
+def preload_models():
+    ensure_models()
+
+    workloads = set(WORKLOAD_ALIAS.values()) | {"matrix_mul"}
+
+    for workload in workloads:
+        for target in TARGETS:
+            model_path = os.path.join(
+                MODEL_DIR, f"model_{workload}_{target}.pkl"
+            )
+            payload = joblib.load(model_path)
+            MODEL_CACHE[(workload, target)] = (
+                payload["model"],
+                payload["log_target"],
+            )
+
+    return "✅ Models loaded successfully."
 
 # -------------------------------------------------
 # Inference Core
@@ -85,8 +108,12 @@ def run_inference(df: pd.DataFrame) -> pd.DataFrame:
 
     # Feature engineering
     for col in [
-        "l1d_size", "l1i_size", "l2_size",
-        "l1d_assoc", "l1i_assoc", "l2_assoc",
+        "l1d_size",
+        "l1i_size",
+        "l2_size",
+        "l1d_assoc",
+        "l1i_assoc",
+        "l2_assoc",
     ]:
         df[f"{col}_log2"] = np.log2(df[col])
 
@@ -127,8 +154,6 @@ def run_inference(df: pd.DataFrame) -> pd.DataFrame:
 # Gradio Wrapper
 # -------------------------------------------------
 def infer_from_csv(file):
-    ensure_models()
-
     df = pd.read_csv(file.name)
     out_df = run_inference(df)
 
@@ -148,6 +173,9 @@ def infer_from_csv(file):
 # UI
 # -------------------------------------------------
 with gr.Blocks(title="AIDE Chip Surrogate Inference") as demo:
+    loading_md = gr.Markdown("## ⏳ Loading surrogate models…", visible=True)
+    ready_md = gr.Markdown("## ✅ Models ready", visible=False)
+
     gr.Markdown(
         """
         # AIDE Chip Surrogate Inference
@@ -160,16 +188,29 @@ with gr.Blocks(title="AIDE Chip Surrogate Inference") as demo:
     )
 
     csv_input = gr.File(label="Input CSV", file_types=[".csv"])
-    run_btn = gr.Button("Run Inference")
+    run_btn = gr.Button("Run Inference", interactive=False)
 
     preview = gr.Dataframe(label="Preview (first 20 rows)")
     output_csv = gr.File(label="Download Full Output CSV")
-    warnings = gr.Textbox(label="Sanity Check Summary")
+    warnings_box = gr.Textbox(label="Sanity Check Summary")
+
+    demo.load(
+        preload_models,
+        inputs=None,
+        outputs=ready_md,
+    ).then(
+        lambda: (
+            gr.update(visible=False),
+            gr.update(visible=True),
+            gr.update(interactive=True),
+        ),
+        outputs=[loading_md, ready_md, run_btn],
+    )
 
     run_btn.click(
         infer_from_csv,
         inputs=csv_input,
-        outputs=[preview, output_csv, warnings],
+        outputs=[preview, output_csv, warnings_box],
     )
 
 if __name__ == "__main__":