feat: submit only 2 models, add failed model tracking to avoid re-submission
This commit is contained in:
43
main.py
43
main.py
@@ -82,15 +82,33 @@ db_conn = None
|
|||||||
def init_db():
|
def init_db():
|
||||||
global db_conn
|
global db_conn
|
||||||
db_conn = sqlite3.connect(':memory:', check_same_thread=False)
|
db_conn = sqlite3.connect(':memory:', check_same_thread=False)
|
||||||
db_conn.execute('''CREATE TABLE IF NOT EXISTS queue (
|
|
||||||
model_id TEXT, gpu TEXT, url TEXT, downloads INTEGER,
|
|
||||||
params TEXT, category TEXT, score REAL,
|
|
||||||
PRIMARY KEY(model_id, gpu)
|
|
||||||
)''')
|
|
||||||
db_conn.execute('''CREATE TABLE IF NOT EXISTS submitted (
|
db_conn.execute('''CREATE TABLE IF NOT EXISTS submitted (
|
||||||
model_id TEXT, gpu TEXT, task_id TEXT, submitted_at TEXT,
|
model_id TEXT, gpu TEXT, task_id TEXT, status TEXT,
|
||||||
|
submitted_at TEXT, checked_at TEXT,
|
||||||
PRIMARY KEY(model_id, gpu)
|
PRIMARY KEY(model_id, gpu)
|
||||||
)''')
|
)''')
|
||||||
|
db_conn.execute('''CREATE TABLE IF NOT EXISTS failed (
|
||||||
|
model_id TEXT, gpu TEXT, reason TEXT, failed_at TEXT,
|
||||||
|
PRIMARY KEY(model_id, gpu)
|
||||||
|
)''')
|
||||||
|
db_conn.commit()
|
||||||
|
def is_model_failed(model_id: str) -> bool:
|
||||||
|
"""检查模型是否已知失败"""
|
||||||
|
if db_conn:
|
||||||
|
row = db_conn.execute(
|
||||||
|
'SELECT 1 FROM failed WHERE model_id=? AND gpu=?',
|
||||||
|
(model_id, TARGET_GPU)
|
||||||
|
).fetchone()
|
||||||
|
return row is not None
|
||||||
|
return False
|
||||||
|
|
||||||
|
def record_failed(model_id: str, reason: str):
|
||||||
|
"""记录失败的模型"""
|
||||||
|
if db_conn:
|
||||||
|
db_conn.execute(
|
||||||
|
'INSERT OR REPLACE INTO failed VALUES (?,?,?,?)',
|
||||||
|
(model_id, TARGET_GPU, reason, datetime.now().isoformat())
|
||||||
|
)
|
||||||
db_conn.commit()
|
db_conn.commit()
|
||||||
|
|
||||||
|
|
||||||
@@ -302,7 +320,7 @@ def submit_model(model_url: str) -> tuple:
|
|||||||
# 主流程
|
# 主流程
|
||||||
# ============================================================
|
# ============================================================
|
||||||
|
|
||||||
def run_pipeline(submit_limit: int = 30):
|
def run_pipeline(submit_limit: int = 2):
|
||||||
"""完整流程:搜索→筛选→提交(只针对目标GPU)"""
|
"""完整流程:搜索→筛选→提交(只针对目标GPU)"""
|
||||||
init_db()
|
init_db()
|
||||||
log("=" * 50)
|
log("=" * 50)
|
||||||
@@ -390,6 +408,10 @@ def run_pipeline(submit_limit: int = 30):
|
|||||||
if check_my_submitted(model_id):
|
if check_my_submitted(model_id):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
# 检查是否已知失败(避免重复提交)
|
||||||
|
if is_model_failed(model_id):
|
||||||
|
continue
|
||||||
|
|
||||||
to_submit.append(m)
|
to_submit.append(m)
|
||||||
if len(to_submit) >= min(submit_limit, available):
|
if len(to_submit) >= min(submit_limit, available):
|
||||||
break
|
break
|
||||||
@@ -410,6 +432,9 @@ def run_pipeline(submit_limit: int = 30):
|
|||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
log(f" ❌ {m['model_id']}: {msg}")
|
log(f" ❌ {m['model_id']}: {msg}")
|
||||||
|
# 永久失败类型记录到 failed 表
|
||||||
|
if any(kw in str(msg) for kw in ['保护期', '白名单', '唯一性']):
|
||||||
|
record_failed(m['model_id'], msg)
|
||||||
time.sleep(0.5)
|
time.sleep(0.5)
|
||||||
|
|
||||||
log(f" 提交完成: {submitted}/{len(to_submit)}")
|
log(f" 提交完成: {submitted}/{len(to_submit)}")
|
||||||
@@ -470,7 +495,7 @@ class AgentHandler(BaseHTTPRequestHandler):
|
|||||||
if content_len > 0:
|
if content_len > 0:
|
||||||
body = json.loads(self.rfile.read(content_len))
|
body = json.loads(self.rfile.read(content_len))
|
||||||
|
|
||||||
limit = body.get('limit', 30)
|
limit = body.get('limit', 2)
|
||||||
|
|
||||||
self._json({'status': 'started', 'gpu': TARGET_GPU, 'limit': limit})
|
self._json({'status': 'started', 'gpu': TARGET_GPU, 'limit': limit})
|
||||||
|
|
||||||
@@ -687,7 +712,7 @@ def main():
|
|||||||
try:
|
try:
|
||||||
state['running'] = True
|
state['running'] = True
|
||||||
state['last_run'] = datetime.now().isoformat()
|
state['last_run'] = datetime.now().isoformat()
|
||||||
count = run_pipeline(submit_limit=30)
|
count = run_pipeline(submit_limit=2)
|
||||||
state['last_result'] = {'submitted': count, 'success': True}
|
state['last_result'] = {'submitted': count, 'success': True}
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log(f"流程异常: {traceback.format_exc()}")
|
log(f"流程异常: {traceback.format_exc()}")
|
||||||
|
|||||||
Reference in New Issue
Block a user