396 lines
16 KiB
Python
396 lines
16 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""模型模板注册 / 加载 / 版本机制（PRD 5.3 ③ 模型框架）。
|
||
|
||
对应 issue #41（父 EPIC #5「③ AI 模型框架 配置化重构」、PRD 5.3
|
||
「③ 模型模板注册 / 加载 / 版本机制」）。
|
||
|
||
PRD 5.3 的核心诉求
|
||
------------------
|
||
|
||
模板化之后，模型 / 配方 / 估计器是「资产」，需要**注册表**统一管理：
|
||
- **注册**：把一个模型模板（含主干、超参、特征列、指标、版本）登记入库；
|
||
- **加载**：按名字 + 版本（或别名）取出可用的模板实例；
|
||
- **版本**：同一模板多版本共存，可回滚、可审计；
|
||
- **阶段**：版本带 stage 标签（dev / staging / prod），灰度发布可控；
|
||
- **校验**：注册时校验模板完整性，防止脏数据进库。
|
||
|
||
本模块交付一个独立的 ``TemplateRegistry``，自包含、不依赖未合并分支，是
|
||
issue #40 ``ModelRegistry`` 的深化（完整版本 + 阶段 + 回滚 + 审计 + 持久化）。
|
||
|
||
本模块交付什么
|
||
--------------
|
||
|
||
1. **``ModelTemplate``**：模型模板数据对象（name / version / backbone /
|
||
hyperparams / feature_columns / metrics / stage / extra），不可变、可序列化。
|
||
2. **``TemplateRegistry``**：模板注册表核心 API：
|
||
- ``register``：注册一个版本（校验完整性 + 同版本号拒重复）；
|
||
- ``get``：按 name + version/别名加载；
|
||
- ``promote``：把版本提升到下一 stage（dev→staging→prod）；
|
||
- ``rollback``：把某 stage 回滚到指定版本；
|
||
- ``list_versions`` / ``list_by_stage`` / ``history``：查询；
|
||
- ``save`` / ``load``：JSON 持久化（文件 / 目录），便于审计与重启恢复。
|
||
3. **``Stage``**：阶段枚举（DEV / STAGING / PROD），PRD 灰度发布三段制。
|
||
4. **校验**：注册时校验 name/version/backbone 非空、版本号语义合法、
|
||
stage 合法，拒绝脏模板。
|
||
|
||
零外部强依赖
|
||
------------
|
||
|
||
纯 Python 实现，无 numpy/sklearn 依赖，CI 可加载与校验。
|
||
|
||
与 issue #34 / #36 / #38 / #40 的关系
|
||
--------------------------------------
|
||
|
||
接口风格对齐 #34 声明式数据对象、#36 / #38 ``Recipe``、#40 ``ModelRegistry``。
|
||
本模块是 #40 ``ModelRegistry`` 的完整版（多阶段 + 回滚 + 持久化），#40 的
|
||
``RegisterStep`` 未来可直接对接本注册表，业务侧零改动。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import re
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from enum import Enum
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
__all__ = [
|
||
"ModelTemplate",
|
||
"TemplateRegistry",
|
||
"Stage",
|
||
"TemplateRegistryError",
|
||
"is_valid_version",
|
||
"next_stage",
|
||
]
|
||
|
||
|
||
class TemplateRegistryError(Exception):
|
||
"""模板注册表统一异常（模板非法 / 版本冲突 / 版本不存在 / 阶段非法）。"""
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 阶段（Stage）：灰度发布三段制
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class Stage(str, Enum):
|
||
"""模型版本阶段：DEV（开发）→ STAGING（预发）→ PROD（生产）。"""
|
||
|
||
DEV = "dev"
|
||
STAGING = "staging"
|
||
PROD = "prod"
|
||
|
||
@classmethod
|
||
def from_str(cls, s: str) -> "Stage":
|
||
try:
|
||
return cls(s.lower())
|
||
except ValueError:
|
||
raise TemplateRegistryError(
|
||
f"非法阶段 {s!r}，允许：{[s.value for s in Stage]}")
|
||
|
||
|
||
# 阶段提升顺序：dev → staging → prod
|
||
_STAGE_ORDER = [Stage.DEV, Stage.STAGING, Stage.PROD]
|
||
|
||
|
||
def next_stage(stage: Stage) -> Optional[Stage]:
|
||
"""返回下一阶段；PROD 已是最高返回 None。"""
|
||
try:
|
||
idx = _STAGE_ORDER.index(stage)
|
||
except ValueError:
|
||
return None
|
||
if idx + 1 >= len(_STAGE_ORDER):
|
||
return None
|
||
return _STAGE_ORDER[idx + 1]
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 版本号语义校验（语义化版本 vMAJOR.MINOR.PATCH 或简单的 vN / N）
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_VERSION_RE = re.compile(r"^v?\d+(\.\d+)*([\-+][0-9A-Za-z.\-]+)?$")
|
||
|
||
|
||
def is_valid_version(version: str) -> bool:
|
||
"""校验版本号是否合法（v1 / 1.0 / v1.2.3 / v1.0-rc1 等）。"""
|
||
if not isinstance(version, str) or not version.strip():
|
||
return False
|
||
return bool(_VERSION_RE.match(version.strip()))
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 模型模板数据对象
|
||
# ---------------------------------------------------------------------------
|
||
|
||
#: 允许的主干类型（对齐 PRD 5.3 四类模型模板 + 通用）
|
||
ALLOWED_BACKBONES = (
|
||
"quality_forecast", # ① 质量预测（issue #36）
|
||
"anomaly_detection", # 异常检测（issue #37）
|
||
"cross_process_opt", # ③ 跨工序寻优（issue #38）
|
||
"recipe_opt", # ② 配方优化
|
||
"generic", # 通用
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ModelTemplate:
|
||
"""模型模板：描述一个可加载的模型资产（主干 + 超参 + 特征 + 指标 + 版本）。
|
||
|
||
不可变数据对象，``to_dict`` / ``from_dict`` 可序列化往返，便于持久化与
|
||
配置台展示。注册表以 (name, version) 为主键管理多个模板实例。
|
||
"""
|
||
|
||
name: str
|
||
version: str
|
||
backbone: str = "generic"
|
||
hyperparams: Dict[str, Any] = field(default_factory=dict)
|
||
feature_columns: Tuple[str, ...] = field(default_factory=tuple)
|
||
target_column: str = ""
|
||
metrics: Dict[str, float] = field(default_factory=dict)
|
||
stage: Stage = Stage.DEV
|
||
registered_at: float = field(default_factory=time.time)
|
||
description: str = ""
|
||
extra: Dict[str, Any] = field(default_factory=dict)
|
||
|
||
def __post_init__(self) -> None:
|
||
if not self.name:
|
||
raise TemplateRegistryError("ModelTemplate 缺少 name")
|
||
if not is_valid_version(self.version):
|
||
raise TemplateRegistryError(
|
||
f"非法版本号 {self.version!r}（例：v1 / 1.0.0 / v1.2-rc1）")
|
||
if self.backbone not in ALLOWED_BACKBONES:
|
||
raise TemplateRegistryError(
|
||
f"非法主干 {self.backbone!r}，允许：{ALLOWED_BACKBONES}")
|
||
if not isinstance(self.stage, Stage):
|
||
# 允许从字符串构造（dataclass frozen 用 object.__setattr__）
|
||
object.__setattr__(self, "stage", Stage.from_str(str(self.stage)))
|
||
|
||
@property
|
||
def key(self) -> Tuple[str, str]:
|
||
return (self.name, self.version)
|
||
|
||
def to_dict(self) -> Dict[str, Any]:
|
||
return {
|
||
"name": self.name,
|
||
"version": self.version,
|
||
"backbone": self.backbone,
|
||
"hyperparams": dict(self.hyperparams),
|
||
"feature_columns": list(self.feature_columns),
|
||
"target_column": self.target_column,
|
||
"metrics": dict(self.metrics),
|
||
"stage": self.stage.value,
|
||
"registered_at": self.registered_at,
|
||
"description": self.description,
|
||
"extra": dict(self.extra),
|
||
}
|
||
|
||
@classmethod
|
||
def from_dict(cls, d: Dict[str, Any]) -> "ModelTemplate":
|
||
try:
|
||
return cls(
|
||
name=d["name"],
|
||
version=d["version"],
|
||
backbone=d.get("backbone", "generic"),
|
||
hyperparams=dict(d.get("hyperparams", {})),
|
||
feature_columns=tuple(d.get("feature_columns", [])),
|
||
target_column=d.get("target_column", ""),
|
||
metrics={k: float(v) for k, v in d.get("metrics", {}).items()},
|
||
stage=Stage.from_str(d.get("stage", "dev")),
|
||
registered_at=float(d.get("registered_at", time.time())),
|
||
description=d.get("description", ""),
|
||
extra=dict(d.get("extra", {})),
|
||
)
|
||
except KeyError as exc: # pragma: no cover
|
||
raise TemplateRegistryError(f"模板缺少必填字段：{exc}") from exc
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 模板注册表
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class TemplateRegistry:
|
||
"""模型模板注册表：注册 / 加载 / 版本 / 阶段提升 / 回滚 / 持久化。
|
||
|
||
* 以 ``name`` 维护多个 ``version``；
|
||
* 每个版本带 ``Stage``（dev/staging/prod），``promote`` 逐级提升；
|
||
* ``rollback`` 把某 stage 指针回退到指定版本（保留历史，可审计）；
|
||
* ``save`` / ``load`` JSON 持久化（单文件或目录每模板一文件）。
|
||
"""
|
||
|
||
def __init__(self) -> None:
|
||
self._templates: Dict[str, Dict[str, ModelTemplate]] = {}
|
||
# 每个 name 的阶段指针：{stage: version}
|
||
self._stage_pointers: Dict[str, Dict[Stage, str]] = {}
|
||
# 审计日志：[(ts, action, name, version, detail)]
|
||
self._history: List[Tuple[float, str, str, str, str]] = []
|
||
|
||
# -- 注册 / 校验 --------------------------------------------------------
|
||
|
||
def register(self, template: ModelTemplate, *, force: bool = False) -> ModelTemplate:
|
||
"""注册一个模板版本。
|
||
|
||
* 同 (name, version) 默认拒绝重复（``force=True`` 可覆盖）；
|
||
* 注册后自动成为该模板 dev 阶段指针（若该 stage 无指针）；
|
||
* 记入审计日志。
|
||
"""
|
||
versions = self._templates.setdefault(template.name, {})
|
||
if template.version in versions and not force:
|
||
raise TemplateRegistryError(
|
||
f"{template.name}@{template.version} 已存在（force=True 可覆盖）")
|
||
versions[template.version] = template
|
||
pointers = self._stage_pointers.setdefault(template.name, {})
|
||
# 新注册默认进 dev；若 dev 无指针则指向它
|
||
if Stage.DEV not in pointers:
|
||
pointers[Stage.DEV] = template.version
|
||
self._log("register", template.name, template.version,
|
||
f"stage={template.stage.value}")
|
||
return template
|
||
|
||
# -- 加载 ---------------------------------------------------------------
|
||
|
||
def get(self, name: str,
|
||
version: Optional[str] = None,
|
||
stage: Optional[Stage] = None) -> ModelTemplate:
|
||
"""按 name + version 或 name + stage 加载模板。
|
||
|
||
* ``version`` 优先；其次 ``stage``（取该阶段指针）；
|
||
* 都不提供则取该模板最新注册版本。
|
||
"""
|
||
versions = self._templates.get(name)
|
||
if not versions:
|
||
raise TemplateRegistryError(f"模板 {name!r} 未注册")
|
||
if version is not None:
|
||
if version not in versions:
|
||
raise TemplateRegistryError(
|
||
f"{name!r} 无版本 {version!r}（可用：{sorted(versions)}）")
|
||
return versions[version]
|
||
if stage is not None:
|
||
ptr = self._stage_pointers.get(name, {}).get(stage)
|
||
if ptr is None:
|
||
raise TemplateRegistryError(
|
||
f"{name!r} 无 {stage.value} 阶段指针")
|
||
return versions[ptr]
|
||
# 默认：按注册时间最新；时间相同时按版本号字典序最新（确定性）
|
||
latest = max(versions.values(),
|
||
key=lambda t: (t.registered_at, t.version))
|
||
return latest
|
||
|
||
# -- 阶段提升 / 回滚 -----------------------------------------------------
|
||
|
||
def promote(self, name: str, version: str) -> ModelTemplate:
|
||
"""把指定版本提升到下一 stage（dev→staging→prod）。
|
||
|
||
提升后该版本成为新 stage 的指针版本。
|
||
"""
|
||
tpl = self.get(name, version)
|
||
nxt = next_stage(tpl.stage)
|
||
if nxt is None:
|
||
raise TemplateRegistryError(
|
||
f"{name}@{version} 已在 PROD，无法继续提升")
|
||
# 更新该模板版本的 stage（需重建不可变对象）
|
||
new_tpl = ModelTemplate(**{**tpl.to_dict(), "stage": nxt.value})
|
||
self._templates[name][version] = new_tpl
|
||
self._stage_pointers.setdefault(name, {})[nxt] = version
|
||
self._log("promote", name, version, f"{tpl.stage.value}->{nxt.value}")
|
||
return new_tpl
|
||
|
||
def rollback(self, name: str, stage: Stage, version: str) -> ModelTemplate:
|
||
"""把某 stage 的指针回退到指定版本（保留历史版本，可审计）。"""
|
||
tpl = self.get(name, version)
|
||
if not isinstance(stage, Stage):
|
||
stage = Stage.from_str(str(stage))
|
||
self._stage_pointers.setdefault(name, {})[stage] = version
|
||
self._log("rollback", name, version, f"stage={stage.value}")
|
||
return tpl
|
||
|
||
def set_stage(self, name: str, version: str, stage: Stage) -> ModelTemplate:
|
||
"""直接把某版本设到指定 stage（覆盖 promote 的逐级约束，用于紧急回滚）。"""
|
||
tpl = self.get(name, version)
|
||
if not isinstance(stage, Stage):
|
||
stage = Stage.from_str(str(stage))
|
||
new_tpl = ModelTemplate(**{**tpl.to_dict(), "stage": stage.value})
|
||
self._templates[name][version] = new_tpl
|
||
self._stage_pointers.setdefault(name, {})[stage] = version
|
||
self._log("set_stage", name, version, f"stage={stage.value}")
|
||
return new_tpl
|
||
|
||
# -- 查询 ---------------------------------------------------------------
|
||
|
||
def list_names(self) -> List[str]:
|
||
return sorted(self._templates.keys())
|
||
|
||
def list_versions(self, name: str) -> List[str]:
|
||
return sorted(self._templates.get(name, {}).keys())
|
||
|
||
def list_by_stage(self, name: str, stage: Stage) -> List[str]:
|
||
"""列出某模板处于指定 stage 的所有版本号。"""
|
||
if not isinstance(stage, Stage):
|
||
stage = Stage.from_str(str(stage))
|
||
return sorted(v for v, t in self._templates.get(name, {}).items()
|
||
if t.stage == stage)
|
||
|
||
def stage_pointer(self, name: str, stage: Stage) -> Optional[str]:
|
||
if not isinstance(stage, Stage):
|
||
stage = Stage.from_str(str(stage))
|
||
return self._stage_pointers.get(name, {}).get(stage)
|
||
|
||
def history(self, name: Optional[str] = None) -> List[Dict[str, Any]]:
|
||
"""返回审计日志（可按 name 过滤）。"""
|
||
out = []
|
||
for ts, action, n, ver, detail in self._history:
|
||
if name is not None and n != name:
|
||
continue
|
||
out.append({"time": ts, "action": action, "name": n,
|
||
"version": ver, "detail": detail})
|
||
return out
|
||
|
||
# -- 持久化 --------------------------------------------------------------
|
||
|
||
def save(self, path: str) -> None:
|
||
"""把整个注册表序列化到 JSON 文件（含模板 + 阶段指针 + 审计日志）。"""
|
||
data = {
|
||
"templates": {
|
||
name: {ver: t.to_dict() for ver, t in vers.items()}
|
||
for name, vers in self._templates.items()
|
||
},
|
||
"stage_pointers": {
|
||
name: {s.value: v for s, v in ptrs.items()}
|
||
for name, ptrs in self._stage_pointers.items()
|
||
},
|
||
"history": [
|
||
{"time": ts, "action": a, "name": n, "version": v, "detail": d}
|
||
for ts, a, n, v, d in self._history
|
||
],
|
||
}
|
||
with open(path, "w", encoding="utf-8") as fh:
|
||
json.dump(data, fh, ensure_ascii=False, indent=2)
|
||
|
||
@classmethod
|
||
def load(cls, path: str) -> "TemplateRegistry":
|
||
"""从 JSON 文件恢复注册表。"""
|
||
with open(path, "r", encoding="utf-8") as fh:
|
||
data = json.load(fh)
|
||
reg = cls()
|
||
for name, vers in data.get("templates", {}).items():
|
||
for ver, td in vers.items():
|
||
reg._templates.setdefault(name, {})[ver] = ModelTemplate.from_dict(td)
|
||
for name, ptrs in data.get("stage_pointers", {}).items():
|
||
for s_str, v in ptrs.items():
|
||
reg._stage_pointers.setdefault(name, {})[Stage.from_str(s_str)] = v
|
||
for h in data.get("history", []):
|
||
reg._history.append(
|
||
(h["time"], h["action"], h["name"], h["version"], h["detail"]))
|
||
return reg
|
||
|
||
# -- 内部 ----------------------------------------------------------------
|
||
|
||
def _log(self, action: str, name: str, version: str, detail: str) -> None:
|
||
self._history.append((time.time(), action, name, version, detail))
|
||
|
||
def __len__(self) -> int:
|
||
return sum(len(v) for v in self._templates.values())
|
||
|
||
def __contains__(self, name: str) -> bool:
|
||
return name in self._templates
|