openstef-foundation-models 4.1.0__py3-none-any.whl

openstef_foundation_models/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+"""Foundation model support for OpenSTEF."""
+
+import logging
+
+root_logger = logging.getLogger(name=__name__)
+if not root_logger.handlers:
+    root_logger.addHandler(logging.NullHandler())
+
+__all__: list[str] = []

openstef_foundation_models/inference/__init__.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""Inference backends for foundation-model forecasters.
+
+An :class:`InferenceBackend` isolates how a checkpoint is executed behind a
+single named-tensor contract; forecasters compose a backend rather than
+inheriting execution behaviour. Only the dependency-free surface (the protocol
+and the execution-provider configs) is re-exported here; a concrete backend
+lives in its own submodule and imports its heavy runtime at module top level.
+"""
+
+from openstef_foundation_models.inference.backend import InferenceBackend
+from openstef_foundation_models.inference.providers import (
+    CoreMLProvider,
+    CpuProvider,
+    CudaProvider,
+    ExecutionProvider,
+    SessionOptionsConfig,
+    TensorRTProvider,
+)
+
+__all__ = [
+    "CoreMLProvider",
+    "CpuProvider",
+    "CudaProvider",
+    "ExecutionProvider",
+    "InferenceBackend",
+    "SessionOptionsConfig",
+    "TensorRTProvider",
+]

openstef_foundation_models/inference/backend.py

@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""The :class:`InferenceBackend` contract shared by all execution backends."""
+
+from collections.abc import Mapping
+from typing import Protocol, runtime_checkable
+
+import numpy as np
+
+from openstef_foundation_models.models.checkpoint import CheckpointMetadata
+
+
+@runtime_checkable
+class InferenceBackend(Protocol):
+    """A model-agnostic execution backend.
+
+    A backend takes a mapping of named input tensors to a mapping of named
+    output tensors. It owns whatever runtime resources are needed (e.g. an ONNX
+    Runtime session) and is loaded once, then reused across an entire backtest.
+    Model-family specifics live in :attr:`metadata`, not in the backend itself.
+    """
+
+    @property
+    def metadata(self) -> CheckpointMetadata:
+        """Metadata describing the checkpoint this backend executes."""
+        ...
+
+    def run(self, inputs: Mapping[str, np.ndarray]) -> Mapping[str, np.ndarray]:
+        """Execute the model on a batch of named input tensors.
+
+        Args:
+            inputs: Named input tensors. Keys must match ``metadata.input_names``.
+
+        Returns:
+            Named output tensors, including ``metadata.output_name``.
+        """
+        ...
+
+    def close(self) -> None:
+        """Release any runtime resources held by the backend."""
+        ...

openstef_foundation_models/inference/onnx_backend.py

@@ -0,0 +1,237 @@
+# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""ONNX Runtime execution backend.
+
+Importing this module requires ONNX Runtime (the ``[cpu]`` or ``[gpu]`` extra)
+and raises :class:`MissingExtraError` if it is missing.
+"""
+
+import logging
+from collections.abc import Mapping, Sequence
+from typing import Self
+
+import numpy as np
+
+from openstef_core.exceptions import MissingExtraError
+from openstef_foundation_models.inference.provider_selection import (
+    DefaultProviderPolicy,
+    HostCapabilities,
+    ProviderPolicy,
+)
+from openstef_foundation_models.inference.providers import (
+    ExecutionProvider,
+    SessionOptionsConfig,
+)
+from openstef_foundation_models.models.checkpoint import CheckpointMetadata, ResolvedCheckpoint
+
+try:
+    import onnxruntime as ort
+except ImportError as e:
+    raise MissingExtraError("onnxruntime", "openstef-foundation-models", install_extra="cpu") from e
+
+# onnxruntime-gpu ships the CUDA execution-provider plugin but loads the CUDA/cuDNN
+# runtime (the nvidia-*-cu12 wheels the [gpu] extra pulls) lazily at session creation.
+# preload_dlls() loads them from the nvidia site-packages so the CUDA provider can be
+# realized without a system CUDA install or LD_LIBRARY_PATH. It is a no-op on the CPU
+# runtime; the guard covers onnxruntime < 1.21, which predates the API.
+if hasattr(ort, "preload_dlls"):
+    ort.preload_dlls()
+
+logger = logging.getLogger(__name__)
+
+
+class OnnxBackend:
+    """An :class:`~openstef_foundation_models.inference.backend.InferenceBackend` backed by ONNX Runtime.
+
+    The session is built once on construction and reused for every
+    :meth:`run` call, so a single backend instance can be shared across an
+    entire backtest. Users may either let the backend build a session from a
+    resolved checkpoint and provider chain, or pass a pre-built session they own.
+    """
+
+    def __init__(
+        self,
+        metadata: CheckpointMetadata,
+        session: ort.InferenceSession,
+    ) -> None:
+        """Wrap a pre-built ONNX Runtime session.
+
+        Prefer :meth:`from_checkpoint` unless you need to own the session
+        lifecycle yourself.
+
+        Args:
+            metadata: Metadata describing the checkpoint the session executes.
+            session: A pre-built ONNX Runtime inference session.
+        """
+        self._metadata = metadata
+        self._session: ort.InferenceSession | None = session
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        checkpoint: ResolvedCheckpoint,
+        providers: Sequence[ExecutionProvider] | None = None,
+        session_options: SessionOptionsConfig | None = None,
+        *,
+        policy: ProviderPolicy | None = None,
+    ) -> Self:
+        """Build a backend by loading a checkpoint into a new ONNX Runtime session.
+
+        With ``providers=None`` the *policy* selects a chain from the checkpoint
+        and host; an explicit ``providers`` list is used as given and *policy* is
+        ignored. See :class:`~openstef_foundation_models.inference.provider_selection.ProviderPolicy`
+        for how a chain is chosen and how strictly its realization is enforced.
+
+        Args:
+            checkpoint: The resolved checkpoint (weights + metadata) to load.
+            providers: Ordered execution providers to try. ``None`` lets the policy
+                pick a host-appropriate chain from the checkpoint metadata.
+            session_options: Optional ONNX Runtime session options.
+            policy: Selection policy used when ``providers is None``. Defaults to
+                :class:`DefaultProviderPolicy`.
+
+        Returns:
+            A backend wrapping the newly built session.
+        """
+        metadata = checkpoint.metadata
+        if providers is not None:
+            provider_configs = list(providers)
+            explicit = True
+            logger.info(
+                "Using explicit execution-provider chain %s for checkpoint '%s'.",
+                [config.to_ort()[0] for config in provider_configs],
+                metadata.model_family,
+            )
+        else:
+            selector = policy or DefaultProviderPolicy()
+            host = HostCapabilities.detect()
+            provider_configs = selector.select(metadata, host)
+            explicit = False
+            logger.debug(
+                "Detected host: platform=%s, available_providers=%s",
+                host.platform,
+                sorted(host.available_providers),
+            )
+            logger.info(
+                "%s selected execution-provider chain %s for checkpoint '%s' (precision=%s, static_shapes=%s) on %s.",
+                type(selector).__name__,
+                [config.to_ort()[0] for config in provider_configs],
+                metadata.model_family,
+                metadata.precision,
+                metadata.static_shapes,
+                host.platform,
+            )
+        ort_providers = [config.to_ort() for config in provider_configs]
+        so = _build_session_options(session_options) if session_options else None
+
+        session = ort.InferenceSession(
+            str(checkpoint.weights_path),
+            sess_options=so,
+            providers=ort_providers,
+        )
+        logger.info("ONNX Runtime session built; realized providers: %s.", session.get_providers())
+        _check_provider_fallback(
+            requested=provider_configs,
+            realized=session.get_providers(),
+            strict=explicit,
+        )
+        return cls(metadata=checkpoint.metadata, session=session)
+
+    @property
+    def metadata(self) -> CheckpointMetadata:
+        """Metadata describing the checkpoint this backend executes."""
+        return self._metadata
+
+    def run(self, inputs: Mapping[str, np.ndarray]) -> Mapping[str, np.ndarray]:
+        """Execute the ONNX model on a batch of named input tensors.
+
+        Args:
+            inputs: Named input tensors. Keys must match ``metadata.input_names``.
+
+        Returns:
+            Named output tensors keyed by the model's output names.
+
+        Raises:
+            RuntimeError: If the backend has been closed.
+        """
+        if self._session is None:
+            msg = "OnnxBackend has been closed."
+            raise RuntimeError(msg)
+        output_names = [out.name for out in self._session.get_outputs()]
+        results = self._session.run(output_names, dict(inputs))
+        return {name: np.asarray(result) for name, result in zip(output_names, results, strict=True)}
+
+    def close(self) -> None:
+        """Release the underlying ONNX Runtime session.
+
+        ONNX Runtime frees native resources on garbage collection, so dropping
+        the reference is the supported way to release them.
+        """
+        self._session = None
+
+
+def _build_session_options(config: SessionOptionsConfig) -> ort.SessionOptions:
+    """Translate a :class:`SessionOptionsConfig` into ONNX Runtime session options.
+
+    Args:
+        config: The typed session-options configuration.
+
+    Returns:
+        The corresponding ONNX Runtime ``SessionOptions``.
+    """
+    so = ort.SessionOptions()
+    so.graph_optimization_level = getattr(
+        ort.GraphOptimizationLevel,
+        f"ORT_{config.graph_optimization_level}",
+    )
+    if config.intra_op_num_threads is not None:
+        so.intra_op_num_threads = config.intra_op_num_threads
+    if config.inter_op_num_threads is not None:
+        so.inter_op_num_threads = config.inter_op_num_threads
+    return so
+
+
+def _check_provider_fallback(
+    requested: Sequence[ExecutionProvider],
+    realized: Sequence[str],
+    *,
+    strict: bool,
+) -> None:
+    """Detect and report a silent fallback to the CPU execution provider.
+
+    Compares the requested chain against what ONNX Runtime actually realized. See
+    :class:`~openstef_foundation_models.inference.provider_selection.ProviderPolicy`
+    for the strict-vs-graceful contract this enforces.
+
+    Args:
+        requested: The execution providers that were requested.
+        realized: The provider names ONNX Runtime actually loaded.
+        strict: When ``True``, raise on any missing accelerator; otherwise warn
+            only on a full fallback to CPU.
+
+    Raises:
+        RuntimeError: If ``strict`` is set and a requested accelerator is missing.
+    """
+    requested_names = {config.to_ort()[0] for config in requested}
+    accelerators = requested_names - {"CPUExecutionProvider"}
+    if not accelerators:
+        return
+    realized_set = set(realized)
+    missing = accelerators - realized_set
+    if strict:
+        if missing:
+            msg = (
+                f"Requested execution provider(s) {sorted(missing)} were not realized; "
+                f"ONNX Runtime fell back to {realized}."
+            )
+            raise RuntimeError(msg)
+        return
+    if accelerators & realized_set:
+        return
+    logger.warning(
+        "No requested accelerator (%s) was realized; ONNX Runtime fell back to %s. Inference will run on CPU.",
+        sorted(accelerators),
+        realized,
+    )

openstef_foundation_models/inference/provider_selection.py

@@ -0,0 +1,142 @@
+# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""Metadata-driven execution-provider selection.
+
+Which ONNX Runtime execution provider is fastest — and even which is *usable* —
+depends on both the host (Apple CoreML vs NVIDIA CUDA/TensorRT vs CPU) and on
+properties of the checkpoint itself (precision, whether its graph has static
+shapes). This module keeps that knowledge in **one replaceable component** rather
+than scattered platform ``if``-ladders:
+
+* :class:`HostCapabilities` carries the host facts as injectable data, with a
+  single impure :meth:`HostCapabilities.detect` classmethod.
+* :class:`ProviderPolicy` is the port; :class:`DefaultProviderPolicy` is the
+  adapter that maps ``(checkpoint, host)`` to an ordered provider chain. Users
+  with exotic hardware implement their own policy.
+
+Importing this module requires ONNX Runtime (the ``[cpu]`` or ``[gpu]`` extra)
+and raises :class:`MissingExtraError` if it is missing.
+"""
+
+import platform
+from typing import Literal, Protocol, Self
+
+from pydantic import Field
+
+from openstef_core.base_model import BaseConfig
+from openstef_core.exceptions import MissingExtraError
+from openstef_foundation_models.inference.providers import (
+    CoreMLProvider,
+    CpuProvider,
+    CudaProvider,
+    ExecutionProvider,
+)
+from openstef_foundation_models.models.checkpoint import CheckpointMetadata
+
+try:
+    import onnxruntime as ort
+except ImportError as e:
+    raise MissingExtraError("onnxruntime", "openstef-foundation-models", install_extra="cpu") from e
+
+
+class HostCapabilities(BaseConfig):
+    """Execution-relevant facts about the host, captured as injectable data.
+
+    Passing host facts into a policy (rather than having the policy call
+    ``platform.system()`` itself) keeps selection a pure function of its inputs,
+    so it can be unit-tested by constructing a fake host.
+    """
+
+    model_config = BaseConfig.model_config | {"frozen": True}
+
+    platform: str = Field(
+        description="OS identifier, lower-cased (e.g. 'darwin', 'linux', 'windows').",
+    )
+    available_providers: frozenset[str] = Field(
+        description="Execution provider names ONNX Runtime reports as available on this host.",
+    )
+
+    @classmethod
+    def detect(cls) -> Self:
+        """Detect the host's capabilities from the platform and ONNX Runtime.
+
+        This is the one impure call in the selection path; it is isolated here so
+        the policy stays a pure function of injected facts.
+
+        Returns:
+            The detected host capabilities.
+        """
+        return cls(
+            platform=platform.system().lower(),
+            available_providers=frozenset(ort.get_available_providers()),
+        )
+
+
+class ProviderPolicy(Protocol):
+    """Port mapping a checkpoint and host to an ordered execution-provider chain.
+
+    Implement this to encode selection rules for hardware the default policy does
+    not cover; pass the implementation to the backend or
+    :class:`~openstef_foundation_models.presets.forecasting_workflow.OnnxBackendConfig`.
+
+    A policy-selected chain is enforced *gracefully*: ONNX Runtime silently drops
+    accelerators it cannot initialize and falls back to CPU, and a policy chain
+    such as ``[CoreML, CPU]`` realizing CoreML is the intended outcome, so a
+    warning is logged only if it falls all the way to CPU. A chain the caller
+    passes explicitly is enforced *strictly* instead: any requested accelerator
+    that is not realized raises.
+    """
+
+    def select(self, metadata: CheckpointMetadata, host: HostCapabilities) -> list[ExecutionProvider]:
+        """Return the ordered provider chain to try for *metadata* on *host*."""
+        ...
+
+
+class DefaultProviderPolicy(BaseConfig):
+    """Default policy mapping ``(checkpoint precision/shape, host)`` to a provider chain.
+
+    Each rule encodes a *measured* hardware conclusion; see the design doc
+    ``design-docs/0001`` and the provider benchmark for the rationale. The chain
+    is ordered preferred-first with CPU as the final fallback.
+    """
+
+    kind: Literal["default"] = Field(default="default", description="Discriminator tag for the policy type.")
+
+    def select(  # instance method to satisfy the ProviderPolicy protocol, though stateless here
+        self, metadata: CheckpointMetadata, host: HostCapabilities
+    ) -> list[ExecutionProvider]:
+        """Select an ordered provider chain for *metadata* on *host*.
+
+        Args:
+            metadata: The checkpoint's metadata (precision, static-shape-ness).
+            host: The detected host capabilities.
+
+        Returns:
+            An ordered execution-provider chain, preferred-first, CPU last.
+        """
+        cuda_ok = "CUDAExecutionProvider" in host.available_providers
+        coreml_ok = "CoreMLExecutionProvider" in host.available_providers and metadata.static_shapes
+
+        # int8 (QDQ) runs fast on CPU; CoreML cannot accelerate the quantized ops,
+        # so it is skipped entirely. CUDA int8 is fine when a GPU is present.
+        if metadata.precision == "int8":
+            return [CudaProvider(), CpuProvider()] if cuda_ok else [CpuProvider()]
+        # macOS: a static-shape fp16/fp32 graph runs on CoreML, but only on the GPU
+        # (MLComputeUnits=ALL/ANE triggers a multi-minute Neural-Engine compile for
+        # no inference win — measured).
+        if host.platform == "darwin" and coreml_ok:
+            return [CoreMLProvider(compute_units="CPUAndGPU"), CpuProvider()]
+        # NVIDIA: CUDA with a CPU fallback. TensorRT stays opt-in (engine-build cost
+        # and fp16 caveats), so the default never selects it.
+        if cuda_ok:
+            return [CudaProvider(), CpuProvider()]
+        return [CpuProvider()]
+
+
+__all__ = [
+    "DefaultProviderPolicy",
+    "HostCapabilities",
+    "ProviderPolicy",
+]

openstef_foundation_models/inference/providers.py

@@ -0,0 +1,158 @@
+# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""Typed ONNX Runtime execution-provider configuration.
+
+Each provider is a small pydantic config that compiles to an ONNX Runtime
+``(name, options)`` tuple via :meth:`ExecutionProviderConfig.to_ort`. Keeping
+providers as typed configs (rather than raw strings) lets users opt into
+hardware acceleration — CUDA, TensorRT FP16, CoreML/ANE — without touching
+model code, and keeps the options validated and discoverable.
+"""
+
+from pathlib import Path
+from typing import Annotated, Literal
+
+from pydantic import Field
+
+from openstef_core.base_model import BaseConfig
+
+#: An ONNX Runtime provider specification: ``(provider_name, provider_options)``.
+type OrtProvider = tuple[str, dict[str, object]]
+
+
+class CpuProvider(BaseConfig):
+    """The default CPU execution provider."""
+
+    kind: Literal["cpu"] = Field(default="cpu", description="Discriminator tag for execution-provider type.")
+
+    def to_ort(self) -> OrtProvider:
+        """Compile to an ONNX Runtime provider tuple.
+
+        Returns:
+            The ``CPUExecutionProvider`` with no options.
+        """
+        return ("CPUExecutionProvider", {})
+
+
+class CudaProvider(BaseConfig):
+    """The CUDA (NVIDIA GPU) execution provider."""
+
+    kind: Literal["cuda"] = Field(default="cuda", description="Discriminator tag for execution-provider type.")
+    device_id: int = Field(default=0, ge=0, description="CUDA device index to run on.")
+
+    def to_ort(self) -> OrtProvider:
+        """Compile to an ONNX Runtime provider tuple.
+
+        Returns:
+            The ``CUDAExecutionProvider`` with the configured device id.
+        """
+        return ("CUDAExecutionProvider", {"device_id": self.device_id})
+
+
+class TensorRTProvider(BaseConfig):
+    """The TensorRT execution provider (NVIDIA, ahead-of-time engine build).
+
+    FP16 with a persistent engine cache is the recommended production path on
+    NVIDIA hardware: the first run pays the engine-build cost, subsequent runs
+    load the cached engine.
+    """
+
+    kind: Literal["tensorrt"] = Field(default="tensorrt", description="Discriminator tag for execution-provider type.")
+    device_id: int = Field(default=0, ge=0, description="CUDA device index to run on.")
+    fp16: bool = Field(default=True, description="Enable FP16 precision for faster inference.")
+    engine_cache_dir: Path | None = Field(
+        default=None,
+        description="Directory to persist built TensorRT engines. When set, engine caching is enabled.",
+    )
+
+    def to_ort(self) -> OrtProvider:
+        """Compile to an ONNX Runtime provider tuple.
+
+        Returns:
+            The ``TensorrtExecutionProvider`` with precision and engine-cache options.
+        """
+        options: dict[str, object] = {
+            "device_id": self.device_id,
+            "trt_fp16_enable": self.fp16,
+        }
+        if self.engine_cache_dir is not None:
+            options["trt_engine_cache_enable"] = True
+            options["trt_engine_cache_path"] = str(self.engine_cache_dir)
+        return ("TensorrtExecutionProvider", options)
+
+
+class CoreMLProvider(BaseConfig):
+    """The CoreML execution provider (Apple, GPU/Neural Engine).
+
+    ``ModelFormat=MLProgram`` is required for modern CoreML: the legacy
+    ``NeuralNetwork`` format fragments the graph and silently falls back to CPU
+    for many ops.
+    """
+
+    kind: Literal["coreml"] = Field(default="coreml", description="Discriminator tag for execution-provider type.")
+    model_format: Literal["MLProgram", "NeuralNetwork"] = Field(
+        default="MLProgram",
+        description="CoreML model format. MLProgram is required for modern op coverage.",
+    )
+    compute_units: Literal["CPUOnly", "CPUAndGPU", "CPUAndNeuralEngine", "ALL"] = Field(
+        default="ALL",
+        description="Which compute units CoreML may dispatch to. Prefer 'CPUAndGPU' for large transformer "
+        "graphs: allowing the Neural Engine ('ALL'/'CPUAndNeuralEngine') can make CoreML's ahead-of-time "
+        "compile run for many minutes for no inference win.",
+    )
+    cache_dir: Path | None = Field(
+        default=None,
+        description="Directory for CoreML's compiled-model cache (ORT 'ModelCacheDirectory'). CoreML compiles "
+        "the graph ahead of time on session build, which is slow; caching it cuts a warm rebuild from tens of "
+        "seconds to a few. The cache is ORT-version/OS/hardware-specific — a local speedup, not a portable "
+        "artifact. Requires 'MLProgram' format.",
+    )
+
+    def to_ort(self) -> OrtProvider:
+        """Compile to an ONNX Runtime provider tuple.
+
+        Returns:
+            The ``CoreMLExecutionProvider`` with format, compute-unit and (optional) cache options.
+        """
+        options: dict[str, object] = {"ModelFormat": self.model_format, "MLComputeUnits": self.compute_units}
+        if self.cache_dir is not None:
+            options["ModelCacheDirectory"] = str(self.cache_dir)
+        return ("CoreMLExecutionProvider", options)
+
+
+#: An execution-provider config, discriminated by its ``kind`` tag.
+ExecutionProvider = Annotated[
+    CpuProvider | CudaProvider | TensorRTProvider | CoreMLProvider,
+    Field(discriminator="kind"),
+]
+
+
+class SessionOptionsConfig(BaseConfig):
+    """A subset of ONNX Runtime ``SessionOptions`` exposed as typed config."""
+
+    graph_optimization_level: Literal["DISABLE_ALL", "ENABLE_BASIC", "ENABLE_EXTENDED", "ENABLE_ALL"] = Field(
+        default="ENABLE_ALL",
+        description="Graph optimization level applied when loading the model.",
+    )
+    intra_op_num_threads: int | None = Field(
+        default=None,
+        ge=0,
+        description="Threads used within a single operator. None uses the ONNX Runtime default.",
+    )
+    inter_op_num_threads: int | None = Field(
+        default=None,
+        ge=0,
+        description="Threads used across operators. None uses the ONNX Runtime default.",
+    )
+
+
+__all__ = [
+    "CoreMLProvider",
+    "CpuProvider",
+    "CudaProvider",
+    "ExecutionProvider",
+    "SessionOptionsConfig",
+    "TensorRTProvider",
+]

openstef_foundation_models/integrations/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-FileCopyrightText: 2025 Contributors to the OpenSTEF project <openstef@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""Integrations with neighbouring OpenSTEF frameworks, one module per framework."""