feat: Add MLflow tracing, trace Assessment Sync, and Trace Export - #6405
feat: Add MLflow tracing, trace Assessment Sync, and Trace Export#6405Vperiodt wants to merge 9 commits into
Conversation
There was a problem hiding this comment.
🔴 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)
Was this helpful? React with 👍 or 👎 to provide feedback.
|
linter is failing |
| @@ -0,0 +1,130 @@ | |||
| # Feast MLflow Tracing | |||
jyejare
left a comment
There was a problem hiding this comment.
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.
| 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: |
There was a problem hiding this comment.
[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:
| 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) |
| 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", |
There was a problem hiding this comment.
[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:
| 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", | |
| ] |
There was a problem hiding this comment.
Added that before it was causing grpcio conflicts and
opentelemetry-exporter-otlp-proto-http is already a transitive dep of mlflow>=2.14.0
| 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" | ||
| ) | ||
|
|
There was a problem hiding this comment.
[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:
| 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}") |
| # async exporter must send the span before that happens. | ||
| if has_traceparent: | ||
| try: | ||
| _mlflow_mod.flush_trace_async_logging() | ||
| except Exception: |
There was a problem hiding this comment.
[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.
| 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() |
There was a problem hiding this comment.
[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:
| 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 |
| 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. """ | ||
|
|
There was a problem hiding this comment.
[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:
| 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 | |||
There was a problem hiding this comment.
better to include this in existing mlflow.md
| * [\[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) |
| _logger.debug("Failed to set attribute %s on LLM span", key) | ||
|
|
||
|
|
||
| def install_feast_span_processor() -> None: |
There was a problem hiding this comment.
is it being called somewhere other than tests ?
There was a problem hiding this comment.
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()
There was a problem hiding this comment.
I updated it in the docs
| @@ -0,0 +1,222 @@ | |||
| """CLI commands for fine-tuning dataset export.""" | |||
There was a problem hiding this comment.
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
| if dry_run: | ||
| click.echo(" DRY RUN — no data will be written.") | ||
|
|
||
| result = sync_mlflow_dataset_to_feast( |
There was a problem hiding this comment.
I think better integration will be if it's MlflowDatasetSource. The dataset sync should be an MlflowDatasetSource (proper Feast primitive).
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ 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
... and 2 files with indirect coverage changes Continue to review full report in Codecov by Harness.
🚀 New features to boost your workflow:
|
|
@Vperiodt Please rebase and solve conflicts, also please fix the linting |
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
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>
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)get-online-features,search,write-to-online-store) creates an MLflow span with feature names, entity count, and project metadata.traceparentheader propagation — Feast spans become children of the caller's trace. MCP server instruments its internal httpx client and forwardsmcp-session-idfor session grouping.feast_span_processor) automatically stamps LLM/CHAT_MODEL spans withfeast.context_featureswhen the SDK is used directly viafeast_trace_scope().redact_entity_piiscrubs entity values from span inputs.trace_sampling_ratiocontrols what fraction of requests produce spans.2. Trace Assessment Sync (
feast mlflow sync-assessments)trace_id,assessment_name,assessment_type,value,source_id,rationale,event_timestamp.--pivot): Groups assessments bytrace_idinto LabelView-compatible rows with--assessment-mappingto control column names and--labeler-columnfor the assessment source.--assessment-namesfilters which assessments to ingest.store.get_label_view()whenstore.get_feature_view()fails, enabling correct offline writes via PushSource.3. Trace Export (
feast mlflow export-traces)CHAT_MODELspans, extracts prompt/completion pairs withfeast.*context attributes and assessments (expectations + feedback).ConflictPolicyor onlineLAST_WRITE_WINS) or promote MLflow expectations tocorrected_response.--datasetlimits export to traces curated into a named MLflow GenAI Dataset.--format openai(OpenAI-compatible JSONL) or--format enriched(includes trace provenance, feature refs, entity values, label info).--registersaves exported JSONL as an MLflow artifact;--register-experiment/--register-runcontrol where.Which issue(s) this PR fixes:
Checks
git commit -s)Testing Strategy
Misc