++
This commit is contained in:
jack
@@ -0,0 +1,396 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
import os
|
||||
import time
|
||||
import json
|
||||
import httpx
|
||||
import psycopg2
|
||||
from httpx import BasicAuth
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
def __init__(self):
|
||||
self.connection = None
|
||||
self.database_name = "alpha"
|
||||
self.init_database()
|
||||
|
||||
def create_database(self):
|
||||
"""创建数据库（如果不存在）"""
|
||||
try:
|
||||
# 先连接到默认的postgres数据库来创建alpha数据库
|
||||
conn = psycopg2.connect(
|
||||
host="localhost",
|
||||
port="5432",
|
||||
database="postgres",
|
||||
user="jack",
|
||||
password="aaaAAA111"
|
||||
)
|
||||
conn.autocommit = True
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 检查数据库是否存在
|
||||
cursor.execute("SELECT 1 FROM pg_catalog.pg_database WHERE datname = %s", (self.database_name,))
|
||||
exists = cursor.fetchone()
|
||||
|
||||
if not exists:
|
||||
cursor.execute(f"CREATE DATABASE {self.database_name}")
|
||||
print(f"数据库 {self.database_name} 创建成功")
|
||||
else:
|
||||
print(f"数据库 {self.database_name} 已存在")
|
||||
|
||||
cursor.close()
|
||||
conn.close()
|
||||
|
||||
except Exception as e:
|
||||
print(f"创建数据库时出错: {e}")
|
||||
raise
|
||||
|
||||
def get_connection(self):
|
||||
"""获取数据库连接"""
|
||||
if self.connection is None or self.connection.closed:
|
||||
self.connection = psycopg2.connect(
|
||||
host="localhost",
|
||||
port="5432",
|
||||
database=self.database_name,
|
||||
user="jack",
|
||||
password="aaaAAA111"
|
||||
)
|
||||
return self.connection
|
||||
|
||||
def init_database(self):
|
||||
"""初始化数据库和表结构"""
|
||||
# 先创建数据库
|
||||
self.create_database()
|
||||
|
||||
# 然后连接到此数据库创建表
|
||||
conn = self.get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
# 创建 alpha_prepare 表
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS alpha_prepare
|
||||
(
|
||||
id
|
||||
SERIAL
|
||||
PRIMARY
|
||||
KEY,
|
||||
alpha
|
||||
TEXT
|
||||
NOT
|
||||
NULL
|
||||
UNIQUE,
|
||||
unused
|
||||
BOOLEAN
|
||||
NOT
|
||||
NULL
|
||||
DEFAULT
|
||||
TRUE,
|
||||
created_time
|
||||
TIMESTAMP
|
||||
DEFAULT
|
||||
CURRENT_TIMESTAMP
|
||||
)
|
||||
''')
|
||||
|
||||
# 创建 simulation 表
|
||||
cursor.execute('''
|
||||
CREATE TABLE IF NOT EXISTS simulation
|
||||
(
|
||||
id
|
||||
SERIAL
|
||||
PRIMARY
|
||||
KEY,
|
||||
expression
|
||||
TEXT
|
||||
NOT
|
||||
NULL,
|
||||
time_consuming
|
||||
REAL
|
||||
NOT
|
||||
NULL,
|
||||
status
|
||||
TEXT
|
||||
NOT
|
||||
NULL,
|
||||
timestamp
|
||||
TEXT
|
||||
NOT
|
||||
NULL,
|
||||
alpha_id
|
||||
TEXT,
|
||||
message
|
||||
TEXT,
|
||||
created_time
|
||||
TIMESTAMP
|
||||
DEFAULT
|
||||
CURRENT_TIMESTAMP
|
||||
)
|
||||
''')
|
||||
|
||||
conn.commit()
|
||||
print(f"数据库 {self.database_name} 表结构初始化完成")
|
||||
|
||||
def get_unused_alpha(self):
|
||||
"""获取所有未使用的alpha表达式"""
|
||||
conn = self.get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("SELECT alpha FROM alpha_prepare WHERE unused = TRUE")
|
||||
results = cursor.fetchall()
|
||||
|
||||
alpha_list = [result[0] for result in results]
|
||||
return alpha_list
|
||||
|
||||
def mark_alpha_used(self, alpha):
|
||||
"""将alpha标记为已使用"""
|
||||
conn = self.get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute("UPDATE alpha_prepare SET unused = FALSE WHERE alpha = %s", (alpha,))
|
||||
conn.commit()
|
||||
|
||||
def insert_simulation_result(self, result_data):
|
||||
"""插入模拟结果到simulation表"""
|
||||
conn = self.get_connection()
|
||||
cursor = conn.cursor()
|
||||
|
||||
cursor.execute('''
|
||||
INSERT INTO simulation
|
||||
(expression, time_consuming, status, timestamp, alpha_id, message)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
''', (
|
||||
result_data["expression"],
|
||||
result_data["time_consuming"],
|
||||
result_data["status"],
|
||||
result_data["timestamp"],
|
||||
result_data.get("alpha_id"),
|
||||
result_data.get("message", "")
|
||||
))
|
||||
|
||||
conn.commit()
|
||||
|
||||
def close_connection(self):
|
||||
"""关闭数据库连接"""
|
||||
if self.connection and not self.connection.closed:
|
||||
self.connection.close()
|
||||
|
||||
|
||||
class AlphaSimulator:
|
||||
def __init__(self, credentials_file='account.txt'):
|
||||
self.credentials_file = credentials_file
|
||||
self.client = None
|
||||
self.brain_api_url = 'https://api.worldquantbrain.com'
|
||||
self.db_manager = DatabaseManager()
|
||||
|
||||
def __del__(self):
|
||||
"""析构函数，确保数据库连接被关闭"""
|
||||
if hasattr(self, 'db_manager'):
|
||||
self.db_manager.close_connection()
|
||||
|
||||
def load_credentials(self):
|
||||
if not os.path.exists(self.credentials_file):
|
||||
print("未找到 account.txt 文件")
|
||||
with open(self.credentials_file, 'w') as f:
|
||||
f.write("")
|
||||
print("account.txt 文件已创建，请填写账号密码, 格式: ['username', 'password]")
|
||||
exit(1)
|
||||
|
||||
with open(self.credentials_file) as f:
|
||||
credentials = eval(f.read())
|
||||
return credentials[0], credentials[1]
|
||||
|
||||
def login(self):
|
||||
username, password = self.load_credentials()
|
||||
self.client = httpx.Client(auth=BasicAuth(username, password))
|
||||
|
||||
response = self.client.post(f'{self.brain_api_url}/authentication')
|
||||
print(f"登录状态: {response.status_code}")
|
||||
|
||||
if response.status_code == 201:
|
||||
print(f"登录成功!:{response.json()}")
|
||||
return response.json()
|
||||
else:
|
||||
print(f"登录失败: {response.json()}")
|
||||
return {}
|
||||
|
||||
def load_alpha_list(self):
|
||||
"""从数据库加载未使用的alpha表达式"""
|
||||
alpha_list = self.db_manager.get_unused_alpha()
|
||||
return alpha_list
|
||||
|
||||
def simulate_alpha(self, expression):
|
||||
if self.client is None:
|
||||
raise Exception("请先登录")
|
||||
|
||||
settings = {
|
||||
'instrumentType': 'EQUITY',
|
||||
'region': 'USA',
|
||||
'universe': 'TOP3000',
|
||||
'delay': 1,
|
||||
'decay': 0,
|
||||
'neutralization': 'INDUSTRY',
|
||||
'truncation': 0.08,
|
||||
'pasteurization': 'ON',
|
||||
'unitHandling': 'VERIFY',
|
||||
'nanHandling': 'OFF',
|
||||
'language': 'FASTEXPR',
|
||||
'visualization': False,
|
||||
}
|
||||
|
||||
simulation_data = {
|
||||
'type': 'REGULAR',
|
||||
'settings': settings,
|
||||
'regular': expression
|
||||
}
|
||||
|
||||
sim_resp = self.client.post(f'{self.brain_api_url}/simulations', json=simulation_data)
|
||||
print(f"模拟提交状态: {sim_resp.status_code}")
|
||||
|
||||
sim_progress_url = sim_resp.headers['location']
|
||||
|
||||
while True:
|
||||
sim_progress_resp = self.client.get(sim_progress_url)
|
||||
retry_after_sec = float(sim_progress_resp.headers.get("Retry-After", 0))
|
||||
|
||||
if retry_after_sec == 0:
|
||||
break
|
||||
|
||||
if sim_progress_resp.json():
|
||||
result = sim_progress_resp.json()
|
||||
progress = result.get('progress', 0)
|
||||
print(f"模拟进度: {float(progress) * 100}%")
|
||||
|
||||
print(f"等待 {retry_after_sec} 秒...")
|
||||
time.sleep(retry_after_sec)
|
||||
|
||||
if sim_progress_resp.json().get("status") == "ERROR":
|
||||
result = sim_progress_resp.json().get("message", "未知错误")
|
||||
print(f"因子模拟失败: {result}")
|
||||
return {"status": "err", "message": result}
|
||||
|
||||
alpha_id = sim_progress_resp.json().get("alpha")
|
||||
print(f"生成的Alpha ID: {alpha_id}")
|
||||
|
||||
return {"status": "ok", "alpha_id": alpha_id}
|
||||
|
||||
def run_batch_simulation(self, alpha_list, batch_size=3):
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
for i in range(0, len(alpha_list), batch_size):
|
||||
batch = alpha_list[i:i + batch_size]
|
||||
print(f"\n开始处理第 {i // batch_size + 1} 批因子，共 {len(batch)} 个")
|
||||
|
||||
for expression in batch:
|
||||
print(f"\n模拟因子: {expression}")
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
result = self.simulate_alpha(expression)
|
||||
end_time = time.time()
|
||||
time_consuming = round(end_time - start_time, 2)
|
||||
|
||||
result_data = {
|
||||
"expression": expression,
|
||||
"time_consuming": time_consuming,
|
||||
"status": result["status"],
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
|
||||
if result["status"] == "ok":
|
||||
result_data["alpha_id"] = result["alpha_id"]
|
||||
success_count += 1
|
||||
print(f"✅ 模拟成功 - Alpha ID: {result['alpha_id']} - 耗时: {time_consuming}秒")
|
||||
else:
|
||||
result_data["message"] = result.get("message", "")
|
||||
fail_count += 1
|
||||
print(f"❌ 模拟失败 - {result.get('message', '未知错误')}")
|
||||
|
||||
# 无论成功失败，都标记为已使用并保存结果到数据库
|
||||
self.db_manager.mark_alpha_used(expression)
|
||||
self.db_manager.insert_simulation_result(result_data)
|
||||
|
||||
except Exception as e:
|
||||
end_time = time.time()
|
||||
time_consuming = round(end_time - start_time, 2)
|
||||
error_result = {
|
||||
"expression": expression,
|
||||
"time_consuming": time_consuming,
|
||||
"status": "err",
|
||||
"message": str(e),
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
fail_count += 1
|
||||
|
||||
# 异常情况也标记为已使用并保存结果到数据库
|
||||
self.db_manager.mark_alpha_used(expression)
|
||||
self.db_manager.insert_simulation_result(error_result)
|
||||
|
||||
print(f"❌ 模拟异常 - {str(e)}")
|
||||
|
||||
print(f"第 {i // batch_size + 1} 批处理完成")
|
||||
|
||||
self._print_summary(success_count, fail_count)
|
||||
return success_count, fail_count
|
||||
|
||||
def _print_summary(self, success_count, fail_count):
|
||||
"""打印总结信息"""
|
||||
now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||||
print_result = f"\n总计: 成功 {success_count} 个, 失败 {fail_count} 个\n\n 完成时间: {now}"
|
||||
print(print_result)
|
||||
print(f"所有结果已保存到 PostgreSQL 数据库 {self.db_manager.database_name} 的 simulation 表中")
|
||||
_send_to_gotify(print_result, now)
|
||||
|
||||
def _send_to_gotify(self, text, now):
|
||||
"""将结果推送到 gotify """
|
||||
title = f"alpha模拟结果 时间: {now}"
|
||||
try:
|
||||
resp = httpx.post("https://gotify.erhe.top/message?token=AvKJCJwQKU6yLP8", json={'title': title, 'message': text})
|
||||
except Exception as e:
|
||||
print(e)
|
||||
|
||||
def main():
|
||||
simulator = AlphaSimulator()
|
||||
|
||||
def refresh_login():
|
||||
"""刷新登录token"""
|
||||
login_result = simulator.login()
|
||||
if not login_result:
|
||||
print("登录失败，程序退出")
|
||||
exit(1)
|
||||
return login_result
|
||||
|
||||
# 初始登录
|
||||
login_result = refresh_login()
|
||||
|
||||
while True:
|
||||
try:
|
||||
expiry = int(login_result['token']['expiry'])
|
||||
if expiry < 1800:
|
||||
print(f"Token剩余 {expiry}秒，重新登录...")
|
||||
login_result = refresh_login()
|
||||
|
||||
alpha_list = simulator.load_alpha_list()
|
||||
if not alpha_list:
|
||||
print("暂无待处理的alpha表达式，10分钟后重新检查...")
|
||||
time.sleep(600)
|
||||
continue
|
||||
|
||||
print(f"共加载 {len(alpha_list)} 个需要模拟的因子表达式")
|
||||
|
||||
success_count, fail_count = simulator.run_batch_simulation(alpha_list, batch_size=3)
|
||||
print(f"本轮处理完成: 成功 {success_count} 个, 失败 {fail_count} 个")
|
||||
|
||||
print("等待5分钟后继续检查...")
|
||||
time.sleep(300)
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n程序被用户中断")
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"程序执行异常: {e}")
|
||||
print("5秒后重试...")
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user