Base application

src/submission/submit.py

@@ -1,9 +1,9 @@
 import json
 import os
-import queue
 import re
 import sys
 import threading
+import time
 import traceback
 from datetime import datetime, timezone
 
@@ -143,13 +143,12 @@ def _save_cert(model_id: str, revision: str, cert: dict) -> str:
     return out_path
 
 
+# Pure, pickleable function signature for ZeroGPU multiprocessing
 @gpu_decorator(duration=120)
-def execute_direct_xray_audit(model_id: str, revision: str = "main", trust_remote_code: bool = False, progress_callback=None) -> dict:
+def execute_direct_xray_audit(model_id: str, revision: str = "main", trust_remote_code: bool = False) -> dict:
     now_str = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
 
-    if progress_callback:
-        progress_callback(f"Verifying '{model_id}' on Hugging Face Hub...")
-    print(f"[LLM-X-RAY] Verifying model '{model_id}'...", flush=True)
+    print(f"[LLM-X-RAY] Verifying repository metadata for '{model_id}'...", flush=True)
 
     try:
         feas = check_feasibility(model_id, revision, TOKEN, MAX_AUDIT_PARAMS_BILLION, trust_remote_code=trust_remote_code)
@@ -163,13 +162,14 @@ def execute_direct_xray_audit(model_id: str, revision: str = "main", trust_remot
             _save_cert(model_id, revision, cert)
             return cert
 
+        # Run multi-layer audit (logs automatically stream to console with flush=True)
         result = run_full_audit(
             model_id=model_id,
             revision=revision,
-            device=AUDIT_DEVICE,
+            device="cuda" if has_spaces else AUDIT_DEVICE,
             token=TOKEN,
             trust_remote_code=trust_remote_code,
-            progress_callback=progress_callback,
+            progress_callback=None,
         )
 
         cert = {
@@ -179,7 +179,7 @@ def execute_direct_xray_audit(model_id: str, revision: str = "main", trust_remot
                 "model_sha": revision,
                 "architecture": feas.architecture or "CausalLM",
                 "params": feas.param_count_b,
-                "precision": "bfloat16" if AUDIT_DEVICE == "cuda" else "float32",
+                "precision": "bfloat16" if (has_spaces or AUDIT_DEVICE == "cuda") else "float32",
                 "license": "open-source",
             },
             "audited_at": now_str,
@@ -216,7 +216,7 @@ def audit_or_search_model(url_or_id: str, trust_remote_code: bool = False):
         )
         return
 
-    # Check for existing cert first
+    # Check for existing cert first (instant load)
     existing_cert = get_certificate_by_model_name(clean_model)
     if existing_cert and existing_cert.get("status") == "ok":
         yield (
@@ -227,21 +227,22 @@ def audit_or_search_model(url_or_id: str, trust_remote_code: bool = False):
         )
         return
 
-    latest_msg = f"Connecting to HF Hub for <b>{clean_model}</b>..."
+    # Initial loading feedback
     yield (
-        styled_loading(latest_msg, "Verifying repository and downloading model tensors..."),
+        styled_loading(f"Connecting to HF Hub for <b>{clean_model}</b>...", "Requesting ZeroGPU slice & downloading model weights..."),
         current_df,
         current_top3,
         current_panel,
     )
 
-    progress_q: "queue.Queue" = queue.Queue()
-    result_holder: dict = {}
+    result_holder = {}
+    done_event = threading.Event()
 
     def _worker():
         try:
+            # Arguments are pure primitives (str, str, bool) -> 100% pickleable on ZeroGPU
             result_holder["cert"] = execute_direct_xray_audit(
-                clean_model, "main", trust_remote_code=trust_remote_code, progress_callback=progress_q.put
+                clean_model, "main", trust_remote_code=trust_remote_code
             )
         except Exception as err:
             traceback.print_exc()
@@ -252,32 +253,31 @@ def audit_or_search_model(url_or_id: str, trust_remote_code: bool = False):
                 "error_message": str(err),
             }
         finally:
-            progress_q.put(None)
+            done_event.set()
 
     thread = threading.Thread(target=_worker, daemon=True)
     thread.start()
 
-    while True:
-        try:
-            msg = progress_q.get(timeout=0.2)
-            if msg is None:
-                break
-            latest_msg = msg
-            yield (
-                styled_loading(latest_msg, "Executing multi-layer operator audit. Progress streams live below."),
-                current_df,
-                current_top3,
-                current_panel,
-            )
-        except queue.Empty:
-            if not thread.is_alive():
-                break
-            yield (
-                styled_loading(latest_msg, "Audit in progress... Please keep this page open."),
-                current_df,
-                current_top3,
-                current_panel,
-            )
+    # Informative UI stage cycle while ZeroGPU finishes audit
+    stages = [
+        "Downloading model weights and tokenizer tensors...",
+        "Layer A: Performing SVD Spectral Tomography across weight tensors...",
+        "Layer B: Streaming activations and calculating operator covariance...",
+        "Layer C: Running empirical factual probe battery...",
+        "Finalizing operator risk and registering certificate...",
+    ]
+    stage_idx = 0
+
+    while not done_event.is_set():
+        current_stage = stages[min(stage_idx, len(stages) - 1)]
+        yield (
+            styled_loading(f"🔬 Auditing <b>{clean_model}</b>...", current_stage),
+            current_df,
+            current_top3,
+            current_panel,
+        )
+        done_event.wait(timeout=2.5)
+        stage_idx += 1
 
     thread.join()
     cert = result_holder.get("cert", {})