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