37 lines
1.2 KiB
Python
37 lines
1.2 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Iterable
|
|
|
|
from app.ml.training import TrainedArtifact
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ModelRegistry:
|
|
def __init__(self, root: str | Path) -> None:
|
|
self._root = Path(root)
|
|
self._root.mkdir(parents=True, exist_ok=True)
|
|
self._artifacts: dict[str, TrainedArtifact] = {}
|
|
|
|
def register(self, artifact: TrainedArtifact) -> TrainedArtifact:
|
|
self._artifacts[artifact.artifact_id] = artifact
|
|
self._persist(artifact)
|
|
return artifact
|
|
|
|
def load_artifact(self, artifact_id: str) -> TrainedArtifact:
|
|
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 _persist(self, artifact: TrainedArtifact) -> None:
|
|
target = self._root / f"{artifact.artifact_id}.json"
|
|
target.write_text(
|
|
f"{artifact.artifact_id}\t{','.join(artifact.supported_sensors)}\n",
|
|
encoding="utf-8",
|
|
)
|
|
logger.info("Modell gespeichert: %s", target) |