3 Commits

60
main.py
View File

@@ -31,7 +31,7 @@ PORT = 8080
STRATEGY_ID = os.getenv("STRATEGY_ID", "") STRATEGY_ID = os.getenv("STRATEGY_ID", "")
# 目标GPU # 目标GPU
TARGET_GPU = "ppu_zw_810e" TARGET_GPU = "Iluvatar_bi-150"
# 账号Token # 账号Token
TARGET_TOKEN = "f45f1aae2c094426be237c88b1085015" TARGET_TOKEN = "f45f1aae2c094426be237c88b1085015"
@@ -98,20 +98,28 @@ def init_db():
# ModelScope 搜索 # ModelScope 搜索
# ============================================================ # ============================================================
MODELSCOPE_API = "https://modelscope.cn/api/v1"
def search_models(keyword: str, limit: int = 50) -> list: def search_models(keyword: str, limit: int = 50) -> list:
"""HuggingFace 搜索模型""" """ModelScope 搜索模型"""
url = "https://huggingface.co/api/models" url = "https://modelscope.cn/openapi/v1/models"
params = { params = {
'search': keyword, 'search': keyword,
'limit': limit, 'page_size': min(limit, 50),
'page_number': 1,
'sort': 'downloads', 'sort': 'downloads',
'direction': -1,
} }
try: try:
resp = requests.get(url, params=params, timeout=20) resp = requests.get(url, params=params, timeout=20,
headers={'User-Agent': 'Mozilla/5.0'})
data = resp.json() data = resp.json()
log(f" [{keyword}]: {len(data)} 个结果") if data.get('success'):
return [{'id': m.get('id'), 'downloads': m.get('downloads', 0)} for m in data] models = data.get('data', {}).get('models', [])
log(f" [{keyword}]: {len(models)} 个结果")
return [{'id': m.get('id'), 'downloads': m.get('downloads', 0)} for m in models]
else:
log(f" [{keyword}]: success=false")
except Exception as e: except Exception as e:
log(f" [{keyword}]: 失败 {e}") log(f" [{keyword}]: 失败 {e}")
return [] return []
@@ -120,7 +128,7 @@ def search_models(keyword: str, limit: int = 50) -> list:
def check_architecture(model_id: str) -> tuple: def check_architecture(model_id: str) -> tuple:
"""检查模型架构""" """检查模型架构"""
try: try:
cfg_url = f"https://huggingface.co/{model_id}/raw/main/config.json" cfg_url = f"{MODELSCOPE_API}/models/{model_id}/repo?Revision=master&FilePath=config.json"
resp = requests.get(cfg_url, timeout=10) resp = requests.get(cfg_url, timeout=10)
if resp.status_code == 200: if resp.status_code == 200:
cfg = resp.json() cfg = resp.json()
@@ -145,7 +153,13 @@ def check_architecture(model_id: str) -> tuple:
def normalize_model_url(model_url: str) -> str: def normalize_model_url(model_url: str) -> str:
"""标准化 URL 格式 - 直接返回 HuggingFace URL""" """标准化 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 return model_url
@@ -285,7 +299,7 @@ def run_pipeline(submit_limit: int = 30):
log(f"提交限制: {submit_limit}") log(f"提交限制: {submit_limit}")
# 1. 搜索 # 1. 搜索
log("\n--- 阶段1: 搜索 HuggingFace ---") log("\n--- 阶段1: 搜索 ModelScope ---")
seen = set() seen = set()
all_models = [] all_models = []
for kw in SEARCH_KEYWORDS: for kw in SEARCH_KEYWORDS:
@@ -298,7 +312,7 @@ def run_pipeline(submit_limit: int = 30):
if downloads >= 50: if downloads >= 50:
all_models.append({ all_models.append({
'model_id': mid, 'model_id': mid,
'url': f"https://huggingface.co/{mid}", 'url': f"https://modelscope.cn/{mid}",
'downloads': downloads, 'downloads': downloads,
}) })
time.sleep(0.3) time.sleep(0.3)
@@ -371,7 +385,7 @@ def run_pipeline(submit_limit: int = 30):
log(f"{m['model_id']}") log(f"{m['model_id']}")
db_conn.execute( db_conn.execute(
'INSERT OR REPLACE INTO submitted VALUES (?,?,?,?)', 'INSERT OR REPLACE INTO submitted VALUES (?,?,?,?)',
(m['model_id'], gpu, str(task_id), datetime.now().isoformat()) (m['model_id'], TARGET_GPU, str(task_id), datetime.now().isoformat())
) )
else: else:
log(f"{m['model_id']}: {msg}") log(f"{m['model_id']}: {msg}")
@@ -492,10 +506,9 @@ class AgentHandler(BaseHTTPRequestHandler):
# 3. ModelHub 查询 API # 3. ModelHub 查询 API
try: try:
token = list(ACCOUNTS.values())[0]
resp = requests.get( resp = requests.get(
'https://modelhub.org.cn/api/adapt/task/page', 'https://modelhub.org.cn/api/adapt/task/page',
headers={'Xc-Token': token, 'Accept': 'application/json'}, headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'},
params={'current': 1, 'pageSize': 1, 'onlyMine': 'true'}, params={'current': 1, 'pageSize': 1, 'onlyMine': 'true'},
timeout=10, timeout=10,
) )
@@ -510,17 +523,16 @@ class AgentHandler(BaseHTTPRequestHandler):
# 4. ModelHub 提交 API (dry test) # 4. ModelHub 提交 API (dry test)
try: try:
token = list(ACCOUNTS.values())[0]
resp = requests.post( resp = requests.post(
'https://modelhub.org.cn/api/adapt/task/add', 'https://modelhub.org.cn/api/adapt/task/add',
headers={'Xc-Token': token, 'Accept': 'application/json', 'Content-Type': 'application/json'}, headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json', 'Content-Type': 'application/json'},
json={ json={
'modelAddress': 'https://www.modelscope.cn/models/Qwen/Qwen3-8B', 'modelAddress': 'https://www.modelscope.cn/models/Qwen/Qwen3-8B',
'taskType': 'text-generation', 'taskType': 'text-generation',
'targetGpu': 'Kunlunxin_p-800', 'targetGpu': TARGET_GPU,
'framework': 'vllm', 'framework': 'vllm',
'strategyId': STRATEGY_ID, 'strategyId': STRATEGY_ID,
'configParams': 'framework: vllm\n', 'configParams': build_config_params(),
}, },
timeout=10, timeout=10,
) )
@@ -614,10 +626,9 @@ def main():
# 3. ModelHub 查询 # 3. ModelHub 查询
try: try:
token = list(ACCOUNTS.values())[0]
resp = requests.get( resp = requests.get(
'https://modelhub.org.cn/api/adapt/task/page', 'https://modelhub.org.cn/api/adapt/task/page',
headers={'Xc-Token': token, 'Accept': 'application/json'}, headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json'},
params={'current': 1, 'pageSize': 1, 'onlyMine': 'true'}, params={'current': 1, 'pageSize': 1, 'onlyMine': 'true'},
timeout=10 timeout=10
) )
@@ -631,15 +642,14 @@ def main():
# 4. ModelHub 提交 # 4. ModelHub 提交
try: try:
token = list(ACCOUNTS.values())[0]
resp = requests.post( resp = requests.post(
'https://modelhub.org.cn/api/adapt/task/add', 'https://modelhub.org.cn/api/adapt/task/add',
headers={'Xc-Token': token, 'Accept': 'application/json', 'Content-Type': 'application/json'}, headers={'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json', 'Content-Type': 'application/json'},
json={ json={
'modelAddress': 'https://www.modelscope.cn/models/Qwen/Qwen3-8B', 'modelAddress': 'https://www.modelscope.cn/models/Qwen/Qwen3-8B',
'taskType': 'text-generation', 'targetGpu': 'Kunlunxin_p-800', 'taskType': 'text-generation', 'targetGpu': TARGET_GPU,
'framework': 'vllm', 'strategyId': STRATEGY_ID, 'framework': 'vllm', 'strategyId': STRATEGY_ID,
'configParams': 'framework: vllm\n', 'configParams': build_config_params(),
}, timeout=10 }, timeout=10
) )
data = resp.json() data = resp.json()