diff --git a/tradingagents/dataflows/polymarket.py b/tradingagents/dataflows/polymarket.py index 147ba95..1b139d9 100644 --- a/tradingagents/dataflows/polymarket.py +++ b/tradingagents/dataflows/polymarket.py @@ -13,11 +13,9 @@ import json import logging from datetime import datetime, timezone -import subprocess - import requests -from .stockstats_utils import _get_ssh_config +from .ssh_proxy_pool import ssh_proxy_fetch_json logger = logging.getLogger(__name__) @@ -35,35 +33,10 @@ DEFAULT_LIMIT = 6 def _request(path: str, params: dict) -> dict: url = f"{GAMMA_BASE}/{path}" - try: - response = requests.get(url, params=params, timeout=DIRECT_TIMEOUT) - response.raise_for_status() - return response.json() - except requests.RequestException: - # Fallback: fetch through SSH proxy - logger.info("Polymarket direct failed, trying SSH proxy for %s...", path) - query = "&".join(f"{k}={v}" for k, v in params.items()) - full_url = f"{url}?{query}" - host, port, password = _get_ssh_config() - cmd = [ - "sshpass", "-p", password, - "ssh", "-o", "StrictHostKeyChecking=no", - "-o", "ConnectTimeout=15", - "-p", str(port), f"root@{host}", - f'curl -s "{full_url}"' - ] - try: - result = subprocess.run(cmd, capture_output=True, text=True, timeout=REQUEST_TIMEOUT) - except subprocess.TimeoutExpired: - raise requests.RequestException("SSH proxy timed out") - if result.returncode != 0: - raise requests.RequestException(f"SSH proxy failed: {result.stderr}") - if not result.stdout or not result.stdout.strip(): - raise requests.RequestException("SSH proxy returned empty response") - try: - return json.loads(result.stdout) - except json.JSONDecodeError as e: - raise requests.RequestException(f"SSH proxy returned invalid JSON: {e}") + # Skip local request, go directly through SSH proxy + query = "&".join(f"{k}={v}" for k, v in params.items()) + full_url = f"{url}?{query}" + return ssh_proxy_fetch_json(full_url, timeout=REQUEST_TIMEOUT) def _parse_json_list(value) -> list: diff --git a/tradingagents/dataflows/ssh_proxy_pool.py b/tradingagents/dataflows/ssh_proxy_pool.py new file mode 100644 index 0000000..73f1641 --- /dev/null +++ b/tradingagents/dataflows/ssh_proxy_pool.py @@ -0,0 +1,123 @@ +"""SSH Proxy Connection Pool — 串行化所有SSH中转请求，避免并发资源竞争。 + +所有需要走中转的数据源统一调用 ssh_proxy_fetch(url)， +内部通过队列串行化SSH连接，避免多线程/多进程并发导致的连接冲突。 +""" +import json +import logging +import subprocess +import threading +import time +from queue import Queue, Empty +from typing import Optional + +logger = logging.getLogger(__name__) + +# SSH配置缓存 +_ssh_config: Optional[tuple] = None +_config_lock = threading.Lock() + + +def _get_ssh_config() -> tuple: + """获取SSH配置（host, port, password），带缓存。""" + global _ssh_config + if _ssh_config is not None: + return _ssh_config + with _config_lock: + if _ssh_config is not None: + return _ssh_config + from .stockstats_utils import _get_ssh_config as _load_config + _ssh_config = _load_config() + return _ssh_config + + +class SSHProxyPool: + """SSH代理连接池 — 串行化所有SSH请求。""" + + def __init__(self): + self._queue = Queue() + self._lock = threading.Lock() + self._initialized = False + + def _ensure_initialized(self): + """确保连接池已初始化。""" + if self._initialized: + return + with self._lock: + if self._initialized: + return + # 放一个占位符，表示可以发起请求 + self._queue.put(True) + self._initialized = True + + def fetch(self, url: str, timeout: int = 30) -> str: + """通过SSH代理获取URL内容，串行化执行。 + + Args: + url: 要获取的URL + timeout: 超时时间（秒） + + Returns: + URL响应内容字符串 + + Raises: + RuntimeError: SSH代理失败时抛出 + """ + self._ensure_initialized() + + # 等待轮到自己（最多等timeout秒） + try: + self._queue.get(timeout=timeout) + except Empty: + raise RuntimeError(f"SSH proxy queue timeout after {timeout}s") + + try: + host, port, password = _get_ssh_config() + cmd = [ + "sshpass", "-p", password, + "ssh", "-o", "StrictHostKeyChecking=no", + "-o", "ConnectTimeout=15", + "-o", "ServerAliveInterval=10", + "-p", str(port), f"root@{host}", + f'curl -sL --max-time {timeout} "{url}"' + ] + result = subprocess.run( + cmd, capture_output=True, text=True, + timeout=timeout + 10 # 给SSH额外时间 + ) + + if result.returncode != 0: + raise RuntimeError(f"SSH proxy failed (rc={result.returncode}): {result.stderr[:200]}") + + if not result.stdout or not result.stdout.strip(): + raise RuntimeError("SSH proxy returned empty response") + + return result.stdout + + except subprocess.TimeoutExpired: + raise RuntimeError(f"SSH proxy subprocess timeout after {timeout}s") + finally: + # 无论成功失败，都放回占位符，让下一个请求执行 + self._queue.put(True) + + def fetch_json(self, url: str, timeout: int = 30) -> dict: + """通过SSH代理获取JSON响应。""" + text = self.fetch(url, timeout) + try: + return json.loads(text) + except json.JSONDecodeError as e: + raise RuntimeError(f"SSH proxy returned invalid JSON: {e}") + + +# 全局单例 +_pool = SSHProxyPool() + + +def ssh_proxy_fetch(url: str, timeout: int = 30) -> str: + """通过SSH代理获取URL内容（串行化）。""" + return _pool.fetch(url, timeout) + + +def ssh_proxy_fetch_json(url: str, timeout: int = 30) -> dict: + """通过SSH代理获取JSON响应（串行化）。""" + return _pool.fetch_json(url, timeout) diff --git a/tradingagents/dataflows/stockstats_utils.py b/tradingagents/dataflows/stockstats_utils.py index 681ddef..018e77b 100644 --- a/tradingagents/dataflows/stockstats_utils.py +++ b/tradingagents/dataflows/stockstats_utils.py @@ -43,19 +43,9 @@ def _get_ssh_config(): def _fetch_via_ssh(url: str, timeout: int = 30) -> str: - """Fetch a URL through the foreign server via SSH tunnel.""" - host, port, password = _get_ssh_config() - cmd = [ - "sshpass", "-p", password, - "ssh", "-o", "StrictHostKeyChecking=no", - "-o", "ConnectTimeout=15", - "-p", str(port), f"root@{host}", - f'curl -s "{url}" -H "User-Agent: Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"' - ] - result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) - if result.returncode != 0: - raise RuntimeError(f"SSH fetch failed: {result.stderr}") - return result.stdout + """Fetch a URL through the foreign server via SSH tunnel (uses connection pool).""" + from .ssh_proxy_pool import ssh_proxy_fetch + return ssh_proxy_fetch(url, timeout) def _yf_download_via_proxy(symbol: str, start: str, end: str) -> pd.DataFrame: