import argparse import os import re import json import random import threading import queue from typing import Dict, Tuple, List, Optional, Set import requests from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry # ============================================================ # 1. 基础配置 # ============================================================ BASE_URL = "https://modelhub.org.cn" LOGIN_ENDPOINT = "/adminApi/user/login" ADD_ADAPT_TASK_ENDPOINT = "/api/adapt/task/add" MODELHUB_DB_CHECK_ENDPOINT = "/api/computility/models/list/page/vo" TASK_TYPE = "text-generation" FRAMEWORK = "llamacpp" # 按当前策略部署约定，ModelHub 登录凭据固定写入策略代码。 USER_ACCOUNT = "1326356020@qq.com" USER_PASSWORD = "4pdpassword" # 不使用手动 token，统一用账号密码登录获取 token MODELHUB_TOKEN = "" # 平台环境变量契约：提交任务时必须携带当前策略 ID。 STRATEGY_ID = os.environ["STRATEGY_ID"] # ============================================================ # 2. Hugging Face / hf-mirror 筛选配置 # ============================================================ ORG_NAME = "mradermacher" HF_API_URL = "https://hf-mirror.com/api/models" PER_PAGE = int(os.getenv("HF_PER_PAGE", "100")) START_PAGE = int(os.getenv("HF_START_PAGE", "1")) END_PAGE = int(os.getenv("HF_END_PAGE", "0")) # 0 表示持续翻页直到末页 REQUEST_DELAY = 1 # 每次请求的基础延迟（秒），可根据情况调整 RANDOM_DELAY_RANGE = (0.5, 2) # 随机延迟范围（秒），避免固定间隔被识别 RETRY_TIMES = 3 # 429错误重试次数 RETRY_DELAY = 5 # 每次重试的等待时间（秒） MAX_FILE_SIZE_GB = 12 # ============================================================ # 3. 提交与轮询配置 # ============================================================ # 如果所有模型都想强制用某种 GPU，就填： # FORCE_TARGET_GPU = "Ascend_910-b4" # 如果保留 None，则 i1/i2/... 模型走 Ascend_910-b4，其他走 Mthreads_s4000 FORCE_TARGET_GPU = None SUBMIT_INTERVAL_SECONDS = 2 TASK_LIMIT_CODE = 60007 POLL_INTERVAL_SECONDS = 120 MAX_LIMIT_RETRY_TIMES = None # None 表示一直等到有空位 PIPELINE_MODE = os.getenv("PIPELINE_MODE", "continuous").strip().lower() SCAN_INTERVAL_SECONDS = int(os.getenv("PIPELINE_SCAN_INTERVAL_SECONDS", "1800")) MAX_SCAN_PAGES = int(os.getenv("HF_MAX_PAGES", "0")) # 0 表示不限制 RETRY_FAILED = os.getenv("PIPELINE_RETRY_FAILED", "true").strip().lower() in { "1", "true", "yes", "on" } SHUTDOWN_TIMEOUT_SECONDS = 25 DEBUG_PRINT_PAYLOAD = False # 是否开启断点续跑：已经成功提交过的模型，下次运行自动跳过 RESUME = True # ============================================================ # 断点续跑 / 重跑配置 # ============================================================ # run_mode 说明： # resume：优先复用已有步骤产物；缺失时才运行该步骤。适合 Step 4 报错后接着跑。 # fresh ：从 start_step 开始重跑，并覆盖对应步骤产物；start_step 之前的步骤读取已有产物。 # 命令行参数优先级高于这里的默认值，例如： # python modelhub_pipeline_resumable.py --start-step 4 --run-mode resume DEFAULT_RUN_MODE = os.getenv("PIPELINE_RUN_MODE", "resume") DEFAULT_START_STEP = int(os.getenv("PIPELINE_START_STEP", "1")) DEFAULT_FORCE_RERUN_STEPS = os.getenv("PIPELINE_FORCE_RERUN_STEPS", "") OUTPUT_DIR = "pipeline_outputs" ALL_MODELS_FILE = os.path.join(OUTPUT_DIR, "01_all_models.txt") SIZE_FILTERED_FILE = os.path.join(OUTPUT_DIR, "02_size_filtered_models.txt") NOT_IN_DB_FILE = os.path.join(OUTPUT_DIR, "03_not_in_db_models.txt") SUBMITTED_FILE = os.path.join(OUTPUT_DIR, "04_submitted_models.txt") FAILED_FILE = os.path.join(OUTPUT_DIR, "05_failed_models.txt") SUMMARY_FILE = os.path.join(OUTPUT_DIR, "summary.json") BASE_HEADERS = { "Content-Type": "application/json", "Accept": "application/json, text/plain, */*", "Origin": BASE_URL, "Referer": BASE_URL + "/", "User-Agent": "Mozilla/5.0", } _shutdown_event = threading.Event() class ShutdownRequested(Exception): """平台请求策略进程优雅停止。""" def request_shutdown() -> None: _shutdown_event.set() def reset_shutdown() -> None: _shutdown_event.clear() def check_shutdown() -> None: if _shutdown_event.is_set(): raise ShutdownRequested("收到停机信号") def interruptible_sleep(seconds: float) -> None: if _shutdown_event.wait(max(0, seconds)): raise ShutdownRequested("收到停机信号") # ============================================================ # 4. 通用工具 # ============================================================ def ensure_output_dir() -> None: os.makedirs(OUTPUT_DIR, exist_ok=True) def write_model_list(path: str, models: List[str]) -> None: with open(path, "w", encoding="utf-8") as f: for model in models: f.write(f'"{model}",\n') def append_model(path: str, model: str) -> None: with open(path, "a", encoding="utf-8") as f: f.write(f'"{model}",\n') def load_model_set(path: str) -> Set[str]: if not os.path.exists(path): return set() with open(path, "r", encoding="utf-8") as f: content = f.read() quoted_items = re.findall(r'["\']([^"\']+)["\']', content) if quoted_items: return set(x.strip() for x in quoted_items if x.strip()) result = set() for line in content.splitlines(): item = line.strip().strip(",").strip('"').strip("'").strip() if item: result.add(item) return result def load_model_list(path: str) -> List[str]: """读取模型列表，保持顺序并去重。""" if not os.path.exists(path): raise FileNotFoundError( f"找不到断点文件: {path}。如果要从中间步骤继续，请先确认前置步骤已成功生成该文件。" ) with open(path, "r", encoding="utf-8") as f: content = f.read() quoted_items = re.findall(r'["\']([^"\']+)["\']', content) if quoted_items: raw_items = [x.strip() for x in quoted_items if x.strip()] else: raw_items = [] for line in content.splitlines(): item = line.strip().strip(",").strip('"').strip("'").strip() if item: raw_items.append(item) seen = set() result = [] for item in raw_items: if item not in seen: result.append(item) seen.add(item) return result def parse_bool_env(value: str, default: bool = False) -> bool: if value is None or value == "": return default return value.strip().lower() in {"1", "true", "yes", "y", "on"} def parse_step_set(value: str) -> Set[int]: result: Set[int] = set() if not value: return result for item in value.split(","): item = item.strip() if not item: continue try: step = int(item) except ValueError: raise ValueError(f"PIPELINE_FORCE_RERUN_STEPS 里包含非法步骤: {item}") if step not in {1, 2, 3}: raise ValueError("可强制重跑的步骤只能是 1、2、3；Step 4 每次都会按断点提交剩余模型。") result.add(step) return result def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="HF / hf-mirror -> ModelHub 断点续跑脚本" ) parser.add_argument( "--run-mode", choices=["resume", "fresh"], default=DEFAULT_RUN_MODE, help="resume=复用已有步骤结果；fresh=从 start-step 开始重跑。默认读取 PIPELINE_RUN_MODE 或 resume。", ) parser.add_argument( "--start-step", type=int, choices=[1, 2, 3, 4], default=DEFAULT_START_STEP, help="从第几步开始。Step 4 报错后继续跑用 --start-step 4。默认读取 PIPELINE_START_STEP 或 1。", ) parser.add_argument( "--force-rerun-steps", default=DEFAULT_FORCE_RERUN_STEPS, help="在 resume 模式下强制重跑指定前置步骤，例如 '2,3'。Step 4 不需要写入这里。", ) parser.add_argument( "--reset-submitted", action="store_true", help="清空 04_submitted_models.txt。谨慎使用：会导致已成功提交过的模型不再被跳过。", ) parser.add_argument( "--reset-failed", action="store_true", help="清空 05_failed_models.txt。适合修复参数后重新统计失败模型。", ) parser.add_argument( "--skip-failed", action="store_true", help="跳过 05_failed_models.txt 中记录过失败的模型。默认不跳过，便于修复问题后自动重试失败模型。", ) args = parser.parse_args() args.force_rerun_steps = parse_step_set(args.force_rerun_steps) return args def should_run_step(step: int, output_file: str, args: argparse.Namespace) -> bool: """ 判断 Step 1-3 是否需要执行。 resume：如果产物存在就读取，不存在才执行；force-rerun-steps 可强制执行。 fresh ：start_step 之前读取已有产物，从 start_step 开始重新执行并覆盖产物。 """ if step in args.force_rerun_steps: return True if args.run_mode == "fresh": return step >= args.start_step # resume mode if step < args.start_step: return False return not os.path.exists(output_file) def print_step_loaded(step: int, path: str, models: List[str]) -> None: print(f"Step {step}/4: 跳过执行，读取已有断点文件: {path}") print(f"读取到 {len(models)} 个模型") def remove_file_if_exists(path: str) -> None: if os.path.exists(path): os.remove(path) print(f"已清空: {path}") def create_hf_session() -> requests.Session: session = requests.Session() retry_strategy = Retry( total=RETRY_TIMES, backoff_factor=RETRY_DELAY, status_forcelist=[429, 500, 502, 503, 504], allowed_methods=["GET"], ) adapter = HTTPAdapter(max_retries=retry_strategy) session.mount("https://", adapter) session.mount("http://", adapter) session.headers.update({ "User-Agent": ( "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " "AppleWebKit/537.36 (KHTML, like Gecko) " "Chrome/120.0.0.0 Safari/537.36" ) }) return session def get_model_filename(model_id: str) -> str: """ 从 model_id 生成 GGUF 文件名。 mradermacher/Qwen3-8B-xxx-i1-GGUF -> Qwen3-8B-xxx.i1-Q4_0.gguf QuantFactory/Apollo2-9B-GGUF -> Apollo2-9B.Q8_0.gguf """ base_name = model_id.split("/")[-1] if "_-_" in base_name: base_name = base_name.split("_-_")[-1] if base_name.lower().endswith("-gguf"): base_name = base_name[:-5] match = re.search(r"-i(\d+)$", base_name) if match: number = match.group(1) return base_name[:match.start()] + f".i{number}-Q4_0.gguf" return base_name + ".Q8_0.gguf" # ============================================================ # 5. Step 1：从 hf-mirror 获取组织模型 # ============================================================ def get_org_models(session: requests.Session, org_name: str) -> List[str]: models: List[str] = [] page = START_PAGE pages_seen = 0 while END_PAGE == 0 or page <= END_PAGE: if MAX_SCAN_PAGES and pages_seen >= MAX_SCAN_PAGES: break try: interruptible_sleep(random.uniform(*RANDOM_DELAY_RANGE)) params = { "author": org_name, "page": page, "perPage": PER_PAGE, "sort": "lastModified", "direction": "-1", } response = session.get(HF_API_URL, params=params, timeout=20) response.raise_for_status() data = response.json() if not data: print(f"第 {page} 页无数据，停止获取") break page_models = [] for item in data: model_id = item.get("modelId") if model_id: page_models.append(model_id) models.append(model_id) print(f"成功获取第 {page} 页，共 {len(page_models)} 个模型") pages_seen += 1 page += 1 except ShutdownRequested: raise except Exception as e: print(f"获取第 {page} 页失败: {e}") pages_seen += 1 page += 1 # 去重但保持顺序 seen = set() unique_models = [] for model in models: if model not in seen: unique_models.append(model) seen.add(model) return unique_models # ============================================================ # 6. Step 2：按目标 GGUF 文件大小筛选 # ============================================================ def extract_gguf_file_size( session: requests.Session, model_id: str, target_filename: str, ) -> float: """ 返回 GB。 找不到或失败返回 -1。 """ interruptible_sleep(REQUEST_DELAY + random.uniform(*RANDOM_DELAY_RANGE)) files_api_url = f"https://hf-mirror.com/api/models/{model_id}/tree/main" for retry in range(RETRY_TIMES + 1): try: response = session.get(files_api_url, timeout=20) response.raise_for_status() files_data = response.json() for file_item in files_data: if ( file_item.get("type") == "file" and file_item.get("path") == target_filename ): return file_item.get("size", 0) / (1024 ** 3) print(f"未找到文件: {model_id} -> {target_filename}") return -1 except requests.exceptions.HTTPError as e: status_code = getattr(e.response, "status_code", None) if status_code == 429 and retry < RETRY_TIMES: wait_time = RETRY_DELAY * (retry + 1) print( f"触发 HF 限流: {model_id}，等待 {wait_time} 秒后重试 " f"({retry + 1}/{RETRY_TIMES})" ) interruptible_sleep(wait_time) continue print(f"提取文件大小失败: {model_id} -> {target_filename}: {e}") return -1 except Exception as e: print(f"提取文件大小失败: {model_id} -> {target_filename}: {e}") return -1 return -1 def filter_models_by_gguf_size( session: requests.Session, models: List[str], max_size_gb: float, ) -> List[str]: filtered: List[str] = [] for index, model in enumerate(models, start=1): target_file = get_model_filename(model) file_size = extract_gguf_file_size(session, model, target_file) prefix = f"[{index}/{len(models)}]" if file_size != -1 and file_size <= max_size_gb: filtered.append(model) print( f"{prefix} 符合大小条件: {model} -> {target_file} " f"({file_size:.2f} GB)" ) else: size_str = f"{file_size:.2f} GB" if file_size != -1 else "未知" print( f"{prefix} 不符合大小条件: {model} -> {target_file} " f"({size_str})" ) return filtered # ============================================================ # 7. Step 3：ModelHub 查重，过滤已入库模型 # ============================================================ def check_model_in_modelhub_db(model_id: str) -> Optional[bool]: """ 返回： True = 已入库 False = 未入库 None = 查询异常 """ url = BASE_URL + MODELHUB_DB_CHECK_ENDPOINT payload = { "current": 1, "pageSize": 20, "searchText": model_id, } try: response = requests.post( url, json=payload, headers=BASE_HEADERS, timeout=15, ) if response.status_code != 200: print(f"查重 HTTP 异常: {model_id} -> {response.status_code}") return None resp_json = response.json() code = resp_json.get("code") records = resp_json.get("data", {}).get("records", []) if code != 0: print( f"查重接口业务异常: {model_id} -> " f"{resp_json.get('message', 'unknown')}" ) return None if len(records) == 0: print(f"未入库: {model_id}") return False first_id = records[0].get("id") print(f"已入库: {model_id}，记录 ID: {first_id}") return True except requests.exceptions.Timeout: print(f"查重超时: {model_id}") return None except requests.exceptions.RequestException as e: print(f"查重网络错误: {model_id}: {e}") return None except Exception as e: print(f"查重未知错误: {model_id}: {e}") return None def filter_models_not_in_db(models: List[str]) -> List[str]: not_in_db: List[str] = [] for index, model in enumerate(models, start=1): print(f"[{index}/{len(models)}] 查重: {model}") result = check_model_in_modelhub_db(model) if result is False: not_in_db.append(model) elif result is True: pass else: # 查询异常时，默认不提交，避免重复入库。 print(f"查重异常，跳过提交: {model}") interruptible_sleep(0.3) return not_in_db # ============================================================ # 8. 登录 # ============================================================ def login() -> str: check_shutdown() if not USER_ACCOUNT or not USER_PASSWORD: raise RuntimeError( "没有设置账号密码，请在脚本中填写 USER_ACCOUNT / USER_PASSWORD。" ) payload = { "userAccount": USER_ACCOUNT, "userPassword": USER_PASSWORD, } print("正在登录 ModelHub...") resp = requests.post( BASE_URL + LOGIN_ENDPOINT, headers=BASE_HEADERS, json=payload, timeout=20, ) if resp.status_code != 200: raise RuntimeError(f"HTTP 登录失败: {resp.status_code} - {resp.text}") data = resp.json() if data.get("code") != 0: raise RuntimeError(f"业务登录失败: {data.get('message') or data}") token = data["data"]["token"] print("登录成功") return token def get_token() -> str: return login() # ============================================================ # 9. Step 4：提交验证任务 # ============================================================ def is_i_variant_model(model_filename: str) -> bool: return re.search(r"\.i\d+-Q4_0\.gguf$", model_filename) is not None def build_config_params(model_id: str) -> Tuple[str, str]: model_filename = get_model_filename(model_id) if FORCE_TARGET_GPU: target_gpu = FORCE_TARGET_GPU elif is_i_variant_model(model_filename): target_gpu = "Ascend_910-b4" else: target_gpu = "Mthreads_s4000" if target_gpu == "Ascend_910-b4": config_params = f"""framework: llamacpp api: chat lang: zh max_model_len: 4096 max_tokens: 1024 temperature: 0.7 repetition_penalty: 1.1 top_p: 0.9 sut_config: gpu_num: 1 values: command: ['/workspace/llama.cpp/build_ascend/bin/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '16', '--n-gpu-layers', '128', '--prio', '3', '--min_p', '0.01', '--ctx-size', '4096', '--host', '0.0.0.0', '--port', '3316', '--jinja', '--flash-attn', 'off'] ref_config: gpu_num: 1 values: command: ['/workspace/llama.cpp/build/bin/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '16', '--n-gpu-layers', '128', '--prio', '3', '--min_p', '0.01', '--ctx-size', '4096', '--host', '0.0.0.0', '--port', '80', '--jinja'] """ elif target_gpu == "Mthreads_s4000": config_params = f"""framework: llamacpp api: completion max_tokens: 1024 temperature: 0 repetition_penalty: 1.1 top_p: 0.9 max_model_len: 4096 sut_config: gpu_num: 1 values: command: ['/app/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '20', '--n-gpu-layers', '999', '--prio', '3', '--min_p', '0.01', '--ctx-size', '2048', '--host', '0.0.0.0', '--port', '8000', '--jinja', '--flash-attn', 'off'] ref_config: gpu_num: 1 values: command: ['/workspace/llama.cpp/build/bin/llama-server', '--model', '/model/{model_filename}', '--alias', 'llm', '--threads', '20', '--n-gpu-layers', '999', '--prio', '3', '--min_p', '0.01', '--ctx-size', '2048', '--host', '0.0.0.0', '--port', '8000', '--jinja', '--flash-attn', 'off'] """ else: raise ValueError(f"暂不支持的 targetGpu: {target_gpu}") return target_gpu, config_params def is_business_success(data: Dict) -> bool: if not isinstance(data, dict): return False code = data.get("code") success = data.get("success") if code in (0, "0", 200, "200"): return True if success is True: return True return False def submit_adapt_task(token: str, model_id: str) -> str: """ 返回： SUCCESS LIMIT FAILED AUTH_FAILED """ check_shutdown() target_gpu, config_params = build_config_params(model_id) headers = { **BASE_HEADERS, "Authorization": f"Bearer {token}", } payload = { "modelAddress": model_id, "strategyId": STRATEGY_ID, "taskType": TASK_TYPE, "targetGpu": target_gpu, "framework": FRAMEWORK, "configParams": config_params, } print(f"\n提交模型: {model_id}") print(f"GPU: {target_gpu}") print(f"模型文件名: {get_model_filename(model_id)}") if DEBUG_PRINT_PAYLOAD: print(json.dumps(payload, ensure_ascii=False, indent=2)) try: resp = requests.post( BASE_URL + ADD_ADAPT_TASK_ENDPOINT, headers=headers, json=payload, timeout=20, ) print(f"HTTP 状态码: {resp.status_code}") if resp.status_code in (401, 403): print(f"鉴权失败: {resp.text}") return "AUTH_FAILED" try: data = resp.json() except Exception: data = None if resp.status_code != 200: print(f"提交失败，HTTP 错误: {resp.status_code} - {resp.text}") return "FAILED" if data is None: print(f"提交完成，但返回不是 JSON: {resp.text}") return "SUCCESS" if is_business_success(data): task_id = ( data.get("data", {}).get("taskId") or data.get("data", {}).get("id") or data.get("taskId") or data.get("id") ) if task_id: print(f"提交成功，Task ID: {task_id}") else: print(f"提交成功，返回: {json.dumps(data, ensure_ascii=False)}") return "SUCCESS" code = data.get("code") message = data.get("message", "") if str(code) == str(TASK_LIMIT_CODE): print(f"达到任务上限，稍后重试: {json.dumps(data, ensure_ascii=False)}") return "LIMIT" if "数量已达上限" in message or "当前等待中或运行中" in message: print(f"达到任务上限，稍后重试: {json.dumps(data, ensure_ascii=False)}") return "LIMIT" print(f"提交失败，业务返回: {json.dumps(data, ensure_ascii=False)}") return "FAILED" except requests.RequestException as e: print(f"请求异常: {e}") return "FAILED" def submit_adapt_task_with_polling(token: str, model_id: str) -> Tuple[bool, str]: retry_times = 0 current_token = token while True: result = submit_adapt_task(current_token, model_id) if result == "SUCCESS": return True, current_token if result == "AUTH_FAILED": print("尝试重新登录后重试当前模型...") current_token = login() interruptible_sleep(2) continue if result == "FAILED": return False, current_token if result == "LIMIT": retry_times += 1 if MAX_LIMIT_RETRY_TIMES is not None and retry_times > MAX_LIMIT_RETRY_TIMES: print(f"达到最大等待次数，放弃当前模型: {model_id}") return False, current_token print( f"平台等待中/运行中的异步验证任务已满。" f"等待 {POLL_INTERVAL_SECONDS} 秒后重试当前模型。" f"模型: {model_id}，等待次数: {retry_times}" ) interruptible_sleep(POLL_INTERVAL_SECONDS) continue print(f"未知提交状态: {result}") return False, current_token def iter_org_models(session: requests.Session, org_name: str): """逐页获取模型，避免等待全量列表完成后才开始处理。""" page = START_PAGE pages_seen = 0 while END_PAGE == 0 or page <= END_PAGE: if MAX_SCAN_PAGES and pages_seen >= MAX_SCAN_PAGES: return check_shutdown() try: interruptible_sleep(random.uniform(*RANDOM_DELAY_RANGE)) response = session.get( HF_API_URL, params={ "author": org_name, "page": page, "perPage": PER_PAGE, "sort": "lastModified", "direction": "-1", }, timeout=20, ) response.raise_for_status() data = response.json() if not data: return page_models = [item.get("modelId") for item in data if item.get("modelId")] print(f"扫描第 {page} 页，共 {len(page_models)} 个模型") for model_id in page_models: yield model_id pages_seen += 1 page += 1 except ShutdownRequested: raise except Exception as exc: print(f"扫描第 {page} 页失败: {exc}") pages_seen += 1 page += 1 def continuous_main() -> None: """持续扫描新模型，并由独立消费者逐个提交适配任务。""" ensure_output_dir() pending: "queue.Queue[str]" = queue.Queue() submitted = load_model_set(SUBMITTED_FILE) failed = load_model_set(FAILED_FILE) scheduled: Set[str] = set() stats = {"scanned": 0, "queued": 0, "submitted": 0, "failed": 0} stats_lock = threading.Lock() def submit_worker() -> None: token: Optional[str] = None while not stop_event_is_set(): try: model_id = pending.get(timeout=1) except queue.Empty: continue try: if token is None: token = get_token() ok, token = submit_adapt_task_with_polling(token, model_id) if ok: submitted.add(model_id) append_model(SUBMITTED_FILE, model_id) with stats_lock: stats["submitted"] += 1 else: failed.add(model_id) append_model(FAILED_FILE, model_id) with stats_lock: stats["failed"] += 1 except ShutdownRequested: return except Exception as exc: print(f"提交线程异常 {model_id}: {exc}") failed.add(model_id) append_model(FAILED_FILE, model_id) with stats_lock: stats["failed"] += 1 finally: scheduled.discard(model_id) pending.task_done() worker = threading.Thread(target=submit_worker, name="modelhub-submitter", daemon=True) worker.start() try: while not stop_event_is_set(): session = create_hf_session() try: for model_id in iter_org_models(session, ORG_NAME): check_shutdown() with stats_lock: stats["scanned"] += 1 if model_id in submitted or model_id in scheduled: continue if model_id in failed and not RETRY_FAILED: continue target_filename = get_model_filename(model_id) size_gb = extract_gguf_file_size(session, model_id, target_filename) if size_gb < 0 or size_gb > MAX_FILE_SIZE_GB: continue in_db = check_model_in_modelhub_db(model_id) if in_db is not False: continue scheduled.add(model_id) pending.put(model_id) with stats_lock: stats["queued"] += 1 print(f"已加入提交队列: {model_id} (队列长度 {pending.qsize()})") finally: session.close() print(f"扫描周期完成: {stats}; 下一次扫描等待 {SCAN_INTERVAL_SECONDS} 秒") interruptible_sleep(SCAN_INTERVAL_SECONDS) except ShutdownRequested: pass finally: request_shutdown() worker.join(timeout=SHUTDOWN_TIMEOUT_SECONDS) def stop_event_is_set() -> bool: return _shutdown_event.is_set() # ============================================================ # 10. 主流程：筛选 -> 查重 -> 提交 # ============================================================ def main() -> None: if PIPELINE_MODE not in {"continuous", "batch"}: raise ValueError("PIPELINE_MODE 必须是 continuous 或 batch") if PIPELINE_MODE == "continuous": continuous_main() return ensure_output_dir() args = parse_args() check_shutdown() if args.reset_submitted: remove_file_if_exists(SUBMITTED_FILE) if args.reset_failed: remove_file_if_exists(FAILED_FILE) print("=" * 80) print("Auto HF -> ModelHub Pipeline") print("=" * 80) print(f"组织: {ORG_NAME}") print(f"页码范围: {START_PAGE} - {END_PAGE}") print(f"每页数量: {PER_PAGE}") print(f"最大 GGUF 文件大小: {MAX_FILE_SIZE_GB} GB") print(f"输出目录: {OUTPUT_DIR}") print(f"运行模式: {args.run_mode}") print(f"起始步骤: Step {args.start_step}") print(f"强制重跑步骤: {sorted(args.force_rerun_steps) if args.force_rerun_steps else '无'}") print(f"断点续跑跳过已提交模型: {RESUME}") print("=" * 80) hf_session: Optional[requests.Session] = None try: # Step 1: 拉取 HF 模型 print("\nStep 1/4: 获取 HF / hf-mirror 模型列表") if should_run_step(1, ALL_MODELS_FILE, args): hf_session = hf_session or create_hf_session() all_models = get_org_models(hf_session, ORG_NAME) write_model_list(ALL_MODELS_FILE, all_models) print(f"共获取到 {len(all_models)} 个模型") print(f"已保存: {ALL_MODELS_FILE}") else: all_models = load_model_list(ALL_MODELS_FILE) print_step_loaded(1, ALL_MODELS_FILE, all_models) if not all_models: print("没有获取到模型，结束") return # Step 2: 按文件大小筛选 print("\nStep 2/4: 按 GGUF 文件大小筛选") if should_run_step(2, SIZE_FILTERED_FILE, args): hf_session = hf_session or create_hf_session() size_filtered_models = filter_models_by_gguf_size( hf_session, all_models, MAX_FILE_SIZE_GB, ) write_model_list(SIZE_FILTERED_FILE, size_filtered_models) print(f"大小筛选后剩余 {len(size_filtered_models)} 个模型") print(f"已保存: {SIZE_FILTERED_FILE}") else: size_filtered_models = load_model_list(SIZE_FILTERED_FILE) print_step_loaded(2, SIZE_FILTERED_FILE, size_filtered_models) if not size_filtered_models: print("没有符合大小条件的模型，结束") return # Step 3: ModelHub 查重 print("\nStep 3/4: 查询 ModelHub 是否已入库") if should_run_step(3, NOT_IN_DB_FILE, args): not_in_db_models = filter_models_not_in_db(size_filtered_models) write_model_list(NOT_IN_DB_FILE, not_in_db_models) print(f"未入库模型共 {len(not_in_db_models)} 个") print(f"已保存: {NOT_IN_DB_FILE}") else: not_in_db_models = load_model_list(NOT_IN_DB_FILE) print_step_loaded(3, NOT_IN_DB_FILE, not_in_db_models) if not not_in_db_models: print("没有需要提交的未入库模型，结束") return # Step 4: 登录并提交 print("\nStep 4/4: 登录并提交验证任务") token = get_token() already_submitted = load_model_set(SUBMITTED_FILE) if RESUME else set() already_failed = load_model_set(FAILED_FILE) if (RESUME and args.skip_failed) else set() skip_models = set(already_submitted) | set(already_failed) submit_candidates = [ model for model in not_in_db_models if model not in skip_models ] if RESUME: print(f"断点续跑开启，已成功提交过 {len(already_submitted)} 个模型") if args.skip_failed: print(f"本次会额外跳过历史失败模型 {len(already_failed)} 个") print(f"本次待提交 {len(submit_candidates)} 个模型") success_models: List[str] = [] failed_models: List[str] = [] for index, model_id in enumerate(submit_candidates, start=1): print("\n" + "=" * 80) print(f"[{index}/{len(submit_candidates)}] 准备提交: {model_id}") ok, token = submit_adapt_task_with_polling(token, model_id) if ok: success_models.append(model_id) append_model(SUBMITTED_FILE, model_id) else: failed_models.append(model_id) append_model(FAILED_FILE, model_id) if index < len(submit_candidates): interruptible_sleep(SUBMIT_INTERVAL_SECONDS) summary = { "org_name": ORG_NAME, "start_page": START_PAGE, "end_page": END_PAGE, "max_file_size_gb": MAX_FILE_SIZE_GB, "run_mode": args.run_mode, "start_step": args.start_step, "force_rerun_steps": sorted(args.force_rerun_steps), "all_models_count": len(all_models), "size_filtered_count": len(size_filtered_models), "not_in_db_count": len(not_in_db_models), "already_submitted_count": len(already_submitted), "submitted_this_run_count": len(success_models), "failed_this_run_count": len(failed_models), "success_models": success_models, "failed_models": failed_models, } with open(SUMMARY_FILE, "w", encoding="utf-8") as f: json.dump(summary, f, ensure_ascii=False, indent=2) print("\n" + "=" * 80) print("流程完成") print(f"HF 获取模型数: {len(all_models)}") print(f"大小筛选后: {len(size_filtered_models)}") print(f"未入库模型数: {len(not_in_db_models)}") print(f"历史已提交跳过: {len(already_submitted)}") print(f"本次提交成功: {len(success_models)}") print(f"本次提交失败: {len(failed_models)}") print(f"汇总文件: {SUMMARY_FILE}") if failed_models: print("\n失败模型:") for model in failed_models: print(model) finally: if hf_session is not None: hf_session.close() if __name__ == "__main__": main()