Skip to content

Commit 3667772

Browse files
VperiodtVperiodt
authored
feat: Add MLflow as a first-class offline DataSource for Feast (#6702)
* add mlflow datasource adaptor Signed-off-by: Vanshika Vanshika <[email protected]> rh-pre-commit.version: 2.3.2 rh-pre-commit.check-secrets: ENABLED * fix: Correct lint Signed-off-by: Vanshika Vanshika <[email protected]> rh-pre-commit.version: 2.3.2 rh-pre-commit.check-secrets: ENABLED * address reviews and fix error handling Signed-off-by: Vanshika Vanshika <[email protected]> * feat: Close MLflow DataSource scope gaps and add RHOAI E2E suite Register feast mlflow CLI commands, add SavedDataset round-trip and auth coexistence tests, document the TP workflow, and provide cluster E2E scripts after rebasing onto upstream master. Signed-off-by: Vanshika Vanshika <[email protected]> * chore: Drop local RHOAI E2E scripts from branch Cluster-only validation scripts are not part of upstream Feast; keep gap fixes, tests, and how-to guide in PR #6702. Signed-off-by: $(git config user.name 2>/dev/null || echo 'User') <$(git config user.email 2>/dev/null || echo 'user@local')> Signed-off-by: Vanshika Vanshika <[email protected]> * chore: Restore protos from master Signed-off-by: Vanshika Vanshika <[email protected]> * chore: Update secrets baseline Signed-off-by: Vanshika Vanshika <[email protected]> x Signed-off-by: Vanshika Vanshika <[email protected]> * chore: Refresh secrets baseline and pixi lock for CI Signed-off-by: Vanshika Vanshika <[email protected]> --------- Signed-off-by: Vanshika Vanshika <[email protected]> Signed-off-by: $(git config user.name 2>/dev/null || echo 'User') <$(git config user.email 2>/dev/null || echo 'user@local')>
1 parent f9e91fa commit 3667772

24 files changed

Lines changed: 7805 additions & 4314 deletions

File tree

‎.secrets.baseline‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎docs/SUMMARY.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@
9898
* [Adding or reusing tests](how-to-guides/adding-or-reusing-tests.md)
9999
* [Starting Feast servers in TLS(SSL) Mode](how-to-guides/starting-feast-servers-tls-mode.md)
100100
* [Importing Features from dbt](how-to-guides/dbt-integration.md)
101+
* [Using MLflow as a Data Source](how-to-guides/mlflow-datasource.md)
101102
* [Entity Key Serialization (v2 to v3)](how-to-guides/entity-reserialization-of-from-v2-to-v3.md)
102103
* [Feature Quality Monitoring](how-to-guides/feature-monitoring.md)
103104

‎docs/how-to-guides/mlflow-datasource.md‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
# Using MLflow as a Data Source in Feast
2+
3+
This guide walks through the end-to-end workflow for using MLflow as a **first-class offline DataSource** in Feast. You define `FeatureView` objects backed by `MlflowDatasetSource`, then use `get_historical_features()` and `create_saved_dataset()` without manual export or ETL.
4+
5+
For API reference, see [MLflow Integration](../reference/mlflow.md#mlflow-as-a-datasource). For operator auto-discovery of MLflow on RHOAI, see [Guide 8 — MLflow Integration](feast-operator/08-mlflow-integration.md).
6+
7+
## Overview
8+
9+
| Mode | Use case | Configuration |
10+
|------|----------|---------------|
11+
| **GenAI Dataset** | Curated evaluation / trace records from MLflow GenAI Datasets | `dataset_name` or `dataset_id` |
12+
| **Artifact** | Parquet or CSV artifacts logged on MLflow runs | `run_id`, `artifact_path`, `artifact_format` |
13+
14+
**Scope (Tech Preview):** offline/historical retrieval only; Parquet and CSV tabular artifacts; reads go through the MLflow Tracking API (no direct object-storage bypass).
15+
16+
## Prerequisites
17+
18+
- Feast with MLflow support: `pip install 'feast[mlflow]'` (MLflow ≥ 2.10 for GenAI datasets).
19+
- A reachable MLflow tracking server.
20+
- On RHOAI: Feast and MLflow operators deployed; FeatureStore CR in Ready state; cluster CA bundle available for TLS if using HTTPS.
21+
22+
## Step 1: Configure `feature_store.yaml`
23+
24+
```yaml
25+
project: my_project
26+
registry: data/registry.db
27+
provider: local
28+
offline_store: duckdb
29+
online_store:
30+
type: sqlite
31+
path: data/online_store.db
32+
entity_key_serialization_version: 3
33+
34+
mlflow:
35+
enabled: true
36+
tracking_uri: https://mlflow.example.com:8443
37+
ca_bundle: /etc/pki/tls/certs/ca-bundle.crt
38+
supported_artifact_formats:
39+
- parquet
40+
- csv
41+
request_timeout: 30 # seconds per HTTP request to MLflow (default: 30)
42+
max_retries: 3 # retry attempts for transient failures (default: 3)
43+
retry_backoff_factor: 1.0 # exponential backoff base (default: 1.0 → 1s, 2s, 4s)
44+
```
45+
46+
On operator-managed clusters, `tracking_uri` and `ui_url` are often injected automatically when an `MLflow` CR is present. You can still set `ca_bundle` explicitly for DataSource reads from a workbench or CI job.
47+
48+
| Field | Default | Description |
49+
|---|---|---|
50+
| `tracking_uri` | `MLFLOW_TRACKING_URI` env | MLflow tracking server URI |
51+
| `ca_bundle` | `None` | CA bundle path for TLS (falls back to `REQUESTS_CA_BUNDLE` env) |
52+
| `supported_artifact_formats` | `["parquet", "csv"]` | Accepted artifact formats |
53+
| `request_timeout` | `30` | Seconds per HTTP request to MLflow |
54+
| `max_retries` | `3` | Retry attempts for transient failures (timeouts, 5xx) |
55+
| `retry_backoff_factor` | `1.0` | Exponential backoff base (wait = factor × 2^attempt) |
56+
57+
## Step 2: Define `MlflowDatasetSource` FeatureViews
58+
59+
`MlflowDatasetSource` is read-only. You must set a `batch_source` (for example `FileSource`) for writeback paths such as `create_saved_dataset()` and materialization.
60+
61+
### GenAI Dataset mode
62+
63+
```python
64+
from feast import Entity, FeatureView, Field, FileSource
65+
from feast.infra.data_sources.mlflow import MlflowDatasetSource
66+
from feast.types import String
67+
68+
batch_sink = FileSource(
69+
path="data/eval_records.parquet",
70+
timestamp_field="event_timestamp",
71+
)
72+
73+
source = MlflowDatasetSource(
74+
name="eval_dataset",
75+
dataset_name="production_validation_set",
76+
batch_source=batch_sink,
77+
timestamp_field="event_timestamp",
78+
field_mapping={"expectations.expected_response": "expected_response"},
79+
)
80+
81+
eval_records = FeatureView(
82+
name="eval_records",
83+
entities=[Entity(name="record_id", join_keys=["record_id"])],
84+
schema=[Field(name="expected_response", dtype=String)],
85+
source=source,
86+
)
87+
```
88+
89+
### Artifact mode (Parquet)
90+
91+
```python
92+
source = MlflowDatasetSource(
93+
name="training_features",
94+
run_id="<mlflow-run-id>",
95+
artifact_path="outputs/features.parquet",
96+
artifact_format="parquet",
97+
batch_source=FileSource(
98+
path="data/features.parquet",
99+
timestamp_field="event_timestamp",
100+
),
101+
timestamp_field="event_timestamp",
102+
)
103+
104+
training_features = FeatureView(
105+
name="training_features",
106+
entities=[Entity(name="record_id", join_keys=["record_id"])],
107+
schema=[...],
108+
source=source,
109+
)
110+
```
111+
112+
## Step 3: Apply and validate
113+
114+
```bash
115+
feast apply
116+
feast mlflow list-sources
117+
feast mlflow validate-source eval_records
118+
```
119+
120+
`validate-source` checks configuration, MLflow reachability, schema introspection, and schema match against the FeatureView.
121+
122+
## Step 4: Retrieve historical features and save datasets
123+
124+
```python
125+
from feast import FeatureStore
126+
127+
store = FeatureStore(".")
128+
129+
job = store.get_historical_features(
130+
entity_df=entity_df,
131+
features=["eval_records:expected_response"],
132+
)
133+
training_df = job.to_df()
134+
135+
from feast.data_format import ParquetFormat
136+
from feast.infra.offline_stores.file_source import SavedDatasetFileStorage
137+
138+
storage = SavedDatasetFileStorage(
139+
path="data/saved_eval.parquet",
140+
file_format=ParquetFormat(),
141+
)
142+
143+
saved = store.create_saved_dataset(
144+
from_=job,
145+
name="eval_training_snapshot",
146+
storage=storage,
147+
)
148+
149+
loaded = store.get_saved_dataset("eval_training_snapshot")
150+
assert loaded.to_df().shape == training_df.shape
151+
```
152+
153+
Feast primitives in scope: **`get_historical_features()`** and **`SavedDataset`** (via `create_saved_dataset()` / `get_saved_dataset()`).
154+
155+
## Step 5: Dataset sync (Tier 3 offline stores)
156+
157+
For BigQuery, Snowflake, or Redshift, sync MLflow GenAI data into the offline store first:
158+
159+
```bash
160+
feast mlflow sync-dataset --feature-view eval_records
161+
feast mlflow sync-dataset --feature-view eval_records --dry-run
162+
feast mlflow preview-dataset --source production_validation_set --limit 10
163+
```
164+
165+
## Authentication
166+
167+
Feast resolves tokens in this order for DataSource reads:
168+
169+
1. User Bearer token from the current Feast API request (REST / Arrow Flight).
170+
2. `MLFLOW_TRACKING_TOKEN` environment variable.
171+
3. Kubernetes ServiceAccount token file (when present).
172+
173+
Feast registers `FeastMLflowHeaderProvider` as an MLflow `request_header_provider` plugin. When a token is active in context, it sets `Authorization: Bearer <token>` on MLflow REST calls.
174+
175+
### Coexistence with `MLFLOW_TRACKING_AUTH`
176+
177+
The Feast operator may inject `MLFLOW_TRACKING_AUTH=kubernetes-namespaced` for pod-level MLflow access. That uses MLflow’s **auth provider** plugin (ServiceAccount token). Feast’s **header provider** runs in parallel; when a ContextVar token is set (user request or `MLFLOW_TRACKING_TOKEN`), the header provider’s `Authorization` value is used for DataSource reads. Automated jobs without a user token continue to rely on kubernetes-namespaced auth.
178+
179+
## Troubleshooting
180+
181+
| Symptom | What to check |
182+
|---------|----------------|
183+
| `ImportError` for `mlflow` | `pip install 'feast[mlflow]'` |
184+
| Connection / timeout errors | `tracking_uri`, network, `ca_bundle`, MLflow pod health |
185+
| HTTP 401 / 403 | Token validity and MLflow RBAC on the target experiment or artifact |
186+
| Unsupported format | Only `parquet` and `csv` are supported in TP |
187+
| Missing artifact / dataset | Run ID, artifact path, or dataset name; 404 errors include context in the message |
188+
189+
## Limitations (Tech Preview)
190+
191+
- No online store integration with MLflow.
192+
- No pickle or other non-tabular artifact formats.
193+
- No bypass of MLflow’s auth layer via direct object storage.
194+
- Operator FeatureStore CR fields for DataSource-specific MLflow config may evolve in follow-up releases.

‎docs/reference/feast-cli-commands.md‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -498,6 +498,24 @@ feast registry create-schema
498498

499499
This command only applies to SQL-based registries (`registry_type: sql`). It is safe to run multiple times — existing tables are not modified.
500500

501+
## MLflow
502+
503+
MLflow integration utilities for GenAI dataset sync and DataSource validation. Requires `pip install 'feast[mlflow]'`. See [MLflow Integration](mlflow.md).
504+
505+
```bash
506+
# Sync MLflow GenAI EvaluationDataset → FeatureView (via MlflowDatasetSource)
507+
feast mlflow sync-dataset --feature-view mlflow_labels
508+
509+
# Preview flattened dataset records
510+
feast mlflow preview-dataset --source agent-feedback-v3 --limit 10
511+
512+
# Validate an MlflowDatasetSource-backed FeatureView
513+
feast mlflow validate-source eval_records
514+
515+
# List all FeatureViews backed by MlflowDatasetSource
516+
feast mlflow list-sources
517+
```
518+
501519
## Teardown
502520

503521
Tear down deployed feature store infrastructure

‎docs/reference/mlflow.md‎

Lines changed: 160 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,3 +345,163 @@ Start the Feast UI with:
345345
```bash
346346
feast ui --host 127.0.0.1 --port 8888
347347
```
348+
349+
## MLflow as a DataSource
350+
351+
Feast supports MLflow as a first-class **offline DataSource** via `MlflowDatasetSource`. Users define `FeatureView(source=MlflowDatasetSource(...))` and call `get_historical_features()` / `create_saved_dataset()` to retrieve tabular data from MLflow-tracked artifacts or GenAI Datasets.
352+
353+
### Dual-mode design
354+
355+
`MlflowDatasetSource` supports two modes:
356+
357+
| Mode | When to use | Config fields |
358+
|---|---|---|
359+
| **GenAI Dataset** | Flywheel path — curated trace/assessment records from MLflow GenAI Datasets | `dataset_name` or `dataset_id` |
360+
| **Artifact** | General purpose — download tabular artifacts (Parquet/CSV) from MLflow runs | `run_id` + `artifact_path` + `artifact_format` |
361+
362+
### GenAI Dataset mode
363+
364+
Reads from MLflow GenAI EvaluationDatasets via `get_dataset().to_df()`. This is the primary path for the tracing → assessment → training data flywheel.
365+
366+
```python
367+
from feast import Entity, FeatureView, Field, FileSource
368+
from feast.infra.data_sources.mlflow import MlflowDatasetSource
369+
from feast.types import String
370+
371+
batch_sink = FileSource(path="data/eval_records.parquet", timestamp_field="event_timestamp")
372+
373+
source = MlflowDatasetSource(
374+
name="eval_dataset",
375+
dataset_name="production_validation_set",
376+
batch_source=batch_sink,
377+
timestamp_field="event_timestamp",
378+
field_mapping={"expectations.expected_response": "expected_response"},
379+
)
380+
381+
fv = FeatureView(
382+
name="eval_records",
383+
entities=[Entity(name="record_id", join_keys=["record_id"])],
384+
schema=[Field(name="expected_response", dtype=String)],
385+
source=source,
386+
)
387+
```
388+
389+
### Artifact mode
390+
391+
Downloads Parquet or CSV artifacts from MLflow runs via `mlflow.artifacts.download_artifacts()`.
392+
393+
```python
394+
source = MlflowDatasetSource(
395+
name="training_features",
396+
run_id="abc123def456",
397+
artifact_path="outputs/features.parquet",
398+
artifact_format="parquet", # or "csv"
399+
batch_source=FileSource(path="data/features.parquet", timestamp_field="ts"),
400+
timestamp_field="ts",
401+
)
402+
```
403+
404+
### Retrieving historical features
405+
406+
Once registered, use `get_historical_features()` as with any other DataSource:
407+
408+
```python
409+
from feast import FeatureStore
410+
411+
store = FeatureStore(".")
412+
features = store.get_historical_features(
413+
entity_df=entity_df,
414+
features=["eval_records:expected_response"],
415+
)
416+
training_df = features.to_df()
417+
```
418+
419+
### `batch_source` for writeback
420+
421+
`MlflowDatasetSource` is read-only — it does not support writing back to MLflow. A `batch_source` (e.g., `FileSource`) is always required for `create_saved_dataset()` and materialization writeback.
422+
423+
### Authentication
424+
425+
For RHOAI deployments, auth tokens are propagated automatically:
426+
427+
1. **User-initiated requests**: The Bearer token from the Feast API request is forwarded to MLflow
428+
2. **Automated jobs** (materialization CronJobs): The pod's ServiceAccount token is used
429+
3. **Local dev**: Falls back to `MLFLOW_TRACKING_TOKEN` env var or no auth
430+
431+
### Graceful degradation
432+
433+
- When `mlflow` is not installed, attempting to use `MlflowDatasetSource` raises a clear `ImportError` with install instructions
434+
- Non-MLflow FeatureViews are completely unaffected by MLflow availability
435+
- The feature server starts normally without MLflow
436+
437+
### Configuration
438+
439+
The `mlflow:` section in `feature_store.yaml` supports DataSource-specific fields:
440+
441+
```yaml
442+
mlflow:
443+
enabled: true
444+
tracking_uri: https://mlflow.example.com:8443
445+
ca_bundle: /etc/pki/tls/odh-trusted-ca-bundle.crt
446+
supported_artifact_formats:
447+
- parquet
448+
- csv
449+
request_timeout: 30
450+
max_retries: 3
451+
retry_backoff_factor: 1.0
452+
```
453+
454+
| Field | Type | Default | Description |
455+
|---|---|---|---|
456+
| `ca_bundle` | `str` | `None` | Path to CA bundle for TLS verification (falls back to `REQUESTS_CA_BUNDLE` env var) |
457+
| `supported_artifact_formats` | `list[str]` | `["parquet", "csv"]` | Artifact formats the adapter will accept |
458+
| `request_timeout` | `int` | `30` | Timeout in seconds for individual HTTP requests to the MLflow tracking server |
459+
| `max_retries` | `int` | `3` | Maximum retry attempts for transient MLflow API failures (timeouts, 5xx, connection errors) |
460+
| `retry_backoff_factor` | `float` | `1.0` | Base factor for exponential backoff between retries (wait = factor × 2^attempt) |
461+
462+
### CLI commands
463+
464+
#### `feast mlflow validate-source`
465+
466+
Validates that an MlflowDatasetSource-backed FeatureView is reachable and schema-compatible:
467+
468+
```bash
469+
feast mlflow validate-source eval_records
470+
```
471+
472+
Checks performed:
473+
- Config validation (mode, batch_source, artifact_format)
474+
- Schema introspection (column names and types from MLflow)
475+
- Schema match (FeatureView features vs source columns)
476+
477+
#### `feast mlflow list-sources`
478+
479+
Lists all FeatureViews backed by MlflowDatasetSource:
480+
481+
```bash
482+
feast mlflow list-sources
483+
```
484+
485+
#### `feast mlflow sync-dataset`
486+
487+
Syncs an MLflow GenAI Dataset into a Feast FeatureView (for Tier 3 offline stores like BigQuery, Snowflake, Redshift):
488+
489+
```bash
490+
feast mlflow sync-dataset --feature-view eval_records
491+
feast mlflow sync-dataset --feature-view eval_records --full-refresh
492+
feast mlflow sync-dataset --feature-view eval_records --dry-run
493+
feast mlflow sync-dataset --feature-view eval_records --batch-size 5000
494+
feast mlflow sync-dataset --feature-view eval_records --field-mapping mapping.json
495+
```
496+
497+
When `--feature-view` is omitted, syncs all FeatureViews whose source is `MlflowDatasetSource`. Supports incremental sync (default) via watermark tags and full refresh via `--full-refresh`.
498+
499+
#### `feast mlflow preview-dataset`
500+
501+
Preview flattened records from an MLflow GenAI Dataset:
502+
503+
```bash
504+
feast mlflow preview-dataset --source agent-feedback-v3 --limit 10
505+
```
506+
507+
Fetches the dataset, flattens nested fields, and displays a preview. Useful for inspecting data shape before defining a FeatureView.

0 commit comments

Comments
 (0)