Add `set_destination` API #14249

B-Step62 · 2025-01-15T05:32:24Z

🛠 DevTools 🛠

Install mlflow from this PR

# Use `%sh` to run this command on Databricks
OPTIONS=$(if pip freeze | grep -q 'mlflow @ git+https://github.com/mlflow/mlflow.git'; then echo '--force-reinstall --no-deps'; fi)
pip install $OPTIONS git+https://github.com/mlflow/mlflow.git@refs/pull/14249/merge

Checkout with GitHub CLI

gh pr checkout 14249

What changes are proposed in this pull request?

Add a new set_destination experimental API to support setting the external span exporter provided by the databricks-agents package.

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

Interface

area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
area/windows: Windows support

Language

language/r: R APIs and clients
language/java: Java APIs and clients
language/new: Proposals for new client languages

Integrations

integrations/azure: Azure and Azure ML integrations
integrations/sagemaker: SageMaker integrations
integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

Signed-off-by: B-Step62 <[email protected]>

github-actions · 2025-01-15T05:32:55Z

Documentation preview for 15a442a will be available when this CircleCI job
completes successfully.

More info

Ignore this comment if this PR does not change the documentation.
It takes a few minutes for the preview to be available.
The preview is updated when a new commit is pushed to this PR.
This comment was created by https://github.com/mlflow/mlflow/actions/runs/12823951771.

B-Step62 · 2025-01-15T05:36:45Z

mlflow/tracing/processor/local.py

+_logger = logging.getLogger(__name__)
+
+
+class LocalSpanProcessor(SimpleSpanProcessor):


note: This class has large logic overlap with InferenceTableSpanProcessor. I will refactor it in a follow-up to minimize the blast radius of shipping this change in between RC and stable release.

dbczumar · 2025-01-15T06:39:32Z

mlflow/tracing/provider.py

+            The exporter is responsible for implementing the logic to send generated
+            traces to the desired destination, such as a trace collector endpoint.
+    """
+    # The destination needs to be persisted because the tracer setup can be re-initialized sometimes


Can we be more specific about when it's reinitialized?

dbczumar · 2025-01-15T06:45:40Z

mlflow/tracing/provider.py

@@ -116,6 +122,28 @@ def detach_span_from_context(token: contextvars.Token):
    context_api.detach(token)


+@experimental
+def set_destination(destination: SpanExporter):


What's the UX if a user wants to set a particular experiment as the destination for their traces? They should just be able to pass the experiment info in, right?

I'm wondering if making users reason about constructing SpanExporter instances is the right API here.

It would be nice if developers defined a "destination" object or URI that users passed in here. Developers could then define the appropriate exporter to be used for sending the traces to that destination. This would also resolve a small terminology gripe I have: an exporter is different from a destination. An exporter is a tool used to send traces to a destination.

As discussed offline, I've updated the logic as follows:

Add TraceDestination base class ("destination" object) and added MlflowExperiment destination that takes experiment ID and tracking URI.

Updated set_destination API to take the desintation object.

The span processor/exporter is hardcoded right now, but we will update it to pluggable registry after the 2.20 release. (Not doing it here to avoid regression).

dbczumar · 2025-01-15T06:46:52Z

mlflow/tracing/provider.py

+        from mlflow.tracing.processor.local import LocalSpanProcessor
+
+        processor = LocalSpanProcessor(_MLFLOW_TRACE_CUSTOM_EXPORTER)


Following from https://github.com/mlflow/mlflow/pull/14249/files#r1916019408, what happens if we want to support setting a particular experiment as a destination? Is LocalSpanProcessor still appropriate for that?

With the change mentioned above, when a user passes MlflowExperiment desintation, MLflow selects theMlflowSpanProcessor and the corresponding exporter.

Signed-off-by: B-Step62 <[email protected]>

mlflow/tracing/destination.py

dbczumar · 2025-01-17T00:26:21Z

mlflow/tracing/provider.py

+        if isinstance(_MLFLOW_TRACE_USER_DESTINATION, MlflowExperiment):
+            from mlflow import MlflowClient
+            from mlflow.tracing.export.mlflow import MlflowSpanExporter
+            from mlflow.tracing.processor.mlflow import MlflowSpanProcessor
+
+            client = MlflowClient(tracking_uri=_MLFLOW_TRACE_USER_DESTINATION.tracking_uri)
+            exporter = MlflowSpanExporter(client)
+            processor = MlflowSpanProcessor(
+                exporter, client, _MLFLOW_TRACE_USER_DESTINATION.experiment_id
+            )
+
+        else:
+            from mlflow.tracing.export.databricks_agent import DatabricksAgentSpanExporter
+            from mlflow.tracing.processor.databricks_agent import DatabricksAgentSpanProcessor
+
+            exporter = DatabricksAgentSpanExporter(_MLFLOW_TRACE_USER_DESTINATION)
+            processor = DatabricksAgentSpanProcessor(exporter)


It would be nice if the Destination implementation defined this logic within the class, e.g. via get_exporter() and get_processor() developer-facing (not user-facing) methods.

Otherwise, it's hard for other developers to add new destinations without modifying provider.py. Eventually, it would be nice to enable developers to define destinations in other packages (e.g. via a plugin system); having the destination define exporter & processor itself would be a good step towards this goal.

Absolutely! @dbczumar can I do that in a follow-up? This part will be entirely replaced by the new pluggable registry implementation, and I can make sure the destination to have exporter/processor control. The intention in this PR is trying to make blast radius minimum because we are bypassing RC.

Signed-off-by: B-Step62 <[email protected]>

dbczumar

LGTM!

Signed-off-by: B-Step62 <[email protected]>

Signed-off-by: B-Step62 <[email protected]> Signed-off-by: k99kurella <[email protected]>

B-Step62 added 4 commits January 10, 2025 17:16

poc

aba59f5

Signed-off-by: B-Step62 <[email protected]>

update

627973e

Signed-off-by: B-Step62 <[email protected]>

experiental

7e589bb

Signed-off-by: B-Step62 <[email protected]>

update

67b7670

Signed-off-by: B-Step62 <[email protected]>

github-actions bot added area/tracking Tracking service, tracking client APIs, autologging rn/none List under Small Changes in Changelogs. v2.20.0 Release blocker for version 2.20.0 labels Jan 15, 2025

B-Step62 commented Jan 15, 2025

View reviewed changes

dbczumar reviewed Jan 15, 2025

View reviewed changes

exporter -> destination

8a002ed

Signed-off-by: B-Step62 <[email protected]>

B-Step62 requested a review from dbczumar January 15, 2025 16:22

B-Step62 added 2 commits January 16, 2025 12:54

doc update

bbe7bf6

Signed-off-by: B-Step62 <[email protected]>

nit

284a93d

Signed-off-by: B-Step62 <[email protected]>

dbczumar reviewed Jan 17, 2025

View reviewed changes

mlflow/tracing/destination.py Outdated Show resolved Hide resolved

dbczumar reviewed Jan 17, 2025

View reviewed changes

B-Step62 added 2 commits January 17, 2025 09:50

id -> type

9b44db0

Signed-off-by: B-Step62 <[email protected]>

id -> type

15a442a

Signed-off-by: B-Step62 <[email protected]>

B-Step62 requested a review from dbczumar January 17, 2025 06:32

dbczumar approved these changes Jan 21, 2025

View reviewed changes

B-Step62 added this pull request to the merge queue Jan 22, 2025

Merged via the queue into mlflow:master with commit 0161bbf Jan 22, 2025
51 checks passed

B-Step62 deleted the trace-destination branch January 22, 2025 00:06

harupy pushed a commit to harupy/mlflow that referenced this pull request Jan 23, 2025

Add set_destination API (mlflow#14249)

14ac92c

Signed-off-by: B-Step62 <[email protected]>

harupy pushed a commit that referenced this pull request Jan 23, 2025

Add set_destination API (#14249)

3ee9499

Signed-off-by: B-Step62 <[email protected]>

harupy mentioned this pull request Jan 23, 2025

Run python3 dev/update_changelog.py --prev-version 2.2... #14314

Merged

karthikkurella pushed a commit to karthikkurella/mlflow that referenced this pull request Jan 30, 2025

Add set_destination API (mlflow#14249)

c308a4b

Signed-off-by: B-Step62 <[email protected]> Signed-off-by: k99kurella <[email protected]>

karthikkurella pushed a commit to karthikkurella/mlflow that referenced this pull request Jan 30, 2025

Add set_destination API (mlflow#14249)

575fe00

Signed-off-by: B-Step62 <[email protected]> Signed-off-by: k99kurella <[email protected]>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `set_destination` API #14249

Add `set_destination` API #14249

B-Step62 commented Jan 15, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Jan 15, 2025 •

edited

Loading

B-Step62 Jan 15, 2025 •

edited

Loading

dbczumar Jan 15, 2025 •

edited

Loading

dbczumar Jan 15, 2025 •

edited

Loading

B-Step62 Jan 15, 2025

dbczumar Jan 15, 2025

B-Step62 Jan 15, 2025 •

edited

Loading

dbczumar Jan 17, 2025

B-Step62 Jan 17, 2025

dbczumar left a comment

		_logger = logging.getLogger(__name__)


		class LocalSpanProcessor(SimpleSpanProcessor):

		from mlflow.tracing.processor.local import LocalSpanProcessor

		processor = LocalSpanProcessor(_MLFLOW_TRACE_CUSTOM_EXPORTER)

Add set_destination API #14249

Add set_destination API #14249

Conversation

B-Step62 commented Jan 15, 2025 • edited by github-actions bot Loading

Install mlflow from this PR

Checkout with GitHub CLI

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

github-actions bot commented Jan 15, 2025 • edited Loading

B-Step62 Jan 15, 2025 • edited Loading

Choose a reason for hiding this comment

dbczumar Jan 15, 2025 • edited Loading

Choose a reason for hiding this comment

dbczumar Jan 15, 2025 • edited Loading

Choose a reason for hiding this comment

B-Step62 Jan 15, 2025

Choose a reason for hiding this comment

dbczumar Jan 15, 2025

Choose a reason for hiding this comment

B-Step62 Jan 15, 2025 • edited Loading

Choose a reason for hiding this comment

dbczumar Jan 17, 2025

Choose a reason for hiding this comment

B-Step62 Jan 17, 2025

Choose a reason for hiding this comment

dbczumar left a comment

Choose a reason for hiding this comment

Add `set_destination` API #14249

Add `set_destination` API #14249

B-Step62 commented Jan 15, 2025 •

edited by github-actions bot

Loading

github-actions bot commented Jan 15, 2025 •

edited

Loading

B-Step62 Jan 15, 2025 •

edited

Loading

dbczumar Jan 15, 2025 •

edited

Loading

dbczumar Jan 15, 2025 •

edited

Loading

B-Step62 Jan 15, 2025 •

edited

Loading