From dd21ca1be260e99abbefe8c396c2505728970544 Mon Sep 17 00:00:00 2001 From: Codex Date: Mon, 27 Jul 2026 17:29:26 +0800 Subject: [PATCH] Package pipeline as ModelHub strategy --- .dockerignore | 8 + .env.example | 9 + .gitignore | 6 + Dockerfile | 23 + README.md | 61 ++ app.py | 156 ++++++ modelhub_pipeline_resumable.py | 982 +++++++++++++++++++++++++++++++++ requirements.txt | 1 + 8 files changed, 1246 insertions(+) create mode 100644 .dockerignore create mode 100644 .env.example create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 app.py create mode 100644 modelhub_pipeline_resumable.py create mode 100644 requirements.txt diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..34fb3ed --- /dev/null +++ b/.dockerignore @@ -0,0 +1,8 @@ +.git +.gitignore +.idea +__pycache__ +*.py[cod] +*.log +.env +pipeline_outputs diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..835a5d2 --- /dev/null +++ b/.env.example @@ -0,0 +1,9 @@ +STRATEGY_ID=replace-with-platform-strategy-id +MODELHUB_USER_ACCOUNT=replace-with-modelhub-account +MODELHUB_USER_PASSWORD=replace-with-modelhub-password +PORT=8080 +PIPELINE_RUN_MODE=resume +PIPELINE_START_STEP=1 +PIPELINE_FORCE_RERUN_STEPS= +PIPELINE_RUN_INTERVAL_SECONDS=3600 +PIPELINE_RETRY_INTERVAL_SECONDS=60 diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ecf4d8c --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +.idea/ +__pycache__/ +*.py[cod] +*.log +.env +pipeline_outputs/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..de4dcd7 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,23 @@ +FROM modelhubxc-4pd.tencentcloudcr.com/xc_agent_platform/python:3.11-slim + +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PORT=8080 + +WORKDIR /app + +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +COPY app.py modelhub_pipeline_resumable.py ./ +RUN mkdir -p /app/pipeline_outputs \ + && chown -R 10001:0 /app + +USER 10001 + +EXPOSE 8080 + +HEALTHCHECK --interval=30s --timeout=3s --start-period=10s --retries=3 \ + CMD python -c "import urllib.request; urllib.request.urlopen('http://127.0.0.1:8080/health', timeout=2).read()" + +CMD ["python", "app.py"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..83bdc2b --- /dev/null +++ b/README.md @@ -0,0 +1,61 @@ +# new-pipeline + +将 Hugging Face 镜像中的 GGUF 模型筛选、查重后，提交到 ModelHub +执行适配验证。项目已按 ModelHub 智能体策略运行规范进行容器化。 + +## 运行方式 + +平台运行时必须提供： + +- `STRATEGY_ID`：当前策略 ID，提交任务时写入 `strategyId`。 +- `MODELHUB_USER_ACCOUNT`：ModelHub 登录账号。 +- `MODELHUB_USER_PASSWORD`：ModelHub 登录密码。 + +可选变量： + +- `PORT`：健康检查端口，默认 `8080`。 +- `PIPELINE_RUN_MODE`：`resume` 或 `fresh`，默认 `resume`。 +- `PIPELINE_START_STEP`：从第 1–4 步中的哪一步开始，默认 `1`。 +- `PIPELINE_FORCE_RERUN_STEPS`：强制重跑的步骤，例如 `1,2,3`。 +- `PIPELINE_RUN_INTERVAL_SECONDS`：成功完成后再次运行的间隔，默认 `3600`。 +- `PIPELINE_RETRY_INTERVAL_SECONDS`：发生未处理异常后的重试间隔，默认 `60`。 + +本地运行： + +```powershell +$env:STRATEGY_ID="your-strategy-id" +$env:MODELHUB_USER_ACCOUNT="your-account" +$env:MODELHUB_USER_PASSWORD="your-password" +python app.py +``` + +健康检查： + +```text +GET http://127.0.0.1:8080/health +``` + +只要策略服务进程存活，该端点返回 HTTP 200。响应同时包含后台流水线的 +运行状态，但不会返回账号、密码或令牌。 + +## Docker + +```bash +docker build -t new-pipeline . +docker run --rm -p 8080:8080 \ + -e STRATEGY_ID=your-strategy-id \ + -e MODELHUB_USER_ACCOUNT=your-account \ + -e MODELHUB_USER_PASSWORD=your-password \ + new-pipeline +``` + +镜像监听 `8080` 端口，并处理 `SIGTERM`。收到停机信号后，后台等待和轮询 +会立即停止；进行中的 HTTP 请求最长等待 20 秒，随后关闭会话并退出。 + +运行时产生的断点文件位于 `/app/pipeline_outputs`。如需跨 Pod 保留断点， +应为该目录挂载持久卷。 + +平台的 CPU/内存 requests 与 limits 由部署配置设置，不在 Dockerfile 中声明： + +- requests：`100m CPU / 256 Mi` +- limits：`1 CPU / 512 Mi` diff --git a/app.py b/app.py new file mode 100644 index 0000000..3c10c37 --- /dev/null +++ b/app.py @@ -0,0 +1,156 @@ +import json +import os +import signal +import threading +import traceback +from datetime import datetime, timezone +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, Dict + +import modelhub_pipeline_resumable as pipeline + + +HOST = "0.0.0.0" +PORT = int(os.getenv("PORT", "8080")) +RUN_INTERVAL_SECONDS = int(os.getenv("PIPELINE_RUN_INTERVAL_SECONDS", "3600")) +RETRY_INTERVAL_SECONDS = int(os.getenv("PIPELINE_RETRY_INTERVAL_SECONDS", "60")) +SHUTDOWN_TIMEOUT_SECONDS = 25 +STRATEGY_ID = os.environ["STRATEGY_ID"] + +stop_event = threading.Event() + + +def utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +class RuntimeState: + def __init__(self) -> None: + self._lock = threading.Lock() + self._data: Dict[str, Any] = { + "pipeline_status": "starting", + "run_count": 0, + "last_started_at": None, + "last_finished_at": None, + "last_error": None, + } + + def update(self, **values: Any) -> None: + with self._lock: + self._data.update(values) + + def snapshot(self) -> Dict[str, Any]: + with self._lock: + return dict(self._data) + + +runtime_state = RuntimeState() + + +class HealthHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + if self.path.split("?", 1)[0] != "/health": + self._write_json(404, {"status": "not_found"}) + return + + state = runtime_state.snapshot() + self._write_json( + 200, + { + "status": "ok", + "strategy_id": STRATEGY_ID, + "pipeline_status": state["pipeline_status"], + "run_count": state["run_count"], + "last_started_at": state["last_started_at"], + "last_finished_at": state["last_finished_at"], + }, + ) + + def _write_json(self, status_code: int, payload: Dict[str, Any]) -> None: + body = json.dumps(payload, ensure_ascii=False).encode("utf-8") + self.send_response(status_code) + self.send_header("Content-Type", "application/json; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, message_format: str, *args: Any) -> None: + print(f"[health] {self.address_string()} - {message_format % args}") + + +def pipeline_worker() -> None: + pipeline.reset_shutdown() + + while not stop_event.is_set(): + previous = runtime_state.snapshot() + runtime_state.update( + pipeline_status="running", + run_count=previous["run_count"] + 1, + last_started_at=utc_now(), + last_error=None, + ) + + try: + pipeline.main() + except pipeline.ShutdownRequested: + runtime_state.update(pipeline_status="stopping") + break + except Exception as exc: + runtime_state.update( + pipeline_status="error", + last_finished_at=utc_now(), + last_error=f"{type(exc).__name__}: {exc}", + ) + traceback.print_exc() + wait_seconds = RETRY_INTERVAL_SECONDS + else: + runtime_state.update( + pipeline_status="idle", + last_finished_at=utc_now(), + ) + wait_seconds = RUN_INTERVAL_SECONDS + + if stop_event.wait(max(1, wait_seconds)): + break + + runtime_state.update(pipeline_status="stopped") + + +def handle_signal(signum: int, _frame: Any) -> None: + print(f"收到信号 {signum}，开始优雅停机") + runtime_state.update(pipeline_status="stopping") + stop_event.set() + pipeline.request_shutdown() + + +def main() -> None: + signal.signal(signal.SIGTERM, handle_signal) + signal.signal(signal.SIGINT, handle_signal) + + worker = threading.Thread( + target=pipeline_worker, + name="modelhub-pipeline", + daemon=True, + ) + worker.start() + + server = ThreadingHTTPServer((HOST, PORT), HealthHandler) + server.timeout = 0.5 + print(f"健康检查服务已启动: http://{HOST}:{PORT}/health") + + try: + while not stop_event.is_set(): + server.handle_request() + finally: + stop_event.set() + pipeline.request_shutdown() + server.server_close() + worker.join(timeout=SHUTDOWN_TIMEOUT_SECONDS) + if worker.is_alive(): + print("后台任务未在停机窗口内结束，主进程将退出") + else: + print("策略已完成资源清理并停止") + + +if __name__ == "__main__": + main() diff --git a/modelhub_pipeline_resumable.py b/modelhub_pipeline_resumable.py new file mode 100644 index 0000000..bff447c --- /dev/null +++ b/modelhub_pipeline_resumable.py @@ -0,0 +1,982 @@ +import os +import argparse +import re +import json +import random +import threading +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" + +# 账号密码只从环境变量读取，避免凭据进入镜像或 Git 仓库。 +# Linux/macOS: +# export MODELHUB_USER_ACCOUNT="你的账号" +# export MODELHUB_USER_PASSWORD="你的密码" +# Windows PowerShell: +# $env:MODELHUB_USER_ACCOUNT="你的账号" +# $env:MODELHUB_USER_PASSWORD="你的密码" +USER_ACCOUNT = os.getenv("MODELHUB_USER_ACCOUNT", "") +USER_PASSWORD = os.getenv("MODELHUB_USER_PASSWORD", "") + +# 不使用手动 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 = 100 +START_PAGE = 6 # 起始页码 +END_PAGE = 6 # 结束页码 +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 表示一直等到有空位 + +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] = [] + + for page in range(START_PAGE, END_PAGE + 1): + 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)} 个模型") + + except ShutdownRequested: + raise + except Exception as e: + print(f"获取第 {page} 页失败: {e}") + continue + + # 去重但保持顺序 + 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( + "没有设置账号密码。请设置环境变量 MODELHUB_USER_ACCOUNT / " + "MODELHUB_USER_PASSWORD，或在脚本里填写 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 + + +# ============================================================ +# 10. 主流程：筛选 -> 查重 -> 提交 +# ============================================================ + +def main() -> None: + 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() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d51edac --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +requests>=2.31.0,<3.0.0