1 Commits

Author SHA1 Message Date
user_4BTLNC
9acd91b239 fix: 支持 pod 内多种鉴权方式(Xc-Token/Bearer/免鉴权)自动探测 2026-10-25 02:49:52 +08:00

50
main.py
View File

@@ -35,6 +35,9 @@ 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 ""
# 平台在 pod 内可能用不同方式鉴权:依次尝试 Xc-Token → Authorization: Bearer → 无鉴权
AUTH_MODES = ("xctoken", "bearer", "none")
AUTH_MODE = None # 运行期探测到的可用方式
TARGET_TASKS = int(os.getenv("TARGET_TASKS", "250")) TARGET_TASKS = int(os.getenv("TARGET_TASKS", "250"))
QUOTA_PER_MODEL = int(os.getenv("QUOTA_PER_MODEL", "2")) QUOTA_PER_MODEL = int(os.getenv("QUOTA_PER_MODEL", "2"))
INTERVAL_SEC = float(os.getenv("INTERVAL_SEC", "1.5")) INTERVAL_SEC = float(os.getenv("INTERVAL_SEC", "1.5"))
@@ -100,15 +103,23 @@ def _load_state() -> None:
# ---------------------------------------------------------------- HTTP # ---------------------------------------------------------------- HTTP
def _request(method: str, path: str, body: dict | None = None, timeout: int = 30) -> tuple[int, str]: 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 url = API_BASE + path
data = None data = None
headers = {"Accept": "application/json", "User-Agent": "modelhub-xc-agent/1.0"} headers = {"Accept": "application/json", "User-Agent": "modelhub-xc-agent/1.0"}
if body is not None: if body is not None:
data = json.dumps(body, ensure_ascii=False).encode("utf-8") data = json.dumps(body, ensure_ascii=False).encode("utf-8")
headers["Content-Type"] = "application/json" headers["Content-Type"] = "application/json"
if XC_TOKEN: headers.update(_auth_header(auth_mode or AUTH_MODE or "xctoken"))
headers["Xc-Token"] = XC_TOKEN
req = urllib.request.Request(url, data=data, headers=headers, method=method) req = urllib.request.Request(url, data=data, headers=headers, method=method)
try: try:
with urllib.request.urlopen(req, timeout=timeout) as resp: with urllib.request.urlopen(req, timeout=timeout) as resp:
@@ -119,16 +130,16 @@ def _request(method: str, path: str, body: dict | None = None, timeout: int = 30
return 0, str(exc) return 0, str(exc)
def _get_json(path: str, timeout: int = 30) -> dict: def _get_json(path: str, timeout: int = 30, auth_mode: str | None = None) -> dict:
status, text = _request("GET", path, timeout=timeout) status, text = _request("GET", path, timeout=timeout, auth_mode=auth_mode)
try: try:
return json.loads(text) return json.loads(text)
except Exception: except Exception:
return {"code": -1, "message": f"HTTP {status}: {text[:200]}"} return {"code": -1, "message": f"HTTP {status}: {text[:200]}"}
def _post_json(path: str, body: dict, timeout: int = 60) -> dict: 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) status, text = _request("POST", path, body=body, timeout=timeout, auth_mode=auth_mode)
try: try:
return json.loads(text) return json.loads(text)
except Exception: except Exception:
@@ -189,7 +200,7 @@ def load_pool() -> list[dict]:
# ---------------------------------------------------------------- 提交 # ---------------------------------------------------------------- 提交
def submit(model: str, card: dict) -> dict: def submit(model: str, card: dict, auth_mode: str | None = None) -> dict:
body = { body = {
"modelAddress": f"https://huggingface.co/{model}", "modelAddress": f"https://huggingface.co/{model}",
"taskType": "text-generation", "taskType": "text-generation",
@@ -200,7 +211,7 @@ def submit(model: str, card: dict) -> dict:
} }
if DRY_RUN: if DRY_RUN:
return {"code": 0, "data": {"id": "dry-run"}, "dryRun": True} return {"code": 0, "data": {"id": "dry-run"}, "dryRun": True}
return _post_json("/api/adapt/task/add", body) return _post_json("/api/adapt/task/add", body, auth_mode=auth_mode)
def worker() -> None: def worker() -> None:
@@ -209,10 +220,7 @@ def worker() -> None:
STATE["error"] = "STRATEGY_ID 未注入,智能体无法提交(平台运行时自动注入)" STATE["error"] = "STRATEGY_ID 未注入,智能体无法提交(平台运行时自动注入)"
log(STATE["error"]) log(STATE["error"])
return return
if not XC_TOKEN and not DRY_RUN: global AUTH_MODE
STATE["error"] = "XC_TOKEN/EXTERNAL_SERVICE_TOKEN 未注入,请在个人主页生成 xcToken 并交给平台"
log(STATE["error"])
return
pool = load_pool() pool = load_pool()
submitted_pairs = {(s["model"], s["card"]) for s in STATE["submitted"]} submitted_pairs = {(s["model"], s["card"]) for s in STATE["submitted"]}
@@ -241,6 +249,22 @@ def worker() -> None:
res = submit(cand["id"], card) res = submit(cand["id"], card)
code = res.get("code") 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["attempts"] += 1
STATE["codes"][str(code)] = STATE["codes"].get(str(code), 0) + 1 STATE["codes"][str(code)] = STATE["codes"].get(str(code), 0) + 1
STATE["last_event"] = { STATE["last_event"] = {