Upload alpha_factory/infra/wq_client.py with huggingface_hub

alpha_factory/infra/wq_client.py

@@ -0,0 +1,143 @@
+"""
+WorldQuant BRAIN API Client — async wrapper with rate limiting.
+Handles submission, polling, and result harvesting.
+"""
+import asyncio
+import aiohttp
+import hashlib
+from typing import Any, Optional
+from ..config import BrainConfig
+from ..schemas import BrainMetrics
+
+
+class BrainClient:
+    """
+    Async client for WorldQuant BRAIN API.
+    Rate-limited, retry-on-429, circuit-breaker-protected.
+    """
+
+    def __init__(self, session: aiohttp.ClientSession, config: BrainConfig):
+        self.session = session
+        self.config = config
+        self.semaphore = asyncio.Semaphore(config.max_concurrent)
+        self._submissions_today = 0
+        self._consecutive_failures = 0
+        self._circuit_open = False
+
+    async def submit_alpha(
+        self,
+        expression: str,
+        neutralization: str = "sector",
+        decay: int = 5,
+    ) -> dict[str, Any]:
+        """
+        Submit an alpha expression to BRAIN for simulation.
+        Returns the simulation response or raises on failure.
+        """
+        if self._circuit_open:
+            raise RuntimeError("Circuit breaker OPEN — too many consecutive failures")
+
+        async with self.semaphore:
+            # Rate limiting
+            await asyncio.sleep(self.config.submit_interval_sec)
+
+            payload = {
+                "type": "REGULAR",
+                "settings": {
+                    "instrumentType": "EQUITY",
+                    "region": self.config.region,
+                    "universe": self.config.universe,
+                    "delay": self.config.delay,
+                    "neutralization": neutralization,
+                    "decay": decay,
+                    "truncation": self.config.truncation,
+                    "pasteurization": self.config.pasteurization,
+                    "nanHandling": self.config.nan_handling,
+                },
+                "regular": expression,
+            }
+
+            try:
+                async with self.session.post(
+                    f"{self.config.api_url}/simulations",
+                    json=payload,
+                ) as resp:
+                    if resp.status == 429:
+                        # Rate limited — back off
+                        await asyncio.sleep(30)
+                        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
+                        )
+
+                    resp.raise_for_status()
+                    self._consecutive_failures = 0
+                    self._submissions_today += 1
+                    return await resp.json()
+
+            except aiohttp.ClientError as e:
+                self._consecutive_failures += 1
+                if self._consecutive_failures >= 5:
+                    self._circuit_open = True
+                raise
+
+    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
+
+        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')}")
+
+            await asyncio.sleep(interval)
+            elapsed += interval
+
+        raise TimeoutError(f"Simulation {sim_id} timed out after {max_wait}s")
+
+    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", [])
+
+        return BrainMetrics(
+            alpha_id=alpha_id,
+            sharpe_full=stats.get("sharpe", 0.0),
+            sharpe_is=stats.get("sharpe_is", stats.get("sharpe", 0.0)),
+            sharpe_os=stats.get("sharpe_os", stats.get("sharpe", 0.0)),
+            fitness=stats.get("fitness", 0.0),
+            turnover=stats.get("turnover", 0.0),
+            returns=stats.get("returns", 0.0),
+            max_drawdown=stats.get("max_drawdown", 0.0),
+            yearly_sharpe=[y.get("sharpe", 0.0) for y in yearly],
+            yearly_returns=[y.get("returns", 0.0) for y in yearly],
+            margin_pct=stats.get("margin", None),
+            long_count=stats.get("long_count", None),
+            short_count=stats.get("short_count", None),
+        )
+
+    @property
+    def submissions_today(self) -> int:
+        return self._submissions_today
+
+    def reset_circuit(self):
+        """Manually reset the circuit breaker."""
+        self._circuit_open = False
+        self._consecutive_failures = 0
+
+    @staticmethod
+    def alpha_hash(expression: str, neutralization: str, decay: int) -> str:
+        """Deterministic hash for dedup."""
+        key = f"{expression.strip()}|{neutralization}|{decay}"
+        return hashlib.sha256(key.encode()).hexdigest()[:16]