zhaohaoyu/new-pipeline
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

Package pipeline as ModelHub strategy

Browse Source
This commit is contained in:
Codex
2026-07-27 17:29:26 +08:00
commit dd21ca1be2
8 changed files with 1246 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
.dockerignore Normal file
View File

@@ -0,0 +1,8 @@
.git
.gitignore
.idea
__pycache__
*.py[cod]
*.log
.env
pipeline_outputs

9
.env.example Normal file
View File

@@ -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

6
.gitignore vendored Normal file
View File

@@ -0,0 +1,6 @@
.idea/
__pycache__/
*.py[cod]
*.log
.env
pipeline_outputs/

23
Dockerfile Normal file
View File

@@ -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"]

61
README.md Normal file
View File

@@ -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`

156
app.py Normal file
View File

@@ -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()

982
modelhub_pipeline_resumable.py Normal file
View File

@@ -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()

1
requirements.txt Normal file
View File

@@ -0,0 +1 @@
requests>=2.31.0,<3.0.0