fix: 支持 pod 内多种鉴权方式(Xc-Token/Bearer/免鉴权)自动探测
This commit is contained in:
50
main.py
50
main.py
@@ -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"] = {
|
||||||
|
|||||||
Reference in New Issue
Block a user