mlflow-optimization 0.1.0__py3-none-any.whl

mlflow_optimization/integrations/pyscipopt/__init__.py

@@ -0,0 +1,5 @@
+"""The PySCIPOpt solver integration package."""
+
+from .adapter import PyScipOptAdapter
+
+__all__ = ["PyScipOptAdapter"]

mlflow_optimization/integrations/pyscipopt/adapter.py

@@ -0,0 +1,657 @@
+"""The PySCIPOpt solver integration."""
+
+import math
+from collections.abc import Callable
+from functools import wraps
+from inspect import Signature, signature
+from types import ModuleType
+from typing import Any
+from weakref import WeakKeyDictionary
+
+import pyscipopt
+import pyscipopt.scip
+from mlflow_optimization.common.patching import (
+    rebind_module_names,
+    receiver_of,
+    restore_module_names,
+)
+from mlflow_optimization.features.autolog.ports import ObservationError, SolveObserver
+from mlflow_optimization.features.autolog.types import SolveObservation
+from mlflow_optimization.features.metadata.types import (
+    ParameterOrigin,
+    ReportedMetadata,
+    SolverParameter,
+)
+from mlflow_optimization.features.metrics.types import (
+    MissingReason,
+    ObjectiveSense,
+    ReportedMetrics,
+)
+from mlflow_optimization.features.outcomes.types import ReportedOutcome, Termination
+
+_STATUS_TERMINATIONS: dict[str, Termination] = {
+    "unknown": Termination.UNKNOWN,
+    "optimal": Termination.OPTIMAL,
+    "infeasible": Termination.INFEASIBLE,
+    "unbounded": Termination.UNBOUNDED,
+    "inforunbd": Termination.INFEASIBLE_OR_UNBOUNDED,
+    "userinterrupt": Termination.INTERRUPTED,
+    "terminate": Termination.INTERRUPTED,
+    "nodelimit": Termination.NODE_LIMIT,
+    "totalnodelimit": Termination.NODE_LIMIT,
+    "stallnodelimit": Termination.NODE_LIMIT,
+    "timelimit": Termination.TIME_LIMIT,
+    "memlimit": Termination.OTHER_LIMIT,
+    "gaplimit": Termination.OTHER_LIMIT,
+    "primallimit": Termination.OTHER_LIMIT,
+    "duallimit": Termination.OTHER_LIMIT,
+    "sollimit": Termination.SOLUTION_LIMIT,
+    "bestsollimit": Termination.SOLUTION_LIMIT,
+    "restartlimit": Termination.OTHER_LIMIT,
+}
+
+
+class PyScipOptAdapter:
+    """The PySCIPOpt solver integration."""
+
+    api: str = "pyscipopt"
+
+    def __init__(self) -> None:
+        self._installed: bool = False
+        self._observer: SolveObserver | None = None
+        self._original_model: type[Any] | None = None
+        self._logging_model: type[Any] | None = None
+        self._rebound_names: list[tuple[ModuleType, str]] = []
+        self._counter_snapshots: WeakKeyDictionary[
+            Any,
+            tuple[int | MissingReason, dict[str, float | MissingReason]],
+        ] = WeakKeyDictionary()
+
+    def install(self, observer: SolveObserver) -> None:
+        """A patch that replaces model references with an observing subclass."""
+        if self._installed:
+            self._observer = observer
+            return
+
+        original_model: type[Any] = pyscipopt.Model
+        integration: PyScipOptAdapter = self
+        default_parameters: dict[str, bool | int | float | str] | None = (
+            _read_default_parameters(original_model)
+        )
+
+        class LoggingModel(original_model):
+            """A PySCIPOpt model that observes optimization calls."""
+
+        entry_point_names: tuple[str, ...] = (
+            "optimize",
+            "optimizeNogil",
+            "solveConcurrent",
+        )
+        for entry_point_name in entry_point_names:
+            original_entry_point: Any = getattr(original_model, entry_point_name, None)
+            if not callable(original_entry_point):
+                continue
+            observed_entry_point: Callable[..., Any] = _wrap_solve_entry_point(
+                original_entry_point,
+                integration,
+                default_parameters,
+            )
+            setattr(LoggingModel, entry_point_name, observed_entry_point)
+
+        logging_model: type[Any] = LoggingModel
+        rebound_names: list[tuple[ModuleType, str]] = rebind_module_names(
+            original_model,
+            logging_model,
+        )
+
+        self._observer = observer
+        self._original_model = original_model
+        self._logging_model = logging_model
+        self._rebound_names = rebound_names
+        self._installed = True
+
+    def uninstall(self) -> None:
+        """The restoration of original PySCIPOpt model references."""
+        if not self._installed:
+            return
+
+        assert self._original_model is not None
+        assert self._logging_model is not None
+        original_model: type[Any] = self._original_model
+        logging_model: type[Any] = self._logging_model
+        restore_module_names(self._rebound_names, original_model, logging_model)
+        rebind_module_names(logging_model, original_model)
+
+        self._installed = False
+        self._observer = None
+        self._original_model = None
+        self._logging_model = None
+        self._rebound_names = []
+        self._counter_snapshots.clear()
+
+
+def _wrap_solve_entry_point(
+    original_entry_point: Callable[..., Any],
+    integration: PyScipOptAdapter,
+    default_parameters: dict[str, bool | int | float | str] | None,
+) -> Callable[..., Any]:
+    """An observed wrapper for one original PySCIPOpt solve entry point."""
+    original_signature: Signature = signature(original_entry_point)
+
+    @wraps(original_entry_point)
+    def observed_entry_point(
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        observer: SolveObserver | None = integration._observer
+        model: Any = receiver_of(original_entry_point, args, kwargs)
+        if model is None or observer is None:
+            return original_entry_point(*args, **kwargs)
+        original_signature.bind(*args, **kwargs)
+        counter_snapshots: WeakKeyDictionary[
+            Any,
+            tuple[int | MissingReason, dict[str, float | MissingReason]],
+        ] = integration._counter_snapshots
+        stage_before_solve: int | MissingReason = _read_value(
+            lambda: int(model.getStage()),
+            MissingReason.FAILED,
+        )
+        counter_snapshot: (
+            tuple[
+                int | MissingReason,
+                dict[str, float | MissingReason],
+            ]
+            | None
+        ) = None
+        counter_cache_available: bool = True
+        try:
+            counter_snapshot = counter_snapshots.get(model)
+        except Exception:  # noqa: BLE001
+            counter_cache_available = False
+        counters_before_solve: dict[str, float | MissingReason] | None = None
+        if not counter_cache_available or isinstance(
+            stage_before_solve,
+            MissingReason,
+        ):
+            counters_before_solve = _unavailable_solve_counters()
+        elif counter_snapshot is None:
+            if stage_before_solve != pyscipopt.scip.PY_SCIP_STAGE.PROBLEM:
+                counters_before_solve = _unavailable_solve_counters()
+        else:
+            previous_stage: int | MissingReason = counter_snapshot[0]
+            if isinstance(previous_stage, MissingReason):
+                counters_before_solve = _unavailable_solve_counters()
+            elif stage_before_solve == previous_stage:
+                counters_before_solve = counter_snapshot[1]
+            elif stage_before_solve != pyscipopt.scip.PY_SCIP_STAGE.PROBLEM:
+                counters_before_solve = _unavailable_solve_counters()
+
+        # NOTE: SCIP changes parameters while solving, so capture configured values first.
+        parameters_before_solve: tuple[SolverParameter, ...] = _read_parameters(
+            model,
+            default_parameters,
+        )
+
+        def call() -> Any:
+            return original_entry_point(*args, **kwargs)
+
+        def extract(
+            result: Any | None, error: BaseException | None
+        ) -> SolveObservation:
+            try:
+                counters_after_solve: dict[str, float | MissingReason] | None = None
+                if error is None:
+                    counters_after_solve = _read_solve_counters(model)
+                    stage_after_solve: int | MissingReason = _read_value(
+                        lambda: int(model.getStage()),
+                        MissingReason.FAILED,
+                    )
+                    try:
+                        counter_snapshots[model] = (
+                            stage_after_solve,
+                            counters_after_solve,
+                        )
+                    except Exception:  # noqa: BLE001
+                        counter_snapshots.clear()
+                else:
+                    try:
+                        counter_snapshots[model] = (
+                            MissingReason.FAILED,
+                            _unavailable_solve_counters(),
+                        )
+                    except Exception:  # noqa: BLE001
+                        counter_snapshots.clear()
+                return _extract_observation(
+                    model,
+                    error,
+                    parameters_before_solve,
+                    counters_before_solve,
+                    counters_after_solve,
+                )
+            except Exception as extraction_error:
+                raise ObservationError(
+                    "Could not build a PySCIPOpt solve observation."
+                ) from extraction_error
+
+        return observer.observe(integration.api, call, extract)
+
+    return observed_entry_point
+
+
+def _extract_observation(
+    model: Any,
+    error: BaseException | None,
+    parameters_before_solve: tuple[SolverParameter, ...],
+    counters_before_solve: dict[str, float | MissingReason] | None,
+    counters_after_solve: dict[str, float | MissingReason] | None,
+) -> SolveObservation:
+    """The SDK-free observation of one PySCIPOpt optimization call."""
+    if error is not None:
+        error_termination: Termination = (
+            Termination.INTERRUPTED
+            if isinstance(error, KeyboardInterrupt)
+            else Termination.ERROR
+        )
+        error_metrics: ReportedMetrics = ReportedMetrics(
+            has_objective=None,
+            sense=None,
+            objective_value=MissingReason.FAILED,
+            dual_bound=MissingReason.FAILED,
+            solver_time_seconds=MissingReason.FAILED,
+            statistics={},
+            native={},
+        )
+        error_metadata: ReportedMetadata = _read_metadata(
+            model, parameters_before_solve
+        )
+        error_outcome: ReportedOutcome = ReportedOutcome(
+            raw_status="unknown",
+            termination=error_termination,
+            has_solution=None,
+        )
+        return SolveObservation(
+            api="pyscipopt",
+            outcome=error_outcome,
+            metrics=error_metrics,
+            metadata=error_metadata,
+        )
+
+    raw_status: str = _read_value(lambda: str(model.getStatus()), "unknown")
+    has_solution: bool | None = _read_value(lambda: model.getNSols() > 0, None)
+    termination: Termination = _STATUS_TERMINATIONS.get(
+        raw_status.lower(), Termination.UNKNOWN
+    )
+
+    has_objective: bool | None = _read_has_objective(model)
+    sense_name: str | MissingReason = _read_value(
+        lambda: str(model.getObjectiveSense()).lower(),
+        MissingReason.FAILED,
+    )
+    sense: ObjectiveSense | None = (
+        {
+            "minimize": ObjectiveSense.MINIMIZE,
+            "maximize": ObjectiveSense.MAXIMIZE,
+        }.get(sense_name)
+        if isinstance(sense_name, str)
+        else None
+    )
+
+    gap: float | MissingReason | None = None
+    if has_objective is False:
+        objective_value: float | MissingReason = MissingReason.UNDEFINED
+        dual_bound: float | MissingReason = MissingReason.UNDEFINED
+    else:
+        solver_infinity: float | MissingReason = _read_number(
+            lambda: float(model.infinity())
+        )
+        dual_bound = _read_scip_number(
+            lambda: float(model.getDualbound()),
+            solver_infinity,
+        )
+        if has_solution is False:
+            objective_value = MissingReason.NO_SOLUTION
+        else:
+            objective_value = _read_scip_number(
+                lambda: float(model.getObjVal()),
+                solver_infinity,
+            )
+        if has_solution is not False:
+            gap = _read_scip_number(
+                lambda: float(model.getGap()),
+                solver_infinity,
+            )
+
+    assert counters_after_solve is not None
+    solver_time_seconds: float | MissingReason = _counter_difference(
+        counters_after_solve["solver_time_seconds"],
+        counters_before_solve,
+        "solver_time_seconds",
+    )
+    node_count: float | MissingReason = _counter_difference(
+        counters_after_solve["node_count"],
+        counters_before_solve,
+        "node_count",
+    )
+    statistics: dict[str, float | MissingReason] = {"node_count": node_count}
+    native: dict[str, float] = {}
+
+    _add_native_counter(
+        native,
+        "solution_count",
+        counters_after_solve["solution_count"],
+        counters_before_solve,
+    )
+    _add_native_counter(
+        native,
+        "solutions_found",
+        counters_after_solve["solutions_found"],
+        counters_before_solve,
+    )
+    _add_native_counter(
+        native,
+        "lp_count",
+        counters_after_solve["lp_count"],
+        counters_before_solve,
+    )
+    _add_native_counter(
+        native,
+        "lp_iterations",
+        counters_after_solve["lp_iterations"],
+        counters_before_solve,
+    )
+    _add_native_counter(
+        native,
+        "total_node_count",
+        counters_after_solve["total_node_count"],
+        counters_before_solve,
+    )
+    _add_native_counter(
+        native,
+        "presolving_time_seconds",
+        counters_after_solve["presolving_time_seconds"],
+        counters_before_solve,
+    )
+    if gap is not None and not isinstance(gap, MissingReason):
+        native["gap"] = gap
+
+    metrics: ReportedMetrics = ReportedMetrics(
+        has_objective=has_objective,
+        sense=sense,
+        objective_value=objective_value,
+        dual_bound=dual_bound,
+        solver_time_seconds=solver_time_seconds,
+        statistics=statistics,
+        native=native,
+    )
+    metadata: ReportedMetadata = _read_metadata(model, parameters_before_solve)
+    outcome: ReportedOutcome = ReportedOutcome(
+        raw_status=raw_status,
+        termination=termination,
+        has_solution=has_solution,
+    )
+
+    return SolveObservation(
+        api="pyscipopt", outcome=outcome, metrics=metrics, metadata=metadata
+    )
+
+
+def _read_has_objective(model: Any) -> bool | None:
+    """Whether the model has a nonzero objective function."""
+    objective_terms_present: bool | MissingReason = _read_value(
+        lambda: bool(model.getObjective().terms),
+        MissingReason.FAILED,
+    )
+    objective_offset: float | MissingReason = _read_value(
+        lambda: float(model.getObjoffset()),
+        MissingReason.FAILED,
+    )
+
+    if objective_terms_present is True:
+        return True
+    if isinstance(objective_offset, float) and objective_offset != 0.0:
+        return True
+    if objective_terms_present is False and objective_offset == 0.0:
+        return False
+    return None
+
+
+def _read_metadata(
+    model: Any,
+    parameters: tuple[SolverParameter, ...],
+) -> ReportedMetadata:
+    """The solver and problem metadata reported by PySCIPOpt."""
+    api_version: str | None = _read_value(
+        lambda: str(
+            pyscipopt.__version__  # pyright: ignore[reportPrivateImportUsage] -- Exposed at runtime.
+        ),
+        None,
+    )
+    major_version: int | None = _read_value(lambda: int(model.getMajorVersion()), None)
+    minor_version: int | None = _read_value(lambda: int(model.getMinorVersion()), None)
+    technical_version: int | None = _read_value(
+        lambda: int(model.getTechVersion()), None
+    )
+    solver_version: str | None = None
+
+    if (
+        major_version is not None
+        and minor_version is not None
+        and technical_version is not None
+    ):
+        solver_version = f"{major_version}.{minor_version}.{technical_version}"
+
+    problem: dict[str, int] = {}
+    variable_count: int | None = _read_value(
+        lambda: int(model.getNVars(transformed=False)), None
+    )
+    constraint_count: int | None = _read_value(
+        lambda: int(model.getNConss(transformed=False)), None
+    )
+
+    if variable_count is not None:
+        problem["num_variables"] = variable_count
+    if constraint_count is not None:
+        problem["num_constraints"] = constraint_count
+
+    variables: list[Any] | MissingReason = _read_value(
+        lambda: model.getVars(transformed=False),
+        MissingReason.FAILED,
+    )
+    if not isinstance(variables, MissingReason):
+        integer_variable_count: int = 0
+        integer_variable_count_available: bool = True
+        for variable in variables:
+            variable_type: str | MissingReason = _read_value(
+                lambda variable=variable: str(variable.vtype()).upper(),
+                MissingReason.FAILED,
+            )
+            if isinstance(variable_type, MissingReason):
+                integer_variable_count_available = False
+                break
+            if variable_type in {"BINARY", "INTEGER", "IMPLINT"}:
+                integer_variable_count += 1
+        if integer_variable_count_available:
+            problem["num_integer_variables"] = integer_variable_count
+
+    return ReportedMetadata(
+        solver="scip",
+        api_version=api_version,
+        solver_version=solver_version,
+        problem=problem,
+        parameters=parameters,
+        wall_time_includes_problem_build=False,
+    )
+
+
+def _read_default_parameters(
+    model_class: type[Any],
+) -> dict[str, bool | int | float | str] | None:
+    """The PySCIPOpt parameter values of a new model."""
+    defaults_model: Any | MissingReason = _read_value(
+        lambda: model_class(),
+        MissingReason.FAILED,
+    )
+    if isinstance(defaults_model, MissingReason):
+        return None
+
+    default_values: dict[str, Any] | MissingReason
+    try:
+        default_values = _read_value(
+            lambda: defaults_model.getParams(),
+            MissingReason.FAILED,
+        )
+    finally:
+        _read_value(lambda: defaults_model.free(), None)
+
+    if isinstance(default_values, MissingReason) or not isinstance(
+        default_values, dict
+    ):
+        return None
+
+    default_parameters: dict[str, bool | int | float | str] = {}
+    for parameter_name, parameter_value in default_values.items():
+        if isinstance(parameter_name, str) and isinstance(
+            parameter_value,
+            (bool, int, float, str),
+        ):
+            default_parameters[parameter_name] = parameter_value
+    return default_parameters
+
+
+def _read_parameters(
+    model: Any,
+    default_parameters: dict[str, bool | int | float | str] | None,
+) -> tuple[SolverParameter, ...]:
+    """The PySCIPOpt parameters that differ from their defaults."""
+    if default_parameters is None:
+        return ()
+
+    parameters: list[SolverParameter] = []
+    for parameter_name, default_value in default_parameters.items():
+        parameter_value: Any | MissingReason = _read_value(
+            lambda parameter_name=parameter_name: model.getParam(parameter_name),
+            MissingReason.FAILED,
+        )
+        if isinstance(parameter_value, MissingReason):
+            continue
+        if not isinstance(parameter_value, (bool, int, float, str)):
+            continue
+        if parameter_value == default_value:
+            continue
+        parameters.append(
+            SolverParameter(
+                name=parameter_name,
+                value=parameter_value,
+                origin=ParameterOrigin.EXPLICIT,
+            )
+        )
+
+    return tuple(parameters)
+
+
+def _read_solve_counters(model: Any) -> dict[str, float | MissingReason]:
+    """The current values of PySCIPOpt counters and timers."""
+    return {
+        "solver_time_seconds": _read_number(lambda: float(model.getSolvingTime())),
+        "node_count": _read_number(lambda: float(model.getNNodes())),
+        "solution_count": _read_number(lambda: float(model.getNSols())),
+        "solutions_found": _read_number(lambda: float(model.getNSolsFound())),
+        "lp_count": _read_number(lambda: float(model.getNLPs())),
+        "lp_iterations": _read_number(lambda: float(model.getNLPIterations())),
+        "total_node_count": _read_number(lambda: float(model.getNTotalNodes())),
+        "presolving_time_seconds": _read_number(
+            lambda: float(model.getPresolvingTime())
+        ),
+    }
+
+
+def _unavailable_solve_counters() -> dict[str, float | MissingReason]:
+    """Unavailable PySCIPOpt counters for a per-solve baseline."""
+    return {
+        "solver_time_seconds": MissingReason.FAILED,
+        "node_count": MissingReason.FAILED,
+        "solution_count": MissingReason.FAILED,
+        "solutions_found": MissingReason.FAILED,
+        "lp_count": MissingReason.FAILED,
+        "lp_iterations": MissingReason.FAILED,
+        "total_node_count": MissingReason.FAILED,
+        "presolving_time_seconds": MissingReason.FAILED,
+    }
+
+
+def _add_native_counter(
+    native: dict[str, float],
+    metric_name: str,
+    current_value: float | MissingReason,
+    counters_before_solve: dict[str, float | MissingReason] | None,
+) -> None:
+    """One available PySCIPOpt counter difference in native metrics."""
+    value: float | MissingReason = _counter_difference(
+        current_value,
+        counters_before_solve,
+        metric_name,
+    )
+    if isinstance(value, MissingReason):
+        return
+    native[metric_name] = value
+
+
+def _read_number(
+    reader: Callable[[], float],
+) -> float | MissingReason:
+    """A guarded PySCIPOpt numeric value."""
+    return _read_value(reader, MissingReason.FAILED)
+
+
+def _counter_difference(
+    current_value: float | MissingReason,
+    counters_before_solve: dict[str, float | MissingReason] | None,
+    counter_name: str,
+) -> float | MissingReason:
+    """The nonnegative difference between two PySCIPOpt counter values."""
+    if isinstance(current_value, MissingReason):
+        return current_value
+    if counters_before_solve is None:
+        if counter_name == "solution_count":
+            return MissingReason.UNSUPPORTED
+        return current_value
+    previous_value: float | MissingReason = counters_before_solve.get(
+        counter_name,
+        MissingReason.FAILED,
+    )
+    if isinstance(previous_value, MissingReason):
+        return previous_value
+    difference: float = current_value - previous_value
+    if difference < 0.0:
+        return MissingReason.FAILED
+    return difference
+
+
+def _read_scip_number(
+    reader: Callable[[], float],
+    solver_infinity: float | MissingReason,
+) -> float | MissingReason:
+    """A PySCIPOpt numeric value with its infinity sentinel normalized."""
+    value: float | MissingReason = _read_number(reader)
+    if isinstance(value, MissingReason):
+        return value
+    if math.isinf(value):
+        return value
+    if isinstance(solver_infinity, MissingReason) or solver_infinity <= 0.0:
+        return MissingReason.FAILED
+    if value >= solver_infinity:
+        return math.inf
+    if value <= -solver_infinity:
+        return -math.inf
+    return value
+
+
+def _read_value[Value](
+    reader: Callable[[], Value],
+    unavailable: Value,
+) -> Value:
+    """A PySCIPOpt value or its fallback after a supported read error."""
+    try:
+        return reader()
+    except Exception:  # noqa: BLE001
+        return unavailable

mlflow_optimization/integrations/routing/__init__.py

@@ -0,0 +1,5 @@
+"""The OR-Tools Routing integration."""
+
+from .adapter import RoutingAdapter
+
+__all__ = ["RoutingAdapter"]