koios-model-utils

Annotate ONNX models with the metadata the Koios IoT platform needs to run them.

This library is the canonical writer (and reader) of the koios.training and koios.bindings metadata_props blocks Koios consumes when a model is uploaded to its predict engine. Bring your own ONNX file from any framework (Stable-Baselines3, PyTorch, TensorFlow/Keras, scikit-learn); the library only touches metadata.

The wire format is fully documented in docs/CONTRACT.md.

Installation

pip install koios-model-utils

Only runtime dependency is onnx>=1.14. No torch, no SB3, no CLI — this is a metadata library.

Quick start

import onnx
from koios_model_utils import (
    Algorithm,
    InputBinding,
    NormalizationSource,
    NormalizationType,
    OutputBinding,
    TrainingMeta,
    embed_koios_metadata,
)

model = onnx.load("my_model.onnx")

embed_koios_metadata(
    model,
    inputs=[
        InputBinding(name="tank_temperature", description="Tank temperature (C)"),
        InputBinding(name="pressure", description="Vessel pressure (kPa)"),
    ],
    outputs=[
        OutputBinding(
            name="valve_position",
            range_min=0.0,
            range_max=100.0,
            normalization_type=NormalizationType.SYMMETRIC,
            normalization_source=NormalizationSource.CUSTOM,
            custom_minimum=0.0,
            custom_maximum=100.0,
            clamp_output=True,
        ),
    ],
    training=TrainingMeta(
        scenario_name="tank_temperature",
        algorithm=Algorithm.PPO,
        obs_depth=5,
        sample_rate=1.0,
    ),
)

onnx.save(model, "my_model_koios.onnx")

The Koios webapp reads both metadata blocks on upload and uses them to configure model.sample_rate, the input/output bindings, normalization rules, and (for DRL classifiers) the action map.

See examples/ for runnable scripts.

Reading metadata back

The mirror of embed_koios_metadata:

from koios_model_utils import parse_koios_metadata

parsed = parse_koios_metadata(onnx.load("my_model_koios.onnx"))
print(parsed.training)        # TrainingMeta dataclass (or None)
print(parsed.inputs)          # list[InputBinding]
print(parsed.outputs)         # list[OutputBinding]
print(parsed.has_metadata)    # True if any koios.* prop was present

If the JSON has already been pulled out of the ONNX file (e.g. persisted in a database column), use parse_koios_metadata_from_dict and pass {"koios.training": {...}, "koios.bindings": {...}} directly.

Public API

Dataclasses

All dataclasses run their validators in __post_init__, so bad combinations (e.g. Z_SCORE + custom_minimum, or a minimum that is not below its maximum) raise at construction.

Enums

NormalizationType, NormalizationSource, ModelType, OutputMode, Algorithm, FailureRangeMode — all enum.StrEnum, all uppercase wire values. See enums.py for the members.

Pass enum members rather than raw strings — the dataclasses accept either, but enums catch typos at type-check time.

Functions

Errors

KoiosMetadataError, UnsupportedSchemaVersionError, KoiosMetadataDecodeError — all ValueError subclasses raised by the parser.

The parser reads output_denormalized, clamp_output and flatten_observation only as a JSON true or false; any other value, the string "false" included, raises KoiosMetadataDecodeError. So does a binding whose minimum is not below its maximum.

Clamping an output

clamp_output=True holds an output's prediction inside a range before Koios writes it. A scaled output is held to its normalization range. An output with no scaling is held to the custom failure bounds its OutputBinding declares instead; with output_denormalized=True, those come from range_min and range_max. If the file declares no custom failure bounds for an unscaled output, there is nothing to clamp to, and Koios stores clamping off for it.

sample_rate vs scan_rate

sample_rate is the interval (seconds) the model was trained at. The Koios predict engine resamples historical inputs to this rate.

scan_rate is how often the predict engine executes the model. When omitted (the common case), it inherits sample_rate — forecasting models train and run at the same cadence. Set it explicitly for RL controllers that need to act faster than the simulator step they were trained against (e.g. sample_rate=1.0 history lookback, scan_rate=0.1 control loop).

Discrete (classification / DQN) heads

For models that emit an action index instead of a continuous value, set output_mode and pass an action_map:

from koios_model_utils import (
    ActionMapEntry, OutputMode, TrainingMeta,
)

training = TrainingMeta(
    algorithm="DQN",
    output_mode=OutputMode.DISCRETE,
    action_map=[
        ActionMapEntry(value=-1.0, label="Decrease setpoint"),
        ActionMapEntry(value=0.0, label="Hold"),
        ActionMapEntry(value=1.0, label="Increase setpoint"),
    ],
)

action_map also accepts raw {"value": ..., "label": ...} dicts and coerces them to ActionMapEntry at construction — extra keys are dropped, missing keys raise.

Contract & versioning

The wire format is at schema_version 1. docs/CONTRACT.md is the source of truth for every field, every default, and every server-side policy applied on upload.

The library and the Koios webapp move in lockstep: any wire-format change here requires a coordinated webapp release. The CONTRACT lays out exactly which keys the webapp reads, the forward-compat rules within v1, and when a v2 bump would be needed.

Development

pip install -e ".[dev]"
make test
make lint

License

Apache License 2.0 — see LICENSE.