"""

``feast.mlflow`` — drop-in replacement for ``import mlflow`` with Feast superpowers.

Any function or attribute available on the ``mlflow`` module can be accessed

via ``feast.mlflow.*``. A subset of calls are **Feast-enhanced** with

automatic tagging, lineage tracking, and feature resolution:

- ``start_run()`` — auto-tags runs with ``feast.project``

- ``log_model()`` — auto-saves ``feast_features.json``

- ``register_model()`` — auto-tags model versions with ``feast.feature_service``

- ``load_model()`` — auto-links prediction runs to training runs

- ``resolve_features()`` — Feast-only: model URI → feature service name

- ``get_training_entity_df()`` — Feast-only: recover training entity data

All other calls (``log_params``, ``log_metrics``, ``set_tag``,

``log_artifact``, ``MlflowClient``, etc.) pass through to the raw

``mlflow`` module unchanged.

**Store resolution order** (first match wins):

1. Explicit ``feast.mlflow.init(store)`` call

2. Most recently created ``FeatureStore`` (auto-registered)

3. ``FeatureStore(".")`` from the current working directory

"""

from __future__ import annotations

import logging

from typing import TYPE_CHECKING, Any, Optional

if TYPE_CHECKING:

from feast import FeatureStore

_logger = logging.getLogger(__name__)

_MISSING = object()

_client: Optional[Any] = None

_registered_store: Optional["FeatureStore"] = None

def _register_store(store: "FeatureStore") -> None:

"""Called by ``FeatureStore.__init__`` to auto-register itself.

This is an internal API — end users should call :func:`init` instead.

"""

global _registered_store

_registered_store = store

def _build_client() -> Any:

"""Create a ``FeastMlflowClient`` using the best available store."""

from feast.mlflow_integration.client import FeastMlflowClient

store = _registered_store

if store is None:

try:

from feast import FeatureStore

store = FeatureStore(".")

except Exception as exc:

raise RuntimeError(

"feast.mlflow could not auto-discover a FeatureStore. "

"Either call feast.mlflow.init(store), create a FeatureStore "

"before using feast.mlflow, or ensure feature_store.yaml "

"exists in the current directory."

) from exc

mlflow_cfg = getattr(store.config, "mlflow", None)

if mlflow_cfg is None or not mlflow_cfg.enabled:

raise RuntimeError(

"MLflow integration is not enabled. "

"Set mlflow.enabled=true in feature_store.yaml, or call "

"feast.mlflow.init(store) with a store whose config has "

"mlflow.enabled=true."

)

try:

return FeastMlflowClient(store)

except ImportError:

raise ImportError(

"mlflow package is not installed. "

"Install it with: pip install feast[mlflow]"

)

def _ensure_client() -> Any:

"""Return the cached client, creating it on first call."""

global _client

if _client is None:

_client = _build_client()

return _client

def init(store: "FeatureStore") -> None:

"""Bind ``feast.mlflow`` to a specific :class:`~feast.FeatureStore`.

Call this once at the start of a notebook or script. All subsequent

``feast.mlflow.*`` calls will use this store's MLflow configuration.

This is **optional** — if you skip it, ``feast.mlflow`` will use the

most recently created ``FeatureStore`` automatically.

"""

global _client, _registered_store

_client = None

_registered_store = store

def get_active_run_id() -> Optional[str]:

"""Return the active MLflow run ID, or ``None``."""

return _ensure_client().active_run_id

def __getattr__(name: str) -> Any:

"""Open delegation: Feast-enhanced client first, raw mlflow fallback.

Lookup order for ``feast.mlflow.<name>``:

1. If ``FeastMlflowClient`` has a public attribute *name*, return it.

This gives Feast-enhanced versions of ``start_run``, ``log_model``,

``register_model``, ``load_model``, etc.

2. Otherwise, fall back to the raw ``mlflow`` module. This makes

``feast.mlflow.log_params``, ``feast.mlflow.set_tag``,

``feast.mlflow.MlflowClient``, etc. work without any wrappers.

"""

if name.startswith("_"):

raise AttributeError(f"module 'feast.mlflow' has no attribute {name!r}")

client = _ensure_client()

client_attr = getattr(client, name, _MISSING)

if client_attr is not _MISSING:

return client_attr

mlflow_attr = getattr(client._mlflow, name, _MISSING)

if mlflow_attr is not _MISSING:

return mlflow_attr

raise AttributeError(f"module 'feast.mlflow' has no attribute {name!r}")