Skip to content

feat: Add MLflow tracing, trace Assessment Sync, and Trace Export - #6405

Open
Vperiodt wants to merge 9 commits into
feast-dev:masterfrom
Vperiodt:mcp-trace
Open

feat: Add MLflow tracing, trace Assessment Sync, and Trace Export#6405
Vperiodt wants to merge 9 commits into
feast-dev:masterfrom
Vperiodt:mcp-trace

Conversation

@Vperiodt

@Vperiodt Vperiodt commented May 14, 2026

Copy link
Copy Markdown
Contributor

What this PR does / why we need it

Adds end-to-end MLflow integration for Feast: distributed tracing, trace export, assessment sync, and MLflow GenAI Dataset sync — enabling ML teams to trace agent-feature interactions, produce training data from production traces, and serve curated evaluation datasets through Feast's online/offline stores.

1. Distributed Tracing (mlflow.enable_distributed_tracing: true)

  • Server-side spans: Every feature retrieval endpoint (get-online-features, search, write-to-online-store) creates an MLflow span with feature names, entity count, and project metadata.
  • Cross-process trace linking: W3C traceparent header propagation — Feast spans become children of the caller's trace. MCP server instruments its internal httpx client and forwards mcp-session-id for session grouping.
  • In-process LLM span tagging: A span processor (feast_span_processor) automatically stamps LLM/CHAT_MODEL spans with feast.context_features when the SDK is used directly via feast_trace_scope().
  • PII redaction: Optional redact_entity_pii scrubs entity values from span inputs.
  • Sampling: trace_sampling_ratio controls what fraction of requests produce spans.
  • Graceful degradation: Missing MLflow, incompatible versions, or span creation failures never block feature serving.

2. Trace Assessment Sync (feast mlflow sync-assessments)

  • Assessment ingestion: Scans MLflow traces in an experiment for expectations and feedback, writes them to a Feast FeatureView or LabelView.
  • Flat mode (default): Each assessment = one row with trace_id, assessment_name, assessment_type, value, source_id, rationale, event_timestamp.
  • Pivot mode (--pivot): Groups assessments by trace_id into LabelView-compatible rows with --assessment-mapping to control column names and --labeler-column for the assessment source.
  • Selective sync: --assessment-names filters which assessments to ingest.
  • LabelView support: Falls back to store.get_label_view() when store.get_feature_view() fails, enabling correct offline writes via PushSource.

3. Trace Export (feast mlflow export-traces)

  • Trace extraction: Queries MLflow for CHAT_MODEL spans, extracts prompt/completion pairs with feast.* context attributes and assessments (expectations + feedback).
  • Label resolution: Two paths — join with a Feast LabelView (offline ConflictPolicy or online LAST_WRITE_WINS) or promote MLflow expectations to corrected_response.
  • Dataset filtering: --dataset limits export to traces curated into a named MLflow GenAI Dataset.
  • Export formats: --format openai (OpenAI-compatible JSONL) or --format enriched (includes trace provenance, feature refs, entity values, label info).
  • Artifact registration: --register saves exported JSONL as an MLflow artifact; --register-experiment / --register-run control where.
Configuration:
   mlflow:
  enabled: true
  tracking_uri: "http://mlflow:5000"
  enable_distributed_tracing: true
  trace_sampling_ratio: 1.0
  redact_entity_pii: false

Which issue(s) this PR fixes:

Checks

  • I've made sure the tests are passing.
  • My commits are signed off (git commit -s)
  • My PR title follows conventional commits format

Testing Strategy

  • Unit tests
  • Integration tests
  • Manual tests
  • Testing is not required for this change

Misc


Open in Devin Review

@Vperiodt Vperiodt changed the title mlflow-tracing [WIP] mlflow-tracing May 14, 2026
@Vperiodt Vperiodt changed the title [WIP] mlflow-tracing [WIP] : mlflow-tracing May 24, 2026
@Vperiodt Vperiodt changed the title [WIP] : mlflow-tracing feat: MLflow native distributed tracing for Feast feature server May 26, 2026
@Vperiodt
Vperiodt marked this pull request as ready for review May 26, 2026 14:59
@Vperiodt
Vperiodt requested a review from a team as a code owner May 26, 2026 14:59
devin-ai-integration[bot]

This comment was marked as resolved.

@Vperiodt Vperiodt changed the title feat: MLflow native distributed tracing for Feast feature server feat: MLflow distributed tracing for Feast feature server May 27, 2026

@devin-ai-integration devin-ai-integration Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Devin Review found 1 new potential issue.

View 7 additional findings in Devin Review.

Open in Devin Review

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔴 Removing detect_types from SQLite connection breaks event_ts in retrieve_online_documents_v2

The PR removes detect_types=sqlite3.PARSE_DECLTYPES | sqlite3.PARSE_COLNAMES from _initialize_conn at sqlite.py:705. This flag was responsible for automatically converting stored timestamps back to datetime objects via the registered converters (lines 99-101). Since datetime objects are stored as epoch integers (via the adapt_datetime_epoch adapter at line 81), removing detect_types means timestamp columns now return raw int values instead of datetime objects.

The online_read method was correctly updated (lines 278-285) to handle int timestamps. However, retrieve_online_documents_v2 was NOT updated. At line 677, the check isinstance(entity_dict[entity_key_value]["event_ts"], datetime) will always be False because event_ts is now an int, causing res_event_ts to always be None. This loses all timestamp information for document retrieval results, which flows into _retrieve_from_online_store_v2 (feature_store.py:3575) and ultimately into the API response.

(Refers to lines 677-678)

Open in Devin Review

Was this helpful? React with 👍 or 👎 to provide feedback.

@franciscojavierarceo

Copy link
Copy Markdown
Member

linter is failing

Comment thread feast-tracing-demo/TRACING.md Outdated
@@ -0,0 +1,130 @@
# Feast MLflow Tracing

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add this official docs

@jyejare jyejare left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR adds comprehensive MLflow distributed tracing support to Feast, enabling automatic trace collection for feature server endpoints, cross-process trace linking via traceparent headers, and in-process feature context tagging. The implementation is well-structured with proper separation of concerns, comprehensive test coverage, and graceful degradation when dependencies are missing.

Comment on lines +278 to +284
if isinstance(ts, (int, float)):
res_ts = datetime.fromtimestamp(ts, tz=timezone.utc)
else:
res_ts = ts.replace(tzinfo=timezone.utc)
ts = cast(datetime, ts)
if ts.tzinfo is not None:
res_ts = ts.astimezone(timezone.utc)
else:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Warning] Inconsistent timestamp handling may cause runtime errors

The code handles both int/float and datetime timestamp types, but the original code assumed datetime objects. This change could break existing code that expects consistent datetime handling and may cause timezone issues with numeric timestamps.

Suggested:

Suggested change
if isinstance(ts, (int, float)):
res_ts = datetime.fromtimestamp(ts, tz=timezone.utc)
else:
res_ts = ts.replace(tzinfo=timezone.utc)
ts = cast(datetime, ts)
if ts.tzinfo is not None:
res_ts = ts.astimezone(timezone.utc)
else:
ts = cast(datetime, ts)
if ts.tzinfo is not None:
res_ts = ts.astimezone(timezone.utc)
else:
res_ts = ts.replace(tzinfo=timezone.utc)

Comment thread pyproject.toml
Comment on lines 105 to +113
oracle = ["ibis-framework[oracle]>=10.0.0"]
mysql = ["pymysql", "types-PyMySQL"]
openlineage = ["openlineage-python>=1.40.0"]
mlflow = [
"mlflow>=2.14.0",
"opentelemetry-api>=1.28.0",
"opentelemetry-sdk>=1.28.0",
"opentelemetry-instrumentation-fastapi>=0.49b0",
"opentelemetry-instrumentation-httpx>=0.49b0",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Suggestion] Consider adding opentelemetry-exporter-otlp-proto-http dependency

The mlflow extra includes several OpenTelemetry packages but is missing opentelemetry-exporter-otlp-proto-http which was mentioned in the pixi.lock file. This could cause issues if OTLP export is needed.

Suggested:

Suggested change
oracle = ["ibis-framework[oracle]>=10.0.0"]
mysql = ["pymysql", "types-PyMySQL"]
openlineage = ["openlineage-python>=1.40.0"]
mlflow = [
"mlflow>=2.14.0",
"opentelemetry-api>=1.28.0",
"opentelemetry-sdk>=1.28.0",
"opentelemetry-instrumentation-fastapi>=0.49b0",
"opentelemetry-instrumentation-httpx>=0.49b0",
mlflow = [
"mlflow>=2.14.0",
"opentelemetry-api>=1.28.0",
"opentelemetry-sdk>=1.28.0",
"opentelemetry-instrumentation-fastapi>=0.49b0",
"opentelemetry-instrumentation-httpx>=0.49b0",
"opentelemetry-exporter-otlp-proto-http>=1.28.0",
]

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added that before it was causing grpcio conflicts and
opentelemetry-exporter-otlp-proto-http is already a transitive dep of mlflow>=2.14.0

Comment on lines +313 to +325
return

try:
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

FastAPIInstrumentor.instrument_app(app)
logger.info("FastAPI OTEL instrumentation enabled for trace propagation")
except ImportError:
logger.debug(
"opentelemetry-instrumentation-fastapi not installed; "
"cross-process trace linking disabled"
)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Suggestion] Improve error handling in FastAPI instrumentation

The ImportError handling is good, but it would be helpful to also catch other potential exceptions during instrumentation and provide more specific error messages.

Suggested:

Suggested change
return
try:
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
FastAPIInstrumentor.instrument_app(app)
logger.info("FastAPI OTEL instrumentation enabled for trace propagation")
except ImportError:
logger.debug(
"opentelemetry-instrumentation-fastapi not installed; "
"cross-process trace linking disabled"
)
try:
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
FastAPIInstrumentor.instrument_app(app)
logger.info("FastAPI OTEL instrumentation enabled for trace propagation")
except ImportError:
logger.debug(
"opentelemetry-instrumentation-fastapi not installed; "
"cross-process trace linking disabled"
)
except Exception as e:
logger.warning(f"Failed to instrument FastAPI app: {e}")

Comment on lines +139 to +143
# async exporter must send the span before that happens.
if has_traceparent:
try:
_mlflow_mod.flush_trace_async_logging()
except Exception:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Warning] Potential performance impact from synchronous flush in async context

Calling flush_trace_async_logging() synchronously in an async endpoint could block the event loop and impact performance. Consider making this operation truly asynchronous or moving it to a background task.

Comment on lines +75 to +82
def feast_trace_scope() -> Iterator[FeastTraceContext]:
"""Context manager that creates and cleans up a ``FeastTraceContext``."""
ctx = FeastTraceContext()
_thread_local.feast_ctx = ctx
try:
yield ctx
finally:
ctx.clear()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Suggestion] Consider adding context validation

The feast_trace_scope context manager could benefit from validation to ensure the context is properly cleaned up even if exceptions occur during setup.

Suggested:

Suggested change
def feast_trace_scope() -> Iterator[FeastTraceContext]:
"""Context manager that creates and cleans up a ``FeastTraceContext``."""
ctx = FeastTraceContext()
_thread_local.feast_ctx = ctx
try:
yield ctx
finally:
ctx.clear()
@contextmanager
def feast_trace_scope() -> Iterator[FeastTraceContext]:
"""Context manager that creates and cleans up a ``FeastTraceContext``."""
old_ctx = getattr(_thread_local, 'feast_ctx', None)
ctx = FeastTraceContext()
_thread_local.feast_ctx = ctx
try:
yield ctx
finally:
ctx.clear()
_thread_local.feast_ctx = old_ctx

Comment on lines +63 to +68
enable_tracing: StrictBool = True
""" bool: When True and mlflow.enabled=True, server-side API calls
create MLflow trace spans via mlflow.start_span(). Spans appear
in the MLflow UI Traces tab and support parent-child linking via
traceparent headers. Defaults to True. """

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Nitpick] Consider more descriptive field name

The field name 'enable_tracing' could be more specific since it only controls MLflow tracing, not all tracing in Feast.

Suggested:

Suggested change
enable_tracing: StrictBool = True
""" bool: When True and mlflow.enabled=True, server-side API calls
create MLflow trace spans via mlflow.start_span(). Spans appear
in the MLflow UI Traces tab and support parent-child linking via
traceparent headers. Defaults to True. """
enable_distributed_tracing: StrictBool = True
""" bool: When True and mlflow.enabled=True, server-side API calls
create MLflow trace spans via mlflow.start_span(). Spans appear
in the MLflow UI Traces tab and support parent-child linking via
traceparent headers. Defaults to True. """

@@ -0,0 +1,132 @@
# MLflow Distributed Tracing

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

better to include this in existing mlflow.md

Comment thread docs/SUMMARY.md Outdated
* [\[Alpha\] Feature View Versioning](reference/alpha-feature-view-versioning.md)
* [OpenLineage Integration](reference/openlineage.md)
* [MLflow Integration](reference/mlflow.md)
* [MLflow Distributed Tracing](reference/mlflow-distributed-tracing.md)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no need of different section

Comment thread sdk/python/feast/tracing_hooks.py Outdated
_logger.debug("Failed to set attribute %s on LLM span", key)


def install_feast_span_processor() -> None:

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is it being called somewhere other than tests ?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It is expected to be called by the users in there agent code which helps to tag LLM spans with feast context, either i can document it or we can add it in the _lazy_init()

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I updated it in the docs

@Vperiodt Vperiodt changed the title feat: MLflow distributed tracing for Feast feature server feat: Add MLflow tracing, fine-tuning export, and GenAI dataset sync Jul 13, 2026
Comment thread sdk/python/feast/cli/finetuning.py Outdated
@@ -0,0 +1,222 @@
"""CLI commands for fine-tuning dataset export."""

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

feast finetuning is not correct naming here.

The feast finetuning name implies Feast has an opinion about LLM training workflows, which is a much larger product decision than what's in this PR.

I'ts better if command is like feast mlflow export-traces or similar - making it clear this is an MLflow integration subcommand

Comment thread sdk/python/feast/cli/datasets.py Outdated
if dry_run:
click.echo(" DRY RUN — no data will be written.")

result = sync_mlflow_dataset_to_feast(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think better integration will be if it's MlflowDatasetSource. The dataset sync should be an MlflowDatasetSource (proper Feast primitive).

@codecov-commenter

codecov-commenter commented Jul 21, 2026

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 42.47076% with 787 lines in your changes missing coverage. Please review.
✅ Project coverage is 46.75%. Comparing base (037c4cd) to head (be79ec5).

Files with missing lines Patch % Lines
sdk/python/feast/trace_export/trace_extractor.py 29.72% 171 Missing and 11 partials ⚠️
...dk/python/feast/mlflow_integration/dataset_sync.py 37.71% 124 Missing and 23 partials ⚠️
sdk/python/feast/feature_server.py 32.55% 115 Missing and 1 partial ⚠️
sdk/python/feast/trace_export/label_resolver.py 25.58% 94 Missing and 2 partials ⚠️
sdk/python/feast/cli/mlflow_cmd.py 30.82% 92 Missing ⚠️
sdk/python/feast/trace_export/dataset_filter.py 0.00% 45 Missing ⚠️
sdk/python/feast/tracing.py 72.80% 27 Missing and 7 partials ⚠️
sdk/python/feast/trace_export/exporters.py 74.41% 14 Missing and 8 partials ⚠️
sdk/python/feast/mlflow_integration/__init__.py 13.33% 13 Missing ⚠️
sdk/python/feast/trace_export/__init__.py 18.75% 13 Missing ⚠️
... and 5 more
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##           master    #6405      +/-   ##
==========================================
- Coverage   46.76%   46.75%   -0.01%     
==========================================
  Files         415      425      +10     
  Lines       50392    51706    +1314     
  Branches     7214     7480     +266     
==========================================
+ Hits        23567    24177     +610     
- Misses      25171    25820     +649     
- Partials     1654     1709      +55     
Flag Coverage Δ
go-feature-server 30.58% <ø> (ø)
python-unit 48.04% <42.47%> (-0.05%) ⬇️
Files with missing lines Coverage Δ
sdk/python/feast/cli/cli.py 55.51% <100.00%> (+0.34%) ⬆️
sdk/python/feast/mlflow_integration/config.py 91.83% <100.00%> (+91.83%) ⬆️
sdk/python/feast/tracing_context.py 95.91% <95.91%> (ø)
sdk/python/feast/tracing_hooks.py 90.00% <90.00%> (ø)
sdk/python/feast/infra/mcp_servers/mcp_server.py 81.70% <66.66%> (-3.59%) ⬇️
sdk/python/feast/feature_store.py 43.03% <52.94%> (+0.06%) ⬆️
sdk/python/feast/infra/online_stores/sqlite.py 59.33% <27.27%> (-0.97%) ⬇️
sdk/python/feast/mlflow_integration/__init__.py 31.57% <13.33%> (+31.57%) ⬆️
sdk/python/feast/trace_export/__init__.py 18.75% <18.75%> (ø)
sdk/python/feast/trace_export/exporters.py 74.41% <74.41%> (ø)
... and 7 more

... and 2 files with indirect coverage changes


Continue to review full report in Codecov by Harness.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 037c4cd...be79ec5. Read the comment docs.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@ntkathole

Copy link
Copy Markdown
Member

@Vperiodt Please rebase and solve conflicts, also please fix the linting

Vperiodt added 3 commits July 31, 2026 15:36
Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED
Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED
Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED
Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED
Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED
@Vperiodt Vperiodt changed the title feat: Add MLflow tracing, fine-tuning export, and GenAI dataset sync feat: Add MLflow tracing, trace Assessment Sync, and Trace Export Aug 4, 2026
@Vperiodt Vperiodt changed the title feat: Add MLflow tracing, trace Assessment Sync, and Trace Export feat:Add MLflow tracing, trace Assessment Sync, and Trace Export Aug 5, 2026
@Vperiodt Vperiodt changed the title feat:Add MLflow tracing, trace Assessment Sync, and Trace Export feat: Add MLflow tracing, trace Assessment Sync, and Trace Export Aug 5, 2026
Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED
Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>

rh-pre-commit.version: 2.3.2
rh-pre-commit.check-secrets: ENABLED

Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants