Files
modelhub-submit-agent/main.py

698 lines
24 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""
ModelHub 全云端提交智能体
部署在平台容器中,直接从内网提交验证任务
功能:
- POST /run → 触发一次完整流程(搜索→筛选→提交)
- GET /status → 查看当前状态
- GET /health → 健康检查
"""
import json
import os
import signal
import sqlite3
import threading
import time
import traceback
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse, parse_qs
import requests
import yaml
# ============================================================
# 配置
# ============================================================
HOST = "0.0.0.0"
PORT = 8080
STRATEGY_ID = os.getenv("STRATEGY_ID", "")
# 目标GPU
TARGET_GPU = "Iluvatar_bi-150"
# 账号Token
TARGET_TOKEN = "f45f1aae2c094426be237c88b1085015"
# 架构白名单
SUPPORTED_ARCH_KEYWORDS = ['Qwen', 'Qwen2', 'Qwen3']
SUPPORTED_MODEL_TYPES = [
'qwen', 'qwen2', 'qwen2_vl', 'qwen2_5_vl', 'qwen2_audio',
'qwen3', 'qwen3_vl', 'qwen3_5', 'qwen3_5_moe',
]
SUPPORTED_SPECIAL_ARCHS = ['Eagle3Speculator', 'LlamaForCausalLMEagle3']
MODELHUB_API = "https://modelhub.org.cn/api"
# 搜索关键词
SEARCH_KEYWORDS = ['qwen', 'Qwen2', 'Qwen3', 'Qwen3.5', 'Qwen1.5', 'Qwen-']
# ============================================================
# 全局状态
# ============================================================
state = {
'running': False,
'last_run': None,
'last_result': None,
'logs': [],
}
state_lock = threading.Lock()
def log(msg: str):
ts = datetime.now().strftime('%H:%M:%S')
line = f"[{ts}] {msg}"
print(line, flush=True)
with state_lock:
state['logs'].append(line)
if len(state['logs']) > 500:
state['logs'] = state['logs'][-300:]
# ============================================================
# 数据库(内存 SQLite
# ============================================================
db_conn = None
def init_db():
global db_conn
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 (
model_id TEXT, gpu TEXT, task_id TEXT, submitted_at TEXT,
PRIMARY KEY(model_id, gpu)
)''')
db_conn.commit()
# ============================================================
# ModelScope 搜索
# ============================================================
MODELSCOPE_API = "https://modelscope.cn/api/v1"
DOWNLOAD_MIN = 50
DOWNLOAD_MAX = 5000
SEARCH_PAGES = 5 # 每个关键词搜5页50*5=250个结果
def search_models(keyword: str) -> list:
"""从 ModelScope 搜索模型多页筛选下载量50-5000的冷门模型"""
url = "https://modelscope.cn/openapi/v1/models"
models = []
for page in range(1, SEARCH_PAGES + 1):
params = {
'search': keyword,
'page_size': 50,
'page_number': page,
'sort': 'downloads',
}
try:
resp = requests.get(url, params=params, timeout=20,
headers={'User-Agent': 'Mozilla/5.0'})
data = resp.json()
if data.get('success'):
page_models = data.get('data', {}).get('models', [])
for m in page_models:
dl = m.get('downloads', 0)
if DOWNLOAD_MIN <= dl <= DOWNLOAD_MAX:
models.append({'id': m.get('id'), 'downloads': dl})
if len(page_models) < 50:
break # 最后一页,不继续
else:
break
except Exception as e:
log(f" [{keyword}] page={page}: {e}")
break
time.sleep(0.3)
log(f" [{keyword}]: {len(models)} 个 (50<={DOWNLOAD_MAX})")
return models
def check_architecture(model_id: str) -> tuple:
"""检查模型架构"""
try:
cfg_url = f"{MODELSCOPE_API}/models/{model_id}/repo?Revision=master&FilePath=config.json"
resp = requests.get(cfg_url, timeout=10)
if resp.status_code == 200:
cfg = resp.json()
archs = cfg.get('architectures', [])
mtype = cfg.get('model_type', '')
arch_str = str(archs)
for special in SUPPORTED_SPECIAL_ARCHS:
if special in arch_str:
return True, special
for kw in SUPPORTED_ARCH_KEYWORDS:
if kw in arch_str:
return True, kw
if mtype in SUPPORTED_MODEL_TYPES:
return True, mtype
return False, f"arch={archs} type={mtype}"
mid_upper = model_id.upper()
if 'QWEN3' in mid_upper or 'QWEN2' in mid_upper:
return True, "GGUF"
return False, "no config, not Qwen"
except Exception as e:
return True, f"check error: {e}"
def normalize_model_url(model_url: str) -> str:
"""标准化 URL 格式"""
if '/models/' in model_url:
return model_url
if 'modelscope.cn/' in model_url:
parts = model_url.split('modelscope.cn/')
if len(parts) == 2:
return f"https://www.modelscope.cn/models/{parts[1]}"
return model_url
# ============================================================
# 平台 API
# ============================================================
def check_platform_verify(model_id: str) -> dict:
"""查询全平台验证状态"""
headers = {'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'}
url = f"{MODELHUB_API}/computility/models/search-by-model-id"
try:
resp = requests.get(url, headers=headers, params={'modelId': model_id}, timeout=10)
data = resp.json()
if data.get('code') == 0:
result = data.get('data', {}).get('verifyResult', {})
return result if isinstance(result, dict) else {}
except Exception:
pass
return {}
def check_my_submitted(model_id: str) -> bool:
"""检查自己是否已提交"""
headers = {'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'}
url = f"{MODELHUB_API}/adapt/task/page"
try:
resp = requests.get(url, headers=headers, params={
'current': 1, 'pageSize': 100, 'onlyMine': 'true',
'gpuType': TARGET_GPU, 'modelId': model_id,
}, timeout=10)
data = resp.json()
if data.get('code') == 0:
records = data.get('data', {}).get('records', [])
for r in records:
if r.get('modelId') == model_id and r.get('gpuType') == gpu:
status = r.get('status', '')
if status in ('waiting', 'running', 'success'):
return True
except Exception:
pass
return False
def check_queue_available() -> int:
"""查询队列可用位置"""
headers = {'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'}
url = f"{MODELHUB_API}/adapt/task/page"
try:
resp = requests.get(url, headers=headers, params={
'current': 1, 'pageSize': 1, 'onlyMine': 'true',
'gpuType': TARGET_GPU, 'status': 'waiting',
}, timeout=10)
data = resp.json()
if data.get('code') == 0:
waiting = int(data['data'].get('total', 0))
return max(0, 100 - waiting)
except Exception:
pass
return -1
def build_config_params() -> str:
"""构建 YAML 配置 - 完全匹配平台自动生成的格式只支持vllm"""
params = {
'framework': 'vllm',
'nv_framework': 'vllm',
'api': 'completion',
'max_tokens': 1024,
'temperature': 0.7,
'repetition_penalty': 1.2,
'top_p': 0.9,
'lang': 'zh',
'max_model_len': 2048,
'sut_config': {
'gpu_num': 1,
'values': {
'command': [
'vllm', 'serve', '/model', '--port', '8000',
'--served-model-name', 'llm', '--max-model-len', '2048',
'--dtype', 'auto', '--gpu-memory-utilization', '0.95',
'-tp', '1', '--enforce-eager', '--trust-remote-code',
]
}
},
'ref_config': {
'gpu_num': 1,
'values': {
'command': [
'vllm', 'serve', '/model', '--port', '80',
'--served-model-name', 'llm', '--max-model-len', '4096',
'--enforce-eager', '--trust-remote-code', '-tp', '1',
]
}
},
}
return yaml.dump(params, default_flow_style=False, allow_unicode=True, width=1000)
def submit_model(model_url: str) -> tuple:
"""提交单个模型"""
headers = {
'Xc-Token': TARGET_TOKEN,
'Accept': 'application/json',
'Content-Type': 'application/json',
}
url = f"{MODELHUB_API}/adapt/task/add"
payload = {
'modelAddress': normalize_model_url(model_url),
'taskType': 'text-generation',
'targetGpu': TARGET_GPU,
'framework': 'vllm',
'strategyId': STRATEGY_ID,
'configParams': build_config_params(),
}
try:
resp = requests.post(url, headers=headers, json=payload, timeout=30)
data = resp.json()
if data.get('code') == 0:
task_id = data.get('data', {}).get('id')
return True, task_id, 'success'
return False, None, data.get('message', 'unknown error')
except Exception as e:
return False, None, str(e)
# ============================================================
# 主流程
# ============================================================
def run_pipeline(submit_limit: int = 30):
"""完整流程搜索→筛选→提交只针对目标GPU"""
init_db()
log("=" * 50)
log("开始执行流程")
log(f"目标GPU: {TARGET_GPU}")
log(f"提交限制: {submit_limit}")
# 1. 搜索
log("\n--- 阶段1: 搜索 ModelScope ---")
seen = set()
all_models = []
for kw in SEARCH_KEYWORDS:
models = search_models(kw)
for m in models:
mid = m.get('id', '')
if mid and mid not in seen:
seen.add(mid)
all_models.append({
'model_id': mid,
'url': f"https://modelscope.cn/{mid}",
'downloads': m.get('downloads', 0),
})
time.sleep(0.3)
log(f"搜索完成: {len(seen)} 个唯一模型, {len(all_models)} 个下载量{DOWNLOAD_MIN}-{DOWNLOAD_MAX}")
# 2. 格式筛选只保留HuggingFace格式排除GGUF
log("\n--- 阶段2: 格式筛选 ---")
hf_models = []
gguf_skipped = 0
for m in all_models:
mid = m['model_id'].upper()
if 'GGUF' in mid:
gguf_skipped += 1
log(f"{m['model_id']}: GGUF格式跳过")
else:
hf_models.append(m)
log(f"格式筛选: {len(hf_models)} 通过 (HuggingFace), {gguf_skipped} 跳过 (GGUF)")
# 3. 架构筛选
log("\n--- 阶段3: 架构筛选 ---")
arch_passed = []
arch_rejected = 0
for m in hf_models:
ok, reason = check_architecture(m['model_id'])
if ok:
arch_passed.append(m)
else:
arch_rejected += 1
log(f"{m['model_id']}: {reason}")
time.sleep(0.15)
log(f"架构筛选: {len(arch_passed)} 通过, {arch_rejected} 拒绝")
# 4. 筛选并提交只针对目标GPU
log(f"\n--- 阶段4: 筛选并提交 [{TARGET_GPU}] ---")
# 检查队列
available = check_queue_available()
if available <= 0:
log(f" 队列满,跳过")
return 0
log(f" 队列可用: {available}")
# 筛选
to_submit = []
for m in arch_passed:
model_id = m['model_id']
# 检查全平台验证状态
verify = check_platform_verify(model_id)
if TARGET_GPU in verify:
continue # 已有记录,跳过
# 检查自己是否已提交
if check_my_submitted(model_id):
continue
to_submit.append(m)
if len(to_submit) >= min(submit_limit, available):
break
time.sleep(0.2)
log(f" 待提交: {len(to_submit)}")
# 提交
submitted = 0
for m in to_submit:
ok, task_id, msg = submit_model(m['url'])
if ok:
submitted += 1
log(f"{m['model_id']}")
db_conn.execute(
'INSERT OR REPLACE INTO submitted VALUES (?,?,?,?)',
(m['model_id'], TARGET_GPU, str(task_id), datetime.now().isoformat())
)
else:
log(f"{m['model_id']}: {msg}")
time.sleep(0.5)
log(f" 提交完成: {submitted}/{len(to_submit)}")
db_conn.commit()
log(f"\n{'=' * 50}")
log(f"流程完成,共提交 {submitted} 个模型")
return submitted
# ============================================================
# HTTP 服务
# ============================================================
class AgentHandler(BaseHTTPRequestHandler):
def do_GET(self):
parsed = urlparse(self.path)
path = parsed.path
if path == '/health':
self._json({'status': 'ok'})
elif path == '/':
self._json({
'name': 'modelhub-submit-agent',
'strategy_id': STRATEGY_ID,
'status': 'running' if state['running'] else 'idle',
'last_run': state['last_run'],
'last_result': state['last_result'],
})
elif path == '/test':
self._json(self._run_connectivity_test())
elif path == '/status':
self._json({
'running': state['running'],
'last_run': state['last_run'],
'last_result': state['last_result'],
'queue_count': db_conn.execute('SELECT COUNT(*) FROM queue').fetchone()[0] if db_conn else 0,
'submitted_count': db_conn.execute('SELECT COUNT(*) FROM submitted').fetchone()[0] if db_conn else 0,
})
elif path == '/logs':
lines = int(parse_qs(parsed.query).get('lines', ['50'])[0])
self._json({'logs': state['logs'][-lines:]})
else:
self._json({'error': 'not found'}, 404)
def do_POST(self):
parsed = urlparse(self.path)
path = parsed.path
if path == '/run':
if state['running']:
self._json({'error': 'already running'}, 409)
return
# 读取请求体
content_len = int(self.headers.get('Content-Length', 0))
body = {}
if content_len > 0:
body = json.loads(self.rfile.read(content_len))
limit = body.get('limit', 30)
self._json({'status': 'started', 'gpu': TARGET_GPU, 'limit': limit})
# 后台运行
def _run():
try:
state['running'] = True
state['last_run'] = datetime.now().isoformat()
count = run_pipeline(submit_limit=limit)
state['last_result'] = {'submitted': count, 'success': True}
except Exception as e:
log(f"流程异常: {traceback.format_exc()}")
state['last_result'] = {'error': str(e), 'success': False}
finally:
state['running'] = False
threading.Thread(target=_run, daemon=True).start()
else:
self._json({'error': 'not found'}, 404)
def _run_connectivity_test(self) -> dict:
"""测试各 API 连通性"""
results = {}
# 1. ModelScope 搜索 API
try:
resp = requests.get(
'https://modelscope.cn/openapi/v1/models',
params={'search': 'qwen', 'page_size': 2, 'sort': 'downloads'},
timeout=10,
headers={'User-Agent': 'Mozilla/5.0'}
)
results['modelscope_search'] = {
'status': resp.status_code,
'ok': resp.status_code == 200,
'body_preview': resp.text[:200] if resp.status_code == 200 else resp.text[:100],
}
except Exception as e:
results['modelscope_search'] = {'ok': False, 'error': str(e)}
# 2. ModelScope config.json API
try:
resp = requests.get(
'https://modelscope.cn/api/v1/models/Qwen/Qwen3-8B/repo?Revision=master&FilePath=config.json',
timeout=10,
)
results['modelscope_config'] = {
'status': resp.status_code,
'ok': resp.status_code == 200,
'body_preview': resp.text[:200] if resp.status_code == 200 else resp.text[:100],
}
except Exception as e:
results['modelscope_config'] = {'ok': False, 'error': str(e)}
# 3. ModelHub 查询 API
try:
resp = requests.get(
'https://modelhub.org.cn/api/adapt/task/page',
headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'},
params={'current': 1, 'pageSize': 1, 'onlyMine': 'true'},
timeout=10,
)
data = resp.json()
results['modelhub_query'] = {
'ok': data.get('code') == 0,
'code': data.get('code'),
'total': data.get('data', {}).get('total'),
}
except Exception as e:
results['modelhub_query'] = {'ok': False, 'error': str(e)}
# 4. ModelHub 提交 API (dry test)
try:
resp = requests.post(
'https://modelhub.org.cn/api/adapt/task/add',
headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json', 'Content-Type': 'application/json'},
json={
'modelAddress': 'https://www.modelscope.cn/models/Qwen/Qwen3-8B',
'taskType': 'text-generation',
'targetGpu': TARGET_GPU,
'framework': 'vllm',
'strategyId': STRATEGY_ID,
'configParams': build_config_params(),
},
timeout=10,
)
data = resp.json()
results['modelhub_submit'] = {
'ok': data.get('code') == 0,
'code': data.get('code'),
'message': data.get('message', '')[:100],
}
except Exception as e:
results['modelhub_submit'] = {'ok': False, 'error': str(e)}
return results
def _json(self, body: dict, status: int = 200):
payload = json.dumps(body, ensure_ascii=False).encode()
self.send_response(status)
self.send_header('Content-Type', 'application/json; charset=utf-8')
self.send_header('Content-Length', str(len(payload)))
self.end_headers()
self.wfile.write(payload)
def log_message(self, fmt, *args):
pass # 静默 HTTP 日志
# ============================================================
# 启动
# ============================================================
shutdown_requested = False
def _handle_signal(signum, _frame):
global shutdown_requested
shutdown_requested = True
log(f"收到信号 {signum},准备关闭")
def main():
signal.signal(signal.SIGTERM, _handle_signal)
signal.signal(signal.SIGINT, _handle_signal)
init_db()
server = ThreadingHTTPServer((HOST, PORT), AgentHandler)
server.timeout = 1
log(f"智能体启动 | {HOST}:{PORT}")
log(f"STRATEGY_ID: {STRATEGY_ID}")
log(f"目标GPU: {TARGET_GPU}")
# 启动后自动运行连通性测试
def _startup_test():
time.sleep(2)
log("=" * 50)
log("连通性测试开始")
log("=" * 50)
# 1. ModelScope 搜索
try:
resp = requests.get(
'https://modelscope.cn/openapi/v1/models',
params={'search': 'qwen', 'page_size': 2, 'sort': 'downloads'},
timeout=10, headers={'User-Agent': 'Mozilla/5.0'}
)
if resp.status_code == 200:
data = resp.json()
models = data.get('data', {}).get('models', [])
log(f"[ModelScope搜索] ✅ status={resp.status_code} models={len(models)}")
for m in models[:2]:
log(f" {m.get('id')} downloads={m.get('downloads')}")
else:
log(f"[ModelScope搜索] ❌ status={resp.status_code} body={resp.text[:100]}")
except Exception as e:
log(f"[ModelScope搜索] ❌ error={e}")
# 2. ModelScope config.json
try:
resp = requests.get(
'https://modelscope.cn/api/v1/models/Qwen/Qwen3-8B/repo?Revision=master&FilePath=config.json',
timeout=10
)
if resp.status_code == 200:
cfg = resp.json()
log(f"[ModelScope配置] ✅ arch={cfg.get('architectures')} type={cfg.get('model_type')}")
else:
log(f"[ModelScope配置] ❌ status={resp.status_code}")
except Exception as e:
log(f"[ModelScope配置] ❌ error={e}")
# 3. ModelHub 查询
try:
resp = requests.get(
'https://modelhub.org.cn/api/adapt/task/page',
headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'},
params={'current': 1, 'pageSize': 1, 'onlyMine': 'true'},
timeout=10
)
data = resp.json()
if data.get('code') == 0:
log(f"[ModelHub查询] ✅ total={data['data'].get('total')}")
else:
log(f"[ModelHub查询] ❌ code={data.get('code')} msg={data.get('message','')[:80]}")
except Exception as e:
log(f"[ModelHub查询] ❌ error={e}")
# 4. ModelHub 提交
try:
resp = requests.post(
'https://modelhub.org.cn/api/adapt/task/add',
headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json', 'Content-Type': 'application/json'},
json={
'modelAddress': 'https://www.modelscope.cn/models/Qwen/Qwen3-8B',
'taskType': 'text-generation', 'targetGpu': TARGET_GPU,
'framework': 'vllm', 'strategyId': STRATEGY_ID,
'configParams': build_config_params(),
}, timeout=10
)
data = resp.json()
log(f"[ModelHub提交] code={data.get('code')} msg={data.get('message','')[:80]}")
except Exception as e:
log(f"[ModelHub提交] ❌ error={e}")
log("=" * 50)
log("连通性测试完成")
log("=" * 50)
# 开始正式提交流程
log("\n自动触发提交流程...")
try:
state['running'] = True
state['last_run'] = datetime.now().isoformat()
count = run_pipeline(submit_limit=30)
state['last_result'] = {'submitted': count, 'success': True}
except Exception as e:
log(f"流程异常: {traceback.format_exc()}")
state['last_result'] = {'error': str(e), 'success': False}
finally:
state['running'] = False
threading.Thread(target=_startup_test, daemon=True).start()
while not shutdown_requested:
server.handle_request()
server.server_close()
log("智能体已关闭")
if __name__ == '__main__':
main()