mlflow-optimization 0.1.0__py3-none-any.whl

mlflow_optimization/features/metadata/collect.py

@@ -0,0 +1,52 @@
+"""Common rules for tracking the API, engine, problem and parameters."""
+
+from .types import (
+    CollectedMetadata,
+    ParameterOrigin,
+    ReportedMetadata,
+)
+
+
+def _format_value(value: float | str | bool) -> str:
+    """A string representation of a recorded value."""
+    if isinstance(value, bool):
+        return str(value).lower()
+    return str(value)
+
+
+def collect_metadata(
+    api: str, reported: ReportedMetadata | None, *, include_parameters: bool
+) -> CollectedMetadata:
+    """The recordable experiment conditions of a solve.
+
+    ``reported`` is None when extraction failed.
+    """
+    tags: dict[str, str] = {"optimization.api": api}
+    collected_parameters: dict[str, str] = {}
+    if reported is None:
+        tags["optimization.solver"] = "unknown"
+        return CollectedMetadata(tags=tags, params=collected_parameters)
+
+    tags["optimization.solver"] = reported.solver or "unknown"
+    if reported.api_version is not None:
+        tags["optimization.api_version"] = reported.api_version
+    if reported.solver_version is not None:
+        tags["optimization.solver_version"] = reported.solver_version
+    tags["optimization.wall_time_includes_problem_build"] = _format_value(
+        reported.wall_time_includes_problem_build
+    )
+
+    for problem_name, problem_value in reported.problem.items():
+        collected_parameters[f"problem.{problem_name}"] = _format_value(problem_value)
+
+    if include_parameters:
+        parameter_prefixes: dict[ParameterOrigin, str] = {
+            ParameterOrigin.EXPLICIT: "solver_param.",
+            ParameterOrigin.DEFAULT: "solver_param_default.",
+            ParameterOrigin.EFFECTIVE: "solver_param_effective.",
+        }
+        for parameter in reported.parameters:
+            parameter_key: str = parameter_prefixes[parameter.origin] + parameter.name
+            collected_parameters[parameter_key] = _format_value(parameter.value)
+
+    return CollectedMetadata(tags=tags, params=collected_parameters)

mlflow_optimization/features/metadata/types.py

@@ -0,0 +1,42 @@
+"""Experiment conditions of a solve."""
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field
+from enum import StrEnum
+
+
+class ParameterOrigin(StrEnum):
+    """Where a solver parameter value comes from."""
+
+    EXPLICIT = "explicit"
+    DEFAULT = "default"
+    EFFECTIVE = "effective"
+
+
+@dataclass(frozen=True)
+class SolverParameter:
+    """One solver parameter value and its origin."""
+
+    name: str
+    value: int | float | str | bool
+    origin: ParameterOrigin = ParameterOrigin.EXPLICIT
+
+
+@dataclass(frozen=True)
+class ReportedMetadata:
+    """The experiment conditions as an adapter read them from the solver."""
+
+    solver: str | None = None
+    api_version: str | None = None
+    solver_version: str | None = None
+    problem: Mapping[str, float | str | bool] = field(default_factory=dict)
+    parameters: Sequence[SolverParameter] = ()
+    wall_time_includes_problem_build: bool = False
+
+
+@dataclass(frozen=True)
+class CollectedMetadata:
+    """The tags and params recorded for the experiment conditions, under final key names."""
+
+    tags: Mapping[str, str]
+    params: Mapping[str, str]

mlflow_optimization/features/metrics/normalize.py

@@ -0,0 +1,121 @@
+"""Common rules for metric names, missing values and gaps."""
+
+from math import isfinite
+
+from .types import (
+    COMMON_METRICS,
+    COMMON_STATISTICS,
+    MissingReason,
+    NormalizedMetrics,
+    ReportedMetrics,
+)
+
+
+def _record_reported_metric(
+    values: dict[str, float],
+    missing: dict[str, MissingReason],
+    metric_name: str,
+    reported_value: float | MissingReason,
+) -> None:
+    """Record a finite reported metric or its missing reason."""
+    if isinstance(reported_value, MissingReason):
+        missing[metric_name] = reported_value
+        return
+    if isfinite(reported_value):
+        values[metric_name] = reported_value
+        return
+    missing[metric_name] = MissingReason.UNDEFINED
+
+
+def normalize_metrics(
+    reported: ReportedMetrics | None,
+    *,
+    api: str,
+    has_solution: bool | None,
+    wall_time_seconds: float,
+    include_native: bool,
+) -> NormalizedMetrics:
+    """The recordable metrics of a solve."""
+    values: dict[str, float] = {}
+    missing: dict[str, MissingReason] = {}
+    tags: dict[str, str] = {}
+
+    _record_reported_metric(values, missing, "wall_time_seconds", wall_time_seconds)
+
+    if reported is None:
+        tags["optimization.has_objective"] = "unknown"
+        for metric_name in COMMON_METRICS:
+            if metric_name != "wall_time_seconds":
+                missing[metric_name] = MissingReason.FAILED
+    else:
+        if reported.has_objective is None:
+            tags["optimization.has_objective"] = "unknown"
+        else:
+            tags["optimization.has_objective"] = str(reported.has_objective).lower()
+        if reported.sense is not None:
+            tags["optimization.objective_sense"] = reported.sense.value
+
+        if reported.has_objective is False:
+            missing["objective_value"] = MissingReason.UNDEFINED
+            missing["dual_bound"] = MissingReason.UNDEFINED
+        elif has_solution is False:
+            missing["objective_value"] = MissingReason.NO_SOLUTION
+            _record_reported_metric(values, missing, "dual_bound", reported.dual_bound)
+        else:
+            _record_reported_metric(
+                values, missing, "objective_value", reported.objective_value
+            )
+            _record_reported_metric(values, missing, "dual_bound", reported.dual_bound)
+
+        _record_reported_metric(
+            values, missing, "solver_time_seconds", reported.solver_time_seconds
+        )
+
+        for metric_name, reported_value in reported.statistics.items():
+            if metric_name in COMMON_STATISTICS:
+                _record_reported_metric(values, missing, metric_name, reported_value)
+
+        if reported.has_objective is False:
+            missing["absolute_gap"] = MissingReason.UNDEFINED
+            missing["relative_gap"] = MissingReason.UNDEFINED
+        elif has_solution is False:
+            missing["absolute_gap"] = MissingReason.NO_SOLUTION
+            missing["relative_gap"] = MissingReason.NO_SOLUTION
+        elif "objective_value" in missing:
+            missing["absolute_gap"] = missing["objective_value"]
+            missing["relative_gap"] = missing["objective_value"]
+        elif "dual_bound" in missing:
+            missing["absolute_gap"] = missing["dual_bound"]
+            missing["relative_gap"] = missing["dual_bound"]
+        else:
+            absolute_gap: float = abs(values["objective_value"] - values["dual_bound"])
+            if isfinite(absolute_gap):
+                values["absolute_gap"] = absolute_gap
+                if values["objective_value"] == 0.0:
+                    if absolute_gap == 0.0:
+                        values["relative_gap"] = 0.0
+                    else:
+                        missing["relative_gap"] = MissingReason.UNDEFINED
+                else:
+                    relative_gap: float = absolute_gap / abs(values["objective_value"])
+                    if isfinite(relative_gap):
+                        values["relative_gap"] = relative_gap
+                    else:
+                        missing["relative_gap"] = MissingReason.UNDEFINED
+            else:
+                missing["absolute_gap"] = MissingReason.UNDEFINED
+                missing["relative_gap"] = MissingReason.UNDEFINED
+
+    for metric_name in COMMON_METRICS:
+        if metric_name not in values and metric_name not in missing:
+            missing[metric_name] = MissingReason.UNSUPPORTED
+        if metric_name in missing:
+            tags[f"optimization.missing.{metric_name}"] = missing[metric_name].value
+
+    if reported is not None and include_native:
+        for native_name, native_value in reported.native.items():
+            api_metric_name: str = f"{api}.{native_name}"
+            if isfinite(native_value):
+                values[api_metric_name] = native_value
+
+    return NormalizedMetrics(values=values, tags=tags)

mlflow_optimization/features/metrics/types.py

@@ -0,0 +1,54 @@
+"""Objective, bound, time and search statistics of a solve."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from enum import StrEnum
+
+
+class MissingReason(StrEnum):
+    """Why a value is absent."""
+
+    UNSUPPORTED = "unsupported"
+    UNDEFINED = "undefined"
+    NO_SOLUTION = "no_solution"
+    FAILED = "failed"
+
+
+class ObjectiveSense(StrEnum):
+    """The optimization direction."""
+
+    MINIMIZE = "minimize"
+    MAXIMIZE = "maximize"
+
+
+COMMON_STATISTICS = ("node_count", "simplex_iterations", "barrier_iterations")
+COMMON_METRICS = (
+    "objective_value",
+    "dual_bound",
+    "absolute_gap",
+    "relative_gap",
+    "wall_time_seconds",
+    "solver_time_seconds",
+    *COMMON_STATISTICS,
+)
+
+
+@dataclass(frozen=True)
+class ReportedMetrics:
+    """The numbers of a solve as an adapter read them from the solver."""
+
+    has_objective: bool | None = None
+    sense: ObjectiveSense | None = None
+    objective_value: float | MissingReason = MissingReason.UNSUPPORTED
+    dual_bound: float | MissingReason = MissingReason.UNSUPPORTED
+    solver_time_seconds: float | MissingReason = MissingReason.UNSUPPORTED
+    statistics: Mapping[str, float | MissingReason] = field(default_factory=dict)
+    native: Mapping[str, float] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class NormalizedMetrics:
+    """The metrics and metric-related tags recorded for a solve, under final key names."""
+
+    values: Mapping[str, float]
+    tags: Mapping[str, str]

mlflow_optimization/features/outcomes/classify.py

@@ -0,0 +1,53 @@
+"""Common rules for termination and solution availability."""
+
+from mlflow_optimization.common.diagnostics import warn
+
+from .types import ReportedOutcome, SolveOutcome, Termination
+
+
+def classify_outcome(
+    reported: ReportedOutcome | None,
+    error: BaseException | None,
+    *,
+    has_objective: bool | None = None,
+) -> SolveOutcome:
+    """The validated outcome of a solve."""
+    if reported is None:
+        if isinstance(error, KeyboardInterrupt):
+            return SolveOutcome("unknown", Termination.INTERRUPTED, None)
+        if error is not None:
+            return SolveOutcome("unknown", Termination.ERROR, None)
+        return SolveOutcome("unknown", Termination.UNKNOWN, None)
+
+    termination: Termination = reported.termination
+    if isinstance(error, KeyboardInterrupt):
+        termination = Termination.INTERRUPTED
+    elif error is not None:
+        termination = Termination.ERROR
+    else:
+        if termination is Termination.OPTIMAL and has_objective is False:
+            termination = Termination.FEASIBLE
+        if termination is Termination.FEASIBLE and reported.has_solution is False:
+            warn("Reported feasible termination conflicts with has_solution=False.")
+            termination = Termination.UNKNOWN
+        elif termination is Termination.OPTIMAL and reported.has_solution is False:
+            warn("Reported optimal termination conflicts with has_solution=False.")
+        elif termination is Termination.INFEASIBLE and reported.has_solution is True:
+            warn("Reported infeasible termination conflicts with has_solution=True.")
+
+    return SolveOutcome(reported.raw_status, termination, reported.has_solution)
+
+
+def outcome_tags(outcome: SolveOutcome) -> dict[str, str]:
+    """The tags that describe an outcome."""
+    has_solution: str = "unknown"
+    if outcome.has_solution is True:
+        has_solution = "true"
+    elif outcome.has_solution is False:
+        has_solution = "false"
+
+    return {
+        "optimization.termination": outcome.termination.value,
+        "optimization.raw_status": outcome.raw_status,
+        "optimization.has_solution": has_solution,
+    }

mlflow_optimization/features/outcomes/types.py

@@ -0,0 +1,40 @@
+"""How a solve ended."""
+
+from dataclasses import dataclass
+from enum import StrEnum
+
+
+class Termination(StrEnum):
+    """The solver-independent reason a solve ended."""
+
+    OPTIMAL = "optimal"
+    FEASIBLE = "feasible"
+    INFEASIBLE = "infeasible"
+    UNBOUNDED = "unbounded"
+    INFEASIBLE_OR_UNBOUNDED = "infeasible_or_unbounded"
+    TIME_LIMIT = "time_limit"
+    ITERATION_LIMIT = "iteration_limit"
+    NODE_LIMIT = "node_limit"
+    SOLUTION_LIMIT = "solution_limit"
+    OTHER_LIMIT = "other_limit"
+    INTERRUPTED = "interrupted"
+    ERROR = "error"
+    UNKNOWN = "unknown"
+
+
+@dataclass(frozen=True)
+class ReportedOutcome:
+    """The end state of a solve as an adapter read it from the solver."""
+
+    raw_status: str
+    termination: Termination
+    has_solution: bool | None
+
+
+@dataclass(frozen=True)
+class SolveOutcome:
+    """The validated end state of a solve."""
+
+    raw_status: str
+    termination: Termination
+    has_solution: bool | None

mlflow_optimization/features/runs/ports.py

@@ -0,0 +1,30 @@
+"""Contract of the run storage."""
+
+from collections.abc import Mapping
+from typing import Protocol
+
+from .types import RunReference, SolveRecord
+
+
+class RunStoreError(Exception):
+    """An error raised when run storage fails."""
+
+
+class RunStore(Protocol):
+    """The storage of solve runs in an experiment tracker."""
+
+    def active_run(self) -> RunReference | None:
+        """The user's currently active run."""
+        ...
+
+    def save(self, record: SolveRecord, parent: RunReference | None) -> RunReference:
+        """Create, fill and terminate the run of one finished solve."""
+        ...
+
+    def log_metric(self, run: RunReference, key: str, value: float) -> None:
+        """Add one metric to an existing run."""
+        ...
+
+    def set_tags(self, run: RunReference, tags: Mapping[str, str]) -> None:
+        """Add tags to an existing run."""
+        ...

mlflow_optimization/features/runs/record.py

@@ -0,0 +1,231 @@
+"""Recording of completed solves and later additions to them."""
+
+from collections.abc import Mapping
+from contextvars import ContextVar
+from dataclasses import dataclass
+
+from mlflow_optimization.common.diagnostics import warn
+from mlflow_optimization.features.annotations.annotate import (
+    annotation_tags,
+    current_annotation,
+    user_metric,
+    user_tags,
+)
+from mlflow_optimization.features.annotations.types import Annotation, UserMetric
+from mlflow_optimization.features.artifacts.policy import artifact_refusal
+from mlflow_optimization.features.artifacts.ports import (
+    ArtifactStore,
+    ArtifactStoreError,
+)
+from mlflow_optimization.features.autolog.types import CompletedSolve, SolveObservation
+from mlflow_optimization.features.metadata.collect import collect_metadata
+from mlflow_optimization.features.metadata.types import CollectedMetadata
+from mlflow_optimization.features.metrics.normalize import normalize_metrics
+from mlflow_optimization.features.metrics.types import NormalizedMetrics
+from mlflow_optimization.features.outcomes.classify import (
+    classify_outcome,
+    outcome_tags,
+)
+from mlflow_optimization.features.outcomes.types import SolveOutcome
+
+from .ports import RunStore, RunStoreError
+from .types import RecordOptions, RunReference, RunStatus, SolveRecord
+
+
+@dataclass(frozen=True)
+class _RecordingContext:
+    """The state captured at solve start."""
+
+    parent: RunReference | None
+    annotation: Annotation
+
+
+class SolveRun:
+    """The handle of one recorded solve."""
+
+    def __init__(
+        self,
+        run: RunReference,
+        runs: RunStore,
+        artifacts: ArtifactStore,
+        options: RecordOptions,
+    ) -> None:
+        """Initialize a handle for a recorded solve.
+
+        Parameters
+        ----------
+        run : RunReference
+            Reference to the recorded solve's run.
+        runs : RunStore
+            Store used to add metrics and tags to the run.
+        artifacts : ArtifactStore
+            Store used to attach files to the run.
+        options : RecordOptions
+            Options controlling recorded solver data and artifact logging.
+        """
+        self.run_id: str = run.run_id
+        self._run: RunReference = run
+        self._runs: RunStore = runs
+        self._artifacts: ArtifactStore = artifacts
+        self._options: RecordOptions = options
+
+    def log_metric(
+        self,
+        name: str,
+        value: float,
+        *,
+        unit: str | None = None,
+        description: str | None = None,
+    ) -> None:
+        """Add a user metric to this solve.
+
+        Parameters
+        ----------
+        name : str
+            User metric name; its recorded key is prefixed with ``user.``.
+        value : float
+            Metric value.
+        unit : str, optional
+            Unit tag stored as ``user.{name}.unit``.
+        description : str, optional
+            Description tag stored as ``user.{name}.description``.
+        """
+        try:
+            entry: UserMetric | None = user_metric(
+                name, value, unit=unit, description=description
+            )
+            if entry is not None:
+                self._runs.log_metric(self._run, entry.metric_key, entry.value)
+                if entry.tags:
+                    self._runs.set_tags(self._run, entry.tags)
+        except RunStoreError as error:
+            warn(f"Could not log metric: {error}")
+
+    def set_tags(self, tags: Mapping[str, str]) -> None:
+        """Add user tags to this solve.
+
+        Parameters
+        ----------
+        tags : Mapping[str, str]
+            Tags to add; keys beginning with ``optimization.`` or ``mlflow.``
+            are ignored.
+        """
+        try:
+            accepted: dict[str, str] = user_tags(tags)
+            if accepted:
+                self._runs.set_tags(self._run, accepted)
+        except RunStoreError as error:
+            warn(f"Could not set tags: {error}")
+
+    def log_artifact(self, local_path: str, artifact_path: str | None = None) -> None:
+        """Attach a local file to this solve.
+
+        Parameters
+        ----------
+        local_path : str
+            Path to an existing local file.
+        artifact_path : str, optional
+            Destination directory in run artifacts; if omitted, the file goes
+            to the artifact root.
+        """
+        try:
+            refusal: str | None = artifact_refusal(
+                local_path, enabled=self._options.log_artifacts
+            )
+            if refusal is not None:
+                warn(refusal)
+                return
+            self._artifacts.log_artifact(self.run_id, local_path, artifact_path)
+        except ArtifactStoreError as error:
+            warn(f"Could not log artifact: {error}")
+
+
+_last_solve: ContextVar[SolveRun | None] = ContextVar("last_solve", default=None)
+
+
+class RecordSolve:
+    """The use case that stores each completed solve as one run."""
+
+    def __init__(
+        self, runs: RunStore, artifacts: ArtifactStore, options: RecordOptions
+    ) -> None:
+        self._runs: RunStore = runs
+        self._artifacts: ArtifactStore = artifacts
+        self.options: RecordOptions = options
+
+    def begin(self) -> object:
+        """The parent run and annotations in effect at solve start."""
+        parent: RunReference | None
+        try:
+            parent = self._runs.active_run()
+        except RunStoreError as error:
+            warn(f"Could not get active run: {error}")
+            parent = None
+        return _RecordingContext(parent, current_annotation())
+
+    def record(self, solve: CompletedSolve, context: object) -> None:
+        """Store one completed solve."""
+        if not isinstance(context, _RecordingContext):
+            raise TypeError("Invalid recording context")
+        observation: SolveObservation | None = solve.observation
+        has_objective: bool | None = (
+            observation.metrics.has_objective if observation is not None else None
+        )
+        outcome: SolveOutcome = classify_outcome(
+            observation.outcome if observation else None,
+            solve.error,
+            has_objective=has_objective,
+        )
+        metrics: NormalizedMetrics = normalize_metrics(
+            observation.metrics if observation else None,
+            api=solve.api,
+            has_solution=outcome.has_solution,
+            wall_time_seconds=solve.wall_time_seconds,
+            include_native=self.options.log_solver_metrics,
+        )
+        metadata: CollectedMetadata = collect_metadata(
+            solve.api,
+            observation.metadata if observation else None,
+            include_parameters=self.options.log_solver_parameters,
+        )
+        tags: dict[str, str] = dict(annotation_tags(context.annotation))
+        tags.update(outcome_tags(outcome))
+        tags.update(metrics.tags)
+        tags.update(metadata.tags)
+        tags["optimization.solve_id"] = solve.solve_id
+        tags["optimization.solve_index"] = str(solve.solve_index)
+        if solve.error is not None:
+            tags["optimization.error"] = type(solve.error).__name__
+        status: RunStatus = RunStatus.FINISHED
+        if isinstance(solve.error, KeyboardInterrupt):
+            status = RunStatus.KILLED
+        elif isinstance(solve.error, Exception):
+            status = RunStatus.FAILED
+        solver: str = metadata.tags.get("optimization.solver") or "unknown"
+        record: SolveRecord = SolveRecord(
+            name=f"{solve.api}/{solver} #{solve.solve_index}",
+            status=status,
+            start_time_ms=solve.start_time_ms,
+            end_time_ms=solve.end_time_ms,
+            tags=tags,
+            metrics=metrics.values,
+            params=metadata.params,
+        )
+        run: RunReference
+        try:
+            run = self._runs.save(record, context.parent)
+        except RunStoreError as error:
+            warn(f"Could not record solve: {error}")
+            return
+        _last_solve.set(SolveRun(run, self._runs, self._artifacts, self.options))
+
+
+def last_solve() -> SolveRun | None:
+    """The most recent solve recorded in the current thread or async context.
+
+    Returns
+    -------
+    SolveRun or None
+        None until a solve is recorded in the current context.
+    """
+    return _last_solve.get()

mlflow_optimization/features/runs/types.py

@@ -0,0 +1,43 @@
+"""The stored form of a solve."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from enum import StrEnum
+
+
+@dataclass(frozen=True)
+class RunReference:
+    """The identity of a run in the experiment tracker."""
+
+    run_id: str
+    experiment_id: str
+
+
+class RunStatus(StrEnum):
+    """The lifecycle state a solve's run ends in."""
+
+    FINISHED = "FINISHED"
+    FAILED = "FAILED"
+    KILLED = "KILLED"
+
+
+@dataclass(frozen=True)
+class SolveRecord:
+    """Everything stored for one solve, under final key names."""
+
+    name: str
+    status: RunStatus
+    start_time_ms: int
+    end_time_ms: int
+    tags: Mapping[str, str]
+    metrics: Mapping[str, float]
+    params: Mapping[str, str]
+
+
+@dataclass
+class RecordOptions:
+    """Which optional parts of a solve are recorded."""
+
+    log_solver_parameters: bool = True
+    log_solver_metrics: bool = True
+    log_artifacts: bool = True

mlflow_optimization/integrations/cp_sat/__init__.py

@@ -0,0 +1,5 @@
+"""OR-Tools CP-SAT integration."""
+
+from .adapter import CpSatAdapter
+
+__all__ = ["CpSatAdapter"]