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

Add continuous model discovery and queued submission

Browse Source
This commit is contained in:
Codex
2026-08-19 12:36:14 +08:00
parent 71ea49b4c9
commit 990be72c54
4 changed files with 170 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

157
modelhub_pipeline_resumable.py
View File

@@ -1,9 +1,10 @@
import os
import argparse
import os
import re
import json
import random
import threading
import queue
from typing import Dict, Tuple, List, Optional, Set
import requests
@@ -42,9 +43,9 @@ STRATEGY_ID = os.environ["STRATEGY_ID"]
ORG_NAME = "mradermacher"
HF_API_URL = "https://hf-mirror.com/api/models"
PER_PAGE = 100
START_PAGE = 6 # 起始页码
END_PAGE = 6 # 结束页码
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错误重试次数
@@ -66,6 +67,14 @@ 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
# 是否开启断点续跑：已经成功提交过的模型，下次运行自动跳过
@@ -349,7 +358,11 @@ def get_model_filename(model_id: str) -> str:
def get_org_models(session: requests.Session, org_name: str) -> List[str]:
models: List[str] = []
for page in range(START_PAGE, END_PAGE + 1):
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))
@@ -377,12 +390,15 @@ def get_org_models(session: requests.Session, org_name: str) -> List[str]:
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}")
continue
pages_seen += 1
page += 1
# 去重但保持顺序
seen = set()
@@ -810,11 +826,140 @@ def submit_adapt_task_with_polling(token: str, model_id: str) -> Tuple[bool, str
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()