aps-agent/scripts/check_solver_runtime.py

from __future__ import annotations

import argparse
import subprocess
import sys
from collections.abc import Callable

RUNTIME_PREFIX = "SOLVER_PROBE_RUNTIME"
PROBE_OK_PREFIX = "SOLVER_PROBE_OK"
PROBE = r"""
import importlib.util
import json
import platform
import site
import sys
from pathlib import Path


def module_origin(package):
    spec = importlib.util.find_spec(package)
    if spec is None:
        return None
    if spec.origin and spec.origin not in {"built-in", "frozen"}:
        return str(Path(spec.origin).resolve())
    locations = tuple(spec.submodule_search_locations or ())
    return str(Path(locations[0]).resolve()) if locations else spec.origin


def is_within_runtime(origin, prefix):
    if not origin:
        return False
    try:
        Path(origin).resolve().relative_to(Path(prefix).resolve())
    except ValueError:
        return False
    return True


package_origins = {
    package: module_origin(package) for package in ("ortools", "numpy", "pandas")
}
runtime = {
    "runtimeIdentity": (
        f"{platform.python_implementation()} {platform.python_version()} "
        f"{platform.system()}-{platform.machine()}"
    ),
    "executable": str(Path(sys.executable).resolve()),
    "baseExecutable": str(
        Path(getattr(sys, "_base_executable", None) or sys.executable).resolve()
    ),
    "prefix": str(Path(sys.prefix).resolve()),
    "basePrefix": str(Path(sys.base_prefix).resolve()),
    "ENABLE_USER_SITE": bool(site.ENABLE_USER_SITE),
    "frozen": bool(getattr(sys, "frozen", False)),
    "packageOrigins": package_origins,
}
print("SOLVER_PROBE_RUNTIME " + json.dumps(runtime, ensure_ascii=True, sort_keys=True), flush=True)

safety_errors = []
if not runtime["frozen"]:
    if runtime["ENABLE_USER_SITE"]:
        safety_errors.append("ENABLE_USER_SITE must be disabled")
    for package, origin in package_origins.items():
        if origin is None:
            safety_errors.append(f"{package} is not installed")
        elif not is_within_runtime(origin, runtime["prefix"]):
            safety_errors.append(f"{package} is outside runtime prefix: {origin}")
if safety_errors:
    raise SystemExit("SOLVER_RUNTIME_UNSAFE: " + "; ".join(safety_errors))

import numpy
import ortools
import pandas
from ortools.sat.python import cp_model

model = cp_model.CpModel()
x = model.new_int_var(0, 10, "x")
model.maximize(x)
solver = cp_model.CpSolver()
status = solver.solve(model)
if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE) or solver.value(x) != 10:
    raise SystemExit("CP-SAT smoke solve failed")

runtime["packages"] = {
    "ortools": {"path": str(Path(ortools.__file__).resolve()), "version": ortools.__version__},
    "numpy": {"path": str(Path(numpy.__file__).resolve()), "version": numpy.__version__},
    "pandas": {"path": str(Path(pandas.__file__).resolve()), "version": pandas.__version__},
}
print("SOLVER_PROBE_OK " + json.dumps(runtime, ensure_ascii=True, sort_keys=True), flush=True)
"""

FATAL_MARKERS = (
    "windows fatal exception",
    "fatal python error",
    "0xc0000139",
    "winerror 127",
    "access violation",
    "segmentation fault",
)


def run_probe(
    iterations: int,
    *,
    timeout_seconds: float = 30.0,
    runner: Callable[..., subprocess.CompletedProcess[str]] = subprocess.run,
    executable: str | None = None,
) -> int:
    command = [executable or sys.executable, "-X", "faulthandler", "-c", PROBE]
    for attempt in range(1, iterations + 1):
        try:
            completed = runner(
                command,
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                timeout=timeout_seconds,
                check=False,
            )
        except subprocess.TimeoutExpired:
            print(
                f"solver probe timed out on attempt {attempt}/{iterations} "
                f"after {timeout_seconds:g}s",
                file=sys.stderr,
            )
            return 1
        except OSError as exc:
            print(
                f"solver probe could not start on attempt {attempt}/{iterations}: {exc}",
                file=sys.stderr,
            )
            return 1

        stdout = completed.stdout or ""
        stderr = completed.stderr or ""
        output = f"{stdout}\n{stderr}"
        fatal = next(
            (marker for marker in FATAL_MARKERS if marker in output.lower()), None
        )
        success_line = next(
            (
                line
                for line in stdout.splitlines()
                if line.startswith(f"{PROBE_OK_PREFIX} ")
            ),
            None,
        )
        runtime_line = next(
            (
                line
                for line in stdout.splitlines()
                if line.startswith(f"{RUNTIME_PREFIX} ")
            ),
            None,
        )
        if (
            completed.returncode != 0
            or fatal
            or success_line is None
            or runtime_line is None
        ):
            print(
                f"solver probe failed on attempt {attempt}/{iterations}",
                file=sys.stderr,
            )
            if fatal:
                print(f"fatal marker detected: {fatal}", file=sys.stderr)
            print(output.strip(), file=sys.stderr)
            return 1
        print(success_line)
    print(f"solver runtime probe passed ({iterations} isolated processes)")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Fail closed on solver runtime or native-loader errors."
    )
    parser.add_argument("--iterations", type=int, default=5)
    parser.add_argument("--timeout-seconds", type=float, default=30.0)
    args = parser.parse_args()
    if args.iterations < 1:
        parser.error("--iterations must be at least 1")
    if args.timeout_seconds <= 0:
        parser.error("--timeout-seconds must be greater than 0")
    return run_probe(args.iterations, timeout_seconds=args.timeout_seconds)


if __name__ == "__main__":
    raise SystemExit(main())