81 lines
3.1 KiB
Python
81 lines
3.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
from collections.abc import Iterable
|
|
|
|
from app.ml.training import TrainedArtifact
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_ARTIFACT_ID_PATTERN = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$")
|
|
|
|
|
|
class ModelRegistry:
|
|
def __init__(self, root: str | Path) -> None:
|
|
self._root = Path(root).resolve()
|
|
self._root.mkdir(parents=True, exist_ok=True)
|
|
self._artifacts: dict[str, TrainedArtifact] = {}
|
|
self._load_existing()
|
|
|
|
def register(self, artifact: TrainedArtifact) -> TrainedArtifact:
|
|
self._validate_artifact_id(artifact.artifact_id)
|
|
self._artifacts[artifact.artifact_id] = artifact
|
|
self._persist(artifact)
|
|
return artifact
|
|
|
|
def load_artifact(self, artifact_id: str) -> TrainedArtifact:
|
|
self._validate_artifact_id(artifact_id)
|
|
if artifact_id not in self._artifacts:
|
|
raise KeyError(f"Artifact '{artifact_id}' nicht registriert.")
|
|
return self._artifacts[artifact_id]
|
|
|
|
def list_models(self) -> Iterable[TrainedArtifact]:
|
|
return list(self._artifacts.values())
|
|
|
|
def _load_existing(self) -> None:
|
|
for source in sorted(self._root.glob("*.json")):
|
|
try:
|
|
raw = json.loads(source.read_text(encoding="utf-8"))
|
|
artifact_id = raw["artifact_id"]
|
|
supported_sensors = raw["supported_sensors"]
|
|
if not isinstance(artifact_id, str) or not isinstance(supported_sensors, list):
|
|
raise ValueError("invalid artifact structure")
|
|
self._validate_artifact_id(artifact_id)
|
|
if source.name != f"{artifact_id}.json":
|
|
raise ValueError("artifact id does not match filename")
|
|
if not all(isinstance(sensor, str) for sensor in supported_sensors):
|
|
raise ValueError("supported_sensors must contain strings")
|
|
except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc:
|
|
raise ValueError(f"Ungültiges Modell-Artefakt: {source.name}") from exc
|
|
|
|
self._artifacts[artifact_id] = TrainedArtifact(
|
|
artifact_id=artifact_id,
|
|
supported_sensors=tuple(supported_sensors),
|
|
)
|
|
|
|
def _persist(self, artifact: TrainedArtifact) -> None:
|
|
target = self._root / f"{artifact.artifact_id}.json"
|
|
temporary = target.with_suffix(".json.tmp")
|
|
payload = {
|
|
"artifact_id": artifact.artifact_id,
|
|
"supported_sensors": list(artifact.supported_sensors),
|
|
}
|
|
temporary.write_text(
|
|
json.dumps(payload, ensure_ascii=True, sort_keys=True) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
os.replace(temporary, target)
|
|
logger.info("Modell gespeichert: %s", target)
|
|
|
|
@staticmethod
|
|
def _validate_artifact_id(artifact_id: str) -> None:
|
|
if not _ARTIFACT_ID_PATTERN.fullmatch(artifact_id) or ".." in artifact_id:
|
|
raise ValueError(
|
|
"artifact_id darf nur Buchstaben, Ziffern, Punkt, Unterstrich "
|
|
"und Bindestrich enthalten."
|
|
)
|