feat(#41): 模型模板注册/加载/版本机制(多版本+阶段dev/staging/prod+提升+回滚+审计+持久化,PRD 5.3 ③模型模板注册加载版本机制)

This commit is contained in:
2026-08-05 01:19:01 +08:00
parent 793dd0a3b8
commit b6392525ce
7 changed files with 837 additions and 0 deletions
+71
View File
@@ -0,0 +1,71 @@
# iAOP-Core · 模型框架层(AI Model Framework)
对应 PRD 5.3「③ AI 模型框架」与 EPIC #5「内核平台化改造」。
本层把化工 AI 的「模型资产」统一注册、版本化、阶段化管理:模型 / 配方 /
估计器是资产,需要注册表统一管理——注册、加载、版本、阶段(灰度)、回滚、审计。
## 当前已交付
| 模块 | 对应 issue | PRD 5.3 模型 | 说明 |
|------|-----------|-------------|------|
| `template_registry` | #41 | ③ 模型模板注册/加载/版本机制 | 多版本 + 阶段(dev/staging/prod) + 提升 + 回滚 + 审计 + 持久化 |
## 模板注册表(`template_registry.py`)
模板化之后,模型 / 配方是「资产」,需要注册表统一管理:
- **注册**:登记一个模型模板(含主干、超参、特征列、指标、版本);
- **加载**:按 name + version/别名/stage 取出;
- **版本**:同模板多版本共存,可回滚、可审计;
- **阶段**:版本带 stage 标签(dev/staging/prod),灰度发布可控。
### 核心组件
- **`ModelTemplate`**:模型模板数据对象(name/version/backbone/hyperparams/
feature_columns/metrics/stage),不可变、可序列化往返。
- **`TemplateRegistry`**:注册表核心 API:
- `register`(校验完整性 + 同版本号拒重复,force 可覆盖)
- `get`(按 version / stage / 默认最新加载)
- `promote`(dev→staging→prod 逐级提升)
- `rollback`(stage 指针回退,保留历史可审计)
- `set_stage`(直接设 stage,紧急回滚)
- `list_versions` / `list_by_stage` / `stage_pointer` / `history`
- `save` / `load`(JSON 持久化,重启恢复)
- **`Stage`**:阶段枚举(DEV / STAGING / PROD),PRD 灰度三段制。
- **`is_valid_version`**:语义化版本号校验(v1 / 1.0.0 / v1.2-rc1)。
### 快速开始
```python
from template_registry import ModelTemplate, Stage, TemplateRegistry
reg = TemplateRegistry()
reg.register(ModelTemplate(
name="ti-quality", version="v1.0", backbone="quality_forecast",
feature_columns=("furnace_temp", "cl2_flow"),
metrics={"accuracy": 0.91}, stage="dev"))
reg.register(ModelTemplate(name="ti-quality", version="v1.1",
backbone="quality_forecast", metrics={"accuracy": 0.94}))
reg.promote("ti-quality", "v1.1") # dev -> staging
reg.promote("ti-quality", "v1.1") # staging -> prod
reg.rollback("ti-quality", Stage.PROD, "v1.0") # 紧急回滚
serving = reg.get("ti-quality", stage=Stage.PROD) # 按阶段加载
reg.save("registry.json") # 持久化审计
```
## 测试
```bash
cd core/model-framework
python -m unittest discover -s tests -v
python _sanity_check.py
```
## 与规划模块的关系
接口风格对齐 #34(Model Recipe)、#36(quality_forecast)、#38
(cross_process_optimizer)、#40(pipeline)。本模块是 #40 `ModelRegistry`
的完整版(多阶段 + 回滚 + 持久化),#40 的 `RegisterStep` 未来可直接对接本
注册表,业务侧零改动。
+22
View File
@@ -0,0 +1,22 @@
# -*- coding: utf-8 -*-
"""iAOP-Core · 模型框架层(AI Model Framework)。
对应 PRD 5.3「③ AI 模型框架」与 EPIC #5(内核平台化改造)。
当前已交付(自包含,不依赖未合并分支):
- ``template_registry``:模型模板注册 / 加载 / 版本机制(多版本 + 阶段 dev/
staging/prod + 提升 + 回滚 + 审计 + JSON 持久化),issue #41。
规划(待相关 PR 合入后无缝对接,业务侧零改动):
- ``pipeline``(issue #40,PR #107)、``cross_process_optimizer``(issue #38,
PR #106)等可注册为本注册表的具名模板。
"""
from model_framework.template_registry import ( # noqa: F401
ALLOWED_BACKBONES,
ModelTemplate,
Stage,
TemplateRegistry,
TemplateRegistryError,
is_valid_version,
next_stage,
)
+90
View File
@@ -0,0 +1,90 @@
# -*- coding: utf-8 -*-
"""模型模板注册 / 加载 / 版本机制 sanity 检查(issue #41)。
验证 PRD 5.3 模板资产管理的端到端能力:
1. 注册多版本 → 阶段提升(dev→staging→prod)→ 回滚;
2. 按版本 / 阶段 / 默认加载;
3. JSON 持久化往返一致;
4. 审计日志完整。
"""
import os
import sys
import tempfile
HERE = os.path.dirname(os.path.abspath(__file__))
if HERE not in sys.path:
sys.path.insert(0, HERE)
from template_registry import ( # noqa: E402
ModelTemplate, Stage, TemplateRegistry, is_valid_version, next_stage)
def main() -> int:
failures = []
# 版本号校验
for v in ["v1", "1.0.0", "v2.1-rc3"]:
if not is_valid_version(v):
failures.append(f"版本号 {v!r} 应合法")
# 注册 + 提升 + 回滚
try:
reg = TemplateRegistry()
reg.register(ModelTemplate(
name="ti-quality", version="v1.0", backbone="quality_forecast",
feature_columns=("furnace_temp", "cl2_flow"),
metrics={"accuracy": 0.91}, stage="dev"))
reg.register(ModelTemplate(
name="ti-quality", version="v1.1", backbone="quality_forecast",
feature_columns=("furnace_temp", "cl2_flow", "impurity_fe"),
metrics={"accuracy": 0.94}, stage="dev"))
assert reg.list_versions("ti-quality") == ["v1.0", "v1.1"]
# v1.1 逐级提升到 prod
reg.promote("ti-quality", "v1.1") # dev->staging
reg.promote("ti-quality", "v1.1") # staging->prod
assert reg.stage_pointer("ti-quality", Stage.PROD) == "v1.1"
# 紧急回滚 prod 到 v1.0
reg.rollback("ti-quality", Stage.PROD, "v1.0")
assert reg.stage_pointer("ti-quality", Stage.PROD) == "v1.0"
assert "v1.1" in reg.list_versions("ti-quality") # 历史保留
# 按阶段加载
serving = reg.get("ti-quality", stage=Stage.PROD)
assert serving.version == "v1.0", f"prod 应为 v1.0,实为 {serving.version}"
print(f"[ti-quality] 版本={reg.list_versions('ti-quality')} "
f"prod={reg.stage_pointer('ti-quality', Stage.PROD)}")
print(f" v1.0 metrics={reg.get('ti-quality','v1.0').metrics}")
print(f" v1.1 metrics={reg.get('ti-quality','v1.1').metrics}")
print(f" 审计日志 {len(reg.history('ti-quality'))} 条")
except Exception as exc: # noqa: BLE001
failures.append(f"注册/提升/回滚流程失败:{exc}")
# 持久化往返
try:
with tempfile.NamedTemporaryFile(
mode="w", suffix=".json", delete=False, encoding="utf-8") as fh:
path = fh.name
reg.save(path)
reg2 = TemplateRegistry.load(path)
assert reg2.list_versions("ti-quality") == ["v1.0", "v1.1"]
assert reg2.stage_pointer("ti-quality", Stage.PROD) == "v1.0"
print(f"[persist] save/load 往返一致,恢复 {len(reg2)} 个模板")
os.unlink(path)
except Exception as exc: # noqa: BLE001
failures.append(f"持久化往返失败:{exc}")
if failures:
print("\n失败项:")
for f in failures:
print(f" ✗ {f}")
return 1
print("\n✓ template_registry sanity check 通过")
return 0
if __name__ == "__main__":
sys.exit(main())
+395
View File
@@ -0,0 +1,395 @@
# -*- 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
+1
View File
@@ -0,0 +1 @@
# -*- coding: utf-8 -*-
+23
View File
@@ -0,0 +1,23 @@
# -*- coding: utf-8 -*-
"""测试引导:把连字符目录 ``core/model-framework`` 加载为可导入包
``model_framework``,使测试可 ``from model_framework import ...``。
"""
import importlib.util
import os
import sys
PKG_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
def _load_package(name: str, path: str) -> None:
if name in sys.modules:
return
init_py = os.path.join(path, "__init__.py")
spec = importlib.util.spec_from_file_location(
name, init_py, submodule_search_locations=[path])
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module)
_load_package("model_framework", PKG_DIR)
@@ -0,0 +1,235 @@
# -*- coding: utf-8 -*-
"""模型模板注册 / 加载 / 版本机制单元测试(issue #41)。
覆盖:
- 版本号语义校验 ``is_valid_version``;
- Stage 枚举与 ``next_stage`` 阶段提升顺序;
- ModelTemplate 构造校验(name/version/backbone/stage)+ 序列化往返;
- TemplateRegistry:注册(拒重复 / force 覆盖)、加载(version/stage/默认)、
阶段提升 promote、回滚 rollback、set_stage、查询(list_*)、审计日志、
JSON 持久化 save/load 往返一致性。
"""
import json
import os
import sys
import tempfile
import time
import unittest
HERE = os.path.dirname(os.path.abspath(__file__))
if HERE not in sys.path:
sys.path.insert(0, HERE)
import _bootstrap # noqa: E402 加载 model_framework 包
from model_framework import ( # noqa: E402
ModelTemplate,
Stage,
TemplateRegistry,
TemplateRegistryError,
is_valid_version,
next_stage,
)
class TestVersionValidation(unittest.TestCase):
def test_valid_versions(self):
for v in ["v1", "1.0", "v1.2.3", "1.2.3", "v1.0-rc1", "v2.0.0+build5"]:
self.assertTrue(is_valid_version(v), f"应合法:{v}")
def test_invalid_versions(self):
for v in ["", "v", "abc", "v1.x", "1..2", None, "v 1"]:
self.assertFalse(is_valid_version(v), f"应非法:{v!r}")
class TestStage(unittest.TestCase):
def test_from_str(self):
self.assertEqual(Stage.from_str("dev"), Stage.DEV)
self.assertEqual(Stage.from_str("PROD"), Stage.PROD)
def test_from_str_invalid(self):
with self.assertRaises(TemplateRegistryError):
Stage.from_str("qa")
def test_next_stage(self):
self.assertEqual(next_stage(Stage.DEV), Stage.STAGING)
self.assertEqual(next_stage(Stage.STAGING), Stage.PROD)
self.assertIsNone(next_stage(Stage.PROD))
class TestModelTemplate(unittest.TestCase):
def test_construct_minimal(self):
t = ModelTemplate(name="m", version="v1")
self.assertEqual(t.backbone, "generic")
self.assertEqual(t.stage, Stage.DEV)
def test_rejects_empty_name(self):
with self.assertRaises(TemplateRegistryError):
ModelTemplate(name="", version="v1")
def test_rejects_bad_version(self):
with self.assertRaises(TemplateRegistryError):
ModelTemplate(name="m", version="abc")
def test_rejects_bad_backbone(self):
with self.assertRaises(TemplateRegistryError):
ModelTemplate(name="m", version="v1", backbone="magic")
def test_accepts_known_backbones(self):
for b in ("quality_forecast", "anomaly_detection",
"cross_process_opt", "recipe_opt", "generic"):
ModelTemplate(name="m", version="v1", backbone=b)
def test_roundtrip(self):
t = ModelTemplate(
name="qa-model", version="v1.2.0", backbone="quality_forecast",
hyperparams={"lr": 0.1}, feature_columns=("a", "b"),
target_column="y", metrics={"accuracy": 0.93},
stage="prod", description="d", extra={"k": "v"})
t2 = ModelTemplate.from_dict(json.loads(json.dumps(t.to_dict())))
self.assertEqual(t, t2)
self.assertEqual(t2.stage, Stage.PROD)
self.assertEqual(t2.metrics["accuracy"], 0.93)
class TestRegistryRegisterLoad(unittest.TestCase):
def setUp(self):
self.reg = TemplateRegistry()
self.t1 = ModelTemplate(name="m", version="v1", backbone="quality_forecast")
self.t2 = ModelTemplate(name="m", version="v2", backbone="quality_forecast")
def test_register_and_get_by_version(self):
self.reg.register(self.t1)
self.assertEqual(self.reg.get("m", "v1").version, "v1")
def test_register_duplicate_rejected(self):
self.reg.register(self.t1)
with self.assertRaises(TemplateRegistryError):
self.reg.register(self.t1)
def test_register_force_overwrites(self):
self.reg.register(self.t1)
t1_updated = ModelTemplate(
name="m", version="v1", description="updated")
self.reg.register(t1_updated, force=True)
self.assertEqual(self.reg.get("m", "v1").description, "updated")
def test_get_missing_name(self):
with self.assertRaises(TemplateRegistryError):
self.reg.get("nope")
def test_get_missing_version(self):
self.reg.register(self.t1)
with self.assertRaises(TemplateRegistryError):
self.reg.get("m", "v99")
def test_get_default_latest(self):
self.reg.register(self.t1)
time.sleep(0.01)
self.reg.register(self.t2)
self.assertEqual(self.reg.get("m").version, "v2")
def test_get_by_stage_pointer(self):
self.reg.register(self.t1)
# 新注册默认进 dev
self.assertEqual(self.reg.get("m", stage=Stage.DEV).version, "v1")
with self.assertRaises(TemplateRegistryError):
self.reg.get("m", stage=Stage.PROD)
class TestPromoteRollback(unittest.TestCase):
def setUp(self):
self.reg = TemplateRegistry()
self.reg.register(ModelTemplate(name="m", version="v1"))
self.reg.register(ModelTemplate(name="m", version="v2"))
def test_promote_chain(self):
self.reg.promote("m", "v1") # dev -> staging
self.assertEqual(self.reg.stage_pointer("m", Stage.STAGING), "v1")
self.reg.promote("m", "v1") # staging -> prod
self.assertEqual(self.reg.stage_pointer("m", Stage.PROD), "v1")
def test_promote_prod_raises(self):
self.reg.promote("m", "v1")
self.reg.promote("m", "v1") # 到 prod
with self.assertRaises(TemplateRegistryError):
self.reg.promote("m", "v1") # prod 无法继续
def test_rollback_stage_pointer(self):
self.reg.promote("m", "v2") # v2 -> staging
self.reg.promote("m", "v2") # v2 -> prod
# 回滚 prod 到 v1
self.reg.rollback("m", Stage.PROD, "v1")
self.assertEqual(self.reg.stage_pointer("m", Stage.PROD), "v1")
# v2 版本本身仍在(可审计)
self.assertIn("v2", self.reg.list_versions("m"))
def test_set_stage_direct(self):
self.reg.set_stage("m", "v1", Stage.PROD)
self.assertEqual(self.reg.get("m", "v1").stage, Stage.PROD)
self.assertEqual(self.reg.stage_pointer("m", Stage.PROD), "v1")
class TestQueries(unittest.TestCase):
def test_list_names_and_versions(self):
reg = TemplateRegistry()
reg.register(ModelTemplate(name="a", version="v1"))
reg.register(ModelTemplate(name="a", version="v2"))
reg.register(ModelTemplate(name="b", version="v1"))
self.assertEqual(reg.list_names(), ["a", "b"])
self.assertEqual(reg.list_versions("a"), ["v1", "v2"])
self.assertIn("a", reg)
self.assertNotIn("c", reg)
self.assertEqual(len(reg), 3)
def test_list_by_stage(self):
reg = TemplateRegistry()
reg.register(ModelTemplate(name="m", version="v1"))
reg.register(ModelTemplate(name="m", version="v2"))
reg.promote("m", "v2") # v2 -> staging
self.assertEqual(reg.list_by_stage("m", Stage.DEV), ["v1"])
self.assertEqual(reg.list_by_stage("m", Stage.STAGING), ["v2"])
class TestHistoryAndPersist(unittest.TestCase):
def test_history_logged(self):
reg = TemplateRegistry()
reg.register(ModelTemplate(name="m", version="v1"))
reg.promote("m", "v1")
h = reg.history("m")
actions = [e["action"] for e in h]
self.assertIn("register", actions)
self.assertIn("promote", actions)
def test_history_filter_by_name(self):
reg = TemplateRegistry()
reg.register(ModelTemplate(name="a", version="v1"))
reg.register(ModelTemplate(name="b", version="v1"))
self.assertEqual(len(reg.history("a")), 1)
self.assertEqual(len(reg.history("b")), 1)
def test_save_load_roundtrip(self):
reg = TemplateRegistry()
reg.register(ModelTemplate(
name="m", version="v1", backbone="quality_forecast",
metrics={"accuracy": 0.9}, stage="dev"))
reg.promote("m", "v1")
with tempfile.NamedTemporaryFile(
mode="w", suffix=".json", delete=False, encoding="utf-8") as fh:
path = fh.name
try:
reg.save(path)
reg2 = TemplateRegistry.load(path)
self.assertEqual(reg2.list_names(), ["m"])
self.assertEqual(reg2.get("m", "v1").backbone, "quality_forecast")
self.assertEqual(reg2.get("m", "v1").metrics["accuracy"], 0.9)
# 阶段指针恢复
self.assertEqual(reg2.stage_pointer("m", Stage.STAGING), "v1")
# 审计日志恢复
self.assertTrue(len(reg2.history("m")) >= 2)
finally:
os.unlink(path)
if __name__ == "__main__":
unittest.main()