Skip to content

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

Draft
Vperiodt wants to merge 1 commit into
feast-dev:masterfrom
Vperiodt:mcp-trace
Draft

Vperiodt wants to merge 1 commit 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 49.64937% with 718 lines in your changes missing coverage. Please review.
✅ Project coverage is 48.11%. Comparing base (3667772) to head (3002c82).

Files with missing lines Patch % Lines
sdk/python/feast/trace_export/trace_extractor.py 29.72% 171 Missing and 11 partials ⚠️
sdk/python/feast/feature_server.py 33.14% 115 Missing and 2 partials ⚠️
sdk/python/feast/cli/mlflow_cmd.py 29.05% 83 Missing ⚠️
...dk/python/feast/mlflow_integration/dataset_sync.py 48.34% 51 Missing and 27 partials ⚠️
sdk/python/feast/infra/registry/sql.py 58.66% 42 Missing and 20 partials ⚠️
sdk/python/feast/trace_export/label_resolver.py 62.79% 30 Missing and 18 partials ⚠️
sdk/python/feast/trace_export/dataset_filter.py 0.00% 45 Missing ⚠️
sdk/python/feast/tracing.py 74.40% 25 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 ⚠️
... and 7 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   48.04%   48.11%   +0.06%     
==========================================
  Files         427      435       +8     
  Lines       53591    54954    +1363     
  Branches     7800     8083     +283     
==========================================
+ Hits        25749    26442     +693     
- Misses      25986    26559     +573     
- Partials     1856     1953      +97     
Flag Coverage Δ
go-feature-server 30.58% <ø> (ø)
python-unit 49.42% <49.64%> (+0.03%) ⬆️
Files with missing lines Coverage Δ
sdk/python/feast/cli/cli.py 55.97% <100.00%> (+0.33%) ⬆️
sdk/python/feast/mlflow_integration/config.py 90.00% <100.00%> (+0.93%) ⬆️
sdk/python/feast/cli/registry.py 41.66% <50.00%> (-1.82%) ⬇️
sdk/python/feast/tracing_context.py 95.91% <95.91%> (ø)
sdk/python/feast/saved_dataset.py 64.50% <88.88%> (+4.38%) ⬆️
sdk/python/feast/tracing_hooks.py 90.00% <90.00%> (ø)
sdk/python/feast/infra/mcp_servers/mcp_server.py 83.89% <66.66%> (-2.04%) ⬇️
sdk/python/feast/feature_store.py 43.57% <52.94%> (+0.05%) ⬆️
sdk/python/feast/mlflow_integration/__init__.py 31.57% <13.33%> (-68.43%) ⬇️
sdk/python/feast/trace_export/__init__.py 18.75% <18.75%> (ø)
... and 9 more

... and 3 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 3667772...3002c82. 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
Vperiodt force-pushed the mcp-trace branch 4 times, most recently from 7f46aa2 to 855aeb9 Compare August 4, 2026 18:35
@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
@Vperiodt
Vperiodt force-pushed the mcp-trace branch 2 times, most recently from fef300c to be79ec5 Compare August 10, 2026 09:35
@ntkathole

Copy link
Copy Markdown
Member

@Vperiodt conflicts

@Vperiodt
Vperiodt force-pushed the mcp-trace branch 2 times, most recently from 8386081 to 3bd6540 Compare September 28, 2026 05:47
@ntkathole

Copy link
Copy Markdown
Member

@Vperiodt need to run pixi lock

@github-advanced-security

Copy link
Copy Markdown

You are seeing this message because GitHub Code Scanning has recently been set up for this repository, or this pull request contains the workflow file for the Code Scanning tool.

What Enabling Code Scanning Means:

  • The 'Security' tab will display more code scanning analysis results (e.g., for the default branch).
  • Depending on your configuration and choice of analysis tool, future pull requests will be annotated with code scanning analysis results.
  • You will be able to see the analysis results for the pull request's branch on this overview once the scans have completed and the checks have passed.

For more information about GitHub Code Scanning, check out the documentation.

Signed-off-by: Vanshika Vanshika <vvanshik@redhat.com>
@Vperiodt
Vperiodt marked this pull request as draft September 29, 2026 13:00
@ntkathole

Copy link
Copy Markdown
Member

Trace_export dir should be under mlflow

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants