koios-model-utils 1.0.0__py3-none-any.whl

koios_model_utils/__init__.py

@@ -0,0 +1,60 @@
+"""koios-model-utils — utilities for exporting and annotating ONNX models for the Koios IoT platform."""
+
+__version__ = "1.0.0"
+
+from koios_model_utils.enums import (
+    Algorithm,
+    ModelType,
+    NormalizationSource,
+    NormalizationType,
+    OutputMode,
+)
+from koios_model_utils.metadata import (
+    InputBinding,
+    OutputBinding,
+    TrainingMeta,
+    build_input_bindings,
+    build_output_bindings,
+    embed_koios_metadata,
+)
+
+# SB3-dependent exports — only available with koios-model-utils[sb3]
+try:
+    from koios_model_utils.callbacks import (
+        ObsRangeCallback,
+        VecNormalizeSyncCallback,
+    )
+    from koios_model_utils.export import (
+        OnnxActorWrapper,
+        OnnxDQNWrapper,
+        detect_algorithm,
+        export_onnx,
+        load_vec_normalize_stats,
+    )
+except ImportError:
+    pass
+
+__all__ = [
+    # enums (always available)
+    "Algorithm",
+    "ModelType",
+    "NormalizationSource",
+    "NormalizationType",
+    "OutputMode",
+    # metadata (always available)
+    "InputBinding",
+    "OutputBinding",
+    "TrainingMeta",
+    "embed_koios_metadata",
+    "build_input_bindings",
+    "build_output_bindings",
+    # export (requires sb3 extra)
+    "export_onnx",
+    "detect_algorithm",
+    "load_vec_normalize_stats",
+    "OnnxActorWrapper",
+    "OnnxDQNWrapper",
+    # callbacks (requires sb3 extra)
+    "ObsRangeCallback",
+    "VecNormalizeSyncCallback",
+]

koios_model_utils/callbacks.py

@@ -0,0 +1,147 @@
+"""SB3 training callbacks for Koios model export."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import numpy as np
+from stable_baselines3.common.callbacks import BaseCallback
+
+if TYPE_CHECKING:
+    from stable_baselines3.common.callbacks import EvalCallback
+
+
+class ObsRangeCallback(BaseCallback):
+    """Track per-feature min/max of raw (pre-normalization) observations.
+
+    Saves ``obs_range.json`` to the run directory on training end, containing
+    the observed min/max per feature across all training steps.  These are used
+    as failure bounds in Koios — if a live value falls outside the range the
+    model was trained on, the binding is flagged as untrusted.
+
+    Works with both raw envs and VecNormalize-wrapped envs.  When VecNormalize
+    is present, the original (un-normalized) observations are tracked.
+    """
+
+    def __init__(self, run_dir: str | Path, verbose: int = 0) -> None:
+        super().__init__(verbose)
+        self._run_dir = Path(run_dir)
+        self._obs_min: np.ndarray | None = None
+        self._obs_max: np.ndarray | None = None
+
+    def _on_step(self) -> bool:
+        # Get raw observations — use original obs when normalization is active
+        obs = None
+
+        # Path 1: VecNormalize (z_score)
+        vec_norm = self.model._vec_normalize_env  # type: ignore[union-attr]
+        if vec_norm is not None:
+            obs = vec_norm.get_original_obs()
+
+        # Path 2: MinMaxObsWrapper — walk wrapper chain for _last_original_obs
+        if obs is None:
+            env = self.training_env
+            # Unwrap through VecEnv layers to find the underlying env
+            while hasattr(env, "envs"):
+                env = env.envs[0]  # type: ignore[union-attr]
+            # Walk Gymnasium wrapper chain
+            current = env
+            while current is not None:
+                if hasattr(current, "_last_original_obs"):
+                    obs = current._last_original_obs
+                    break
+                current = getattr(current, "env", None)
+
+        # Path 3: fallback to raw obs
+        if obs is None:
+            obs = self.locals.get("new_obs")
+
+        if obs is None:
+            return True
+
+        obs = np.asarray(obs, dtype=np.float32)
+
+        # Flatten to (n_envs, num_features) — collapse obs_depth if present
+        if obs.ndim == 3:
+            # (n_envs, obs_depth, num_features) → min/max across envs and depth
+            flat = obs.reshape(-1, obs.shape[-1])
+        elif obs.ndim == 2:
+            flat = obs
+        else:
+            flat = obs.reshape(1, -1)
+
+        step_min = flat.min(axis=0)
+        step_max = flat.max(axis=0)
+
+        if self._obs_min is None:
+            self._obs_min = step_min.copy()
+            self._obs_max = step_max.copy()
+        else:
+            np.minimum(self._obs_min, step_min, out=self._obs_min)
+            np.maximum(self._obs_max, step_max, out=self._obs_max)
+
+        return True
+
+    def _on_training_end(self) -> None:
+        if self._obs_min is None or self._obs_max is None:
+            return
+
+        out_path = self._run_dir / "obs_range.json"
+        data = {
+            "obs_min": self._obs_min.tolist(),
+            "obs_max": self._obs_max.tolist(),
+        }
+        with open(out_path, "w") as f:
+            json.dump(data, f, indent=2)
+
+        if self.verbose >= 1:
+            print(f"Observation ranges saved to {out_path}")
+            for i, (lo, hi) in enumerate(zip(self._obs_min, self._obs_max)):
+                print(f"  feature {i}: [{lo:.4f}, {hi:.4f}]")
+
+
+class VecNormalizeSyncCallback(BaseCallback):
+    """Save VecNormalize stats alongside every new best model.
+
+    SB3's ``EvalCallback`` saves ``best_model.zip`` when eval reward improves,
+    but doesn't save the VecNormalize running statistics.  Without the stats
+    file, an early ``koios-export`` on ``best_model.zip`` produces an ONNX with
+    no z-score metadata — the PE then feeds raw observations to a model that
+    expects normalized ones.
+
+    This callback checks after each eval whether a new best model was saved
+    and, if so, writes ``vec_normalize.pkl`` to the same directory.
+    """
+
+    def __init__(self, eval_callback: EvalCallback, verbose: int = 0) -> None:
+        super().__init__(verbose)
+        self._eval_cb = eval_callback
+        self._last_best_timestep: int = -1
+
+    def _on_step(self) -> bool:
+        # EvalCallback sets best_mean_reward and n_calls when a new best is saved
+        best_ts = getattr(self._eval_cb, "best_mean_reward_timestep", None)
+        if best_ts is None:
+            # Fallback: check if best_model.zip was updated by comparing n_calls
+            best_ts = getattr(self._eval_cb, "n_calls", 0)
+
+        # Detect new best model by checking if EvalCallback updated its best
+        vec_norm = self.model._vec_normalize_env  # type: ignore[union-attr]
+        if vec_norm is None:
+            return True
+
+        best_path = Path(self._eval_cb.best_model_save_path) / "best_model.zip"
+        stats_path = Path(self._eval_cb.best_model_save_path) / "vec_normalize.pkl"
+
+        # Save if best_model.zip exists and is newer than our last save
+        if best_path.exists():
+            mtime = best_path.stat().st_mtime_ns
+            if mtime != self._last_best_timestep:
+                self._last_best_timestep = mtime
+                vec_norm.save(str(stats_path))
+                if self.verbose >= 1:
+                    print(f"[VecNormSync] Saved {stats_path}")
+
+        return True

koios_model_utils/enums.py

@@ -0,0 +1,67 @@
+"""String enums for Koios metadata values.
+
+Every enum here is a :class:`enum.StrEnum` — members are real strings and
+serialize transparently through ``json.dumps``, so users may pass either
+the enum member or the raw string literal wherever metadata values are
+expected.  The server consumes strings (case-insensitive for
+normalization keys, exact-match for training keys), so both forms flow
+through identically.
+
+The string values match the names of the corresponding Django
+``IntegerChoices`` members in
+``koios-webapp-py/app/AiOpsInterface/apps/aiops/choices.py`` — renaming
+a member on the server requires updating the matching member here (the
+contract tests in ``tests/test_metadata.py`` pin this).
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+
+class NormalizationType(StrEnum):
+    """How a binding's value is mapped into / out of model space."""
+
+    NONE = "none"
+    MIN_MAX = "min_max"
+    SYMMETRIC = "symmetric"
+    Z_SCORE = "z_score"
+
+
+class NormalizationSource(StrEnum):
+    """Where normalization parameters come from at runtime."""
+
+    # The tag's engineering range (range_min / range_max on the bound tag).
+    TAG_RANGE = "tag_range"
+    # Explicit ``custom_minimum`` / ``custom_maximum`` / ``custom_mean`` /
+    # ``custom_std`` values embedded in the ONNX metadata.
+    CUSTOM = "custom"
+
+
+class ModelType(StrEnum):
+    """How the predict engine applies the model's output to the bound tag.
+
+    Maps to the server's ``OutputApplicationChoices``. The legacy
+    ``POSITIONAL`` / ``VELOCITY`` aliases are still accepted by the
+    server but should not be used in new code.
+    """
+
+    ABSOLUTE = "ABSOLUTE"
+    RELATIVE = "RELATIVE"
+
+
+class OutputMode(StrEnum):
+    """Whether the model emits a continuous value or a categorical action."""
+
+    CONTINUOUS = "CONTINUOUS"
+    DISCRETE = "DISCRETE"
+
+
+class Algorithm(StrEnum):
+    """Stable-Baselines3 algorithm identifiers recognized by the exporter."""
+
+    PPO = "PPO"
+    SAC = "SAC"
+    TD3 = "TD3"
+    A2C = "A2C"
+    DQN = "DQN"