fix: 支持 pod 内多种鉴权方式（Xc-Token/Bearer/免鉴权）自动探测
This commit is contained in:
user_4BTLNC
752
main.py
752
main.py
@@ -1,364 +1,388 @@
|
||||
"""信创模盒（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()
|
||||
"""信创模盒（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 ""
|
||||
# 平台在 pod 内可能用不同方式鉴权：依次尝试 Xc-Token → Authorization: Bearer → 无鉴权
|
||||
AUTH_MODES = ("xctoken", "bearer", "none")
|
||||
AUTH_MODE = None # 运行期探测到的可用方式
|
||||
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 _auth_header(mode: str) -> dict:
|
||||
if not XC_TOKEN or mode == "none":
|
||||
return {}
|
||||
if mode == "bearer":
|
||||
return {"Authorization": "Bearer " + XC_TOKEN}
|
||||
return {"Xc-Token": XC_TOKEN}
|
||||
|
||||
|
||||
def _request(method: str, path: str, body: dict | None = None, timeout: int = 30,
|
||||
auth_mode: str | None = None) -> 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"
|
||||
headers.update(_auth_header(auth_mode or AUTH_MODE or "xctoken"))
|
||||
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, auth_mode: str | None = None) -> dict:
|
||||
status, text = _request("GET", path, timeout=timeout, auth_mode=auth_mode)
|
||||
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, auth_mode: str | None = None) -> dict:
|
||||
status, text = _request("POST", path, body=body, timeout=timeout, auth_mode=auth_mode)
|
||||
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, auth_mode: str | None = None) -> 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, auth_mode=auth_mode)
|
||||
|
||||
|
||||
def worker() -> None:
|
||||
_load_state()
|
||||
if not STRATEGY_ID:
|
||||
STATE["error"] = "STRATEGY_ID 未注入，智能体无法提交（平台运行时自动注入）"
|
||||
log(STATE["error"])
|
||||
return
|
||||
global AUTH_MODE
|
||||
|
||||
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")
|
||||
# 401 时换一种鉴权方式重试（pod 内可能用 Bearer 或免鉴权）
|
||||
if code == 401:
|
||||
for mode in AUTH_MODES:
|
||||
if mode == (AUTH_MODE or "xctoken"):
|
||||
continue
|
||||
res = submit(cand["id"], card, auth_mode=mode)
|
||||
code = res.get("code")
|
||||
if code != 401:
|
||||
AUTH_MODE = mode
|
||||
log(f"auth mode ok: {mode}")
|
||||
break
|
||||
if code == 401:
|
||||
STATE["error"] = "401: 三种鉴权方式均被拒绝"
|
||||
log(STATE["error"])
|
||||
_save_state()
|
||||
STOP.wait(300)
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user