add tracing to evaluate #11911

jessechancy · 2024-05-06T18:00:09Z

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Add tracing to mlflow.evaluate() with support for regular python functions, Pyfunc, and langchain

When a user passes a chain or function to mlflow.evaluate() that has been instrumented with MLflow tracing, MLflow should log traces during evaluation and link them to the evaluation run using the same tagging approach from https://databricks.atlassian.net/browse/ML-40783 .

We should make sure to support LangChain: even if the user hasn’t enabled LangChain autologging, mlflow.evaluate() should enable autologging for the duration of the evaluation so that traces are captured.

We should make sure to only enable the tracing part of LangChain autologging, not the part that logs models

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)

github-actions · 2024-05-06T18:00:46Z

Documentation preview for df3c6ef 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/9229937298.

liangz1

Looks good so far! Shall we add a test?

mlflow/models/evaluation/base.py

BenWilson2 · 2024-05-07T22:03:04Z

The tracing branch is now closed. Can you rebase with a merge target to master?

liangz1

Left some questions about clarifying the supported functionality. Thanks!

liangz1 · 2024-05-09T22:58:17Z

tests/evaluate/test_evaluation.py

@@ -361,6 +363,73 @@ def baseline_model_uri(request):
    return None


+def test_mlflow_evaluate_logs_traces():


Could you update the docs for mlflow.evaluate() to describe this new functionality? I'm curious to see:

When would the traces be generated with an evaluate() call. Is all model type supported?

How to find the traces produced by a certain evaluate run.

Is model = function supported? If so, what does it look like? Is there an example for this case?

All models types should be supported. Anything other than langchain would just trace the input/output. In langchain, it would create a span for each individual section in langchain.

The traces would appear in the traces section as if you were to log a regular trace

I think there isn't a difference between a pyfunc model or function here, because mlflow.evaluate does the neccessary steps to turn it all into a model with a predict method. Here we would be tracking the input/output of this predict method.

Anything other than langchain would just trace the input/output

It's up to the user to decide what level of granularity to use when tracing; if the user instruments each subroutine in their model with @mlflow.trace, then they won't just get overall input / output. We're not limiting / restricting this granularity in this PR, right @jessechancy ?

The traces would appear in the traces section as if you were to log a regular trace

The key point is that we set the run ID as a tag on the trace, which appears in the UI.

I think there isn't a difference between a pyfunc model or function here, because mlflow.evaluate does the neccessary steps to turn it all into a model with a predict method. Here we would be tracking the input/output of this predict method.

Same as above - the user might instrument at a level that's more granular than just input / output. Can we confirm this works? Looks like your tests already cover that.

liangz1 · 2024-05-09T23:00:00Z

tests/evaluate/test_evaluation.py

+        targets=iris_dataset._constructor_args["targets"],
+        evaluators="dummy_evaluator",
+    )
+    assert len(get_traces()) == 1


What does the trace look like with a regressor model? In which scenario would the user want to see a trace for a regressor evaluation?

Hi @jessechancy , is this addressed somewhere? Could you share a pointer if possible?

Besides, I'm wondering if we should provide a config to turn off the tracing for evaluate() runs, in case non-genai users get confused with the auto-logged traces. cc @BenWilson2

tests/evaluate/test_evaluation.py

BenWilson2 · 2024-05-10T15:35:50Z

mlflow/utils/autologging_utils/__init__.py

+        k: {**v, "disable": True} for k, v in AUTOLOGGING_INTEGRATIONS.items()
+    }
+    try:
+        yield None


Suggested change

yield None

yield

@jessechancy can we address this nit?

BenWilson2 · 2024-05-13T19:14:58Z

tests/evaluate/test_evaluation.py

+
+    TRACE_BUFFER.clear()
+
+


We may want to add one final test to this suite.

Call evaluate on a model

Verify traces are recorded for a simple langchain chain

Call mlflow autolog (specifying no overrides to the default)

Invoke the model

Ensure that the default settings for autologging apply and that the expected tracked run and artifacts / metrics / params are logged to the run.

Mostly to ensure that modifications to autologging behavior within a single sessions are locally scoped with the configuration overrides being used in the context handler (it works as expected now, but in the future, if this logic needs to change, having the test will be a guard against a regression in the behavior introduced in this PR)

Great idea. Can we also test:

Call autolog with some specific settings

Call evaluate on a model

Verify traces are recorded for a simple langchain chain

Invoke the model

Ensure that the settings for autologging in (1) apply and that the expected tracked run and artifacts / metrics / params are logged to the run.

Can we also test these cases when evaluate fails with an exception?

BenWilson2

LGTM once the safeguard test is added

liangz1

LGTM overall, just a question about the non-genai eval run case.
@jessechancy Do we have another JIRA ticket or PR for updating the documentation? If so, could you link it? Thanks!

liangz1 · 2024-05-13T19:45:30Z

tests/evaluate/test_evaluation.py

+        targets=iris_dataset._constructor_args["targets"],
+        evaluators="dummy_evaluator",
+    )
+    assert len(get_traces()) == 1


Hi @jessechancy , is this addressed somewhere? Could you share a pointer if possible?

Besides, I'm wondering if we should provide a config to turn off the tracing for evaluate() runs, in case non-genai users get confused with the auto-logged traces. cc @BenWilson2

liangz1

I noticed some langchain deps changes that might interrupt non-langchain users. Can you help make sure the experience is still smooth in that case? Thanks!

liangz1 · 2024-05-13T19:52:07Z

mlflow/models/evaluation/base.py

+
+        def monkey_patch_predict(x):
+            # Disable all autologging except for traces
+            mlflow.langchain.autolog(log_inputs_outputs=False, disable=False)


Shall we check whether langchain is imported? I see this code block is not langchain-specific, so if the customer only uses non-langchain packages, shall we skip this?

+1 we might want to only apply this enablement if langchain core module is available through a validation with importlib

liangz1 · 2024-05-13T19:53:36Z

.github/workflows/master.yml

@@ -240,7 +240,7 @@ jobs:
      - name: Install dependencies
        run: |
          source ./dev/install-common-deps.sh
-          pip install pyspark torch transformers
+          pip install pyspark torch transformers langchain


Do we want to make langchain a common deps? It seems to me that langchain should not be required for non-genai users. (related to my comment in mlflow/models/evaluation/base.py)

Can we test if langchain is not installed, a user should still be able to call mlflow.evaluate()?

It should not be a common dependency - this line change is only for test infra.
That being said, a +1 to the validation check that the patch function with evaluate will work if langchain is not installed. Temporarily removing the install from the REPL context would be good to validate.

BenWilson2 · 2024-05-13T20:36:47Z

Regarding #11911 (comment) - this is valid. I feel like the non-genai use cases (classification, regression) should be opt-in for this behavior as it doesn't make a great deal of sense and could potentially generate a very large dataset that doesn't provide any value, even for debugging.

dbczumar · 2024-05-15T19:44:25Z

mlflow/utils/autologging_utils/__init__.py

@@ -50,6 +51,8 @@

 # Flag indicating whether autologging is globally disabled for all integrations.
 _AUTOLOGGING_GLOBALLY_DISABLED = False
+# Exceptions autologging flavors to the above flag
+_AUTOLOGGING_GLOBALLY_DISABLED_EXCEPTIONS = []


Suggested change

_AUTOLOGGING_GLOBALLY_DISABLED_EXCEPTIONS = []

_AUTOLOGGING_GLOBALLY_DISABLED_EXEMPTIONS = []

dbczumar · 2024-05-15T19:45:06Z

mlflow/models/evaluation/base.py

+MLFLOW_EVALUATE_LANGCHAIN_AUTOLOG_CONFIG = {
+    "log_input_examples": False,
+    "log_model_signatures": False,
+    "log_models": False,
+    "log_datasets": False,
+    "log_inputs_outputs": False,
+    "disable": False,
+    "exclusive": False,
+    "disable_for_unsupported_versions": True,
+    "silent": False,
+    "registered_model_name": None,
+    "extra_tags": None,
+}


This seems to be unused. Can we move it to https://github.com/mlflow/mlflow/pull/11911/files#r1602160113?

dbczumar · 2024-05-15T19:49:30Z

mlflow/utils/autologging_utils/__init__.py

+        mlflow.langchain.autolog(log_inputs_outputs=False)
+        try:
+            yield None
+        finally:
+            mlflow.langchain.autolog(**prev_langchain_params)


Suggested change

mlflow.langchain.autolog(log_inputs_outputs=False)

try:

yield None

finally:

mlflow.langchain.autolog(**prev_langchain_params)

try:

mlflow.langchain.autolog(log_inputs_outputs=False)

yield

finally:

mlflow.langchain.autolog(**prev_langchain_params)

liangz1

LGTM with the non-langchain case. Thanks @jessechancy !

liangz1 · 2024-05-15T20:49:14Z

tests/evaluate/test_evaluation.py

+    assert run.info.run_id == get_traces()[0].info.request_metadata["mlflow.sourceRun"]
+
+
+def test_evaluate_works_with_no_langchain_installed():


Looks good!

dbczumar · 2024-05-22T08:57:36Z

mlflow/models/evaluation/base.py

@@ -1265,6 +1265,7 @@ def _validate_dataset_type_supports_predictions(data, supported_predictions_data
 def _evaluate(
    *,
    model,
+    model_predict_func,


@jessechancy This will break all existing third-party evaluator plugins, including ones developed by Databricks. Can we construct model_predict_func inside the default evaluator instead?

+1 that it seems naturally scoped to _extract_predict_fn except in the case that a user would pass this explicitly into a DefaultEvaluator.evaluate. Is there a clear use case around that we are trying to address?

I think the _extract_predict_fn doesn't cover all cases here, since I see direct usage of model.predict

dbczumar

https://github.com/mlflow/mlflow/pull/11911/files#r1609573086

dbczumar · 2024-05-22T09:00:55Z

mlflow/utils/autologging_utils/__init__.py

+
+@contextlib.contextmanager
+def restrict_langchain_autologging_to_traces_only():
+    if importlib.util.find_spec("langchain") is None:


Is it sufficient to check sys.modules instead? Are there examples that recommend using find_spec() instead?

There are a few examples throughtout the codebase that use importlib.util.find_spec, but I can change it to sys.modules

mlflow/models/evaluation/base.py

mlflow/utils/autologging_utils/__init__.py

# This is the 1st commit message: add tracing to evaluate Signed-off-by: Jesse Chan <jesse.chan@databricks.com> monkeypatch with langchain Signed-off-by: Jesse Chan <jesse.chan@databricks.com> recursion fix Signed-off-by: Jesse Chan <jesse.chan@databricks.com> remove print Signed-off-by: Jesse Chan <jesse.chan@databricks.com> Autologging langchain config and cleanup Signed-off-by: Jesse Chan <jesse.chan@databricks.com> comment fixes Signed-off-by: Jesse Chan <jesse.chan@databricks.com> wip Signed-off-by: Jesse Chan <jesse.chan@databricks.com> wip Signed-off-by: Jesse Chan <jesse.chan@databricks.com> fixes Signed-off-by: Jesse Chan <jesse.chan@databricks.com> test import global Signed-off-by: Jesse Chan <jesse.chan@databricks.com> test fixes Signed-off-by: Jesse Chan <jesse.chan@databricks.com> fixes Signed-off-by: Jesse Chan <jesse.chan@databricks.com> fixes Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#2: fixes Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#3: fixes Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#4: retrigger tests Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#5: remove copy to prevent retriggers Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#6: fixes + opt in langchain Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#7: fixed tests Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#8: added test for langchain not installed Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#9: add langchain-experimental package Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#10: Create copy of model.predict for tracing Signed-off-by: Jesse Chan <jesse.chan@databricks.com> # This is the commit message mlflow#11: fixed tests Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

ian-ack-db · 2024-05-24T17:27:42Z

tests/evaluate/test_evaluation.py

+def test_langchain_autolog_parameters_matches_default_parameters():
+    # get parameters from mlflow.langchain.autolog
+    params = inspect.signature(mlflow.langchain.autolog).parameters
+    for name in params:
+        assert name in MLFLOW_EVALUATE_RESTRICT_LANGCHAIN_AUTOLOG_TO_TRACES_CONFIG
+    for name in MLFLOW_EVALUATE_RESTRICT_LANGCHAIN_AUTOLOG_TO_TRACES_CONFIG:
+        assert name in params


Thanks for adding this test. Can you give guidance to future engineers who will see this test fail without much context? How should they determine the value of any new parameter in MLFLOW_EVALUATE_RESTRICT_LANGCHAIN_AUTOLOG_TO_TRACES_CONFIG to not violate traces only requirement?

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

ian-ack-db

LGTM, though not familiar enough with implications for official approval

dbczumar · 2024-05-25T01:38:54Z

mlflow/models/evaluation/default_evaluator.py

+def _extract_predict_fn(model, raw_model, model_predict_fn=None):
    predict_fn = model.predict if model is not None else None
+    if model_predict_fn is not None:
+        predict_fn = model_predict_fn


Can we add a docstring that indicates the precedence order and restructure this method so that the logic is if / elif / else? It's currently difficult to reason about what happens when model_predict_fn is specified and raw_model is also specified, etc.

Can we also add type hints to the function to indicate which parameters are optional?

dbczumar

LGTM once small comments are addressed. Thanks @jessechancy !

github-actions bot added the rn/none List under Small Changes in Changelogs. label May 6, 2024

jessechancy requested review from BenWilson2, B-Step62, harupy, liangz1 and prithvikannan and removed request for harupy and BenWilson2 May 7, 2024 00:22

liangz1 reviewed May 7, 2024

View reviewed changes

mlflow/models/evaluation/base.py Outdated Show resolved Hide resolved

jessechancy force-pushed the jesse-evaluation-trace branch from 3e8b3f6 to 2a3b878 Compare May 7, 2024 22:54

jessechancy changed the base branch from tracing to master May 7, 2024 22:54

jessechancy force-pushed the jesse-evaluation-trace branch from c82328c to d6f7c38 Compare May 8, 2024 20:09

jessechancy requested a review from liangz1 May 8, 2024 20:10

liangz1 reviewed May 9, 2024

View reviewed changes

jessechancy force-pushed the jesse-evaluation-trace branch from 8ecbb1c to 0c2cdf9 Compare May 10, 2024 06:23

BenWilson2 reviewed May 10, 2024

View reviewed changes

BenWilson2 added the enable-dev-tests Enables cross-version tests for dev versions label May 10, 2024

jessechancy requested review from liangz1 and BenWilson2 May 10, 2024 17:08

jessechancy removed the enable-dev-tests Enables cross-version tests for dev versions label May 10, 2024

BenWilson2 reviewed May 13, 2024

View reviewed changes

BenWilson2 approved these changes May 13, 2024

View reviewed changes

liangz1 approved these changes May 13, 2024

View reviewed changes

liangz1 requested changes May 13, 2024

View reviewed changes

jessechancy requested review from dbczumar and removed request for B-Step62 and prithvikannan May 15, 2024 18:16

dbczumar reviewed May 15, 2024

View reviewed changes

liangz1 reviewed May 15, 2024

View reviewed changes

liangz1 requested review from liangz1 and removed request for liangz1 May 15, 2024 20:53

liangz1 approved these changes May 15, 2024

View reviewed changes

jessechancy force-pushed the jesse-evaluation-trace branch from 8f8478a to 84597ba Compare May 15, 2024 23:34

dbczumar reviewed May 22, 2024

View reviewed changes

dbczumar requested changes May 22, 2024

View reviewed changes

dbczumar requested a review from ian-ack-db May 22, 2024 08:57

dbczumar reviewed May 22, 2024

View reviewed changes

ian-ack-db reviewed May 23, 2024

View reviewed changes

mlflow/models/evaluation/base.py Outdated Show resolved Hide resolved

mlflow/utils/autologging_utils/__init__.py Outdated Show resolved Hide resolved

jessechancy added 6 commits May 24, 2024 10:59

added testing for evaluate with exception

b7ac9c2

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

comment fixes

4d9c9c0

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

rebase master

f8084d0

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

added langchain autologging param

3850e55

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

test case for syncing params

3fff538

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

jessechancy force-pushed the jesse-evaluation-trace branch from 497f28b to 3fff538 Compare May 24, 2024 15:11

jessechancy requested review from ian-ack-db and dbczumar May 24, 2024 15:13

ian-ack-db reviewed May 24, 2024

View reviewed changes

jessechancy added 3 commits May 24, 2024 16:07

fixes

8ee2591

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

added comment to test

ad34fe7

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

fix tests

df3c6ef

Signed-off-by: Jesse Chan <jesse.chan@databricks.com>

ian-ack-db reviewed May 24, 2024

View reviewed changes

dbczumar reviewed May 25, 2024

View reviewed changes

dbczumar approved these changes May 25, 2024

View reviewed changes

		@@ -361,6 +363,73 @@ def baseline_model_uri(request):
		return None


		def test_mlflow_evaluate_logs_traces():

	_AUTOLOGGING_GLOBALLY_DISABLED_EXCEPTIONS = []
	_AUTOLOGGING_GLOBALLY_DISABLED_EXEMPTIONS = []

		assert run.info.run_id == get_traces()[0].info.request_metadata["mlflow.sourceRun"]


		def test_evaluate_works_with_no_langchain_installed():

add tracing to evaluate #11911

Are you sure you want to change the base?

add tracing to evaluate #11911

Conversation

jessechancy commented May 6, 2024 • edited

Related Issues/PRs

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 May 6, 2024 • edited

liangz1 left a comment

Choose a reason for hiding this comment

BenWilson2 commented May 7, 2024

liangz1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dbczumar May 25, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BenWilson2 left a comment

Choose a reason for hiding this comment

liangz1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liangz1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

BenWilson2 commented May 13, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

liangz1 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ian-ack-db May 23, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dbczumar left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ian-ack-db left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dbczumar left a comment

Choose a reason for hiding this comment

jessechancy commented May 6, 2024 •

edited

github-actions bot commented May 6, 2024 •

edited

dbczumar May 25, 2024 •

edited

ian-ack-db May 23, 2024 •

edited