update
This commit is contained in:
jack
+243
-59
@@ -3,14 +3,188 @@ 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):
|
||||
@@ -33,22 +207,14 @@ class AlphaSimulator:
|
||||
|
||||
if response.status_code == 201:
|
||||
print(f"登录成功!:{response.json()}")
|
||||
return True
|
||||
return response.json()
|
||||
else:
|
||||
print(f"登录失败: {response.json()}")
|
||||
return False
|
||||
|
||||
def load_alpha_list(self, file_path='alpha.txt'):
|
||||
if not os.path.exists(file_path):
|
||||
print(f"{file_path} 文件不存在")
|
||||
with open(file_path, 'w', encoding='utf-8') as file:
|
||||
file.write("")
|
||||
print(f"已创建 {file_path} 文件, 请添加因子后重新运行, 一行一个因子")
|
||||
return []
|
||||
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
alpha_list = [line.strip() for line in file if line.strip()]
|
||||
return {}
|
||||
|
||||
def load_alpha_list(self):
|
||||
"""从数据库加载未使用的alpha表达式"""
|
||||
alpha_list = self.db_manager.get_unused_alpha()
|
||||
return alpha_list
|
||||
|
||||
def simulate_alpha(self, expression):
|
||||
@@ -99,16 +265,16 @@ class AlphaSimulator:
|
||||
if sim_progress_resp.json().get("status") == "ERROR":
|
||||
result = sim_progress_resp.json().get("message", "未知错误")
|
||||
print(f"因子模拟失败: {result}")
|
||||
return {"status": "error", "message": result}
|
||||
return {"status": "err", "message": result}
|
||||
|
||||
alpha_id = sim_progress_resp.json().get("alpha")
|
||||
print(f"生成的Alpha ID: {alpha_id}")
|
||||
|
||||
return {"status": "success", "alpha_id": alpha_id}
|
||||
return {"status": "ok", "alpha_id": alpha_id}
|
||||
|
||||
def run_batch_simulation(self, alpha_list, batch_size=3):
|
||||
success_results = []
|
||||
fail_results = []
|
||||
success_count = 0
|
||||
fail_count = 0
|
||||
|
||||
for i in range(0, len(alpha_list), batch_size):
|
||||
batch = alpha_list[i:i + batch_size]
|
||||
@@ -121,81 +287,99 @@ class AlphaSimulator:
|
||||
try:
|
||||
result = self.simulate_alpha(expression)
|
||||
end_time = time.time()
|
||||
time_consuming = end_time - start_time
|
||||
time_consuming = round(end_time - start_time, 2)
|
||||
|
||||
result_data = {
|
||||
"expression": expression,
|
||||
"time_consuming": round(time_consuming, 2),
|
||||
"time_consuming": time_consuming,
|
||||
"status": result["status"],
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
|
||||
if result["status"] == "success":
|
||||
result_data["alpha_id"] = result["alpha_id"]
|
||||
success_results.append(result_data)
|
||||
print(f"✅ 模拟成功 - Alpha ID: {result['alpha_id']} - 耗时: {time_consuming:.2f}秒")
|
||||
success_count += 1
|
||||
print(f"✅ 模拟成功 - Alpha ID: {result['alpha_id']} - 耗时: {time_consuming}秒")
|
||||
else:
|
||||
result_data["message"] = result.get("message", "")
|
||||
fail_results.append(result_data)
|
||||
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 = end_time - start_time
|
||||
time_consuming = round(end_time - start_time, 2)
|
||||
error_result = {
|
||||
"expression": expression,
|
||||
"time_consuming": round(time_consuming, 2),
|
||||
"status": "error",
|
||||
"time_consuming": time_consuming,
|
||||
"status": "err",
|
||||
"message": str(e),
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
|
||||
}
|
||||
fail_results.append(error_result)
|
||||
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._save_results(success_results, fail_results)
|
||||
return success_results, fail_results
|
||||
self._print_summary(success_count, fail_count)
|
||||
return success_count, fail_count
|
||||
|
||||
def _save_results(self, success_results, fail_results):
|
||||
base_dir = 'simulator_result'
|
||||
success_dir = os.path.join(base_dir, 'success')
|
||||
fail_dir = os.path.join(base_dir, 'fail')
|
||||
|
||||
os.makedirs(success_dir, exist_ok=True)
|
||||
os.makedirs(fail_dir, exist_ok=True)
|
||||
|
||||
timestamp = str(int(time.time()))
|
||||
|
||||
if success_results:
|
||||
success_file = os.path.join(success_dir, f"success_results_{timestamp}.json")
|
||||
with open(success_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(success_results, f, ensure_ascii=False, indent=2)
|
||||
print(f"成功结果已保存到 {success_file}")
|
||||
|
||||
if fail_results:
|
||||
fail_file = os.path.join(fail_dir, f"fail_results_{timestamp}.json")
|
||||
with open(fail_file, 'w', encoding='utf-8') as f:
|
||||
json.dump(fail_results, f, ensure_ascii=False, indent=2)
|
||||
print(f"失败结果已保存到 {fail_file}")
|
||||
|
||||
print(f"\n总计: 成功 {len(success_results)} 个, 失败 {len(fail_results)} 个")
|
||||
def _print_summary(self, success_count, fail_count):
|
||||
"""打印总结信息"""
|
||||
print(f"\n总计: 成功 {success_count} 个, 失败 {fail_count} 个")
|
||||
print(f"所有结果已保存到 PostgreSQL 数据库 {self.db_manager.database_name} 的 simulation 表中")
|
||||
|
||||
|
||||
def main():
|
||||
simulator = AlphaSimulator()
|
||||
|
||||
if not simulator.login():
|
||||
return
|
||||
def refresh_login():
|
||||
"""刷新登录token"""
|
||||
login_result = simulator.login()
|
||||
if not login_result:
|
||||
print("登录失败，程序退出")
|
||||
exit(1)
|
||||
return login_result
|
||||
|
||||
alpha_list = simulator.load_alpha_list('alpha.txt')
|
||||
if not alpha_list:
|
||||
print("未找到有效的因子表达式，请检查 alpha.txt 文件")
|
||||
return
|
||||
# 初始登录
|
||||
login_result = refresh_login()
|
||||
|
||||
while True:
|
||||
try:
|
||||
expiry = int(login_result['token']['expiry'])
|
||||
if expiry < 1800:
|
||||
print(f"Token剩余 {expiry}秒，重新登录...")
|
||||
login_result = refresh_login()
|
||||
|
||||
print(f"共加载 {len(alpha_list)} 个因子表达式")
|
||||
alpha_list = simulator.load_alpha_list()
|
||||
if not alpha_list:
|
||||
print("暂无待处理的alpha表达式，10分钟后重新检查...")
|
||||
time.sleep(600)
|
||||
continue
|
||||
|
||||
simulator.run_batch_simulation(alpha_list, batch_size=3)
|
||||
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__":
|
||||
|
||||
Reference in New Issue
Block a user