xc-model-auto-adaptation-agent/main.py

"""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()