"""ModelHub XC model compatibility preflight agent. The service is intentionally read-only: it evaluates supplied model and device facts but never submits adaptation jobs or calls external services. """ from __future__ import annotations import json import os import signal import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from typing import Any AGENT_NAME = "xc-model-compat-agent" AGENT_VERSION = "1.0.0" PORT = int(os.getenv("PORT", "8080")) STRATEGY_ID = os.getenv("STRATEGY_ID", "") MAX_BODY_BYTES = 1_000_000 DTYPE_BYTES = { "fp32": 4.0, "float32": 4.0, "fp16": 2.0, "float16": 2.0, "bf16": 2.0, "bfloat16": 2.0, "int8": 1.0, "8bit": 1.0, "int4": 0.5, "4bit": 0.5, } class ValidationError(ValueError): """Raised when a provided request field is invalid.""" 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_float(value: Any, field: str) -> float | None: if value in (None, ""): return None try: number = float(value) except (TypeError, ValueError) as exc: raise ValidationError(f"{field} 必须是数字") from exc if number <= 0: raise ValidationError(f"{field} 必须大于 0") return number 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 analyze(payload: dict[str, Any]) -> dict[str, Any]: """Return a conservative, evidence-labelled compatibility preflight.""" if not isinstance(payload, dict): raise ValidationError("请求正文必须是 JSON 对象") model_name = str( _first(payload, "model_name", "model", "model_address") or "未指定模型" ) architecture = str(_first(payload, "architecture", "model_architecture") or "") hardware = str(_first(payload, "hardware", "gpu_type", "target_gpu") or "") dtype = str(_first(payload, "dtype", "precision") or "bf16").lower() parameters_b = _positive_float( _first(payload, "parameters_b", "parameter_billions"), "parameters_b" ) device_memory_gb = _positive_float( _first(payload, "device_memory_gb", "memory_gb"), "device_memory_gb" ) cards = _positive_int(_first(payload, "cards", "card_count"), "cards", 1) context_length = _positive_int( _first(payload, "context_length", "max_sequence_length"), "context_length", 4096, ) missing_fields: list[str] = [] if parameters_b is None: missing_fields.append("parameters_b") if device_memory_gb is None: missing_fields.append("device_memory_gb") if not hardware: missing_fields.append("hardware") risks: list[str] = [] recommendations: list[str] = [] estimates: dict[str, float] = {} verdict = "information_required" bytes_per_parameter = DTYPE_BYTES.get(dtype) if bytes_per_parameter is None: risks.append(f"精度 {dtype} 未进入估算表，需确认格式与推理后端支持。") elif dtype in {"int8", "8bit", "int4", "4bit"}: risks.append("量化精度需要目标国产芯片后端提供对应量化算子。") if not architecture: risks.append("模型架构未提供，无法预判注意力、归一化及自定义算子兼容性。") if context_length > 32768: risks.append("长上下文会显著增加 KV Cache；当前权重估算未包含业务并发负载。") if cards > 1: risks.append("多卡方案需额外验证张量并行、集合通信和卡间带宽。") if parameters_b is not None and bytes_per_parameter is not None: weights_gib = parameters_b * 1_000_000_000 * bytes_per_parameter / (1024**3) runtime_budget_gib = weights_gib * 1.20 estimates = { "weights_gib": round(weights_gib, 2), "minimum_runtime_budget_gib": round(runtime_budget_gib, 2), } if device_memory_gb is not None: total_device_memory_gb = device_memory_gb * cards estimates["total_device_memory_gb"] = round(total_device_memory_gb, 2) estimates["headroom_gb"] = round(total_device_memory_gb - runtime_budget_gib, 2) if total_device_memory_gb < runtime_budget_gib: verdict = "insufficient_memory" recommendations.append("降低精度、增加卡数或选择更大显存设备后再做算子验证。") elif runtime_budget_gib > total_device_memory_gb * 0.85: verdict = "tight_memory" recommendations.append("显存余量偏小，先用低并发和较短上下文完成冒烟测试。") else: verdict = "preflight_feasible" recommendations.append("显存预检通过，可进入目标后端算子兼容性测试。") recommendations.extend( [ "确认目标芯片 SDK、PyTorch 后端与推理框架版本组合。", "检查 Attention、RMSNorm/LayerNorm、RoPE 和自定义算子支持情况。", "先跑单请求冒烟测试，再分别测首 Token 延迟、吞吐和峰值显存。", ] ) return { "agent": AGENT_NAME, "version": AGENT_VERSION, "model": model_name, "target": { "hardware": hardware or None, "cards": cards, "device_memory_gb": device_memory_gb, }, "inputs": { "architecture": architecture or None, "parameters_b": parameters_b, "dtype": dtype, "context_length": context_length, }, "verdict": verdict, "missing_fields": missing_fields, "estimates": estimates, "risks": risks, "recommendations": recommendations, "disclaimer": "这是部署前预检，不代表目标芯片已完成真实运行验证。", } class Handler(BaseHTTPRequestHandler): server_version = "ModelHubCompatAgent/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 - BaseHTTPRequestHandler contract if self.path == "/health": self._json({"status": "ok", "agent": AGENT_NAME, "version": AGENT_VERSION}) return if self.path == "/": self._json( { "name": AGENT_NAME, "version": AGENT_VERSION, "description": "信创模型兼容性与显存预检", "strategy_id_present": bool(STRATEGY_ID), "endpoints": ["GET /health", "POST /analyze", "POST /task"], "external_writes": False, } ) 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: self._json({"error": "payload_too_large"}, 413) return payload = json.loads(self.rfile.read(length)) self._json(analyze(payload)) except (json.JSONDecodeError, UnicodeDecodeError): self._json({"error": "invalid_json", "message": "请求正文必须是有效 JSON"}, 400) except ValidationError as exc: self._json({"error": "validation_error", "message": str(exc)}, 422) def log_message(self, _format: str, *_args: Any) -> None: return STOP = threading.Event() def _handle_signal(signum: int, _frame: Any) -> None: print(f"received signal {signum}; shutting down", flush=True) 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 print(f"{AGENT_NAME} {AGENT_VERSION} listening on 0.0.0.0:{PORT}", flush=True) while not STOP.is_set(): server.handle_request() server.server_close() if __name__ == "__main__": main()