feat(ml): add model metadata persistence with version hashing
SHA-256 based version hashing for ONNX artifacts and async metadata persistence that replaces previous active model entries.
This commit is contained in:
parent
0a9cf12212
commit
ee25636437
|
|
@ -0,0 +1,82 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
metadata.py
|
||||
"""
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.models.model_metadata import ModelMetadata
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MODEL_TYPES: dict[str, str] = {
|
||||
"ae.onnx": "autoencoder",
|
||||
"rf.onnx": "random_forest",
|
||||
"if.onnx": "isolation_forest",
|
||||
}
|
||||
|
||||
VERSION_HASH_LENGTH = 12
|
||||
|
||||
|
||||
def compute_model_version(artifact_path: Path) -> str:
|
||||
"""
|
||||
Compute a 12-char hex version string from the SHA-256 hash of a file
|
||||
"""
|
||||
sha = hashlib.sha256(artifact_path.read_bytes())
|
||||
return sha.hexdigest()[:VERSION_HASH_LENGTH]
|
||||
|
||||
|
||||
async def save_model_metadata(
|
||||
session: AsyncSession,
|
||||
model_dir: Path,
|
||||
training_samples: int,
|
||||
metrics: dict[str, object],
|
||||
mlflow_run_id: str | None = None,
|
||||
threshold: float | None = None,
|
||||
) -> list[ModelMetadata]:
|
||||
"""
|
||||
Persist metadata for all 3 model types, deactivating previous active versions
|
||||
"""
|
||||
rows: list[ModelMetadata] = []
|
||||
|
||||
for filename, model_type in MODEL_TYPES.items():
|
||||
artifact_path = model_dir / filename
|
||||
version = compute_model_version(artifact_path)
|
||||
|
||||
result = await session.execute(
|
||||
select(ModelMetadata).where(
|
||||
ModelMetadata.model_type == model_type,
|
||||
ModelMetadata.is_active == True, # noqa: E712
|
||||
)
|
||||
)
|
||||
for old in result.scalars().all():
|
||||
await session.delete(old)
|
||||
await session.flush()
|
||||
|
||||
row = ModelMetadata(
|
||||
model_type=model_type,
|
||||
version=version,
|
||||
training_samples=training_samples,
|
||||
metrics=metrics,
|
||||
artifact_path=str(artifact_path),
|
||||
is_active=True,
|
||||
mlflow_run_id=mlflow_run_id,
|
||||
threshold=threshold,
|
||||
)
|
||||
session.add(row)
|
||||
rows.append(row)
|
||||
|
||||
await session.commit()
|
||||
|
||||
logger.info(
|
||||
"Saved metadata for %d models (samples=%d)",
|
||||
len(rows),
|
||||
training_samples,
|
||||
)
|
||||
|
||||
return rows
|
||||
|
|
@ -0,0 +1,212 @@
|
|||
"""
|
||||
©AngelaMos | 2026
|
||||
test_metadata.py
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.ext.asyncio import (
|
||||
AsyncSession,
|
||||
async_sessionmaker,
|
||||
create_async_engine,
|
||||
)
|
||||
from sqlalchemy.pool import StaticPool
|
||||
from sqlalchemy import select
|
||||
from sqlmodel import SQLModel
|
||||
|
||||
from app.models.model_metadata import ModelMetadata
|
||||
from ml.metadata import compute_model_version, save_model_metadata
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def db_session(tmp_path: Path):
|
||||
"""
|
||||
In-memory SQLite session for metadata tests
|
||||
"""
|
||||
from app.models import model_metadata as _reg # noqa: F401
|
||||
|
||||
engine = create_async_engine(
|
||||
"sqlite+aiosqlite://",
|
||||
connect_args={"check_same_thread": False},
|
||||
poolclass=StaticPool,
|
||||
)
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(SQLModel.metadata.create_all)
|
||||
|
||||
factory = async_sessionmaker(
|
||||
engine, class_=AsyncSession, expire_on_commit=False
|
||||
)
|
||||
async with factory() as session:
|
||||
yield session
|
||||
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_artifacts(tmp_path: Path) -> Path:
|
||||
"""
|
||||
Create fake ONNX model files for version hashing
|
||||
"""
|
||||
(tmp_path / "ae.onnx").write_bytes(b"ae-model-data-123")
|
||||
(tmp_path / "rf.onnx").write_bytes(b"rf-model-data-456")
|
||||
(tmp_path / "if.onnx").write_bytes(b"if-model-data-789")
|
||||
(tmp_path / "scaler.json").write_text(
|
||||
json.dumps({"center": [0.0], "scale": [1.0]})
|
||||
)
|
||||
(tmp_path / "threshold.json").write_text(
|
||||
json.dumps({"threshold": 0.05})
|
||||
)
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestComputeModelVersion:
|
||||
"""
|
||||
Test SHA-256 based model version hashing
|
||||
"""
|
||||
|
||||
def test_returns_12_char_hex(
|
||||
self, model_artifacts: Path
|
||||
) -> None:
|
||||
"""
|
||||
Version string is a 12-character hex digest
|
||||
"""
|
||||
version = compute_model_version(
|
||||
model_artifacts / "ae.onnx"
|
||||
)
|
||||
|
||||
assert len(version) == 12
|
||||
assert all(c in "0123456789abcdef" for c in version)
|
||||
|
||||
def test_same_file_same_version(
|
||||
self, model_artifacts: Path
|
||||
) -> None:
|
||||
"""
|
||||
Same file produces the same version string
|
||||
"""
|
||||
v1 = compute_model_version(
|
||||
model_artifacts / "ae.onnx"
|
||||
)
|
||||
v2 = compute_model_version(
|
||||
model_artifacts / "ae.onnx"
|
||||
)
|
||||
|
||||
assert v1 == v2
|
||||
|
||||
def test_different_files_different_versions(
|
||||
self, model_artifacts: Path
|
||||
) -> None:
|
||||
"""
|
||||
Different files produce different version strings
|
||||
"""
|
||||
v_ae = compute_model_version(
|
||||
model_artifacts / "ae.onnx"
|
||||
)
|
||||
v_rf = compute_model_version(
|
||||
model_artifacts / "rf.onnx"
|
||||
)
|
||||
|
||||
assert v_ae != v_rf
|
||||
|
||||
|
||||
class TestSaveModelMetadata:
|
||||
"""
|
||||
Test model metadata persistence to database
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_creates_three_rows(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
model_artifacts: Path,
|
||||
) -> None:
|
||||
"""
|
||||
save_model_metadata creates one row per model type
|
||||
"""
|
||||
rows = await save_model_metadata(
|
||||
db_session,
|
||||
model_dir=model_artifacts,
|
||||
training_samples=500,
|
||||
metrics={"f1": 0.9, "pr_auc": 0.88},
|
||||
)
|
||||
|
||||
assert len(rows) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_all_rows_active(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
model_artifacts: Path,
|
||||
) -> None:
|
||||
"""
|
||||
All newly saved models are marked as active
|
||||
"""
|
||||
rows = await save_model_metadata(
|
||||
db_session,
|
||||
model_dir=model_artifacts,
|
||||
training_samples=500,
|
||||
metrics={"f1": 0.9},
|
||||
)
|
||||
|
||||
assert all(r.is_active for r in rows)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_model_types_correct(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
model_artifacts: Path,
|
||||
) -> None:
|
||||
"""
|
||||
Row model types are autoencoder, random_forest, isolation_forest
|
||||
"""
|
||||
rows = await save_model_metadata(
|
||||
db_session,
|
||||
model_dir=model_artifacts,
|
||||
training_samples=500,
|
||||
metrics={},
|
||||
)
|
||||
types = {r.model_type for r in rows}
|
||||
|
||||
assert types == {
|
||||
"autoencoder",
|
||||
"random_forest",
|
||||
"isolation_forest",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_previous_active_replaced(
|
||||
self,
|
||||
db_session: AsyncSession,
|
||||
model_artifacts: Path,
|
||||
) -> None:
|
||||
"""
|
||||
Saving new metadata replaces previous active models
|
||||
"""
|
||||
await save_model_metadata(
|
||||
db_session,
|
||||
model_dir=model_artifacts,
|
||||
training_samples=500,
|
||||
metrics={"f1": 0.9},
|
||||
)
|
||||
|
||||
(model_artifacts / "ae.onnx").write_bytes(
|
||||
b"new-ae-data"
|
||||
)
|
||||
await save_model_metadata(
|
||||
db_session,
|
||||
model_dir=model_artifacts,
|
||||
training_samples=600,
|
||||
metrics={"f1": 0.95},
|
||||
)
|
||||
|
||||
result = await db_session.execute(
|
||||
select(ModelMetadata)
|
||||
)
|
||||
all_rows = result.scalars().all()
|
||||
active_rows = [r for r in all_rows if r.is_active]
|
||||
|
||||
assert len(active_rows) == 3
|
||||
assert all(
|
||||
r.training_samples == 600 for r in active_rows
|
||||
)
|
||||
Loading…
Reference in New Issue