561 lines
21 KiB
Python
561 lines
21 KiB
Python
"""Card-specific ModelHub XC deployment preflight agent.
|
||
|
||
It scans public candidates and only writes a deduplicated task when private runtime
|
||
credentials and verified platform configuration are present.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import signal
|
||
import threading
|
||
from datetime import datetime, timezone
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
from typing import Any
|
||
from urllib.parse import quote, urlencode, urlsplit
|
||
from urllib.request import ProxyHandler, Request, build_opener
|
||
|
||
|
||
AGENT_NAME = "xc-tianga100-advisor-agent"
|
||
AGENT_VERSION = "1.2.0"
|
||
CARD_VENDOR = "天数智芯"
|
||
CARD_MODEL = "天垓100"
|
||
DEFAULT_TARGET_GPU = "Iluvatar_bi-100"
|
||
PORT = int(os.getenv("PORT", "8080"))
|
||
STRATEGY_ID = os.getenv("STRATEGY_ID", "")
|
||
XC_TOKEN = os.getenv("MODELHUB_XC_TOKEN", "")
|
||
TARGET_GPU = os.getenv("MODELHUB_TARGET_GPU", DEFAULT_TARGET_GPU)
|
||
AUTO_SCAN_ENABLED = os.getenv("AUTO_SCAN_ENABLED", "true").lower() not in {
|
||
"0", "false", "no", "off"
|
||
}
|
||
AUTO_SUBMIT_ENABLED = os.getenv("AUTO_SUBMIT_ENABLED", "true").lower() not in {
|
||
"0", "false", "no", "off"
|
||
}
|
||
SCAN_INTERVAL_SECONDS = max(int(os.getenv("SCAN_INTERVAL_SECONDS", "3600")), 60)
|
||
CANDIDATE_ENDPOINT = "https://modelhub.org.cn/api/computility/models/top/models"
|
||
TASK_PAGE_ENDPOINT = "https://modelhub.org.cn/api/adapt/task/page"
|
||
TASK_ADD_ENDPOINT = "https://modelhub.org.cn/api/adapt/task/add"
|
||
DEFAULT_CONFIG_PARAMS = "framework: vllm\napi: completion\nlang: zh\nmax_model_len: 2048\nmax_tokens: 256\ntemperature: 0.1\nrepetition_penalty: 1.0\ntop_p: 0.9\nsut_config:\n gpu_num: 1\n values:\n command:\n - python3\n - -m\n - vllm.entrypoints.openai.api_server\n - --host\n - 0.0.0.0\n - --port\n - '20644'\n - --served-model-name\n - llm\n - --model\n - /model\n - --max-model-len\n - '2048'\n - --tensor-parallel-size\n - '1'\n - --max-num-seqs\n - '8'\n - --enforce-eager\n - --disable-log-requests\n - --enable-prefix-caching\n - --trust-remote-code\nref_config:\n gpu_num: 1\n values:\n command:\n - vllm\n - serve\n - /model\n - --port\n - '80'\n - --served-model-name\n - llm\n - --max-model-len\n - '2048'\n - -tp\n - '1'\n - --enforce-eager\n - --trust-remote-code\n"
|
||
CONFIG_PARAMS = os.getenv("MODELHUB_CONFIG_PARAMS", DEFAULT_CONFIG_PARAMS)
|
||
MAX_BODY_BYTES = 1_000_000
|
||
SECRET_ASSIGNMENT_RE = re.compile(
|
||
r"(?i)(token|secret|password|api[_-]?key)=([^&\s]+)"
|
||
)
|
||
HTTP_OPENER = build_opener(ProxyHandler({}))
|
||
SCANNER_LOCK = threading.Lock()
|
||
SCANNER_STATE: dict[str, Any] = {
|
||
"enabled": AUTO_SCAN_ENABLED,
|
||
"running": False,
|
||
"last_scan_at": None,
|
||
"last_candidate": None,
|
||
"last_action": "not_started",
|
||
"last_error_type": None,
|
||
}
|
||
|
||
|
||
class ValidationError(ValueError):
|
||
"""Raised when a request field is invalid."""
|
||
|
||
|
||
def _safe_model_name(value: Any) -> str:
|
||
"""Return a bounded model identifier without URL credentials or query data."""
|
||
|
||
text = str(value or "未指定模型").strip()
|
||
if "://" in text:
|
||
try:
|
||
parsed = urlsplit(text)
|
||
text = f"{parsed.scheme}://{parsed.hostname or ''}{parsed.path}"
|
||
except ValueError:
|
||
text = "无效模型地址"
|
||
text = SECRET_ASSIGNMENT_RE.sub(r"\1=[REDACTED]", text)
|
||
return text[:200] or "未指定模型"
|
||
|
||
|
||
def _log_event(event: str, **fields: Any) -> None:
|
||
"""Write one structured, unbuffered log record to stdout."""
|
||
|
||
record = {
|
||
"timestamp": datetime.now(timezone.utc).isoformat(),
|
||
"event": event,
|
||
"agent": AGENT_NAME,
|
||
"version": AGENT_VERSION,
|
||
**fields,
|
||
}
|
||
print(json.dumps(record, ensure_ascii=False, separators=(",", ":")), flush=True)
|
||
|
||
|
||
def _scanner_state_update(**fields: Any) -> None:
|
||
with SCANNER_LOCK:
|
||
SCANNER_STATE.update(fields)
|
||
|
||
|
||
def scanner_snapshot() -> dict[str, Any]:
|
||
with SCANNER_LOCK:
|
||
return dict(SCANNER_STATE)
|
||
|
||
|
||
def _http_json(
|
||
method: str,
|
||
url: str,
|
||
*,
|
||
payload: dict[str, Any] | None = None,
|
||
token: str = "",
|
||
) -> dict[str, Any]:
|
||
"""Call one ModelHub JSON endpoint without logging headers or response bodies."""
|
||
|
||
body = None
|
||
headers = {"Accept": "application/json"}
|
||
if payload is not None:
|
||
body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
|
||
headers["Content-Type"] = "application/json"
|
||
if token:
|
||
headers["Xc-Token"] = token
|
||
request = Request(url, data=body, headers=headers, method=method)
|
||
with HTTP_OPENER.open(request, timeout=15) as response:
|
||
result = json.loads(response.read().decode("utf-8"))
|
||
if not isinstance(result, dict):
|
||
raise RuntimeError("unexpected_response_shape")
|
||
return result
|
||
|
||
|
||
def _candidate_records(payload: dict[str, Any]) -> list[dict[str, str]]:
|
||
"""Extract text-generation candidates from the public ModelHub response."""
|
||
|
||
data = payload.get("data")
|
||
if not isinstance(data, list):
|
||
return []
|
||
candidates: list[dict[str, str]] = []
|
||
for item in data:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
model_id = str(item.get("modelId") or "").strip()
|
||
info = item.get("taskLevelsInfo")
|
||
info = info if isinstance(info, dict) else {}
|
||
raw_type = str(
|
||
info.get("taskLevelName")
|
||
or info.get("taskLevelCode")
|
||
or info.get("taskLevelEnglishName")
|
||
or ""
|
||
).lower()
|
||
chinese_type = str(info.get("taskLevelChineseName") or "")
|
||
if not model_id or not (
|
||
"text-generation" in raw_type
|
||
or "text_generation" in raw_type
|
||
or chinese_type.strip() == "文本生成"
|
||
):
|
||
continue
|
||
candidates.append({"model_id": model_id, "task_type": "text-generation"})
|
||
return candidates
|
||
|
||
|
||
def _model_size_score(model_id: str) -> float:
|
||
match = re.search(r"(?i)(\d+(?:\.\d+)?)\s*([bm])(?:\b|[_-])", model_id)
|
||
if not match:
|
||
return 1_000_000.0
|
||
size = float(match.group(1))
|
||
return size * (1_000 if match.group(2).lower() == "b" else 1)
|
||
|
||
|
||
def select_candidate(payload: dict[str, Any]) -> dict[str, str] | None:
|
||
"""Prefer a small official-EngineX text model from the public hot list."""
|
||
|
||
candidates = _candidate_records(payload)
|
||
if not candidates:
|
||
return None
|
||
|
||
def score(item: dict[str, str]) -> tuple[int, float, str]:
|
||
model_id = item["model_id"].lower()
|
||
supported_family = any(
|
||
name in model_id for name in ("qwen3", "llama3", "deepseek-r1-distill")
|
||
)
|
||
return (0 if supported_family else 1, _model_size_score(model_id), model_id)
|
||
|
||
return min(candidates, key=score)
|
||
|
||
|
||
def _model_address(model_id: str) -> str:
|
||
safe_path = "/".join(quote(part, safe="") for part in model_id.split("/"))
|
||
return f"https://www.modelscope.cn/models/{safe_path}"
|
||
|
||
|
||
def scan_once(request_json: Any = None) -> dict[str, Any]:
|
||
"""Discover one candidate, deduplicate it, and optionally submit one task."""
|
||
|
||
call = request_json or _http_json
|
||
scan_at = datetime.now(timezone.utc).isoformat()
|
||
_scanner_state_update(running=True, last_scan_at=scan_at, last_error_type=None)
|
||
_log_event("candidate_scan_started", source="modelhub_hot_models")
|
||
candidate_payload = call("GET", CANDIDATE_ENDPOINT)
|
||
candidate = select_candidate(candidate_payload)
|
||
if candidate is None:
|
||
result = {"action": "no_candidate"}
|
||
_scanner_state_update(running=False, last_candidate=None, last_action=result["action"])
|
||
_log_event("candidate_scan_completed", candidate_count=0, action=result["action"])
|
||
return result
|
||
|
||
model_id = _safe_model_name(candidate["model_id"])
|
||
_scanner_state_update(last_candidate=model_id)
|
||
_log_event(
|
||
"candidate_discovered",
|
||
model=model_id,
|
||
task_type=candidate["task_type"],
|
||
target_gpu_configured=bool(TARGET_GPU),
|
||
)
|
||
|
||
missing = []
|
||
if not AUTO_SUBMIT_ENABLED:
|
||
missing.append("auto_submit_disabled")
|
||
if not XC_TOKEN:
|
||
missing.append("xc_token_missing")
|
||
if not STRATEGY_ID:
|
||
missing.append("strategy_id_missing")
|
||
if not TARGET_GPU:
|
||
missing.append("target_gpu_missing")
|
||
if not CONFIG_PARAMS:
|
||
missing.append("task_config_missing")
|
||
if missing:
|
||
result = {"action": "submission_skipped", "reasons": missing, "model": model_id}
|
||
_scanner_state_update(running=False, last_action=result["action"])
|
||
_log_event("task_submission_skipped", model=model_id, reasons=missing)
|
||
return result
|
||
|
||
query = urlencode(
|
||
{
|
||
"current": 1,
|
||
"pageSize": 1,
|
||
"onlyMine": "true",
|
||
"modelId": candidate["model_id"],
|
||
"gpuType": TARGET_GPU,
|
||
}
|
||
)
|
||
history = call("GET", f"{TASK_PAGE_ENDPOINT}?{query}", token=XC_TOKEN)
|
||
history_data = history.get("data") if isinstance(history, dict) else None
|
||
records = history_data.get("records") if isinstance(history_data, dict) else []
|
||
if records:
|
||
result = {"action": "duplicate_skipped", "model": model_id}
|
||
_scanner_state_update(running=False, last_action=result["action"])
|
||
_log_event("task_duplicate_skipped", model=model_id, target_gpu=TARGET_GPU)
|
||
return result
|
||
|
||
task_payload = {
|
||
"modelAddress": _model_address(candidate["model_id"]),
|
||
"taskType": candidate["task_type"],
|
||
"targetGpu": TARGET_GPU,
|
||
"framework": "vllm",
|
||
"strategyId": STRATEGY_ID,
|
||
"configParams": CONFIG_PARAMS,
|
||
}
|
||
response = call("POST", TASK_ADD_ENDPOINT, payload=task_payload, token=XC_TOKEN)
|
||
response_data = response.get("data") if isinstance(response, dict) else None
|
||
task_id = response_data.get("taskId") if isinstance(response_data, dict) else response_data
|
||
result = {"action": "submitted", "model": model_id, "task_id": task_id}
|
||
_scanner_state_update(running=False, last_action=result["action"])
|
||
_log_event(
|
||
"task_submitted",
|
||
model=model_id,
|
||
target_gpu=TARGET_GPU,
|
||
task_id=task_id,
|
||
)
|
||
return result
|
||
|
||
|
||
def _scanner_loop() -> None:
|
||
_log_event(
|
||
"scanner_started",
|
||
interval_seconds=SCAN_INTERVAL_SECONDS,
|
||
auto_submit=AUTO_SUBMIT_ENABLED,
|
||
token_present=bool(XC_TOKEN),
|
||
target_gpu_configured=bool(TARGET_GPU),
|
||
)
|
||
while not STOP.is_set():
|
||
try:
|
||
scan_once()
|
||
except Exception as exc:
|
||
error_type = type(exc).__name__
|
||
_scanner_state_update(
|
||
running=False,
|
||
last_action="scan_error",
|
||
last_error_type=error_type,
|
||
)
|
||
_log_event("scanner_error", error_type=error_type)
|
||
if STOP.wait(SCAN_INTERVAL_SECONDS):
|
||
break
|
||
|
||
|
||
def _first(payload: dict[str, Any], *names: str) -> Any:
|
||
for name in names:
|
||
if name in payload and payload[name] not in (None, ""):
|
||
return payload[name]
|
||
return None
|
||
|
||
|
||
def _positive_int(value: Any, field: str, default: int) -> int:
|
||
if value in (None, ""):
|
||
return default
|
||
try:
|
||
number = int(value)
|
||
except (TypeError, ValueError) as exc:
|
||
raise ValidationError(f"{field} 必须是整数") from exc
|
||
if number <= 0:
|
||
raise ValidationError(f"{field} 必须大于 0")
|
||
return number
|
||
|
||
|
||
def _boolean(value: Any, field: str, default: bool = False) -> bool:
|
||
if value in (None, ""):
|
||
return default
|
||
if isinstance(value, bool):
|
||
return value
|
||
if isinstance(value, str):
|
||
normalized = value.strip().lower()
|
||
if normalized in {"true", "1", "yes", "on"}:
|
||
return True
|
||
if normalized in {"false", "0", "no", "off"}:
|
||
return False
|
||
raise ValidationError(f"{field} 必须是布尔值")
|
||
|
||
|
||
def _normalize(value: str) -> str:
|
||
return re.sub(r"[^0-9a-z一-鿿]+", "", value.lower())
|
||
|
||
|
||
def analyze(payload: dict[str, Any]) -> dict[str, Any]:
|
||
"""Return an evidence-labelled preflight for this repository's card."""
|
||
|
||
if not isinstance(payload, dict):
|
||
raise ValidationError("请求正文必须是 JSON 对象")
|
||
|
||
model_name = str(_first(payload, "model_name", "model", "model_address") or "")
|
||
framework = str(_first(payload, "framework") or "")
|
||
backend = str(_first(payload, "backend", "inference_engine", "engine") or "")
|
||
sdk_version = str(_first(payload, "sdk_version", "sdk") or "")
|
||
driver_version = str(_first(payload, "driver_version", "driver") or "")
|
||
precision = str(_first(payload, "precision", "dtype") or "bf16").lower()
|
||
cards = _positive_int(_first(payload, "cards", "card_count"), "cards", 1)
|
||
context_length = _positive_int(
|
||
_first(payload, "context_length", "max_sequence_length"),
|
||
"context_length",
|
||
4096,
|
||
)
|
||
custom_ops = _boolean(_first(payload, "custom_ops"), "custom_ops")
|
||
dynamic_shapes = _boolean(
|
||
_first(payload, "dynamic_shapes"), "dynamic_shapes"
|
||
)
|
||
|
||
supplied_hardware = str(
|
||
_first(payload, "hardware", "device", "target_card", "target_gpu") or ""
|
||
)
|
||
target_tokens = {_normalize(CARD_VENDOR), _normalize(CARD_MODEL)}
|
||
normalized_hardware = _normalize(supplied_hardware)
|
||
target_matches = not supplied_hardware or any(
|
||
token and token in normalized_hardware for token in target_tokens
|
||
)
|
||
|
||
environment = {
|
||
"model_name": model_name,
|
||
"framework": framework,
|
||
"backend": backend,
|
||
"sdk_version": sdk_version,
|
||
"driver_version": driver_version,
|
||
}
|
||
missing_fields = [name for name, value in environment.items() if not value]
|
||
|
||
risks: list[str] = []
|
||
if precision in {"int4", "4bit", "int8", "8bit"}:
|
||
risks.append("量化方案需要实测目标卡后端是否提供对应权重格式与算子内核。")
|
||
if cards > 1:
|
||
risks.append("多卡运行需要验证集合通信、进程数、拓扑和并行策略。")
|
||
if context_length > 32768:
|
||
risks.append("长上下文会增加 KV Cache 压力，需要按真实并发测峰值内存。")
|
||
if custom_ops:
|
||
risks.append("模型包含自定义算子，需要确认编译链、ABI 与目标后端注册情况。")
|
||
if dynamic_shapes:
|
||
risks.append("动态形状需要验证图编译缓存、回退路径与重复编译开销。")
|
||
|
||
if not target_matches:
|
||
verdict = "target_mismatch"
|
||
risks.insert(0, f"本智能体仅面向 {CARD_VENDOR}｜{CARD_MODEL}。")
|
||
elif missing_fields:
|
||
verdict = "information_required"
|
||
else:
|
||
verdict = "preflight_ready"
|
||
|
||
recommendations = [
|
||
f"确认实际目标设备标识为 {CARD_VENDOR}｜{CARD_MODEL}。",
|
||
"记录驱动、SDK、框架和推理后端的完整版本矩阵。",
|
||
"先用厂商基础样例确认设备可见，再运行模型最小输入。",
|
||
"依次验证模型加载、首个算子、单请求输出和资源峰值。",
|
||
"真实支持性结论必须来自目标设备运行日志与健康检查。",
|
||
]
|
||
|
||
return {
|
||
"agent": AGENT_NAME,
|
||
"version": AGENT_VERSION,
|
||
"official_target": {
|
||
"vendor": CARD_VENDOR,
|
||
"model": CARD_MODEL,
|
||
"source_scope": "信创模盒模型 X 算力页面",
|
||
},
|
||
"supplied_hardware": supplied_hardware or None,
|
||
"target_matches": target_matches,
|
||
"verdict": verdict,
|
||
"inputs": {
|
||
**{key: value or None for key, value in environment.items()},
|
||
"precision": precision,
|
||
"cards": cards,
|
||
"context_length": context_length,
|
||
"custom_ops": custom_ops,
|
||
"dynamic_shapes": dynamic_shapes,
|
||
},
|
||
"missing_fields": missing_fields,
|
||
"risks": risks,
|
||
"recommendations": recommendations,
|
||
"disclaimer": "未提供或未实测的硬件规格与兼容性不会被推断为已支持。",
|
||
}
|
||
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
server_version = "ModelHubCardAdvisor/1.2"
|
||
|
||
def _json(self, payload: dict[str, Any], status: int = 200) -> None:
|
||
body = json.dumps(payload, 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(body)))
|
||
self.end_headers()
|
||
self.wfile.write(body)
|
||
|
||
def do_GET(self) -> None: # noqa: N802 - BaseHTTPRequestHandler contract
|
||
if self.path == "/health":
|
||
self._json(
|
||
{
|
||
"status": "ok",
|
||
"agent": AGENT_NAME,
|
||
"version": AGENT_VERSION,
|
||
"scanner": scanner_snapshot(),
|
||
}
|
||
)
|
||
return
|
||
if self.path == "/scanner":
|
||
self._json(scanner_snapshot())
|
||
return
|
||
if self.path == "/":
|
||
self._json(
|
||
{
|
||
"name": AGENT_NAME,
|
||
"version": AGENT_VERSION,
|
||
"description": f"{CARD_VENDOR} {CARD_MODEL} 专属部署前预检",
|
||
"official_target": {"vendor": CARD_VENDOR, "model": CARD_MODEL},
|
||
"strategy_id_present": bool(STRATEGY_ID),
|
||
"target_gpu": TARGET_GPU or None,
|
||
"endpoints": [
|
||
"GET /health",
|
||
"GET /scanner",
|
||
"POST /analyze",
|
||
"POST /task",
|
||
],
|
||
"external_writes_enabled": bool(
|
||
AUTO_SUBMIT_ENABLED
|
||
and XC_TOKEN
|
||
and STRATEGY_ID
|
||
and TARGET_GPU
|
||
and CONFIG_PARAMS
|
||
),
|
||
}
|
||
)
|
||
return
|
||
self._json({"error": "not_found"}, 404)
|
||
|
||
def do_POST(self) -> None: # noqa: N802 - BaseHTTPRequestHandler contract
|
||
if self.path not in {"/analyze", "/task"}:
|
||
self._json({"error": "not_found"}, 404)
|
||
return
|
||
try:
|
||
length = int(self.headers.get("Content-Length", "0"))
|
||
if length <= 0:
|
||
raise ValidationError("请求正文不能为空")
|
||
if length > MAX_BODY_BYTES:
|
||
_log_event("request_rejected", path=self.path, reason="payload_too_large")
|
||
self._json({"error": "payload_too_large"}, 413)
|
||
return
|
||
payload = json.loads(self.rfile.read(length))
|
||
model_name = _safe_model_name(
|
||
_first(payload, "model_name", "model", "model_address")
|
||
if isinstance(payload, dict)
|
||
else None
|
||
)
|
||
_log_event("analysis_started", path=self.path, model=model_name)
|
||
result = analyze(payload)
|
||
_log_event(
|
||
"analysis_completed",
|
||
path=self.path,
|
||
model=model_name,
|
||
verdict=result["verdict"],
|
||
target_matches=result["target_matches"],
|
||
)
|
||
self._json(result)
|
||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||
_log_event("request_rejected", path=self.path, reason="invalid_json")
|
||
self._json({"error": "invalid_json", "message": "请求正文必须是有效 JSON"}, 400)
|
||
except ValidationError as exc:
|
||
_log_event(
|
||
"request_rejected",
|
||
path=self.path,
|
||
reason="validation_error",
|
||
message=str(exc),
|
||
)
|
||
self._json({"error": "validation_error", "message": str(exc)}, 422)
|
||
|
||
def log_message(self, _format: str, *args: Any) -> None:
|
||
status = str(args[1]) if len(args) > 1 else "unknown"
|
||
_log_event(
|
||
"http_access",
|
||
method=self.command,
|
||
path=urlsplit(self.path).path,
|
||
status=status,
|
||
)
|
||
|
||
|
||
STOP = threading.Event()
|
||
|
||
|
||
def _handle_signal(signum: int, _frame: Any) -> None:
|
||
_log_event("shutdown_requested", signal=signum)
|
||
STOP.set()
|
||
|
||
|
||
def main() -> None:
|
||
signal.signal(signal.SIGTERM, _handle_signal)
|
||
signal.signal(signal.SIGINT, _handle_signal)
|
||
server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
|
||
server.timeout = 1
|
||
_log_event(
|
||
"service_started",
|
||
host="0.0.0.0",
|
||
port=PORT,
|
||
target_vendor=CARD_VENDOR,
|
||
target_model=CARD_MODEL,
|
||
)
|
||
scanner_thread = None
|
||
if AUTO_SCAN_ENABLED:
|
||
scanner_thread = threading.Thread(
|
||
target=_scanner_loop,
|
||
name="modelhub-candidate-scanner",
|
||
daemon=True,
|
||
)
|
||
scanner_thread.start()
|
||
else:
|
||
_scanner_state_update(last_action="disabled")
|
||
_log_event("scanner_disabled", reason="auto_scan_disabled")
|
||
while not STOP.is_set():
|
||
server.handle_request()
|
||
server.server_close()
|
||
if scanner_thread is not None:
|
||
scanner_thread.join(timeout=5)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|