实现模型自动适配任务智能体 v1.0.0
This commit is contained in:
12
Dockerfile
Normal file
12
Dockerfile
Normal file
@@ -0,0 +1,12 @@
|
||||
FROM modelhubxc-4pd.tencentcloudcr.com/xc_agent_platform/python:3.11-slim
|
||||
|
||||
ENV PYTHONDONTWRITEBYTECODE=1 \
|
||||
PYTHONUNBUFFERED=1 \
|
||||
PORT=8080
|
||||
|
||||
WORKDIR /app
|
||||
COPY main.py /app/main.py
|
||||
|
||||
EXPOSE 8080
|
||||
|
||||
CMD ["python", "/app/main.py"]
|
||||
55
README.md
Normal file
55
README.md
Normal file
@@ -0,0 +1,55 @@
|
||||
# xc-model-auto-adaptation-agent
|
||||
|
||||
信创模盒的“模型自动适配任务智能体”。它不是只给建议的兼容性预检工具,而是一个持续运行的任务执行器。
|
||||
|
||||
## 自动执行链路
|
||||
|
||||
1. 启动后立即读取信创模盒公开热门模型。
|
||||
2. 筛选文本生成模型,并优先选择较小的 Qwen3、Llama3 或 DeepSeek-R1-Distill 候选。
|
||||
3. 用模型 ID 与国产卡枚举查询当前账号的历史验证任务。
|
||||
4. 已有相同任务则记录去重结果;没有则生成 EngineX/vLLM 配置并提交验证任务。
|
||||
5. 每次扫描最多提交一个任务,默认每小时重复检查。
|
||||
|
||||
第一版目标卡为 **天数智芯|天垓100**,平台卡型枚举为 `Iluvatar_bi-100`,使用单卡、短上下文的保守验证配置。
|
||||
|
||||
## 平台契约
|
||||
|
||||
- 仓库根目录包含 `Dockerfile`,监听 `8080`。
|
||||
- `GET /health` 始终返回 HTTP 200,并附带扫描器最近状态。
|
||||
- 从平台注入的 `STRATEGY_ID` 获取自身策略 ID。
|
||||
- 优先读取 `MODELHUB_XC_TOKEN`,同时兼容官方参考策略使用的 `EXTERNAL_SERVICE_TOKEN`。
|
||||
- 正确处理 `SIGTERM`,在平台 30 秒窗口内退出。
|
||||
- 仅使用 Python 标准库,符合 1 CPU / 512 MiB 的平台限制。
|
||||
|
||||
## 安全边界
|
||||
|
||||
- 令牌只从运行环境读取,不写入仓库、响应或日志。
|
||||
- stdout 只记录令牌是否存在,不记录令牌值、请求头或完整 API 响应。
|
||||
- 提交前必须完成精确去重;API 返回非成功业务码时不会伪报成功。
|
||||
- 未获得令牌或策略 ID 时会明确记录 `task_submission_blocked`,不会静默假运行。
|
||||
|
||||
## 端点与关键日志
|
||||
|
||||
- `GET /health`:服务与扫描器状态。
|
||||
- `GET /scanner`:最近候选、动作和错误类型。
|
||||
- `POST /analyze`:返回该模型的实际适配执行计划。
|
||||
|
||||
正常启动后的关键事件依次为:
|
||||
|
||||
```text
|
||||
service_started
|
||||
scanner_started
|
||||
candidate_scan_started
|
||||
candidate_discovered
|
||||
task_duplicate_skipped 或 task_submitted
|
||||
```
|
||||
|
||||
如果运行环境配置不完整,会看到 `task_submission_blocked` 和具体缺失项。
|
||||
|
||||
## 本地验证
|
||||
|
||||
```bash
|
||||
python3 -m unittest -v
|
||||
python3 main.py
|
||||
curl http://localhost:8080/health
|
||||
```
|
||||
541
main.py
Normal file
541
main.py
Normal file
@@ -0,0 +1,541 @@
|
||||
"""ModelHub XC automatic model adaptation task agent.
|
||||
|
||||
The worker discovers a public text-generation model, checks whether the same
|
||||
model/card pair already has a task, and submits at most one verified adaptation
|
||||
task per scan. Credentials are read only from the runtime environment.
|
||||
"""
|
||||
|
||||
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, Mapping
|
||||
from urllib.parse import quote, urlencode, urlsplit
|
||||
from urllib.request import ProxyHandler, Request, build_opener
|
||||
|
||||
|
||||
AGENT_NAME = "xc-model-auto-adaptation-agent"
|
||||
AGENT_VERSION = "1.0.0"
|
||||
TARGET_VENDOR = "天数智芯"
|
||||
TARGET_CARD = "天垓100"
|
||||
DEFAULT_TARGET_GPU = "Iluvatar_bi-100"
|
||||
|
||||
PORT = int(os.getenv("PORT", "8080"))
|
||||
STRATEGY_ID = os.getenv("STRATEGY_ID", "")
|
||||
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)
|
||||
MAX_BODY_BYTES = 1_000_000
|
||||
|
||||
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
|
||||
api: completion
|
||||
lang: zh
|
||||
max_model_len: 2048
|
||||
max_tokens: 256
|
||||
temperature: 0.1
|
||||
repetition_penalty: 1.0
|
||||
top_p: 0.9
|
||||
sut_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- python3
|
||||
- -m
|
||||
- vllm.entrypoints.openai.api_server
|
||||
- --host
|
||||
- 0.0.0.0
|
||||
- --port
|
||||
- '20644'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --model
|
||||
- /model
|
||||
- --max-model-len
|
||||
- '2048'
|
||||
- --tensor-parallel-size
|
||||
- '1'
|
||||
- --max-num-seqs
|
||||
- '8'
|
||||
- --enforce-eager
|
||||
- --disable-log-requests
|
||||
- --enable-prefix-caching
|
||||
- --trust-remote-code
|
||||
ref_config:
|
||||
gpu_num: 1
|
||||
values:
|
||||
command:
|
||||
- vllm
|
||||
- serve
|
||||
- /model
|
||||
- --port
|
||||
- '80'
|
||||
- --served-model-name
|
||||
- llm
|
||||
- --max-model-len
|
||||
- '2048'
|
||||
- -tp
|
||||
- '1'
|
||||
- --enforce-eager
|
||||
- --trust-remote-code
|
||||
"""
|
||||
CONFIG_PARAMS = os.getenv("MODELHUB_CONFIG_PARAMS", DEFAULT_CONFIG_PARAMS)
|
||||
|
||||
SECRET_ASSIGNMENT_RE = re.compile(
|
||||
r"(?i)(token|secret|password|api[_-]?key)=([^&\s]+)"
|
||||
)
|
||||
HTTP_OPENER = build_opener(ProxyHandler({}))
|
||||
STATE_LOCK = threading.Lock()
|
||||
STOP = threading.Event()
|
||||
SCANNER_STATE: dict[str, Any] = {
|
||||
"enabled": AUTO_SCAN_ENABLED,
|
||||
"running": False,
|
||||
"ready": False,
|
||||
"last_scan_at": None,
|
||||
"last_candidate": None,
|
||||
"last_action": "not_started",
|
||||
"last_error_type": None,
|
||||
}
|
||||
|
||||
|
||||
class ValidationError(ValueError):
|
||||
"""Raised when an incoming analysis request is invalid."""
|
||||
|
||||
|
||||
class PlatformAPIError(RuntimeError):
|
||||
"""Raised when ModelHub returns a non-success business response."""
|
||||
|
||||
|
||||
def resolve_runtime_token(environ: Mapping[str, str] | None = None) -> str:
|
||||
"""Resolve documented/compatible secret names without ever logging values."""
|
||||
|
||||
source = environ if environ is not None else os.environ
|
||||
return (
|
||||
source.get("MODELHUB_XC_TOKEN", "")
|
||||
or source.get("EXTERNAL_SERVICE_TOKEN", "")
|
||||
or source.get("XC_TOKEN", "")
|
||||
)
|
||||
|
||||
|
||||
XC_TOKEN = resolve_runtime_token()
|
||||
|
||||
|
||||
def _safe_text(value: Any, limit: int = 200) -> str:
|
||||
text = str(value or "").strip()
|
||||
text = SECRET_ASSIGNMENT_RE.sub(r"\1=[REDACTED]", text)
|
||||
return text[:limit]
|
||||
|
||||
|
||||
def _safe_model_name(value: Any) -> str:
|
||||
text = _safe_text(value) or "未指定模型"
|
||||
if "://" in text:
|
||||
try:
|
||||
parsed = urlsplit(text)
|
||||
text = f"{parsed.scheme}://{parsed.hostname or ''}{parsed.path}"
|
||||
except ValueError:
|
||||
text = "无效模型地址"
|
||||
return text[:200]
|
||||
|
||||
|
||||
def _log_event(event: str, **fields: Any) -> None:
|
||||
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 _state_update(**fields: Any) -> None:
|
||||
with STATE_LOCK:
|
||||
SCANNER_STATE.update(fields)
|
||||
|
||||
|
||||
def scanner_snapshot() -> dict[str, Any]:
|
||||
with STATE_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 JSON endpoint without logging headers, bodies, or credentials."""
|
||||
|
||||
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=20) as response:
|
||||
result = json.loads(response.read().decode("utf-8"))
|
||||
if not isinstance(result, dict):
|
||||
raise PlatformAPIError("unexpected_response_shape")
|
||||
return result
|
||||
|
||||
|
||||
def _require_success(payload: dict[str, Any], operation: str) -> None:
|
||||
code = payload.get("code")
|
||||
if code not in (None, 0, "0", 200, "200"):
|
||||
message = _safe_text(payload.get("message") or "platform_rejected", 120)
|
||||
raise PlatformAPIError(f"{operation}_failed:{code}:{message}")
|
||||
|
||||
|
||||
def _candidate_records(payload: dict[str, Any]) -> list[dict[str, str]]:
|
||||
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 "").strip()
|
||||
if model_id and (
|
||||
"text-generation" in raw_type
|
||||
or "text_generation" in raw_type
|
||||
or chinese_type == "文本生成"
|
||||
):
|
||||
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 text-generation model in a verified EngineX family."""
|
||||
|
||||
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 = any(
|
||||
family in model_id for family in ("qwen3", "llama3", "deepseek-r1-distill")
|
||||
)
|
||||
return (0 if supported 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 _history_records(payload: dict[str, Any]) -> list[Any]:
|
||||
data = payload.get("data")
|
||||
if not isinstance(data, dict):
|
||||
return []
|
||||
records = data.get("records")
|
||||
return records if isinstance(records, list) else []
|
||||
|
||||
|
||||
def scan_once(request_json: Any = None) -> dict[str, Any]:
|
||||
"""Discover, deduplicate, and submit at most one adaptation task."""
|
||||
|
||||
call = request_json or _http_json
|
||||
scan_at = datetime.now(timezone.utc).isoformat()
|
||||
ready = bool(
|
||||
AUTO_SUBMIT_ENABLED and XC_TOKEN and STRATEGY_ID and TARGET_GPU and CONFIG_PARAMS
|
||||
)
|
||||
_state_update(
|
||||
running=True,
|
||||
ready=ready,
|
||||
last_scan_at=scan_at,
|
||||
last_error_type=None,
|
||||
)
|
||||
_log_event("candidate_scan_started", source="modelhub_hot_models")
|
||||
|
||||
candidate_payload = call("GET", CANDIDATE_ENDPOINT)
|
||||
_require_success(candidate_payload, "candidate_query")
|
||||
candidate = select_candidate(candidate_payload)
|
||||
if candidate is None:
|
||||
result = {"action": "no_candidate"}
|
||||
_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"])
|
||||
_state_update(last_candidate=model_id)
|
||||
_log_event(
|
||||
"candidate_discovered",
|
||||
model=model_id,
|
||||
task_type=candidate["task_type"],
|
||||
target_gpu=TARGET_GPU,
|
||||
)
|
||||
|
||||
missing: list[str] = []
|
||||
if not AUTO_SUBMIT_ENABLED:
|
||||
missing.append("auto_submit_disabled")
|
||||
if not XC_TOKEN:
|
||||
missing.append("runtime_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_blocked", "reasons": missing, "model": model_id}
|
||||
_state_update(running=False, ready=False, last_action=result["action"])
|
||||
_log_event("task_submission_blocked", model=model_id, reasons=missing)
|
||||
return result
|
||||
|
||||
query = urlencode(
|
||||
{
|
||||
"current": 1,
|
||||
"pageSize": 20,
|
||||
"onlyMine": "true",
|
||||
"modelId": candidate["model_id"],
|
||||
"gpuType": TARGET_GPU,
|
||||
}
|
||||
)
|
||||
history = call("GET", f"{TASK_PAGE_ENDPOINT}?{query}", token=XC_TOKEN)
|
||||
_require_success(history, "history_query")
|
||||
if _history_records(history):
|
||||
result = {"action": "duplicate_skipped", "model": model_id}
|
||||
_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)
|
||||
_require_success(response, "task_create")
|
||||
response_data = response.get("data")
|
||||
task_id = None
|
||||
if isinstance(response_data, dict):
|
||||
task_id = response_data.get("id") or response_data.get("taskId")
|
||||
elif response_data not in (None, ""):
|
||||
task_id = response_data
|
||||
result = {"action": "submitted", "model": model_id, "task_id": task_id}
|
||||
_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 build_adaptation_plan(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return the concrete adaptation plan used by this worker."""
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
raise ValidationError("请求正文必须是 JSON 对象")
|
||||
model = _safe_model_name(
|
||||
payload.get("model_name") or payload.get("model") or payload.get("model_address")
|
||||
)
|
||||
return {
|
||||
"agent": AGENT_NAME,
|
||||
"version": AGENT_VERSION,
|
||||
"model": model,
|
||||
"target": {
|
||||
"vendor": TARGET_VENDOR,
|
||||
"card": TARGET_CARD,
|
||||
"platform_gpu": TARGET_GPU,
|
||||
"cards": 1,
|
||||
},
|
||||
"framework": "vllm",
|
||||
"task_type": "text-generation",
|
||||
"execution": [
|
||||
"查询同模型与卡型的现有验证任务",
|
||||
"生成单卡 EngineX/vLLM 验证配置",
|
||||
"提交平台验证任务",
|
||||
"由平台执行模型加载、推理与指标采集",
|
||||
],
|
||||
"automatic_submission_ready": bool(
|
||||
AUTO_SUBMIT_ENABLED
|
||||
and XC_TOKEN
|
||||
and STRATEGY_ID
|
||||
and TARGET_GPU
|
||||
and CONFIG_PARAMS
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _scanner_loop() -> None:
|
||||
_log_event(
|
||||
"scanner_started",
|
||||
interval_seconds=SCAN_INTERVAL_SECONDS,
|
||||
auto_submit=AUTO_SUBMIT_ENABLED,
|
||||
runtime_token_present=bool(XC_TOKEN),
|
||||
strategy_id_present=bool(STRATEGY_ID),
|
||||
target_gpu=TARGET_GPU,
|
||||
)
|
||||
while not STOP.is_set():
|
||||
try:
|
||||
scan_once()
|
||||
except Exception as exc: # boundary: keep health endpoint alive
|
||||
error_type = type(exc).__name__
|
||||
_state_update(
|
||||
running=False,
|
||||
ready=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
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
server_version = "ModelHubAutoAdaptationAgent/1.0"
|
||||
|
||||
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
|
||||
path = urlsplit(self.path).path
|
||||
if path == "/health":
|
||||
self._json(
|
||||
{
|
||||
"status": "ok",
|
||||
"agent": AGENT_NAME,
|
||||
"version": AGENT_VERSION,
|
||||
"scanner": scanner_snapshot(),
|
||||
}
|
||||
)
|
||||
return
|
||||
if path == "/scanner":
|
||||
self._json(scanner_snapshot())
|
||||
return
|
||||
if path == "/":
|
||||
self._json(
|
||||
{
|
||||
"name": AGENT_NAME,
|
||||
"version": AGENT_VERSION,
|
||||
"description": "自动发现模型并提交国产卡适配验证任务",
|
||||
"target": {"vendor": TARGET_VENDOR, "card": TARGET_CARD},
|
||||
"endpoints": ["GET /health", "GET /scanner", "POST /analyze"],
|
||||
"automatic_submission_ready": 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
|
||||
if urlsplit(self.path).path != "/analyze":
|
||||
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:
|
||||
self._json({"error": "payload_too_large"}, 413)
|
||||
return
|
||||
payload = json.loads(self.rfile.read(length))
|
||||
model = _safe_model_name(payload.get("model_name") if isinstance(payload, dict) else None)
|
||||
_log_event("adaptation_plan_started", model=model)
|
||||
plan = build_adaptation_plan(payload)
|
||||
_log_event("adaptation_plan_completed", model=plan["model"])
|
||||
self._json(plan)
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
self._json({"error": "invalid_json"}, 400)
|
||||
except ValidationError as 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,
|
||||
)
|
||||
|
||||
|
||||
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=TARGET_VENDOR,
|
||||
target_card=TARGET_CARD,
|
||||
)
|
||||
scanner_thread = None
|
||||
if AUTO_SCAN_ENABLED:
|
||||
scanner_thread = threading.Thread(
|
||||
target=_scanner_loop,
|
||||
name="modelhub-auto-adaptation-scanner",
|
||||
daemon=True,
|
||||
)
|
||||
scanner_thread.start()
|
||||
else:
|
||||
_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()
|
||||
163
test_main.py
Normal file
163
test_main.py
Normal file
@@ -0,0 +1,163 @@
|
||||
import io
|
||||
import json
|
||||
import unittest
|
||||
from contextlib import redirect_stdout
|
||||
from unittest.mock import patch
|
||||
|
||||
import main as agent_main
|
||||
from main import (
|
||||
AGENT_NAME,
|
||||
PlatformAPIError,
|
||||
_log_event,
|
||||
_safe_model_name,
|
||||
build_adaptation_plan,
|
||||
resolve_runtime_token,
|
||||
scan_once,
|
||||
select_candidate,
|
||||
)
|
||||
|
||||
|
||||
class AutoAdaptationTests(unittest.TestCase):
|
||||
@staticmethod
|
||||
def candidate_payload():
|
||||
return {
|
||||
"code": 0,
|
||||
"data": [
|
||||
{
|
||||
"modelId": "Qwen/Qwen3-30B-Instruct",
|
||||
"taskLevelsInfo": {"taskLevelChineseName": "文本生成"},
|
||||
},
|
||||
{
|
||||
"modelId": "Qwen/Qwen3-4B-Instruct-2507",
|
||||
"taskLevelsInfo": {"taskLevelChineseName": "文本生成"},
|
||||
},
|
||||
{
|
||||
"modelId": "example/image-model",
|
||||
"taskLevelsInfo": {"taskLevelChineseName": "文本生成图片"},
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
def test_official_runtime_token_name_is_supported(self):
|
||||
token = resolve_runtime_token({"EXTERNAL_SERVICE_TOKEN": "private-value"})
|
||||
self.assertEqual(token, "private-value")
|
||||
|
||||
def test_explicit_modelhub_token_takes_precedence(self):
|
||||
token = resolve_runtime_token(
|
||||
{
|
||||
"MODELHUB_XC_TOKEN": "preferred",
|
||||
"EXTERNAL_SERVICE_TOKEN": "fallback",
|
||||
}
|
||||
)
|
||||
self.assertEqual(token, "preferred")
|
||||
|
||||
def test_candidate_selection_prefers_small_supported_text_model(self):
|
||||
candidate = select_candidate(self.candidate_payload())
|
||||
self.assertIsNotNone(candidate)
|
||||
self.assertEqual(candidate["model_id"], "Qwen/Qwen3-4B-Instruct-2507")
|
||||
|
||||
def test_scan_blocks_submission_without_runtime_token(self):
|
||||
calls = []
|
||||
|
||||
def fake_request(method, url, **kwargs):
|
||||
calls.append((method, url, kwargs))
|
||||
return self.candidate_payload()
|
||||
|
||||
with patch.object(agent_main, "XC_TOKEN", ""):
|
||||
result = scan_once(fake_request)
|
||||
self.assertEqual(result["action"], "submission_blocked")
|
||||
self.assertIn("runtime_token_missing", result["reasons"])
|
||||
self.assertEqual([call[0] for call in calls], ["GET"])
|
||||
|
||||
def test_scan_deduplicates_before_submission(self):
|
||||
calls = []
|
||||
|
||||
def fake_request(method, url, **kwargs):
|
||||
calls.append((method, url, kwargs))
|
||||
if "top/models" in url:
|
||||
return self.candidate_payload()
|
||||
return {"code": 0, "data": {"records": [{"id": 1}]}}
|
||||
|
||||
with (
|
||||
patch.object(agent_main, "XC_TOKEN", "private-token"),
|
||||
patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
|
||||
patch.object(agent_main, "TARGET_GPU", "verified-gpu"),
|
||||
patch.object(agent_main, "CONFIG_PARAMS", "framework: vllm\nsut_config: test"),
|
||||
patch.object(agent_main, "AUTO_SUBMIT_ENABLED", True),
|
||||
):
|
||||
result = scan_once(fake_request)
|
||||
self.assertEqual(result["action"], "duplicate_skipped")
|
||||
self.assertEqual([call[0] for call in calls], ["GET", "GET"])
|
||||
|
||||
def test_scan_submits_exactly_one_task(self):
|
||||
calls = []
|
||||
|
||||
def fake_request(method, url, **kwargs):
|
||||
calls.append((method, url, kwargs))
|
||||
if "top/models" in url:
|
||||
return self.candidate_payload()
|
||||
if method == "GET":
|
||||
return {"code": 0, "data": {"records": []}}
|
||||
return {"code": 0, "data": {"id": 987, "status": "created"}}
|
||||
|
||||
with (
|
||||
patch.object(agent_main, "XC_TOKEN", "private-token"),
|
||||
patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
|
||||
patch.object(agent_main, "TARGET_GPU", "verified-gpu"),
|
||||
patch.object(agent_main, "CONFIG_PARAMS", "framework: vllm\nsut_config: test"),
|
||||
patch.object(agent_main, "AUTO_SUBMIT_ENABLED", True),
|
||||
):
|
||||
result = scan_once(fake_request)
|
||||
self.assertEqual(result, {
|
||||
"action": "submitted",
|
||||
"model": "Qwen/Qwen3-4B-Instruct-2507",
|
||||
"task_id": 987,
|
||||
})
|
||||
self.assertEqual([call[0] for call in calls], ["GET", "GET", "POST"])
|
||||
submitted = calls[-1][2]["payload"]
|
||||
self.assertEqual(submitted["strategyId"], "strategy-id")
|
||||
self.assertEqual(submitted["targetGpu"], "verified-gpu")
|
||||
self.assertIn("sut_config", submitted["configParams"])
|
||||
|
||||
def test_platform_business_error_is_not_reported_as_success(self):
|
||||
def fake_request(method, url, **kwargs):
|
||||
if "top/models" in url:
|
||||
return self.candidate_payload()
|
||||
return {"code": 403, "message": "forbidden"}
|
||||
|
||||
with (
|
||||
patch.object(agent_main, "XC_TOKEN", "private-token"),
|
||||
patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
|
||||
patch.object(agent_main, "TARGET_GPU", "verified-gpu"),
|
||||
patch.object(agent_main, "CONFIG_PARAMS", "framework: vllm"),
|
||||
patch.object(agent_main, "AUTO_SUBMIT_ENABLED", True),
|
||||
):
|
||||
with self.assertRaises(PlatformAPIError):
|
||||
scan_once(fake_request)
|
||||
|
||||
def test_adaptation_plan_is_execution_oriented(self):
|
||||
with (
|
||||
patch.object(agent_main, "XC_TOKEN", "private-token"),
|
||||
patch.object(agent_main, "STRATEGY_ID", "strategy-id"),
|
||||
):
|
||||
plan = build_adaptation_plan({"model_name": "Qwen/Qwen3-4B"})
|
||||
self.assertEqual(plan["task_type"], "text-generation")
|
||||
self.assertIn("提交平台验证任务", plan["execution"])
|
||||
self.assertTrue(plan["automatic_submission_ready"])
|
||||
|
||||
def test_model_identifier_removes_credentials_and_query(self):
|
||||
value = "https://user:pass@example.com/org/model?token=secret#fragment"
|
||||
self.assertEqual(_safe_model_name(value), "https://example.com/org/model")
|
||||
|
||||
def test_structured_log_does_not_need_secret_values(self):
|
||||
stream = io.StringIO()
|
||||
with redirect_stdout(stream):
|
||||
_log_event("scanner_started", runtime_token_present=True)
|
||||
record = json.loads(stream.getvalue())
|
||||
self.assertEqual(record["agent"], AGENT_NAME)
|
||||
self.assertTrue(record["runtime_token_present"])
|
||||
self.assertNotIn("token", record)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user