1121 lines
36 KiB
Python
1121 lines
36 KiB
Python
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()
|