365 lines
14 KiB
Python
365 lines
14 KiB
Python
"""信创模盒（ModelHub XC）适配智能体。
|
||
|
||
平台契约（见 https://modelhub.org.cn/#/docs/6）：
|
||
- 仓库根目录必须有 Dockerfile，且 EXPOSE 一个端口用于健康检查
|
||
- 必须提供 GET /health，返回 200 表示存活
|
||
- 通过环境变量 STRATEGY_ID 获取自身策略 ID
|
||
- 必须正确处理 SIGTERM，30s 内完成退出
|
||
- 资源：requests 100m/256Mi，limits 1C/512Mi
|
||
|
||
业务逻辑：
|
||
1. 拉取社区模型库，筛出「文本生成 + 仓库完整 + 未在该卡适配过」的模型；
|
||
2. 按卡依次调用 /api/adapt/task/add 提交验证任务；
|
||
3. 命中 60027（该卡已验证成功）/ 60028（该卡验证中）就换下一个组合；
|
||
4. 直到提交满 TARGET_TASKS 条或模型池耗尽。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import signal
|
||
import sys
|
||
import threading
|
||
import time
|
||
import urllib.error
|
||
import urllib.parse
|
||
import urllib.request
|
||
from datetime import datetime
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
||
from cards import CARDS
|
||
|
||
# ---------------------------------------------------------------- 配置
|
||
|
||
API_BASE = os.getenv("MODELHUB_API_BASE", "https://www.modelhub.org.cn").rstrip("/")
|
||
STRATEGY_ID = os.getenv("STRATEGY_ID", "")
|
||
XC_TOKEN = os.getenv("XC_TOKEN") or os.getenv("EXTERNAL_SERVICE_TOKEN") or ""
|
||
TARGET_TASKS = int(os.getenv("TARGET_TASKS", "250"))
|
||
QUOTA_PER_MODEL = int(os.getenv("QUOTA_PER_MODEL", "2"))
|
||
INTERVAL_SEC = float(os.getenv("INTERVAL_SEC", "1.5"))
|
||
MAX_ATTEMPTS = int(os.getenv("MAX_ATTEMPTS", "4000"))
|
||
DRY_RUN = os.getenv("DRY_RUN", "").strip().lower() in ("1", "true", "yes", "on")
|
||
STATE_FILE = os.getenv("STATE_FILE", "/tmp/modelhub_agent_state.json")
|
||
POOL_FILE = os.getenv("POOL_FILE", "/tmp/modelhub_agent_pool.json")
|
||
PAGE_SIZE = int(os.getenv("PAGE_SIZE", "50"))
|
||
|
||
SEED_QUERIES = [q for q in os.getenv(
|
||
"SEED_QUERIES",
|
||
"Qwen2.5-0.5B,Qwen2.5-1.5B,Qwen2.5-3B,Qwen2.5-7B,Qwen3-0.6B,Qwen3-1.7B,Qwen3-4B,Qwen3-8B,"
|
||
"Llama-3.2,Llama-3.1,InternLM2.5,MiniCPM,Phi-3,Phi-4,gemma-2,Mistral-7B,Yi-1.5,GLM-4,"
|
||
"ChatGLM3,DeepSeek-R1-Distill,Qwen2-1.5B,Qwen2-7B",
|
||
).split(",") if q.strip()]
|
||
|
||
# 明显不能直接跑文本生成验证的仓库（统一转小写后匹配）
|
||
BAD_TOKENS = ("gensyn", "gguf", "awq", "gptq", "fp8", "fp4", "nvfp4", "adapter", "lora",
|
||
"int4", "int8", "w4a16", "w8a8", "w8a16", "quantized", "bnb", "exl2", "hqq",
|
||
"mlx", "onnx", "openvino", "tflite")
|
||
MIN_MACHINES = int(os.getenv("MIN_MACHINES", "2")) # 至少在一张卡上适配成功过，说明仓库完整
|
||
|
||
# ---------------------------------------------------------------- 全局状态
|
||
|
||
STOP = threading.Event()
|
||
LOCK = threading.Lock()
|
||
STATE = {
|
||
"started_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
"submitted": [], # [{model, card, taskId, at}]
|
||
"attempts": 0,
|
||
"codes": {}, # code -> count
|
||
"last_event": None,
|
||
"finished": False,
|
||
"error": None,
|
||
}
|
||
|
||
|
||
def log(msg: str) -> None:
|
||
print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}", flush=True)
|
||
|
||
|
||
def _save_state() -> None:
|
||
try:
|
||
tmp = STATE_FILE + ".tmp"
|
||
with open(tmp, "w", encoding="utf-8") as f:
|
||
json.dump(STATE, f, ensure_ascii=False, indent=1)
|
||
os.replace(tmp, STATE_FILE)
|
||
except Exception as exc: # 状态落盘失败不影响主流程
|
||
log(f"save state failed: {exc}")
|
||
|
||
|
||
def _load_state() -> None:
|
||
try:
|
||
with open(STATE_FILE, "r", encoding="utf-8") as f:
|
||
saved = json.load(f)
|
||
STATE["submitted"] = saved.get("submitted", [])
|
||
log(f"resume: {len(STATE['submitted'])} submitted")
|
||
except FileNotFoundError:
|
||
pass
|
||
except Exception as exc:
|
||
log(f"load state failed: {exc}")
|
||
|
||
|
||
# ---------------------------------------------------------------- HTTP
|
||
|
||
def _request(method: str, path: str, body: dict | None = None, timeout: int = 30) -> tuple[int, str]:
|
||
url = API_BASE + path
|
||
data = None
|
||
headers = {"Accept": "application/json", "User-Agent": "modelhub-xc-agent/1.0"}
|
||
if body is not None:
|
||
data = json.dumps(body, ensure_ascii=False).encode("utf-8")
|
||
headers["Content-Type"] = "application/json"
|
||
if XC_TOKEN:
|
||
headers["Xc-Token"] = XC_TOKEN
|
||
req = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
||
return resp.status, resp.read().decode("utf-8", "replace")
|
||
except urllib.error.HTTPError as exc:
|
||
return exc.code, exc.read().decode("utf-8", "replace")
|
||
except Exception as exc:
|
||
return 0, str(exc)
|
||
|
||
|
||
def _get_json(path: str, timeout: int = 30) -> dict:
|
||
status, text = _request("GET", path, timeout=timeout)
|
||
try:
|
||
return json.loads(text)
|
||
except Exception:
|
||
return {"code": -1, "message": f"HTTP {status}: {text[:200]}"}
|
||
|
||
|
||
def _post_json(path: str, body: dict, timeout: int = 60) -> dict:
|
||
status, text = _request("POST", path, body=body, timeout=timeout)
|
||
try:
|
||
return json.loads(text)
|
||
except Exception:
|
||
return {"code": -1, "message": f"HTTP {status}: {text[:200]}"}
|
||
|
||
|
||
# ---------------------------------------------------------------- 模型池
|
||
|
||
def build_pool() -> list[dict]:
|
||
"""拉模型库，返回 [{id, machines}]，按「已适配卡数」升序（可选卡更多）。"""
|
||
seen: dict[str, dict] = {}
|
||
for query in SEED_QUERIES:
|
||
for page in range(1, 6):
|
||
payload = {"pageNum": page, "pageSize": PAGE_SIZE, "searchText": query}
|
||
res = _post_json("/api/computility/models/list/page/vo", payload, timeout=40)
|
||
records = (res.get("data") or {}).get("records") or []
|
||
if not records:
|
||
break
|
||
for item in records:
|
||
model_id = item.get("modelId") or ""
|
||
machines = [m.get("machineName") for m in (item.get("machines") or [])]
|
||
level = (item.get("taskLevelsInfo") or {}).get("taskLevelChineseName")
|
||
if not model_id or not machines or level != "文本生成":
|
||
continue
|
||
lower = model_id.lower()
|
||
if any(tok in lower for tok in BAD_TOKENS):
|
||
continue
|
||
if len(machines) < MIN_MACHINES:
|
||
continue
|
||
if model_id in seen:
|
||
continue
|
||
seen[model_id] = {"id": model_id, "machines": machines}
|
||
if len(records) < PAGE_SIZE:
|
||
break
|
||
if STOP.is_set():
|
||
break
|
||
pool = sorted(seen.values(), key=lambda x: len(x["machines"]))
|
||
log(f"pool built: {len(pool)} models")
|
||
try:
|
||
with open(POOL_FILE, "w", encoding="utf-8") as f:
|
||
json.dump(pool, f)
|
||
except Exception:
|
||
pass
|
||
return pool
|
||
|
||
|
||
def load_pool() -> list[dict]:
|
||
try:
|
||
with open(POOL_FILE, "r", encoding="utf-8") as f:
|
||
pool = json.load(f)
|
||
if pool:
|
||
log(f"pool loaded from cache: {len(pool)}")
|
||
return pool
|
||
except Exception:
|
||
pass
|
||
return build_pool()
|
||
|
||
|
||
# ---------------------------------------------------------------- 提交
|
||
|
||
def submit(model: str, card: dict) -> dict:
|
||
body = {
|
||
"modelAddress": f"https://huggingface.co/{model}",
|
||
"taskType": "text-generation",
|
||
"targetGpu": card["gpu"],
|
||
"framework": card["framework"],
|
||
"strategyId": STRATEGY_ID,
|
||
"configParams": card["configParams"],
|
||
}
|
||
if DRY_RUN:
|
||
return {"code": 0, "data": {"id": "dry-run"}, "dryRun": True}
|
||
return _post_json("/api/adapt/task/add", body)
|
||
|
||
|
||
def worker() -> None:
|
||
_load_state()
|
||
if not STRATEGY_ID:
|
||
STATE["error"] = "STRATEGY_ID 未注入，智能体无法提交（平台运行时自动注入）"
|
||
log(STATE["error"])
|
||
return
|
||
if not XC_TOKEN and not DRY_RUN:
|
||
STATE["error"] = "XC_TOKEN/EXTERNAL_SERVICE_TOKEN 未注入，请在个人主页生成 xcToken 并交给平台"
|
||
log(STATE["error"])
|
||
return
|
||
|
||
pool = load_pool()
|
||
submitted_pairs = {(s["model"], s["card"]) for s in STATE["submitted"]}
|
||
per_model: dict[str, int] = {}
|
||
for s in STATE["submitted"]:
|
||
per_model[s["model"]] = per_model.get(s["model"], 0) + 1
|
||
|
||
while not STOP.is_set() and len(STATE["submitted"]) < TARGET_TASKS:
|
||
progress = False
|
||
for card in CARDS:
|
||
if STOP.is_set() or len(STATE["submitted"]) >= TARGET_TASKS:
|
||
break
|
||
for cand in pool:
|
||
if STOP.is_set() or len(STATE["submitted"]) >= TARGET_TASKS:
|
||
break
|
||
if STATE["attempts"] >= MAX_ATTEMPTS:
|
||
log("reach MAX_ATTEMPTS, stop")
|
||
return
|
||
if card["gpu"] in cand["machines"]:
|
||
continue
|
||
if per_model.get(cand["id"], 0) >= QUOTA_PER_MODEL:
|
||
continue
|
||
pair = (cand["id"], card["gpu"])
|
||
if pair in submitted_pairs:
|
||
continue
|
||
|
||
res = submit(cand["id"], card)
|
||
code = res.get("code")
|
||
STATE["attempts"] += 1
|
||
STATE["codes"][str(code)] = STATE["codes"].get(str(code), 0) + 1
|
||
STATE["last_event"] = {
|
||
"at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
"model": cand["id"], "gpu": card["gpu"], "code": code,
|
||
"message": (res.get("message") or "")[:160],
|
||
}
|
||
progress = True
|
||
|
||
if code == 0:
|
||
task_id = (res.get("data") or {}).get("id")
|
||
with LOCK:
|
||
STATE["submitted"].append({
|
||
"model": cand["id"], "card": card["gpu"],
|
||
"taskId": task_id,
|
||
"at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
})
|
||
per_model[cand["id"]] = per_model.get(cand["id"], 0) + 1
|
||
submitted_pairs.add(pair)
|
||
log(f"OK {len(STATE['submitted'])}/{TARGET_TASKS} | {card['label']} | {cand['id']} | {task_id}")
|
||
_save_state()
|
||
elif code == 60027: # 该卡已经验证成功过
|
||
submitted_pairs.add(pair)
|
||
elif code == 60028: # 该组合正在验证中
|
||
submitted_pairs.add(pair)
|
||
elif code == 50000: # 拉不到上游模型信息
|
||
submitted_pairs.add(pair)
|
||
elif code in (401, 403, 60014):
|
||
log(f"fatal {code}: {res.get('message')}")
|
||
STATE["error"] = f"{code}: {res.get('message')}"
|
||
_save_state()
|
||
STOP.wait(120)
|
||
elif code in (-1, 0) and res.get("message", "").startswith("HTTP"):
|
||
log(f"network issue: {res.get('message')}")
|
||
STOP.wait(10)
|
||
time.sleep(INTERVAL_SEC)
|
||
if not progress:
|
||
log("no progress this round; rebuild pool")
|
||
pool = build_pool()
|
||
if not pool:
|
||
log("empty pool, exit")
|
||
break
|
||
time.sleep(30)
|
||
|
||
STATE["finished"] = True
|
||
_save_state()
|
||
log(f"worker done: submitted={len(STATE['submitted'])} attempts={STATE['attempts']}")
|
||
|
||
|
||
# ---------------------------------------------------------------- HTTP 服务
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
server_version = "modelhub-xc-agent/1.0"
|
||
|
||
def do_GET(self) -> None: # noqa: N802
|
||
path = urllib.parse.urlparse(self.path).path
|
||
if path == "/health":
|
||
self._send({"status": "ok"})
|
||
return
|
||
if path == "/stats":
|
||
with LOCK:
|
||
payload = json.loads(json.dumps(STATE, ensure_ascii=False))
|
||
payload["target"] = TARGET_TASKS
|
||
payload["progress"] = f"{len(payload['submitted'])}/{TARGET_TASKS}"
|
||
payload["strategyId"] = STRATEGY_ID or None
|
||
payload["dryRun"] = DRY_RUN
|
||
self._send(payload)
|
||
return
|
||
if path == "/":
|
||
self._send({
|
||
"name": "modelhub-xc-adapt-agent",
|
||
"status": "stopped" if STOP.is_set() else "running",
|
||
"strategyId": STRATEGY_ID or None,
|
||
"target": TARGET_TASKS,
|
||
"submitted": len(STATE["submitted"]),
|
||
"attempts": STATE["attempts"],
|
||
"dryRun": DRY_RUN,
|
||
})
|
||
return
|
||
self._send({"error": "not found"}, status=404)
|
||
|
||
def log_message(self, fmt: str, *args) -> None:
|
||
print(f"{self.address_string()} - {fmt % args}", flush=True)
|
||
|
||
def _send(self, body: dict, status: int = 200) -> None:
|
||
payload = json.dumps(body, ensure_ascii=False).encode("utf-8")
|
||
self.send_response(status)
|
||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||
self.send_header("Content-Length", str(len(payload)))
|
||
self.end_headers()
|
||
self.wfile.write(payload)
|
||
|
||
|
||
def _handle_signal(signum: int, _frame) -> None:
|
||
log(f"received signal {signum}, shutting down")
|
||
STOP.set()
|
||
|
||
|
||
def main() -> None:
|
||
signal.signal(signal.SIGTERM, _handle_signal)
|
||
signal.signal(signal.SIGINT, _handle_signal)
|
||
|
||
thread = threading.Thread(target=worker, name="submit-worker", daemon=True)
|
||
thread.start()
|
||
|
||
server = ThreadingHTTPServer(("0.0.0.0", int(os.getenv("PORT", "8080"))), Handler)
|
||
server.timeout = 1
|
||
log(f"agent listening on 0.0.0.0:{os.getenv('PORT', '8080')} (dryRun={DRY_RUN}, target={TARGET_TASKS})")
|
||
|
||
while not STOP.is_set():
|
||
server.handle_request()
|
||
|
||
thread.join(timeout=25)
|
||
server.server_close()
|
||
_save_state()
|
||
log("bye")
|
||
sys.exit(0)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|