diff --git a/main.py b/main.py index aeb0e60..40ee2fd 100644 --- a/main.py +++ b/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()