Update src/registry.py

src/registry.py

@@ -17,6 +17,11 @@ import requests
 
 logger = logging.getLogger("SniperRegistry")
 
+# Patch __main__ immediately so joblib can resolve SniperModel from
+# any pickle that was trained with __main__.SniperModel
+from src.sniper_model import patch_main as _patch_main, safe_load as _safe_load
+_patch_main()
+
 CACHE_DIR = Path("/tmp/sniper_cache")
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
@@ -193,7 +198,7 @@ def _load_pkl(repo_id: str, path_in_repo: str, token: str = None):
         logger.warning(f"Could not download {path_in_repo}")
         return None
     try:
-        return joblib.load(local)
+        return _safe_load(local)
     except Exception as e:
         logger.error(f"Failed to load {local}: {e}")
         return None