Upload alpha_factory/infra/wq_client.py with huggingface_hub

alpha_factory/infra/wq_client.py

@@ -1,19 +1,47 @@
 """
-WorldQuant BRAIN API Client — async wrapper with rate limiting.
-Handles submission, polling, and result harvesting.
+WorldQuant BRAIN API Client v3 — async wrapper with proper session auth,
+rate limiting, circuit breaker, and comprehensive exception handling.
+
+NOTE: BRAIN uses session-based auth (cookies), not API keys.
+Session tokens expire quickly. The client reads BRAIN_SESSION_TOKEN
+from env and includes it as a Cookie header.
 """
 import asyncio
 import aiohttp
 import hashlib
+import os
 from typing import Any, Optional
 from ..config import BrainConfig
 from ..schemas import BrainMetrics
 
 
+class BrainClientError(Exception):
+    """Base exception for BRAIN client errors."""
+    pass
+
+
+class BrainAuthError(BrainClientError):
+    """Raised when BRAIN authentication fails (401)."""
+    pass
+
+
+class BrainRateLimitError(BrainClientError):
+    """Raised when BRAIN rate limits the request (429)."""
+    pass
+
+
+class BrainServerError(BrainClientError):
+    """Raised when BRAIN returns a 5xx error."""
+    pass
+
+
 class BrainClient:
     """
     Async client for WorldQuant BRAIN API.
     Rate-limited, retry-on-429, circuit-breaker-protected.
+    
+    Auth: reads BRAIN_SESSION_TOKEN from environment.
+    This is a browser session cookie that must be refreshed manually.
     """
 
     def __init__(self, session: aiohttp.ClientSession, config: BrainConfig):
@@ -23,6 +51,46 @@ class BrainClient:
         self._submissions_today = 0
         self._consecutive_failures = 0
         self._circuit_open = False
+        
+        # Auth: session token from env
+        self._session_token = os.getenv("BRAIN_SESSION_TOKEN", "")
+        if not self._session_token:
+            # Also try legacy name
+            self._session_token = os.getenv("BRAIN_AUTH", "")
+        
+        self._headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        }
+        if self._session_token:
+            # BRAIN uses session cookies, NOT Bearer tokens
+            self._headers["Cookie"] = f"session={self._session_token}"
+
+    def _check_auth(self):
+        """Verify we have auth credentials."""
+        if not self._session_token:
+            raise BrainAuthError(
+                "No BRAIN_SESSION_TOKEN set. "
+                "Get it from browser devtools after logging into brain.worldquant.com "
+                "(Network tab → any API request → copy Cookie header). "
+                "Set: export BRAIN_SESSION_TOKEN=your_token"
+            )
+
+    def _handle_response_error(self, resp: aiohttp.ClientResponse) -> None:
+        """Handle HTTP error responses with proper exception types."""
+        if resp.status == 401:
+            self._consecutive_failures += 1
+            raise BrainAuthError(
+                "BRAIN auth failed (401). Session token expired. "
+                "Get a new token from browser and set BRAIN_SESSION_TOKEN."
+            )
+        elif resp.status == 429:
+            raise BrainRateLimitError("BRAIN rate limit hit (429). Backing off.")
+        elif resp.status >= 500:
+            self._consecutive_failures += 1
+            if self._consecutive_failures >= 5:
+                self._circuit_open = True
+            raise BrainServerError(f"BRAIN server error: {resp.status}")
 
     async def submit_alpha(
         self,
@@ -35,7 +103,9 @@ class BrainClient:
         Returns the simulation response or raises on failure.
         """
         if self._circuit_open:
-            raise RuntimeError("Circuit breaker OPEN — too many consecutive failures")
+            raise BrainClientError("Circuit breaker OPEN — too many consecutive failures")
+
+        self._check_auth()
 
         async with self.semaphore:
             # Rate limiting
@@ -61,56 +131,72 @@ class BrainClient:
                 async with self.session.post(
                     f"{self.config.api_url}/simulations",
                     json=payload,
+                    headers=self._headers,
                 ) as resp:
                     if resp.status == 429:
-                        # Rate limited — back off
-                        await asyncio.sleep(30)
+                        # Rate limited — back off with jitter
+                        wait = 30 + (self._consecutive_failures * 10)
+                        await asyncio.sleep(wait)
                         return await self.submit_alpha(expression, neutralization, decay)
-                    elif resp.status >= 500:
-                        self._consecutive_failures += 1
-                        if self._consecutive_failures >= 5:
-                            self._circuit_open = True
-                        raise aiohttp.ClientResponseError(
-                            resp.request_info, resp.history, status=resp.status
-                        )
+                    elif resp.status >= 400:
+                        self._handle_response_error(resp)
 
                     resp.raise_for_status()
                     self._consecutive_failures = 0
                     self._submissions_today += 1
                     return await resp.json()
 
-            except aiohttp.ClientError as e:
+            except (aiohttp.ClientError, asyncio.TimeoutError, ConnectionError) as e:
                 self._consecutive_failures += 1
                 if self._consecutive_failures >= 5:
                     self._circuit_open = True
-                raise
+                raise BrainClientError(f"BRAIN API request failed: {e}") from e
 
     async def poll_simulation(self, sim_id: str, max_wait: int = 300) -> dict[str, Any]:
         """Poll a simulation until completion or timeout."""
         elapsed = 0
         interval = 5
+        last_error = None
 
         while elapsed < max_wait:
-            async with self.session.get(
-                f"{self.config.api_url}/simulations/{sim_id}"
-            ) as resp:
-                if resp.status == 200:
-                    data = await resp.json()
-                    if data.get("status") == "DONE":
-                        return data
-                    elif data.get("status") == "ERROR":
-                        raise RuntimeError(f"Simulation failed: {data.get('error', 'unknown')}")
+            try:
+                async with self.session.get(
+                    f"{self.config.api_url}/simulations/{sim_id}",
+                    headers=self._headers,
+                ) as resp:
+                    if resp.status == 200:
+                        data = await resp.json()
+                        if data.get("status") == "DONE":
+                            return data
+                        elif data.get("status") == "ERROR":
+                            raise BrainClientError(
+                                f"Simulation failed: {data.get('error', 'unknown')}"
+                            )
+                    elif resp.status == 401:
+                        raise BrainAuthError("Auth expired during polling. Refresh BRAIN_SESSION_TOKEN.")
+                    elif resp.status >= 500:
+                        last_error = f"Server error {resp.status}"
+                    # Non-fatal errors: log and retry
+            except (aiohttp.ClientError, asyncio.TimeoutError, ConnectionError) as e:
+                last_error = str(e)
 
             await asyncio.sleep(interval)
             elapsed += interval
 
-        raise TimeoutError(f"Simulation {sim_id} timed out after {max_wait}s")
+        raise BrainClientError(
+            f"Simulation {sim_id} timed out after {max_wait}s. Last error: {last_error}"
+        )
 
     def parse_metrics(self, sim_result: dict, alpha_id: str) -> BrainMetrics:
         """Extract structured metrics from BRAIN simulation result."""
         stats = sim_result.get("stats", {})
         yearly = sim_result.get("yearly", [])
 
+        # BRAIN sometimes nests stats differently
+        if not stats and "result" in sim_result:
+            stats = sim_result["result"].get("stats", {})
+            yearly = sim_result["result"].get("yearly", [])
+
         return BrainMetrics(
             alpha_id=alpha_id,
             sharpe_full=stats.get("sharpe", 0.0),