From 709ea0fd78e7e324b1c966569445c53fa81da53b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B0=8F=E9=92=B0?= Date: Wed, 8 Jul 2026 14:34:41 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20SSH=E8=BF=9E=E6=8E=A5=E6=B1=A0=E4=B8=B2?= =?UTF-8?q?=E8=A1=8C=E5=8C=96=20+=20Polymarket/=E6=95=B0=E6=8D=AE=E6=BA=90?= =?UTF-8?q?=E7=BB=9F=E4=B8=80=E8=B5=B0=E8=BF=9E=E6=8E=A5=E6=B1=A0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增ssh_proxy_pool.py: 全局单例连接池，队列串行化SSH请求 - stockstats_utils.py: _fetch_via_ssh改用连接池 - polymarket.py: 跳过本地直接走SSH连接池 - 消除并发SSH资源竞争问题 --- tradingagents/dataflows/polymarket.py | 37 +----- tradingagents/dataflows/ssh_proxy_pool.py | 123 ++++++++++++++++++++ tradingagents/dataflows/stockstats_utils.py | 16 +-- 3 files changed, 131 insertions(+), 45 deletions(-) create mode 100644 tradingagents/dataflows/ssh_proxy_pool.py 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: