From 727f0f678f1163f9ced49790fb9cd3ba1d2700ac Mon Sep 17 00:00:00 2001 From: z3st Date: Thu, 23 Jul 2026 23:19:20 +0800 Subject: [PATCH] feat: single GPU, HuggingFace only, vllm only - Target single GPU (ppu_zw_810e) instead of 4 GPUs - Search only from HuggingFace API (removed ModelScope) - Framework fixed to vllm (removed llama.cpp) - Filter out GGUF format models - Simplified config and functions --- main.py | 338 +++++++++++++++++++++----------------------------------- 1 file changed, 129 insertions(+), 209 deletions(-) diff --git a/main.py b/main.py index ca4db28..515ff93 100644 --- a/main.py +++ b/main.py @@ -30,33 +30,11 @@ HOST = "0.0.0.0" PORT = 8080 STRATEGY_ID = os.getenv("STRATEGY_ID", "") -# 四个账号的 Token -ACCOUNTS = { - 'MetaX_c-500': 'f8e60d1dac7f4472967e7ca40145747b', - 'Kunlunxin_p-800': 'f45f1aae2c094426be237c88b1085015', - 'Ascend_910-b4': 'f3c05879e7c34bbba92f399f12884183', - 'hygon_k100-ai': 'b88507029b884ad3b4bad8ba09e6546e', -} +# 目标GPU +TARGET_GPU = "ppu_zw_810e" -# GPU 引擎配置 -GPU_CONFIGS = { - 'MetaX_c-500': { - 'framework': 'vllm', - 'docker_image': 'modelhubxc-4pd.tencentcloudcr.com/enginex/enginex-metax/vllm:0.9.1', - }, - 'Kunlunxin_p-800': { - 'framework': 'vllm', - 'docker_image': 'modelhubxc-4pd.tencentcloudcr.com/enginex/sunjichen/xc-llm-kunlun:latest', - }, - 'Ascend_910-b4': { - 'framework': 'vllm', - 'docker_image': 'git.modelhub.org.cn:9443/enginex-ascend/vllm-ascend:v0.11.0rc0', - }, - 'hygon_k100-ai': { - 'framework': 'llama.cpp', - 'docker_image': 'modelhubxc-4pd.tencentcloudcr.com/enginex/enginex-hygon/hygon-llama.cpp:b7516', - }, -} +# 账号Token +TARGET_TOKEN = "f45f1aae2c094426be237c88b1085015" # 架构白名单 SUPPORTED_ARCH_KEYWORDS = ['Qwen', 'Qwen2', 'Qwen3'] @@ -66,7 +44,6 @@ SUPPORTED_MODEL_TYPES = [ ] SUPPORTED_SPECIAL_ARCHS = ['Eagle3Speculator', 'LlamaForCausalLMEagle3'] -MODELSCOPE_API = "https://modelscope.cn/api/v1" MODELHUB_API = "https://modelhub.org.cn/api" # 搜索关键词 @@ -122,33 +99,28 @@ def init_db(): # ============================================================ def search_models(keyword: str, limit: int = 50) -> list: - """从 ModelScope 搜索模型""" - url = "https://modelscope.cn/openapi/v1/models" + """从 HuggingFace 搜索模型""" + url = "https://huggingface.co/api/models" params = { 'search': keyword, - 'page_size': min(limit, 50), # API 上限 50 - 'page_number': 1, + 'limit': limit, 'sort': 'downloads', + 'direction': -1, } try: - resp = requests.get(url, params=params, timeout=20, - headers={'User-Agent': 'Mozilla/5.0'}) + resp = requests.get(url, params=params, timeout=20) data = resp.json() - if data.get('success'): - models = data.get('data', {}).get('models', []) - log(f" 搜索 [{keyword}]: status={resp.status_code} models={len(models)}") - return models - else: - log(f" 搜索 [{keyword}]: success=false, data={str(data)[:200]}") + log(f" [{keyword}]: {len(data)} 个结果") + return [{'id': m.get('id'), 'downloads': m.get('downloads', 0)} for m in data] except Exception as e: - log(f" 搜索失败 [{keyword}]: {e}") + log(f" [{keyword}]: 失败 {e}") return [] def check_architecture(model_id: str) -> tuple: """检查模型架构""" try: - cfg_url = f"{MODELSCOPE_API}/models/{model_id}/repo?Revision=master&FilePath=config.json" + cfg_url = f"https://huggingface.co/{model_id}/raw/main/config.json" resp = requests.get(cfg_url, timeout=10) if resp.status_code == 200: cfg = resp.json() @@ -173,13 +145,7 @@ def check_architecture(model_id: str) -> tuple: 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]}" + """标准化 URL 格式 - 直接返回 HuggingFace URL""" return model_url @@ -189,7 +155,7 @@ def normalize_model_url(model_url: str) -> str: def check_platform_verify(model_id: str) -> dict: """查询全平台验证状态""" - headers = {'Xc-Token': list(ACCOUNTS.values())[0], 'Accept': 'application/json'} + 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) @@ -202,14 +168,14 @@ def check_platform_verify(model_id: str) -> dict: return {} -def check_my_submitted(model_id: str, gpu: str, token: str) -> bool: +def check_my_submitted(model_id: str) -> bool: """检查自己是否已提交""" - headers = {'Xc-Token': token, 'Accept': 'application/json'} + 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': gpu, 'modelId': model_id, + 'gpuType': TARGET_GPU, 'modelId': model_id, }, timeout=10) data = resp.json() if data.get('code') == 0: @@ -224,14 +190,14 @@ def check_my_submitted(model_id: str, gpu: str, token: str) -> bool: return False -def check_queue_available(gpu: str, token: str) -> int: +def check_queue_available() -> int: """查询队列可用位置""" - headers = {'Xc-Token': token, 'Accept': 'application/json'} + 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': gpu, 'status': 'waiting', + 'gpuType': TARGET_GPU, 'status': 'waiting', }, timeout=10) data = resp.json() if data.get('code') == 0: @@ -242,99 +208,58 @@ def check_queue_available(gpu: str, token: str) -> int: return -1 -def build_config_params(gpu: str) -> str: - """构建 YAML 配置 - 完全匹配平台自动生成的格式""" - config = GPU_CONFIGS.get(gpu, {}) - framework = config.get('framework', 'vllm') - - if framework == 'llama.cpp': - # hygon_k100-ai 使用 llama.cpp - params = { - 'framework': 'llama.cpp', - 'nv_framework': 'llama.cpp', - 'api': 'completion', - 'max_tokens': 1024, - 'temperature': 0.7, - 'repetition_penalty': 1.2, - 'top_p': 0.9, - 'lang': 'zh', - 'max_model_len': 4096, - 'sut_config': { - 'gpu_num': 1, - 'values': { - 'command': [ - 'llama-server', '--model', '/model', '--alias', 'llm', - '--threads', '20', '--n-gpu-layers', '999', '--prio', '3', - '--min_p', '0.01', '--ctx-size', '4096', - '--host', '0.0.0.0', '--port', '8000', - '--jinja', '--flash-attn', 'off', - ] - } - }, - 'ref_config': { - 'gpu_num': 1, - 'values': { - 'command': [ - 'llama-server', '--model', '/model', '--alias', 'llm', - '--threads', '20', '--n-gpu-layers', '999', - '--ctx-size', '4096', '--host', '0.0.0.0', '--port', '8000', - ] - } - }, - } - else: - # vllm (P800, Ascend, MetaX) - 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', - ] - } - }, - } +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, gpu: str, token: str) -> tuple: +def submit_model(model_url: str) -> tuple: """提交单个模型""" headers = { - 'Xc-Token': token, + 'Xc-Token': TARGET_TOKEN, 'Accept': 'application/json', 'Content-Type': 'application/json', } url = f"{MODELHUB_API}/adapt/task/add" - config = GPU_CONFIGS.get(gpu, {}) payload = { 'modelAddress': normalize_model_url(model_url), 'taskType': 'text-generation', - 'targetGpu': gpu, - 'framework': config.get('framework', 'vllm'), + 'targetGpu': TARGET_GPU, + 'framework': 'vllm', 'strategyId': STRATEGY_ID, - 'configParams': build_config_params(gpu), + 'configParams': build_config_params(), } try: resp = requests.post(url, headers=headers, json=payload, timeout=30) @@ -351,19 +276,16 @@ def submit_model(model_url: str, gpu: str, token: str) -> tuple: # 主流程 # ============================================================ -def run_pipeline(gpus: list = None, submit_limit: int = 30): - """完整流程：搜索→筛选→提交""" - if gpus is None: - gpus = list(GPU_CONFIGS.keys()) - +def run_pipeline(submit_limit: int = 30): + """完整流程：搜索→筛选→提交（只针对目标GPU）""" init_db() log("=" * 50) log("开始执行流程") - log(f"目标GPU: {', '.join(gpus)}") - log(f"提交限制: 每GPU {submit_limit} 个") + log(f"目标GPU: {TARGET_GPU}") + log(f"提交限制: {submit_limit} 个") # 1. 搜索 - log("\n--- 阶段1: 搜索 ModelScope ---") + log("\n--- 阶段1: 搜索 HuggingFace ---") seen = set() all_models = [] for kw in SEARCH_KEYWORDS: @@ -376,19 +298,30 @@ def run_pipeline(gpus: list = None, submit_limit: int = 30): if downloads >= 50: all_models.append({ 'model_id': mid, - 'url': f"https://modelscope.cn/{mid}", + 'url': f"https://huggingface.co/{mid}", 'downloads': downloads, - 'params': m.get('params', ''), - 'category': 'quantized' if 'GGUF' in mid.upper() else 'standard', }) time.sleep(0.3) log(f"搜索完成: {len(seen)} 个唯一模型, {len(all_models)} 个下载量>=50") - # 2. 架构筛选 - log("\n--- 阶段2: 架构筛选 ---") + # 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 all_models: + for m in hf_models: ok, reason = check_architecture(m['model_id']) if ok: arch_passed.append(m) @@ -398,70 +331,58 @@ def run_pipeline(gpus: list = None, submit_limit: int = 30): time.sleep(0.15) log(f"架构筛选: {len(arch_passed)} 通过, {arch_rejected} 拒绝") - # 3. 按 GPU 筛选并提交 - log("\n--- 阶段3: 筛选并提交 ---") - total_submitted = 0 - for gpu in gpus: - token = ACCOUNTS.get(gpu) - if not token: + # 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 - log(f"\n[{gpu}]") + to_submit.append(m) + if len(to_submit) >= min(submit_limit, available): + break + time.sleep(0.2) - # 检查队列 - available = check_queue_available(gpu, token) - if available <= 0: - log(f" 队列满，跳过") - continue - log(f" 队列可用: {available}") + log(f" 待提交: {len(to_submit)}") - # 筛选 - to_submit = [] - for m in arch_passed: - model_id = m['model_id'] + # 提交 + 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'], gpu, str(task_id), datetime.now().isoformat()) + ) + else: + log(f" ❌ {m['model_id']}: {msg}") + time.sleep(0.5) - # hygon 只接受 GGUF - if gpu == 'hygon_k100-ai' and 'GGUF' not in model_id.upper(): - continue - - # 检查全平台验证状态 - verify = check_platform_verify(model_id) - if gpu in verify: - continue # 已有记录，跳过 - - # 检查自己是否已提交 - if check_my_submitted(model_id, gpu, token): - 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'], gpu, token) - if ok: - submitted += 1 - log(f" ✅ {m['model_id']}") - db_conn.execute( - 'INSERT OR REPLACE INTO submitted VALUES (?,?,?,?)', - (m['model_id'], gpu, str(task_id), datetime.now().isoformat()) - ) - else: - log(f" ❌ {m['model_id']}: {msg}") - time.sleep(0.5) - - log(f" 提交完成: {submitted}/{len(to_submit)}") - total_submitted += submitted + log(f" 提交完成: {submitted}/{len(to_submit)}") db_conn.commit() log(f"\n{'=' * 50}") - log(f"流程完成，共提交 {total_submitted} 个模型") - return total_submitted + log(f"流程完成，共提交 {submitted} 个模型") + return submitted # ============================================================ @@ -514,17 +435,16 @@ class AgentHandler(BaseHTTPRequestHandler): if content_len > 0: body = json.loads(self.rfile.read(content_len)) - gpus = body.get('gpus', list(GPU_CONFIGS.keys())) limit = body.get('limit', 30) - self._json({'status': 'started', 'gpus': gpus, 'limit': limit}) + self._json({'status': 'started', 'gpu': TARGET_GPU, 'limit': limit}) # 后台运行 def _run(): try: state['running'] = True state['last_run'] = datetime.now().isoformat() - count = run_pipeline(gpus=gpus, submit_limit=limit) + count = run_pipeline(submit_limit=limit) state['last_result'] = {'submitted': count, 'success': True} except Exception as e: log(f"流程异常: {traceback.format_exc()}") @@ -651,7 +571,7 @@ def main(): log(f"智能体启动 | {HOST}:{PORT}") log(f"STRATEGY_ID: {STRATEGY_ID}") - log(f"GPU: {', '.join(GPU_CONFIGS.keys())}") + log(f"目标GPU: {TARGET_GPU}") # 启动后自动运行连通性测试 def _startup_test():