mlflow-optimization 0.1.0__py3-none-any.whl

mlflow_optimization/__init__.py

@@ -0,0 +1,50 @@
+"""Automatic MLflow logging of mathematical optimization solves."""
+
+from collections.abc import Iterable
+
+from mlflow_optimization import _bootstrap
+from mlflow_optimization.common.diagnostics import MlflowOptimizationWarning
+from mlflow_optimization.features.annotations.annotate import context
+from mlflow_optimization.features.runs.record import SolveRun, last_solve
+
+__all__ = ["MlflowOptimizationWarning", "SolveRun", "autolog", "context", "last_solve"]
+
+
+def autolog(
+    disable: bool = False,
+    *,
+    libraries: Iterable[str] | None = None,
+    log_solver_parameters: bool = True,
+    log_solver_metrics: bool = True,
+    log_artifacts: bool = True,
+) -> None:
+    """Configure automatic logging of optimization solves.
+
+    Parameters
+    ----------
+    disable : bool, default False
+        Disable automatic logging; other settings are ignored when True.
+    libraries : Iterable[str] or None, optional
+        Supported names: gurobipy, pyscipopt, pyomo, mathopt, cp_sat, routing.
+        None selects all available integrations; an empty iterable selects none.
+    log_solver_parameters : bool, default True
+        Whether solver parameters are logged.
+    log_solver_metrics : bool, default True
+        Whether native solver metrics are logged.
+    log_artifacts : bool, default True
+        Whether files attached through ``SolveRun.log_artifact`` are logged.
+
+    Raises
+    ------
+    ImportError
+        An explicitly selected integration cannot be imported.
+    ValueError
+        An unsupported library name is selected while logging is enabled.
+    """
+    _bootstrap.configure(
+        disable=disable,
+        libraries=libraries,
+        log_solver_parameters=log_solver_parameters,
+        log_solver_metrics=log_solver_metrics,
+        log_artifacts=log_artifacts,
+    )

mlflow_optimization/_bootstrap.py

@@ -0,0 +1,129 @@
+"""Composition root: wires feature ports to concrete adapters."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from importlib import import_module
+from types import ModuleType
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from mlflow_optimization.features.autolog.ports import SolverIntegration
+    from mlflow_optimization.features.autolog.service import AutologManager
+    from mlflow_optimization.features.runs.record import RecordSolve
+    from mlflow_optimization.features.runs.types import RecordOptions
+    from mlflow_optimization.integrations.mlflow.artifacts import MlflowArtifactStore
+    from mlflow_optimization.integrations.mlflow.runs import MlflowRunStore
+
+INTEGRATION_MODULES: dict[str, tuple[str, str]] = {
+    "gurobipy": ("mlflow_optimization.integrations.gurobipy.adapter", "GurobiAdapter"),
+    "pyscipopt": (
+        "mlflow_optimization.integrations.pyscipopt.adapter",
+        "PyScipOptAdapter",
+    ),
+    "pyomo": ("mlflow_optimization.integrations.pyomo.adapter", "PyomoAdapter"),
+    "mathopt": ("mlflow_optimization.integrations.mathopt.adapter", "MathOptAdapter"),
+    "cp_sat": ("mlflow_optimization.integrations.cp_sat.adapter", "CpSatAdapter"),
+    "routing": ("mlflow_optimization.integrations.routing.adapter", "RoutingAdapter"),
+}
+
+_record_solve: RecordSolve | None = None
+_autolog_manager: AutologManager | None = None
+_integration_cache: dict[type[SolverIntegration], SolverIntegration] = {}
+
+
+def configure(
+    *,
+    disable: bool,
+    libraries: Iterable[str] | None,
+    log_solver_parameters: bool,
+    log_solver_metrics: bool,
+    log_artifacts: bool,
+) -> None:
+    """Apply an autolog configuration to the process-wide wiring."""
+    global _autolog_manager
+    global _record_solve
+
+    integrations: list[SolverIntegration] = []
+    if not disable:
+        integrations = _resolve_integrations(libraries)
+
+    if _autolog_manager is None:
+        from mlflow_optimization.features.autolog.service import AutologManager
+        from mlflow_optimization.features.runs.record import RecordSolve
+        from mlflow_optimization.features.runs.types import RecordOptions
+        from mlflow_optimization.integrations.mlflow.artifacts import (
+            MlflowArtifactStore,
+        )
+        from mlflow_optimization.integrations.mlflow.runs import MlflowRunStore
+
+        run_store: MlflowRunStore = MlflowRunStore()
+        artifact_store: MlflowArtifactStore = MlflowArtifactStore()
+        record_options: RecordOptions = RecordOptions()
+        if not disable:
+            record_options = RecordOptions(
+                log_solver_parameters=log_solver_parameters,
+                log_solver_metrics=log_solver_metrics,
+                log_artifacts=log_artifacts,
+            )
+        record_solve: RecordSolve = RecordSolve(
+            run_store, artifact_store, record_options
+        )
+        autolog_manager: AutologManager = AutologManager(record_solve)
+
+        _record_solve = record_solve
+        _autolog_manager = autolog_manager
+    else:
+        assert _record_solve is not None
+        if not disable:
+            _record_solve.options.log_solver_parameters = log_solver_parameters
+            _record_solve.options.log_solver_metrics = log_solver_metrics
+            _record_solve.options.log_artifacts = log_artifacts
+
+    assert _autolog_manager is not None
+    if disable:
+        _autolog_manager.disable()
+        return
+
+    _autolog_manager.enable(integrations)
+
+
+def _resolve_integrations(libraries: Iterable[str] | None) -> list[SolverIntegration]:
+    """The selected solver integrations."""
+    library_modules: dict[str, tuple[str, str]] = dict(INTEGRATION_MODULES)
+    selected_libraries: tuple[str, ...] = (
+        tuple(library_modules) if libraries is None else tuple(dict.fromkeys(libraries))
+    )
+
+    for library_name in selected_libraries:
+        if library_name not in library_modules:
+            valid_library_names: str = ", ".join(library_modules)
+            raise ValueError(
+                f"Unknown library {library_name!r}. Valid names: {valid_library_names}"
+            )
+
+    integrations: list[SolverIntegration] = []
+    for library_name in selected_libraries:
+        module_name: str
+        class_name: str
+        module_name, class_name = library_modules[library_name]
+        try:
+            integration_module: ModuleType = import_module(module_name)
+            integration_class: type[SolverIntegration] = getattr(
+                integration_module, class_name
+            )
+            cached_integration: SolverIntegration | None = _integration_cache.get(
+                integration_class
+            )
+            if cached_integration is not None:
+                integration: SolverIntegration = cached_integration
+            else:
+                integration = integration_class()
+                _integration_cache[integration_class] = integration
+        except ImportError:
+            if libraries is not None:
+                raise
+            continue
+        integrations.append(integration)
+
+    return integrations

mlflow_optimization/common/diagnostics.py

@@ -0,0 +1,19 @@
+"""Notification channel shared by every feature and integration."""
+
+import logging
+import warnings
+
+
+class MlflowOptimizationWarning(UserWarning):
+    """The category of every warning emitted by this library."""
+
+
+def warn(message: str, *, stacklevel: int = 2) -> None:
+    """Emit an ``MlflowOptimizationWarning``."""
+    try:
+        warnings.warn(message, MlflowOptimizationWarning, stacklevel=stacklevel + 1)
+    except BaseException:  # noqa: BLE001
+        try:
+            logging.getLogger("mlflow_optimization").warning(message)
+        except BaseException:  # noqa: BLE001,S110
+            pass

mlflow_optimization/common/patching.py

@@ -0,0 +1,82 @@
+"""Reversible attribute replacement."""
+
+import sys
+from collections.abc import Callable, Mapping
+from inspect import signature
+from types import ModuleType
+from typing import Any
+
+
+def receiver_of(
+    original: Callable[..., Any], args: tuple[Any, ...], kwargs: Mapping[str, Any]
+) -> Any:
+    """The receiver of a method call."""
+    receiver_name: str | None = next(iter(signature(original).parameters), None)
+    if receiver_name is None:
+        return None
+    if args:
+        return args[0]
+    return kwargs.get(receiver_name)
+
+
+def rebind_module_names(
+    original: Any,
+    replacement: Any,
+) -> list[tuple[ModuleType, str]]:
+    """The module-level names rebound from an original object to a replacement."""
+    rebound_names: list[tuple[ModuleType, str]] = []
+    modules: tuple[Any, ...] = tuple(sys.modules.values())
+    for module_value in modules:
+        if not isinstance(module_value, ModuleType):
+            continue
+        module: ModuleType = module_value
+        module_values: dict[str, Any] = vars(module)
+        module_bindings: tuple[tuple[str, Any], ...] = tuple(module_values.items())
+        for name, value in module_bindings:
+            if value is not original:
+                continue
+            module_values[name] = replacement
+            rebound_names.append((module, name))
+    return rebound_names
+
+
+def restore_module_names(
+    rebound_names: list[tuple[ModuleType, str]],
+    original: Any,
+    replacement: Any,
+) -> None:
+    """A conditional restoration of module-level names."""
+    for module, name in reversed(rebound_names):
+        module_values: dict[str, Any] = vars(module)
+        if name not in module_values or module_values[name] is not replacement:
+            continue
+        module_values[name] = original
+
+
+class Patch:
+    """A reversible replacement of one attribute on a class or module."""
+
+    def __init__(self, owner: Any, name: str) -> None:
+        self.owner: Any = owner
+        self.name: str = name
+        self._installed: bool = False
+
+    def install(self, factory: Callable[[Any], Any]) -> None:
+        """Replace the attribute with ``factory(original)``; no-op when installed."""
+        if self._installed:
+            return
+        original: Any = getattr(self.owner, self.name)
+        self._own: bool = self.name in vars(self.owner)
+        self._descriptor: Any = vars(self.owner).get(self.name)
+        setattr(self.owner, self.name, factory(original))
+        self._installed = True
+
+    def uninstall(self) -> None:
+        """Restore the exact original attribute; no-op when not installed."""
+        if not self._installed:
+            return
+        if self._own:
+            setattr(self.owner, self.name, self._descriptor)
+        else:
+            delattr(self.owner, self.name)
+        self._installed = False

mlflow_optimization/features/annotations/annotate.py

@@ -0,0 +1,97 @@
+"""Scoping and naming rules for user annotations."""
+
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from math import isfinite
+
+from mlflow_optimization.common.diagnostics import warn
+
+from .types import Annotation, UserMetric
+
+_annotation_context: ContextVar[Annotation] = ContextVar(
+    "annotation_context",
+    default=Annotation(),  # noqa: B039 -- context() copies tags before merging.
+)
+
+
+@contextmanager
+def context(
+    *,
+    instance: str | None = None,
+    description: str | None = None,
+    tags: Mapping[str, str] | None = None,
+) -> Iterator[None]:
+    """A scope whose annotations apply to every solve started inside it.
+
+    Parameters
+    ----------
+    instance : str, optional
+        Instance name; inherited from the enclosing scope when omitted.
+    description : str, optional
+        Description; inherited from the enclosing scope when omitted.
+    tags : Mapping[str, str], optional
+        User tags merge with enclosing tags; inner values replace outer values for
+        matching keys. Keys prefixed ``optimization.`` or ``mlflow.`` are ignored.
+    """
+    enclosing_annotation: Annotation = current_annotation()
+    merged_tags: dict[str, str] = dict(enclosing_annotation.tags)
+    if tags is not None:
+        merged_tags.update(user_tags(tags))
+    merged_annotation: Annotation = Annotation(
+        instance=instance if instance is not None else enclosing_annotation.instance,
+        description=(
+            description if description is not None else enclosing_annotation.description
+        ),
+        tags=merged_tags,
+    )
+    token: Token[Annotation] = _annotation_context.set(merged_annotation)
+    try:
+        yield
+    finally:
+        _annotation_context.reset(token)
+
+
+def current_annotation() -> Annotation:
+    """The merged annotations of the enclosing ``context`` scopes."""
+    return _annotation_context.get()
+
+
+def annotation_tags(annotation: Annotation) -> dict[str, str]:
+    """The tags that record an annotation."""
+    tags: dict[str, str] = user_tags(annotation.tags)
+    if annotation.instance is not None:
+        tags["optimization.instance"] = annotation.instance
+    if annotation.description is not None:
+        tags["mlflow.note.content"] = annotation.description
+    return tags
+
+
+def user_tags(tags: Mapping[str, str]) -> dict[str, str]:
+    """The subset of user tags that may be recorded."""
+    accepted_tags: dict[str, str] = {}
+    for key, value in tags.items():
+        if key.startswith(("optimization.", "mlflow.")):
+            warn(f"Ignoring user tag with reserved key {key!r}.")
+        else:
+            accepted_tags[key] = value
+    return accepted_tags
+
+
+def user_metric(
+    name: str, value: float, *, unit: str | None = None, description: str | None = None
+) -> UserMetric | None:
+    """The record of one user metric, or None when it may not be recorded."""
+    if name.startswith(("optimization.", "mlflow.")):
+        warn(f"Ignoring user metric with reserved name {name!r}.")
+        return None
+    if not isfinite(value):
+        warn(f"Ignoring user metric {name!r} because its value is not finite.")
+        return None
+    metric_key: str = f"user.{name}"
+    tags: dict[str, str] = {}
+    if unit is not None:
+        tags[f"{metric_key}.unit"] = unit
+    if description is not None:
+        tags[f"{metric_key}.description"] = description
+    return UserMetric(metric_key=metric_key, value=value, tags=tags)

mlflow_optimization/features/annotations/types.py

@@ -0,0 +1,22 @@
+"""User-provided information attached to solves."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+
+
+@dataclass(frozen=True)
+class Annotation:
+    """The user annotations in effect for a solve."""
+
+    instance: str | None = None
+    description: str | None = None
+    tags: Mapping[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class UserMetric:
+    """One user metric under final key names."""
+
+    metric_key: str
+    value: float
+    tags: Mapping[str, str]

mlflow_optimization/features/artifacts/policy.py

@@ -0,0 +1,16 @@
+"""Rules for which explicit artifact requests are honored."""
+
+from pathlib import Path
+
+
+def artifact_refusal(local_path: str, *, enabled: bool) -> str | None:
+    """The reason an explicit artifact request is refused, or None when it is honored."""
+    if not enabled:
+        return "Artifact logging is disabled."
+    try:
+        is_file: bool = Path(local_path).is_file()
+    except OSError:
+        return f"Artifact path cannot be inspected: {local_path}"
+    if not is_file:
+        return f"Artifact path is not an existing file: {local_path}"
+    return None

mlflow_optimization/features/artifacts/ports.py

@@ -0,0 +1,17 @@
+"""Contract of the artifact storage."""
+
+from typing import Protocol
+
+
+class ArtifactStoreError(Exception):
+    """An error raised when artifact storage fails."""
+
+
+class ArtifactStore(Protocol):
+    """The storage of files attached to runs."""
+
+    def log_artifact(
+        self, run_id: str, local_path: str, artifact_path: str | None
+    ) -> None:
+        """Attach a local file to a run."""
+        ...

mlflow_optimization/features/autolog/ports.py

@@ -0,0 +1,49 @@
+"""Contracts between solve observation and its surroundings."""
+
+from collections.abc import Callable
+from typing import Protocol
+
+from .types import CompletedSolve, SolveObservation
+
+
+class ObservationError(Exception):
+    """An error raised when a solve observation cannot be built."""
+
+
+class SolveObserver(Protocol):
+    """The entry through which an adapter routes a solve call."""
+
+    def observe[T](
+        self,
+        api: str,
+        call: Callable[[], T],
+        extract: Callable[[T | None, BaseException | None], SolveObservation],
+    ) -> T:
+        """The result of ``call()``, observed as one solve unless nested or disabled."""
+        ...
+
+
+class SolverIntegration(Protocol):
+    """The hook of one Python solver API."""
+
+    api: str
+
+    def install(self, observer: SolveObserver) -> None:
+        """Route the SDK's solve entry points through ``observer``."""
+        ...
+
+    def uninstall(self) -> None:
+        """Restore the SDK's original solve entry points."""
+        ...
+
+
+class SolveSink(Protocol):
+    """The destination of completed solves."""
+
+    def begin(self) -> object:
+        """The recording context captured at solve start."""
+        ...
+
+    def record(self, solve: CompletedSolve, context: object) -> None:
+        """Store one completed solve."""
+        ...

mlflow_optimization/features/autolog/service.py

@@ -0,0 +1,147 @@
+"""Activation of integrations and arbitration of observed solve calls."""
+
+from collections.abc import Callable, Iterator, Sequence
+from contextvars import ContextVar, Token
+from dataclasses import replace
+from itertools import count
+from threading import Lock
+from time import perf_counter, time
+from typing import cast
+from uuid import uuid4
+
+from mlflow_optimization.common.diagnostics import warn
+
+from .ports import SolverIntegration, SolveSink
+from .types import CompletedSolve, SolveObservation
+
+_solve_indices: Iterator[int] = count(1)
+_solve_index_lock: Lock = Lock()
+
+
+class _SolveScope:
+    """The lifetime of an observed solve call."""
+
+    def __init__(self) -> None:
+        self.active: bool = True
+
+
+_observing: ContextVar[_SolveScope | None] = ContextVar(
+    "autolog_observing", default=None
+)
+
+
+class AutologManager:
+    """The observer that turns outermost solve calls into completed solves."""
+
+    def __init__(self, sink: SolveSink) -> None:
+        self._sink: SolveSink = sink
+        self._integrations: list[SolverIntegration] = []
+        self._enabled: bool = False
+
+    def enable(self, integrations: Sequence[SolverIntegration]) -> None:
+        """Observe exactly the given integrations."""
+        for integration in self._integrations[:]:
+            if not any(integration is candidate for candidate in integrations):
+                integration.uninstall()
+                self._integrations.remove(integration)
+        for integration in integrations:
+            if not any(integration is installed for installed in self._integrations):
+                integration.install(self)
+                self._integrations.append(integration)
+        self._enabled = True
+
+    def disable(self) -> None:
+        """Stop observing and uninstall every integration."""
+        self._enabled = False
+        for integration in self._integrations[:]:
+            integration.uninstall()
+            self._integrations.remove(integration)
+
+    def observe[T](
+        self,
+        api: str,
+        call: Callable[[], T],
+        extract: Callable[[T | None, BaseException | None], SolveObservation],
+    ) -> T:
+        """The result of ``call()``, observed as one solve unless nested or disabled."""
+        parent: _SolveScope | None = _observing.get()
+        if not self._enabled or (parent is not None and parent.active):
+            return call()
+        scope: _SolveScope = _SolveScope()
+        token: Token[_SolveScope | None] = _observing.set(scope)
+        try:
+            failures: list[str] = []
+
+            def guard(step: str, operation: Callable[[], object]) -> object:
+                """The result of a guarded observation operation."""
+                try:
+                    return operation()
+                # NOTE: A defect in logging must never change the result of the user's solve.
+                except Exception as exception:  # noqa: BLE001
+                    failures.append(
+                        f"Solve {step} failed for {api}: "
+                        f"{type(exception).__name__}: {exception}"
+                    )
+                    return None
+
+            context: object = guard("begin", self._sink.begin)
+            with _solve_index_lock:
+                solve_index: int = next(_solve_indices)
+            solve_id: str = uuid4().hex
+            start_time_ms: int = int(time() * 1000)
+            result: T | None = None
+
+            def finish(solve: CompletedSolve) -> None:
+                """Extract and record a completed solve."""
+                observation: SolveObservation | None = None
+                try:
+                    observation = extract(result, solve.error)
+                # NOTE: A defect in extraction must never change the result of the user's solve.
+                except Exception as exception:  # noqa: BLE001
+                    failures.append(
+                        f"Solve extraction failed for {api}: "
+                        f"{type(exception).__name__}: {exception}"
+                    )
+                self._sink.record(replace(solve, observation=observation), context)
+
+            start: float = perf_counter()
+            try:
+                result = call()
+            except BaseException as error:
+                end: float = perf_counter()
+                scope.active = False
+                end_time_ms: int = int(time() * 1000)
+                solve: CompletedSolve = CompletedSolve(
+                    api,
+                    solve_id,
+                    solve_index,
+                    start_time_ms,
+                    end_time_ms,
+                    end - start,
+                    None,
+                    error,
+                )
+                guard("recording", lambda: finish(solve))
+                for failure in failures:
+                    warn(failure)
+                raise
+            end = perf_counter()
+            scope.active = False
+            end_time_ms = int(time() * 1000)
+            solve = CompletedSolve(
+                api,
+                solve_id,
+                solve_index,
+                start_time_ms,
+                end_time_ms,
+                end - start,
+                None,
+                None,
+            )
+            guard("recording", lambda: finish(solve))
+            for failure in failures:
+                warn(failure)
+            return cast(T, result)
+        finally:
+            scope.active = False
+            _observing.reset(token)

mlflow_optimization/features/autolog/types.py

@@ -0,0 +1,31 @@
+"""What is observed about a solve."""
+
+from dataclasses import dataclass, field
+
+from mlflow_optimization.features.metadata.types import ReportedMetadata
+from mlflow_optimization.features.metrics.types import ReportedMetrics
+from mlflow_optimization.features.outcomes.types import ReportedOutcome
+
+
+@dataclass(frozen=True)
+class SolveObservation:
+    """The SDK-free snapshot an adapter takes of one finished solve."""
+
+    api: str
+    outcome: ReportedOutcome
+    metrics: ReportedMetrics = field(default_factory=ReportedMetrics)
+    metadata: ReportedMetadata = field(default_factory=ReportedMetadata)
+
+
+@dataclass(frozen=True)
+class CompletedSolve:
+    """One finished outermost solve call as seen by the observer."""
+
+    api: str
+    solve_id: str
+    solve_index: int
+    start_time_ms: int
+    end_time_ms: int
+    wall_time_seconds: float
+    observation: SolveObservation | None
+    error: BaseException | None