feast/sdk/python/feast/mlflow.py at codex/flink-doc-fix · feast-dev/feast · GitHub

History

139 lines (103 loc) · 4.66 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

"""

``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}")

feast/sdk/python/feast/mlflow.py at codex/flink-doc-fix · feast-dev/feast · GitHub

Uh oh!

Expand file tree

Breadcrumbs

mlflow.py

Latest commit

History

Breadcrumbs

mlflow.py

File metadata and controls