SystemSimulationApp/app/simulation/warmup.py

"""Process-local warm-up for the numerical simulation runtime."""

from __future__ import annotations

from dataclasses import asdict, dataclass
import logging
import os
from threading import Lock
from time import perf_counter
from typing import Literal


LOGGER = logging.getLogger(__name__)
WarmupStatus = Literal["completed", "failed", "disabled"]


@dataclass(frozen=True)
class SimulationWarmupReport:
    status: WarmupStatus
    duration_ms: float
    error: str | None = None

    def as_dict(self) -> dict[str, object]:
        return asdict(self)


_WARMUP_LOCK = Lock()
_WARMUP_REPORT: SimulationWarmupReport | None = None


def simulation_warmup_enabled() -> bool:
    raw_value = os.getenv("SIMULATIONAPP_WARMUP", "on").strip().lower()
    if raw_value in {"", "1", "true", "yes", "on"}:
        return True
    if raw_value in {"0", "false", "no", "off"}:
        return False
    raise ValueError(
        "SIMULATIONAPP_WARMUP must be one of: on, off, true, false, 1, 0."
    )


def _run_numerical_warmup() -> None:
    """Exercise only in-memory SciPy paths used by real simulations."""

    import numpy as np
    from scipy.integrate import BDF, DOP853, LSODA, RK23, RK45, Radau, solve_ivp
    from scipy.optimize import brentq, least_squares
    from scipy.optimize._numdiff import group_columns
    from scipy.sparse import csc_matrix, csr_matrix

    # Importing these classes is intentional even though the micro solve below
    # uses BDF: the stepwise solver selects them dynamically at runtime.
    solver_types = (BDF, DOP853, LSODA, RK23, RK45, Radau)
    if len(solver_types) != 6:
        raise RuntimeError("SciPy solver warm-up did not load every supported method.")

    sparsity = csc_matrix(np.array([[1.0]], dtype=float))
    groups = group_columns(sparsity)
    if groups.shape != (1,):
        raise RuntimeError("SciPy Jacobian grouping warm-up returned an invalid shape.")

    integration = solve_ivp(
        lambda _time, state: -state,
        (0.0, 1.0e-4),
        np.array([1.0], dtype=float),
        method="BDF",
        t_eval=np.array([0.0, 1.0e-4], dtype=float),
        jac_sparsity=sparsity,
        rtol=1.0e-6,
        atol=1.0e-9,
    )
    if not integration.success or not np.isfinite(integration.y).all():
        raise RuntimeError("SciPy integration warm-up did not complete successfully.")

    algebraic_sparsity = csr_matrix(np.eye(2, dtype=bool))
    algebraic = least_squares(
        lambda state: np.array(
            [state[0] - 1.0, state[1] - 2.0],
            dtype=float,
        ),
        np.array([0.5, 0.5], dtype=float),
        bounds=(
            np.array([0.0, 0.0], dtype=float),
            np.array([3.0, 3.0], dtype=float),
        ),
        jac_sparsity=algebraic_sparsity,
        tr_solver="lsmr",
    )
    if (
        not algebraic.success
        or not np.isfinite(algebraic.x).all()
        or not np.allclose(algebraic.x, np.array([1.0, 2.0]), atol=1.0e-8)
    ):
        raise RuntimeError("SciPy algebraic warm-up did not complete successfully.")

    root = brentq(lambda value: value - 0.5, 0.0, 1.0)
    if abs(root - 0.5) > 1.0e-12:
        raise RuntimeError("SciPy scalar root warm-up returned an invalid result.")

    # Compile the cached v3 XSD through the same public validation path.  The
    # intentionally incomplete document is never accepted or persisted.
    from app.system_xml import validate_system_xml_document

    validate_system_xml_document(b"<System/>")


def warm_up_simulation_runtime() -> SimulationWarmupReport:
    """Warm one worker exactly once, returning a startup diagnostic report.

    Ordinary warm-up failures are reported but do not prevent the editor and
    non-simulation APIs from starting.  ``MemoryError`` remains fatal because
    continuing a worker under memory exhaustion is unsafe.
    """

    global _WARMUP_REPORT

    with _WARMUP_LOCK:
        if _WARMUP_REPORT is not None:
            return _WARMUP_REPORT
        if not simulation_warmup_enabled():
            _WARMUP_REPORT = SimulationWarmupReport(
                status="disabled",
                duration_ms=0.0,
            )
            return _WARMUP_REPORT

        started = perf_counter()
        try:
            _run_numerical_warmup()
        except MemoryError:
            raise
        except Exception as exc:
            _WARMUP_REPORT = SimulationWarmupReport(
                status="failed",
                duration_ms=(perf_counter() - started) * 1000.0,
                error=f"{type(exc).__name__}: {exc}",
            )
            LOGGER.exception("Simulation runtime warm-up failed; startup will continue.")
        else:
            _WARMUP_REPORT = SimulationWarmupReport(
                status="completed",
                duration_ms=(perf_counter() - started) * 1000.0,
            )
            LOGGER.info(
                "Simulation runtime warm-up completed in %.1f ms.",
                _WARMUP_REPORT.duration_ms,
            )
        return _WARMUP_REPORT


def _reset_simulation_warmup_for_tests() -> None:
    global _WARMUP_REPORT
    with _WARMUP_LOCK:
        _WARMUP_REPORT = None


__all__ = [
    "SimulationWarmupReport",
    "simulation_warmup_enabled",
    "warm_up_simulation_runtime",
]