import threading import uuid from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from typing import Any, Literal from unittest import mock from unittest.mock import ANY, MagicMock import pandas as pd import pytest import mlflow from mlflow.entities.assessment import Expectation, Feedback from mlflow.entities.assessment_source import AssessmentSource, AssessmentSourceType from mlflow.entities.span import SpanType from mlflow.entities.trace import Trace from mlflow.entities.trace_data import TraceData from mlflow.exceptions import MlflowException from mlflow.genai.datasets import EvaluationDataset, create_dataset from mlflow.genai.evaluation.entities import EvalItem, EvalResult, EvaluationResult from mlflow.genai.evaluation.harness import ( AUTO_INITIAL_RPS, NoOpRateLimiter, _get_new_expectations, _run_predict, _ScoreSubmitter, _should_clone_trace, backpressure_buffer, ) from mlflow.genai.evaluation.rate_limiter import RPSRateLimiter from mlflow.genai.scorers.base import scorer from mlflow.genai.scorers.builtin_scorers import RelevanceToQuery from mlflow.genai.simulators import ConversationSimulator from mlflow.server import handlers from mlflow.server.fastapi_app import app from mlflow.server.handlers import initialize_backend_stores from mlflow.tracing.constant import AssessmentMetadataKey, TraceMetadataKey from tests.helper_functions import get_safe_port from tests.tracing.helper import ( create_test_trace_info, create_test_trace_info_with_uc_table, get_traces, ) from tests.tracking.integration_test_utils import ServerThread @pytest.fixture def mlflow_experiment_trace(): return Trace( info=create_test_trace_info(trace_id="tr-123", experiment_id="exp-123"), data=TraceData(spans=[]), ) _DUMMY_CHAT_RESPONSE = { "id": "1", "object": "text_completion", "created": "2021-10-01T00:00:00.000000Z", "model": "gpt-4o-mini", "choices": [ { "index": 0, "message": { "content": "This is a response", "role": "assistant", }, "finish_reason": "length", } ], "usage": { "prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2, }, } class TestModel: def predict(self, question: str) -> str: return "I don't know" @scorer def exact_match(outputs, expectations): return outputs == expectations["expected_response"] @scorer def is_concise(outputs, expectations): return len(outputs) <= expectations["max_length"] @scorer def relevance(inputs, outputs): return Feedback( name="relevance", value="yes", rationale="The response is relevant to the question", source=AssessmentSource(source_id="gpt", source_type="LLM_JUDGE"), ) @scorer @mlflow.trace(span_type=SpanType.EVALUATOR) def has_trace(trace): return trace is not None class FailingSessionScorer(mlflow.genai.Scorer): def __init__(self): super().__init__(name="failing_session_scorer") @property def is_session_level_scorer(self) -> bool: return True def __call__(self, session=None, **kwargs): raise ValueError("Session scorer error") class WorkingSessionScorer(mlflow.genai.Scorer): def __init__(self): super().__init__(name="working_session_scorer") @property def is_session_level_scorer(self) -> bool: return True def __call__(self, session=None, **kwargs): return len(session or []) def _validate_assessments(traces): """Validate assessments are added to the traces""" for trace in traces: assert len(trace.info.assessments) == 6, ( f"Expected 6 assessments, got {len(trace.info.assessments)}" f"Assessments: {[a.name for a in trace.info.assessments]}" ) # 2 expectations + 4 feedbacks assessments = {a.name: a for a in trace.info.assessments} a_exact_match = assessments["exact_match"] assert isinstance(a_exact_match, Feedback) assert a_exact_match.trace_id == trace.info.trace_id assert isinstance(a_exact_match.value, bool) assert a_exact_match.source.source_type == AssessmentSourceType.CODE # Scorer name is used as source_id assert a_exact_match.source.source_id == "exact_match" assert a_exact_match.metadata[AssessmentMetadataKey.SOURCE_RUN_ID] is not None a_is_concise = assessments["is_concise"] assert isinstance(a_is_concise, Feedback) assert isinstance(a_is_concise.value, bool) assert a_is_concise.metadata[AssessmentMetadataKey.SOURCE_RUN_ID] is not None a_has_trace = assessments["has_trace"] assert isinstance(a_has_trace, Feedback) assert a_has_trace.value is True assert a_has_trace.metadata[AssessmentMetadataKey.SOURCE_RUN_ID] is not None a_relevance = assessments["relevance"] assert isinstance(a_relevance, Feedback) assert a_relevance.value == "yes" assert a_relevance.source.source_id == "gpt" assert a_relevance.source.source_type == "LLM_JUDGE" assert a_relevance.rationale == "The response is relevant to the question" assert a_relevance.metadata[AssessmentMetadataKey.SOURCE_RUN_ID] is not None a_expected_response = assessments["expected_response"] assert isinstance(a_expected_response, Expectation) assert isinstance(a_expected_response.value, str) assert a_expected_response.source.source_type == AssessmentSourceType.HUMAN assert a_expected_response.source.source_id is not None a_max_length = assessments["max_length"] assert isinstance(a_max_length, Expectation) assert isinstance(a_max_length.value, (int, float)) assert a_max_length.source.source_type == AssessmentSourceType.HUMAN def _validate_eval_result_df(result: EvaluationResult): search_traces_df = mlflow.search_traces(run_id=result.run_id) assert result.result_df is not None assert len(result.result_df) == len(search_traces_df) assert set(result.result_df.columns) >= set(search_traces_df.columns) actual = result.result_df.sort_values(by="trace_id").reset_index(drop=True) expected = search_traces_df.sort_values(by="trace_id").reset_index(drop=True) for i in range(len(actual)): assert actual.iloc[i].trace_id == expected.iloc[i].trace_id assert actual.iloc[i].spans == expected.iloc[i].spans assert actual.iloc[i].assessments == expected.iloc[i].assessments assert actual.iloc[i]["exact_match/value"] is not None assert actual.iloc[i]["is_concise/value"] is not None assert actual.iloc[i]["relevance/value"] is not None assert actual.iloc[i]["has_trace/value"] is not None assert actual.iloc[i]["expected_response/value"] is not None assert actual.iloc[i]["max_length/value"] is not None # backwards compatibility assert len(result.tables["eval_results"]) == len(result.result_df) @dataclass class ServerConfig: host_type: Literal["local", "remote", "databricks"] backend_type: Literal["file", "sqlalchemy"] | None = None # Test with different server configurations # 1. local file backend # 2. local sqlalchemy backend # 3. remote server running on file backend # 4. remote server running on sqlalchemy backend @pytest.fixture( params=[ ServerConfig(host_type="local", backend_type="file"), ServerConfig(host_type="local", backend_type="sqlalchemy"), ServerConfig(host_type="remote", backend_type="file"), ServerConfig(host_type="remote", backend_type="sqlalchemy"), ], ids=["local_file", "local_sqlalchemy", "remote_file", "remote_sqlalchemy"], ) def server_config(request, tmp_path: Path, db_uri: str): """Provides an MLflow Tracking API client pointed at the local tracking server.""" config = request.param if config.backend_type == "file": pytest.skip("FileStore is no longer supported.") match config.backend_type: case "file": backend_uri = tmp_path.joinpath("file").as_uri() case "sqlalchemy": backend_uri = db_uri match config.host_type: case "local": mlflow.set_tracking_uri(backend_uri) yield config case "remote": # Force-reset backend stores before each test. handlers._tracking_store = None handlers._model_registry_store = None initialize_backend_stores(backend_uri, default_artifact_root=tmp_path.as_uri()) with ServerThread(app, get_safe_port()) as url: mlflow.set_tracking_uri(url) yield config def test_evaluate_with_static_dataset(server_config): data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", "expectations": { "expected_response": "MLflow is a tool for ML", "max_length": 100, }, }, { "inputs": {"question": "What is Spark?"}, "outputs": "Spark is a fast data processing engine", "expectations": { "expected_response": "Spark is a fast data processing engine", "max_length": 1, }, }, ] result = mlflow.genai.evaluate( data=data, scorers=[exact_match, is_concise, relevance, has_trace], ) # OSS evaluator doesn't support metrics aggregation yet. metrics = result.metrics assert metrics["exact_match/mean"] == 1.0 assert metrics["is_concise/mean"] == 0.5 assert metrics["relevance/mean"] == 1.0 assert metrics["has_trace/mean"] == 1.0 # Exact number of traces should be generated traces = get_traces() assert len(traces) == len(data) # Traces should be associated with the eval run traces = mlflow.search_traces(run_id=result.run_id, return_type="list") assert len(traces) == len(data) # Re-order traces to match with the order of the input data traces = sorted(traces, key=lambda t: t.data.spans[0].inputs["question"]) for i in range(len(traces)): assert len(traces[i].data.spans) == 1 span = traces[i].data.spans[0] assert span.name == "root_span" assert span.inputs == data[i]["inputs"] assert span.outputs == data[i]["outputs"] _validate_assessments(traces) _validate_eval_result_df(result) # Dataset input should be logged to the run run = mlflow.get_run(result.run_id) assert len(run.inputs.dataset_inputs) == 1 assert run.inputs.dataset_inputs[0].dataset.name == "dataset" assert run.inputs.dataset_inputs[0].dataset.source_type == "code" def test_evaluate_with_empty_scorers_logs_expectations(server_config): """Regression test for #23746. When scorers=[] (no scorers), dataset expectations must still be persisted to the trace as Expectation assessments. Before the fix, the no-scorers branch in harness._run_pipeline set EvalResult(assessments=[]) without calling _get_new_expectations / _log_assessments, so expectations were silently dropped. """ data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", "expectations": { "expected_response": "MLflow is a tool for ML", "max_length": 100, }, }, { "inputs": {"question": "What is Spark?"}, "outputs": "Spark is a fast data processing engine", "expectations": { "expected_response": "Spark is a fast data processing engine", "max_length": 1, }, }, ] # Empty scorers list: this is the regressed code path. result = mlflow.genai.evaluate(data=data, scorers=[]) traces = mlflow.search_traces(run_id=result.run_id, return_type="list") assert len(traces) == len(data) traces = sorted(traces, key=lambda t: t.data.spans[0].inputs["question"]) for i in range(len(traces)): trace = traces[i] assessments = {a.name: a for a in trace.info.assessments} # No scorers ran, so exactly the 2 dataset expectations must be present # (and no Feedback assessments). assert len(trace.info.assessments) == 2 assert set(assessments) == {"expected_response", "max_length"}, ( f"Expected only dataset expectations, got {list(assessments)}" ) a_expected_response = assessments["expected_response"] assert isinstance(a_expected_response, Expectation) assert a_expected_response.trace_id == trace.info.trace_id assert a_expected_response.value == data[i]["expectations"]["expected_response"] assert a_expected_response.source.source_type == AssessmentSourceType.HUMAN a_max_length = assessments["max_length"] assert isinstance(a_max_length, Expectation) assert a_max_length.value == data[i]["expectations"]["max_length"] assert a_max_length.source.source_type == AssessmentSourceType.HUMAN def test_evaluate_with_empty_scorers_logs_dataset_tags(server_config): """Regression test for #23746 (tags part). The no-scorers short-circuit in harness._run_pipeline also skipped the eval_item.tags -> set_trace_tag step that _run_score performs, so with scorers=[] dataset tags were silently dropped from the traces. """ data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", "tags": {"dataset_split": "validation", "case_id": "case-1"}, } ] result = mlflow.genai.evaluate(data=data, scorers=[]) traces = mlflow.search_traces(run_id=result.run_id, return_type="list") assert len(traces) == 1 tags = traces[0].info.tags assert tags["dataset_split"] == "validation" assert tags["case_id"] == "case-1" def test_evaluate_passed_respects_scorer_pass_if(server_config): @scorer(pass_if=lambda v: v >= 0.8) def confidence(outputs): return 0.9 if outputs == "good" else 0.5 passing = mlflow.genai.evaluate( data=[{"inputs": {"q": "x"}, "outputs": "good"}], scorers=[confidence], ) assert passing.pass_criteria.get("confidence") is not None assert passing.passed, passing.reason failing = mlflow.genai.evaluate( data=[{"inputs": {"q": "x"}, "outputs": "bad"}], scorers=[confidence], ) assert not failing.passed assert "confidence" in failing.reason def test_evaluate_numeric_value_without_pass_if_fails_loudly(server_config): @scorer def confidence(outputs): return 0.9 result = mlflow.genai.evaluate( data=[{"inputs": {"q": "x"}, "outputs": "good"}], scorers=[confidence], ) # A bare numeric value is not guessed as pass/fail; the user must declare pass_if. assert not result.passed assert "pass_if" in result.reason def test_evaluate_errored_scorer_fails_not_silently_passes(server_config): @scorer def boom(outputs): raise ValueError("kaboom") result = mlflow.genai.evaluate( data=[{"inputs": {"q": "x"}, "outputs": "good"}], scorers=[boom], ) assert not result.passed assert "boom" in result.reason assert "kaboom" in result.reason def test_evaluate_reason_includes_scorer_rationale(server_config): @scorer def judged(outputs): return Feedback(value="no", rationale="answer was wrong") result = mlflow.genai.evaluate( data=[{"inputs": {"q": "x"}, "outputs": "good"}], scorers=[judged], ) assert not result.passed assert "answer was wrong" in result.reason @pytest.mark.parametrize("is_predict_fn_traced", [True, False]) def test_evaluate_with_predict_fn(is_predict_fn_traced, server_config): model_id = mlflow.set_active_model(name="test-model-id").model_id data = [ { "inputs": {"question": "What is MLflow?"}, "expectations": { "expected_response": "MLflow is a tool for ML", "max_length": 100, }, }, { "inputs": {"question": "What is Spark?"}, "expectations": { "expected_response": "Spark is a fast data processing engine", "max_length": 1, }, }, ] model = TestModel() predict_fn = mlflow.trace(model.predict) if is_predict_fn_traced else model.predict result = mlflow.genai.evaluate( predict_fn=predict_fn, data=data, scorers=[exact_match, is_concise, relevance, has_trace], model_id=model_id, ) metrics = result.metrics assert metrics["exact_match/mean"] == 0.0 assert metrics["is_concise/mean"] == 0.5 assert metrics["relevance/mean"] == 1.0 assert metrics["has_trace/mean"] == 1.0 # Metrics should be logged to the model ID as well model = mlflow.get_logged_model(model_id) assert metrics == {m.key: m.value for m in model.metrics} # Exact number of traces should be generated traces = get_traces() assert len(traces) == len(data) # Traces should be associated with the eval run traces = mlflow.search_traces(run_id=result.run_id, return_type="list") assert len(traces) == len(data) # Re-order traces to match with the order of the input data traces = sorted(traces, key=lambda t: t.data.spans[0].inputs["question"]) # Check if the model_id is set in the traces assert traces[0].info.trace_metadata[TraceMetadataKey.MODEL_ID] == model_id assert traces[1].info.trace_metadata[TraceMetadataKey.MODEL_ID] == model_id # Validate assessments are added to the traces for i in range(len(traces)): assert len(traces[i].data.spans) == 1 span = traces[i].data.spans[0] assert span.name == "predict" assert span.inputs == data[i]["inputs"] assert span.outputs == "I don't know" _validate_assessments(traces) _validate_eval_result_df(result) @pytest.mark.parametrize("return_type", ["pandas", "list"]) def test_evaluate_with_traces(monkeypatch: pytest.MonkeyPatch, server_config, return_type): questions = ["What is MLflow?", "What is Spark?"] @mlflow.trace(span_type=SpanType.AGENT) def predict(question: str) -> str: return TestModel().predict(question) predict(questions[0]) trace_id = mlflow.get_last_active_trace_id() mlflow.log_expectation( trace_id=trace_id, name="expected_response", value="MLflow is a tool for ML", source=AssessmentSource(source_id="me", source_type="HUMAN"), ) mlflow.log_expectation( trace_id=trace_id, name="max_length", value=100, source=AssessmentSource(source_id="me", source_type="HUMAN"), ) predict(questions[1]) trace_id = mlflow.get_last_active_trace_id() mlflow.log_expectation( trace_id=trace_id, name="expected_response", value="Spark is a fast data processing engine", source=AssessmentSource(source_id="me", source_type="HUMAN"), ) mlflow.log_expectation( trace_id=trace_id, name="max_length", value=1, source=AssessmentSource(source_id="me", source_type="HUMAN"), ) data = mlflow.search_traces(return_type=return_type) assert len(data) == len(questions) result = mlflow.genai.evaluate( data=data, scorers=[exact_match, is_concise, relevance, has_trace], ) metrics = result.metrics assert metrics["exact_match/mean"] == 0.0 assert metrics["is_concise/mean"] == 0.5 assert metrics["relevance/mean"] == 1.0 assert metrics["has_trace/mean"] == 1.0 if server_config.backend_type == "sqlalchemy": # Assessments should be added to the traces in-place and no new trace should be created traces = get_traces() assert len(traces) == len(questions) else: # File store doesn't support trace linking, so each trace will be cloned to the eval run assert len(get_traces()) == len(questions) * 2 # Traces are associated with the eval run traces = mlflow.search_traces(run_id=result.run_id, return_type="list") assert len(traces) == len(questions) # Re-order traces to match with the order of the input data traces = sorted(traces, key=lambda t: t.data.spans[0].inputs["question"]) # Validate assessments are added to the traces _validate_assessments(traces) _validate_eval_result_df(result) def test_evaluate_with_managed_dataset(is_in_databricks): if is_in_databricks: # Databricks path: Use managed dataset with mocks class MockDatasetClient: def __init__(self): # dataset_id -> list of records self.records = {} def create_dataset(self, uc_table_name: str, experiment_ids: list[str]): from databricks.agents.datasets import Dataset dataset = Dataset( dataset_id=str(uuid.uuid4()), name=uc_table_name, digest=None, source_type="databricks-uc-table", ) self.records[dataset.dataset_id] = [] return dataset def list_dataset_records(self, dataset_id: str): return self.records[dataset_id] def batch_create_dataset_records(self, name: str, dataset_id: str, records): self.records[dataset_id].extend(records) def upsert_dataset_record_expectations( self, name: str, dataset_id: str, record_id: str, expectations: list[dict[str, Any]] ): for record in self.records[dataset_id]: if record.id == record_id: record.expectations.update(expectations) def sync_dataset_to_uc(self, dataset_id: str, uc_table_name: str): pass mock_client = MockDatasetClient() with ( mock.patch("databricks.rag_eval.datasets.api._get_client", return_value=mock_client), mock.patch( "databricks.rag_eval.datasets.entities._get_client", return_value=mock_client ), mock.patch("mlflow.genai.datasets.is_databricks_uri", return_value=True), ): dataset = create_dataset( uc_table_name="mlflow.managed.dataset", experiment_id="exp-123" ) dataset.merge_records([ { "inputs": {"question": "What is MLflow?"}, "expectations": { "expected_response": "MLflow is a tool for ML", "max_length": 100, }, }, { "inputs": {"question": "What is Spark?"}, "expectations": { "expected_response": "Spark is a fast data processing engine", "max_length": 1, }, }, ]) result = mlflow.genai.evaluate( data=dataset, predict_fn=TestModel().predict, scorers=[exact_match, is_concise, relevance, has_trace], ) else: dataset = create_dataset( name="eval_test_dataset", tags={"source": "test", "version": "1.0"} ) dataset.merge_records([ { "inputs": {"question": "What is MLflow?"}, "expectations": { "expected_response": "MLflow is a tool for ML", "max_length": 100, }, }, { "inputs": {"question": "What is Spark?"}, "expectations": { "expected_response": "Spark is a fast data processing engine", "max_length": 1, }, }, ]) result = mlflow.genai.evaluate( data=dataset, predict_fn=TestModel().predict, scorers=[exact_match, is_concise, relevance, has_trace], ) metrics = result.metrics assert metrics["exact_match/mean"] == 0.0 assert metrics["is_concise/mean"] == 0.5 assert metrics["relevance/mean"] == 1.0 assert metrics["has_trace/mean"] == 1.0 run = mlflow.get_run(result.run_id) # Dataset metadata should be added to the run assert len(run.inputs.dataset_inputs) == 1 assert run.inputs.dataset_inputs[0].dataset.name == dataset.name assert run.inputs.dataset_inputs[0].dataset.digest == dataset.digest # Check for the correct source_type based on whether we're in Databricks or OSS expected_source_type = ( "databricks-uc-table" if is_in_databricks else "mlflow_evaluation_dataset" ) assert run.inputs.dataset_inputs[0].dataset.source_type == expected_source_type # Traces are associated with the eval run traces = mlflow.search_traces(run_id=result.run_id, return_type="list") assert len(traces) == 2 _validate_assessments(traces) _validate_eval_result_df(result) def test_evaluate_with_managed_dataset_from_searched_traces(): for i in range(3): with mlflow.start_span(name=f"qa_span_{i}") as span: question = f"What is item {i}?" span.set_inputs({"question": question}) span.set_outputs({"answer": f"Item {i} is something"}) mlflow.log_expectation( trace_id=span.trace_id, name="expected_response", value=f"Item {i} is a detailed answer", ) mlflow.log_expectation( trace_id=span.trace_id, name="max_length", value=50 if i % 2 == 0 else 10, ) traces_df = mlflow.search_traces() dataset = create_dataset( name="traces_eval_dataset", tags={"source": "traces", "evaluation": "test"} ) dataset.merge_records(traces_df) result = mlflow.genai.evaluate( data=dataset, predict_fn=TestModel().predict, scorers=[exact_match, is_concise, has_trace], ) metrics = result.metrics assert "exact_match/mean" in metrics assert "is_concise/mean" in metrics assert "has_trace/mean" in metrics assert metrics["has_trace/mean"] == 1.0 def test_model_from_deployment_endpoint(is_in_databricks): with mock.patch("mlflow.deployments.get_deploy_client") as mock_get_deploy_client: mock_client = mock_get_deploy_client.return_value mock_client.predict.return_value = _DUMMY_CHAT_RESPONSE data = [ { "inputs": { "messages": [ {"content": "You are a helpful assistant.", "role": "system"}, {"content": "What is Spark?", "role": "user"}, ], "max_tokens": 10, } }, { "inputs": { "messages": [ {"content": "What is MLflow?", "role": "user"}, ] } }, ] predict_fn = mlflow.genai.to_predict_fn("endpoints:/chat") result = mlflow.genai.evaluate( data=data, predict_fn=predict_fn, scorers=[has_trace], ) databricks_options = {"databricks_options": {"return_trace": True}} mock_client.predict.assert_has_calls( [ # Test call to check if the function is traced or not mock.call(endpoint="chat", inputs={**data[0]["inputs"], **databricks_options}), # First evaluation call mock.call(endpoint="chat", inputs={**data[0]["inputs"], **databricks_options}), # Second evaluation call mock.call(endpoint="chat", inputs={**data[1]["inputs"], **databricks_options}), ], any_order=True, ) # Validate traces traces = mlflow.search_traces(run_id=result.run_id, return_type="list") assert len(traces) == 2 spans = traces[0].data.spans assert len(spans) == 1 assert spans[0].name == "predict" # Eval harness runs prediction in parallel, so the order is not deterministic assert spans[0].inputs in (data[0]["inputs"], data[1]["inputs"]) assert spans[0].outputs == _DUMMY_CHAT_RESPONSE def test_missing_scorers_argument(): with pytest.raises(TypeError, match=r"evaluate\(\) missing 1 required positional"): mlflow.genai.evaluate(data=[{"inputs": "Hello", "outputs": "Hi"}]) def test_empty_scorers_allowed(): mock_result = EvaluationResult(run_id="test-run", metrics={}, result_df=pd.DataFrame()) data = [{"inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is an ML platform"}] with mock.patch("mlflow.genai.evaluation.base._run_harness") as mock_evaluate_oss: mock_evaluate_oss.return_value = (mock_result, {}) result = mlflow.genai.evaluate(data=data, scorers=[]) assert result is mock_result mock_evaluate_oss.assert_called_once() @pytest.mark.parametrize("pass_full_dataframe", [True, False]) def test_trace_input_can_contain_string_input(pass_full_dataframe, is_in_databricks): """ The `inputs` column must be a dictionary when a static dataset is provided. However, when a trace is provided, it doesn't need to be validated and the harness can handle it nicely. """ with mlflow.start_span() as span: span.set_inputs("What is MLflow?") span.set_outputs("MLflow is a tool for ML") traces = mlflow.search_traces() if not pass_full_dataframe: traces = traces[["trace"]] # Harness should run without an error mlflow.genai.evaluate(data=traces, scorers=[RelevanceToQuery()]) def test_max_workers_env_var(monkeypatch): # Disable rate limits so auto-derivation doesn't override the default monkeypatch.setenv("MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT", "0") monkeypatch.setenv("MLFLOW_GENAI_EVAL_SCORER_RATE_LIMIT", "0") def _validate_max_workers(expected_max_workers): with mock.patch( "mlflow.genai.evaluation.harness.ThreadPoolExecutor", wraps=ThreadPoolExecutor ) as mock_executor: mlflow.genai.evaluate( data=[ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", } ], scorers=[RelevanceToQuery()], ) # ThreadPoolExecutor is called twice in OSS (harness + scorers) first_call = mock_executor.call_args_list[0] assert first_call[1]["max_workers"] == expected_max_workers # default workers is 10 _validate_max_workers(10) # override workers with env var monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_WORKERS", "20") _validate_max_workers(20) # legacy env var for backward compatibility monkeypatch.delenv("MLFLOW_GENAI_EVAL_MAX_WORKERS", raising=False) monkeypatch.setenv("RAG_EVAL_MAX_WORKERS", "30") _validate_max_workers(30) def test_dataset_name_is_logged_correctly(is_in_databricks): data = pd.DataFrame({ "inputs": [{"question": "What is MLflow?"}], "outputs": ["MLflow is a tool for ML"], }) with mlflow.start_run() as run: mlflow.genai.evaluate( data=data, scorers=[RelevanceToQuery()], ) if not is_in_databricks: run_data = mlflow.get_run(run.info.run_id) assert run_data.inputs is not None assert run_data.inputs.dataset_inputs is not None assert len(run_data.inputs.dataset_inputs) > 0 dataset_input = run_data.inputs.dataset_inputs[0] dataset = dataset_input.dataset assert dataset.name == "dataset" def test_evaluate_with_dataset_preserves_name(is_in_databricks): from mlflow.entities import Dataset as DatasetEntity data = pd.DataFrame({ "inputs": [{"question": "What is MLflow?"}], "outputs": ["MLflow is a tool for ML"], }) mock_managed_dataset = MagicMock(spec=EvaluationDataset) type(mock_managed_dataset).name = mock.PropertyMock(return_value="my_managed_dataset") mock_managed_dataset.to_df.return_value = data mock_managed_dataset.digest = "test_digest" mock_managed_dataset.source = MagicMock() mock_managed_dataset.source.to_json.return_value = "{}" mock_managed_dataset.source._get_source_type.return_value = "test" mock_managed_dataset._to_mlflow_entity.return_value = DatasetEntity( name="my_managed_dataset", digest="test_digest", source_type="test", source="{}", schema=None, profile=None, ) if not is_in_databricks: with mlflow.start_run() as run: mlflow.genai.evaluate( data=data, scorers=[RelevanceToQuery()], ) run_data = mlflow.get_run(run.info.run_id) dataset_input = run_data.inputs.dataset_inputs[0] assert dataset_input.dataset.name == "dataset" with mlflow.start_run() as run: mlflow.genai.evaluate( data=mock_managed_dataset, scorers=[RelevanceToQuery()], ) run_data = mlflow.get_run(run.info.run_id) dataset_input = run_data.inputs.dataset_inputs[0] assert dataset_input.dataset.name == "my_managed_dataset" def test_evaluate_with_managed_dataset_preserves_name(): mock_managed_dataset = MagicMock() mock_managed_dataset.dataset_id = "d-1234567890abcdef1234567890abcdef" mock_managed_dataset.name = "test.evaluation.sample_dataset" mock_managed_dataset.digest = "abc123" mock_managed_dataset.schema = None mock_managed_dataset.profile = None mock_managed_dataset.source_type = "databricks-uc-table" mock_managed_dataset.create_time = None mock_managed_dataset.created_by = None mock_managed_dataset.last_update_time = None mock_managed_dataset.last_updated_by = None mock_managed_dataset.to_df.return_value = pd.DataFrame({ "inputs": [{"question": "What is MLflow?"}], "outputs": ["MLflow is a tool for ML"], }) dataset = EvaluationDataset(mock_managed_dataset) with mlflow.start_run() as run: mlflow.genai.evaluate( data=dataset, scorers=[RelevanceToQuery()], ) run_data = mlflow.get_run(run.info.run_id) assert run_data.inputs is not None assert run_data.inputs.dataset_inputs is not None assert len(run_data.inputs.dataset_inputs) > 0 dataset_input = run_data.inputs.dataset_inputs[0] logged_dataset = dataset_input.dataset assert logged_dataset.name == "test.evaluation.sample_dataset" @pytest.mark.parametrize( ("tags_data", "expected_calls"), [ # Regular tags ( [ {"environment": "test", "model_version": "v1.0"}, {"environment": "production", "team": "data-science"}, ], [ ("environment", "test"), ("model_version", "v1.0"), ("environment", "production"), ("team", "data-science"), ], ), # Empty tags dict ( [{}, {}], [], ), # None tags (no tags field) ( [None, None], [], ), # Mix of tags and empty/None ( [{"env": "test"}, {}, None], [("env", "test")], ), ], ) def test_evaluate_with_tags(tags_data, expected_calls): data = [ { "inputs": {"question": f"What is question {i}?"}, "outputs": f"Answer {i}", "expectations": {"expected_response": f"Answer {i}"}, "tags": tags, } for i, tags in enumerate(tags_data) ] with mock.patch("mlflow.set_trace_tag") as mock_set_trace_tag: mlflow.genai.evaluate( data=data, scorers=[exact_match], ) # Check that all expected calls were made (order may vary due to parallel execution) actual_calls = mock_set_trace_tag.call_args_list expected_mock_calls = [ mock.call(trace_id=ANY, key=key, value=value) for key, value in expected_calls ] assert len(actual_calls) == len(expected_mock_calls) for expected_call in expected_mock_calls: assert expected_call in actual_calls def test_evaluate_with_traces_tags_no_warnings(): with mlflow.start_span() as span: span.set_inputs({"question": "Hello?"}) traces = mlflow.search_traces() with mock.patch("mlflow.tracing.client._logger.warning") as mock_warning: mlflow.genai.evaluate( data=traces, scorers=[has_trace], ) assert not any( "immutable and cannot be set on a trace" in call.args[0] for call in mock_warning.call_args_list ) def test_evaluate_with_tags_error_handling(is_in_databricks): data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", "expectations": {"expected_response": "MLflow is a tool for ML"}, "tags": {"invalid_tag": "value"}, } ] # Mock set_trace_tag to raise an exception with mock.patch("mlflow.set_trace_tag", side_effect=Exception("Tag logging failed")): # This should not raise an exception result = mlflow.genai.evaluate( data=data, scorers=[exact_match], ) # Evaluation should still succeed assert "exact_match/mean" in result.metrics def test_evaluate_with_invalid_tags_type(): data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", "expectations": {"expected_response": "MLflow is a tool for ML"}, "tags": "invalid_tags_string", # Should be dict } ] with pytest.raises(MlflowException, match="Tags must be a dictionary"): mlflow.genai.evaluate( data=data, scorers=[exact_match], ) def test_evaluate_without_inputs_in_eval_dataset(): answers = [ "MLflow is an open-source platform for managing ML lifecycle", "Apache Spark is a fast data processing engine", "I don't know", ] for answer in answers: with mlflow.start_span() as span: span.set_outputs(answer) trace_df = mlflow.search_traces() trace_df["inputs"] = None trace_df["expectations"] = pd.Series([ {"expected_response": answer, "max_length": 100} for answer in answers ]) result = mlflow.genai.evaluate( data=trace_df, scorers=[is_concise, exact_match, has_trace], ) assert "is_concise/mean" in result.metrics assert "exact_match/mean" in result.metrics assert "has_trace/mean" in result.metrics @scorer def input_exist(inputs): if inputs is None: return False return True trace_df["outputs"] = None result = mlflow.genai.evaluate( data=trace_df, scorers=[input_exist], ) assert result.metrics["input_exist/mean"] == 0.0 def test_evaluate_with_only_trace_in_eval_dataset(): for _ in range(3): with mlflow.start_span(): pass trace_df = mlflow.search_traces() trace_df = trace_df[["trace"]] result = mlflow.genai.evaluate( data=trace_df, scorers=[has_trace], ) assert result.metrics["has_trace/mean"] == 1.0 @pytest.mark.parametrize("is_enabled", [True, False]) def test_evaluate_with_scorer_tracing(server_config, monkeypatch, is_enabled): monkeypatch.setenv("MLFLOW_GENAI_EVAL_ENABLE_SCORER_TRACING", str(is_enabled).lower()) data = [ { "inputs": {"question": "What is MLflow?"}, "expectations": { "expected_response": "MLflow is a tool for ML", "max_length": 100, }, }, { "inputs": {"question": "What is Spark?"}, "expectations": { "expected_response": "Spark is a fast data processing engine", "max_length": 1, }, }, ] result = mlflow.genai.evaluate( predict_fn=TestModel().predict, data=data, scorers=[exact_match, is_concise, relevance, has_trace], ) metrics = result.metrics assert metrics["exact_match/mean"] == 0.0 assert metrics["is_concise/mean"] == 0.5 assert metrics["relevance/mean"] == 1.0 assert metrics["has_trace/mean"] == 1.0 traces = get_traces() if is_enabled: assert len(traces) == len(data) * 5 # 1 trace for prediction + 4 scorer traces else: assert len(traces) == len(data) # Traces should be associated with the eval run traces = mlflow.search_traces( filter_string="tags.`mlflow.eval.requestId` != 'None'", run_id=result.run_id, return_type="list", ) assert len(traces) == len(data) # Each assessment should have a source trace ID for trace in traces: for a in trace.info.assessments: if isinstance(a, Feedback) and is_enabled: assert a.metadata[AssessmentMetadataKey.SCORER_TRACE_ID] is not None assert a.metadata[AssessmentMetadataKey.SCORER_TRACE_ID] != trace.info.trace_id else: assert AssessmentMetadataKey.SCORER_TRACE_ID not in a.metadata @pytest.mark.parametrize("diff_experiment_id", [True, False]) def test_eval_with_traces_log_spans_correctly(diff_experiment_id): exp_id = mlflow.set_experiment("traces exp").experiment_id with mlflow.start_span() as span: span.set_inputs({"question": "What is MLflow?"}) span.set_outputs({"answer": "MLflow is a tool for ML"}) span.set_attributes({"key": "value"}) with mlflow.start_span() as child_span: child_span.set_inputs("test") # set to a different experiment if diff_experiment_id: mlflow.set_experiment("diff exp") # search traces from the original experiment trace_df = mlflow.search_traces(locations=[exp_id]) result = mlflow.genai.evaluate( data=trace_df, scorers=[has_trace], ) assert result.metrics["has_trace/mean"] == 1.0 traces = get_traces() assert len(traces) == 1 # copied trace should contain all spans assert len(traces[0].data.spans) == 2 span = traces[0].data.spans[0] assert span.get_attribute("key") == "value" assert span.inputs == {"question": "What is MLflow?"} assert span.outputs == {"answer": "MLflow is a tool for ML"} child_span = traces[0].data.spans[1] assert child_span.inputs == "test" def test_evaluate_with_mixed_single_turn_and_multi_turn_scorers(server_config): """Test evaluation with a combination of single-turn and multi-turn scorers. Validates that: - Single-turn scorers are applied to all traces - Multi-turn scorers are only applied to the first trace of each session """ # Define a multi-turn scorer that counts conversation turns class ConversationLengthScorer(mlflow.genai.Scorer): def __init__(self): super().__init__(name="conversation_length") @property def is_session_level_scorer(self) -> bool: return True def __call__(self, session=None, **kwargs): """Return the number of turns in the conversation.""" return len(session or []) # Define a single-turn scorer @scorer def response_length(outputs) -> int: """Return the length of the response.""" return len(outputs) if isinstance(outputs, str) else 0 # Create a traced model function @mlflow.trace(span_type=SpanType.CHAT_MODEL) def model(question, session_id): mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) return f"Answer to: {question}" # Generate traces for 2 sessions (3 turns + 2 turns = 5 total traces) mlflow.set_experiment("multi_turn_test") with mlflow.start_run() as run: # Session 1: 3 turns for q in ["Q1", "Q2", "Q3"]: model(q, session_id="session_1") # Session 2: 2 turns for q in ["Q4", "Q5"]: model(q, session_id="session_2") # Get traces for evaluation traces = mlflow.search_traces( locations=[run.info.experiment_id], filter_string=f'run_id = "{run.info.run_id}"' ) # Evaluate with both single-turn and multi-turn scorers result = mlflow.genai.evaluate( data=traces, scorers=[response_length, ConversationLengthScorer()] ) # Validate results result_df = result.result_df # Should have one row per trace assert len(result_df) == 5, f"Expected 5 traces, got {len(result_df)}" # Single-turn scorer should be applied to all traces single_turn_scores = result_df["response_length/value"].notna().sum() assert single_turn_scores == 5, ( f"Expected single-turn scores for all 5 traces, got {single_turn_scores}" ) # Multi-turn scorer should only be applied to first trace of each session (2 total) multi_turn_scores = result_df["conversation_length/value"].notna().sum() assert multi_turn_scores == 2, ( f"Expected multi-turn scores for 2 sessions (first trace only), got {multi_turn_scores}" ) # Validate the conversation length values # Session 1 should have 3 turns, Session 2 should have 2 turns conv_lengths = result_df["conversation_length/value"].dropna().sort_values().tolist() assert conv_lengths == [2.0, 3.0], ( f"Expected conversation lengths [2.0, 3.0], got {conv_lengths}" ) # Validate that all single-turn scores are the same (based on our dummy response) response_lengths = result_df["response_length/value"].dropna() # All responses should be "Answer to: Qx" format, so lengths should be consistent assert all(length > 0 for length in response_lengths) # Verify multi-turn assessments were persisted to the backend (not just in-memory). # Fetch each trace from the backend and check the first trace of each session has # the conversation_length assessment logged. all_traces = mlflow.search_traces( locations=[run.info.experiment_id], run_id=result.run_id, return_type="list", ) for trace in all_traces: assessment_names = {a.name for a in trace.info.assessments} has_session = "mlflow.trace.session" in trace.info.trace_metadata is_first_in_session = "conversation_length" in assessment_names if has_session and is_first_in_session: a = next(a for a in trace.info.assessments if a.name == "conversation_length") assert isinstance(a, Feedback) assert a.value in (2.0, 3.0) def test_evaluate_with_evaluation_dataset_and_session_level_scorers(): # Define a session-level scorer class ConversationLengthScorer(mlflow.genai.Scorer): def __init__(self): super().__init__(name="conversation_length") @property def is_session_level_scorer(self) -> bool: return True def __call__(self, session=None, **kwargs): return len(session or []) # Create traces with session metadata (2 traces in session_1, 1 in session_2) @mlflow.trace(span_type=mlflow.entities.SpanType.CHAT_MODEL) def model(question, session_id): mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) return f"Answer to {question}" model("Q1", session_id="session_1") trace_1 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model("Q2", session_id="session_1") trace_2 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model("Q3", session_id="session_2") trace_3 = mlflow.get_trace(mlflow.get_last_active_trace_id()) # Create dataset from traces dataset = create_dataset(name="multi_turn_dataset") dataset.merge_records([trace_1, trace_2, trace_3]) # Evaluate with session-level scorer result = mlflow.genai.evaluate(data=dataset, scorers=[ConversationLengthScorer()]) result_df = result.result_df # Session-level scorer should produce 2 scores (one per session) assert "conversation_length/value" in result_df.columns assert result_df["conversation_length/value"].notna().sum() == 2 # Verify conversation lengths: session_1 has 2 traces, session_2 has 1 trace conv_lengths = result_df["conversation_length/value"].dropna().sort_values().tolist() assert conv_lengths == [1.0, 2.0] def test_evaluate_dataset_mixed_traces_with_and_without_sessions(): class SessionScorer(mlflow.genai.Scorer): def __init__(self): super().__init__(name="session_length") @property def is_session_level_scorer(self): return True def __call__(self, session=None, **kwargs): return len(session or []) # Create mixed traces @mlflow.trace(span_type=mlflow.entities.SpanType.CHAT_MODEL) def model_with_session(question, session_id): mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) return "answer" @mlflow.trace(span_type=mlflow.entities.SpanType.CHAT_MODEL) def model_without_session(question): return "answer" model_with_session("Q1", "session_1") trace_1 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model_without_session("Q2") trace_2 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model_with_session("Q3", "session_1") trace_3 = mlflow.get_trace(mlflow.get_last_active_trace_id()) # Create dataset and evaluate dataset = create_dataset(name="mixed_dataset") dataset.merge_records([trace_1, trace_2, trace_3]) result = mlflow.genai.evaluate(data=dataset, scorers=[SessionScorer()]) result_df = result.result_df # Should have 1 session-level score (for session_1 with 2 traces) # The trace without session should not be scored by session-level scorer assert result_df["session_length/value"].notna().sum() == 1 assert result_df["session_length/value"].dropna().iloc[0] == 2.0 def test_max_scorer_workers_env_var(monkeypatch): @scorer def dummy_scorer_1(outputs): return True @scorer def dummy_scorer_2(outputs): return True @scorer def dummy_scorer_3(outputs): return True def _validate_scorer_max_workers(expected_max_workers, num_scorers): scorers_list = [dummy_scorer_1, dummy_scorer_2, dummy_scorer_3][:num_scorers] with mock.patch( "mlflow.genai.evaluation.harness.ThreadPoolExecutor", wraps=ThreadPoolExecutor ) as mock_executor: mlflow.genai.evaluate( data=[ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", } ], scorers=scorers_list, ) # Find the scorer pool call by its thread_name_prefix scorer_call = next( call for call in mock_executor.call_args_list if call[1].get("thread_name_prefix") == "MlflowGenAIEvalScorer" ) assert scorer_call[1]["max_workers"] == expected_max_workers # default scorer workers is 10, but limited by number of scorers (3) _validate_scorer_max_workers(expected_max_workers=3, num_scorers=3) # override scorer workers with env var (limit to 2) monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_SCORER_WORKERS", "2") _validate_scorer_max_workers(expected_max_workers=2, num_scorers=3) # when num_scorers < max_scorer_workers, use num_scorers monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_SCORER_WORKERS", "10") _validate_scorer_max_workers(expected_max_workers=2, num_scorers=2) # set to 1 for sequential execution monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_SCORER_WORKERS", "1") _validate_scorer_max_workers(expected_max_workers=1, num_scorers=3) # ===================== ConversationSimulator Integration Tests ===================== def test_evaluate_with_conversation_simulator_requires_predict_fn(): simulator = ConversationSimulator( test_cases=[{"goal": "Learn about MLflow"}], max_turns=2, ) with pytest.raises(MlflowException, match="predict_fn is required"): mlflow.genai.evaluate( data=simulator, scorers=[has_trace], ) def test_evaluate_with_conversation_simulator_empty_simulation_error(): def failing_predict_fn(input: list[dict[str, Any]], **kwargs): raise Exception("Simulated failure") simulator = ConversationSimulator( test_cases=[{"goal": "Learn about MLflow"}], max_turns=2, ) with mock.patch( "mlflow.genai.simulators.simulator.invoke_model_without_tracing" ) as mock_invoke: # Simulate a failure that produces no traces mock_invoke.side_effect = Exception("LLM call failed") with pytest.raises(MlflowException, match="Simulation produced no traces"): mlflow.genai.evaluate( data=simulator, predict_fn=failing_predict_fn, scorers=[has_trace], ) def test_session_level_evaluation_with_predict_fn_without_simulator(): class SessionScorer(mlflow.genai.Scorer): def __init__(self): super().__init__(name="session_scorer") @property def is_session_level_scorer(self): return True def __call__(self, session=None, **kwargs): return len(session or []) data = [ {"inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool"}, ] with pytest.raises( MlflowException, match=( r"Session-level scorers require traces with session IDs.*" r"session_scorer.*" r"Either pass a ConversationSimulator to `data` with `predict_fn`" ), ): mlflow.genai.evaluate( data=data, predict_fn=TestModel().predict, scorers=[SessionScorer()], ) def test_evaluate_with_conversation_simulator_calls_simulate(): simulator = ConversationSimulator( test_cases=[{"goal": "Learn MLflow"}], max_turns=2, ) def mock_predict_fn(input: list[dict[str, Any]], **kwargs): return {"output": "Mock response"} with mock.patch.object(simulator, "simulate") as mock_simulate: # Return empty list to trigger the "no traces" error mock_simulate.return_value = [] with pytest.raises(MlflowException, match="Simulation produced no traces"): mlflow.genai.evaluate( data=simulator, predict_fn=mock_predict_fn, scorers=[has_trace], ) # Verify simulate was called with predict_fn mock_simulate.assert_called_once_with(mock_predict_fn) @scorer def always_pass(outputs): return True @pytest.mark.parametrize( ("expectations_values", "expected_count"), [ ([None, {}], 2), ([float("nan")], 1), ([{"expected": "value"}, None, {}], 3), ], ) def test_evaluate_handles_empty_expectations(expectations_values, expected_count): data = [ { "inputs": {"question": f"Q{i}"}, "outputs": f"A{i}", "expectations": exp, } for i, exp in enumerate(expectations_values) ] result = mlflow.genai.evaluate(data=data, scorers=[always_pass]) assert result is not None assert result.metrics is not None assert result.result_df is not None assert len(result.result_df) == expected_count assert result.metrics["always_pass/mean"] == 1.0 from mlflow.genai.simulators.simulator import BaseSimulatedUserAgent class MockUserAgent(BaseSimulatedUserAgent): def __init__(self, **kwargs): pass def generate_message(self, context): if context.turn == 0: return f"Hello, I want to {context.goal}" return "[GOAL ACHIEVED]" def test_evaluate_with_simulator_creates_single_run(tmp_path): mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db") mlflow.set_experiment("test-experiment") simulator = ConversationSimulator( test_cases=[{"goal": "get help"}], max_turns=2, user_agent_class=MockUserAgent, ) def mock_predict_fn(input: list[dict[str, Any]], **kwargs): return {"content": "Response"} with mock.patch( "mlflow.genai.simulators.simulator.invoke_model_without_tracing", return_value='{"rationale": "Goal achieved!", "result": "yes"}', ): mlflow.genai.evaluate(data=simulator, predict_fn=mock_predict_fn, scorers=[]) runs = mlflow.search_runs() assert len(runs) == 1 def test_evaluate_with_simulator_within_parent_run(tmp_path): mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db") mlflow.set_experiment("test-experiment") simulator = ConversationSimulator( test_cases=[{"goal": "get help"}], max_turns=2, user_agent_class=MockUserAgent, ) def mock_predict_fn(input: list[dict[str, Any]], **kwargs): return {"content": "Response"} with mock.patch( "mlflow.genai.simulators.simulator.invoke_model_without_tracing", return_value='{"rationale": "Goal achieved!", "result": "yes"}', ): with mlflow.start_run(run_name="parent-run") as parent_run: parent_run_id = parent_run.info.run_id mlflow.genai.evaluate(data=simulator, predict_fn=mock_predict_fn, scorers=[]) assert mlflow.active_run().info.run_id == parent_run_id runs = mlflow.search_runs() assert len(runs) == 1 assert runs.iloc[0]["tags.mlflow.runName"] == "parent-run" # ===================== Rate Limiting & Pipelining Tests ===================== def test_predict_rate_limiter_is_wired_to_predict_fn(monkeypatch): monkeypatch.setenv("MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT", "100") with mock.patch.object( RPSRateLimiter, "acquire", autospec=True, side_effect=lambda self: None ) as mock_acquire: data = [{"inputs": {"q": f"Q{i}"}} for i in range(5)] mlflow.genai.evaluate(data=data, predict_fn=lambda q: "answer", scorers=[]) assert mock_acquire.call_count == 5 def test_scorer_rate_limiter_is_wired_to_scorers(monkeypatch): monkeypatch.setenv("MLFLOW_GENAI_EVAL_SCORER_RATE_LIMIT", "100") with mock.patch.object( RPSRateLimiter, "acquire", autospec=True, side_effect=lambda self: None ) as mock_acquire: data = [{"inputs": {"q": f"Q{i}"}, "outputs": "a"} for i in range(3)] mlflow.genai.evaluate(data=data, scorers=[always_pass, always_pass]) # 3 items x 2 scorers = 6 scorer acquire calls (no predict_fn → no predict acquires) assert mock_acquire.call_count == 6 def test_pipelining_scores_while_predicts_pending(monkeypatch): monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_WORKERS", "2") # Gate that blocks predicts after the first batch completes gate = threading.Event() predict_call_count = 0 predict_lock = threading.Lock() scoring_started_while_predicts_pending = threading.Event() def gated_predict_fn(q): nonlocal predict_call_count with predict_lock: predict_call_count += 1 call_num = predict_call_count # Let the first 2 predictions through immediately, block the rest if call_num > 2: gate.wait(timeout=5) return "answer" @scorer def signaling_scorer(outputs): # If we're scoring while predicts are still pending, signal success with predict_lock: current_predicts = predict_call_count if ( current_predicts > 0 and current_predicts < 6 ): # 6 total items, so some must still be pending scoring_started_while_predicts_pending.set() return True data = [{"inputs": {"q": f"Q{i}"}} for i in range(6)] try: # Run evaluate in a background thread so we can release the gate result_holder = [] def run_eval(): result = mlflow.genai.evaluate( data=data, predict_fn=gated_predict_fn, scorers=[signaling_scorer] ) result_holder.append(result) eval_thread = threading.Thread(name="test-evaluation-eval", target=run_eval) eval_thread.start() # Wait for scoring to signal it started while predicts are pending signaled = scoring_started_while_predicts_pending.wait(timeout=10) # Release all blocked predicts gate.set() eval_thread.join(timeout=30) assert signaled finally: gate.set() # Ensure we don't leave threads blocked def test_backpressure_limits_in_flight_items(monkeypatch): workers = 2 monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_WORKERS", str(workers)) monkeypatch.setenv("MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT", "0") monkeypatch.setenv("MLFLOW_GENAI_EVAL_SCORER_RATE_LIMIT", "0") # Skip pre-flight predict_fn call that runs outside the backpressure semaphore monkeypatch.setenv("MLFLOW_GENAI_EVAL_SKIP_TRACE_VALIDATION", "true") buffer = backpressure_buffer(workers) max_in_flight = 0 in_flight = 0 lock = threading.Lock() predict_done = threading.Semaphore(0) score_gate = threading.Event() def tracking_predict(q): nonlocal in_flight, max_in_flight with lock: in_flight += 1 max_in_flight = max(max_in_flight, in_flight) predict_done.release() return "answer" @scorer def blocking_scorer(outputs): score_gate.wait() nonlocal in_flight with lock: in_flight -= 1 return True num_items = 3 * buffer data = [{"inputs": {"q": f"Q{i}"}} for i in range(num_items)] eval_thread = threading.Thread( name="test-evaluation-blocking", target=lambda: mlflow.genai.evaluate( data=data, predict_fn=tracking_predict, scorers=[blocking_scorer] ), ) try: eval_thread.start() # Wait for exactly `buffer` predicts to complete. The semaphore guarantees # each acquire() returns once a predict runs. After this, the submit thread # is blocked on the backpressure semaphore and no more predicts can run # (scorers are blocked on score_gate). for _ in range(buffer): predict_done.acquire() with lock: observed_max = max_in_flight finally: score_gate.set() eval_thread.join() # The semaphore bounds in-flight items. Without backpressure all # num_items would pile up. assert observed_max == buffer def test_evaluate_logs_scorer_failure_summary(): @scorer def failing_scorer(inputs, outputs): raise ValueError("Model endpoint not found") @scorer def working_scorer(inputs, outputs): return True data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a platform", }, { "inputs": {"question": "What is Python?"}, "outputs": "Python is a language", }, ] with mock.patch("mlflow.genai.evaluation.harness._logger.warning") as mock_warning: result = mlflow.genai.evaluate( data=data, scorers=[failing_scorer, working_scorer], ) # Evaluation should complete without raising an exception assert "working_scorer/mean" in result.metrics # Verify warning was logged with failure summary warning_calls = [call.args[0] for call in mock_warning.call_args_list] failure_warnings = [ msg for msg in warning_calls if "Some scorer invocations failed during evaluation" in msg ] warning_message = failure_warnings[0] assert "'failing_scorer': 2/2 failed" in warning_message assert "Check individual trace assessments for detailed error messages" in warning_message def test_evaluate_no_warning_when_all_scorers_succeed(): @scorer def working_scorer_1(inputs, outputs): return True @scorer def working_scorer_2(inputs, outputs): return False data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a platform", }, ] with mock.patch("mlflow.genai.evaluation.harness._logger.warning") as mock_warning: result = mlflow.genai.evaluate( data=data, scorers=[working_scorer_1, working_scorer_2], ) assert "working_scorer_1/mean" in result.metrics assert "working_scorer_2/mean" in result.metrics warning_calls = [call.args[0] for call in mock_warning.call_args_list] assert not any( "Some scorer invocations failed during evaluation" in msg for msg in warning_calls ) def test_evaluate_logs_scorer_failure_summary_with_multi_turn_scorers(): @mlflow.trace def model(question, session_id): mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) return f"Answer to {question}" model("Q1", session_id="session_1") trace_1 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model("Q2", session_id="session_1") trace_2 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model("Q3", session_id="session_2") trace_3 = mlflow.get_trace(mlflow.get_last_active_trace_id()) dataset = create_dataset(name="multi_turn_failure_test") dataset.merge_records([trace_1, trace_2, trace_3]) with mock.patch("mlflow.genai.evaluation.harness._logger.warning") as mock_warning: result = mlflow.genai.evaluate( data=dataset, scorers=[FailingSessionScorer(), WorkingSessionScorer()], ) assert "working_session_scorer/mean" in result.metrics warning_calls = [call.args[0] for call in mock_warning.call_args_list] failure_warnings = [ msg for msg in warning_calls if "Some scorer invocations failed during evaluation" in msg ] assert len(failure_warnings) > 0 warning_message = failure_warnings[0] assert "'failing_session_scorer': 2/2 failed" in warning_message assert "Check individual trace assessments for detailed error messages" in warning_message def test_evaluate_logs_scorer_failure_summary_with_mixed_scorers(): @scorer def failing_single_turn(inputs, outputs): raise ValueError("Single turn scorer error") @mlflow.trace def model(question, session_id): mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id}) return f"Answer to {question}" model("Q1", session_id="session_1") trace_1 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model("Q2", session_id="session_1") trace_2 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model("Q3", session_id="session_2") trace_3 = mlflow.get_trace(mlflow.get_last_active_trace_id()) model("Q4", session_id="session_2") trace_4 = mlflow.get_trace(mlflow.get_last_active_trace_id()) # Create dataset from traces dataset = create_dataset(name="mixed_scorer_failure_test") dataset.merge_records([trace_1, trace_2, trace_3, trace_4]) with mock.patch("mlflow.genai.evaluation.harness._logger.warning") as mock_warning: result = mlflow.genai.evaluate( data=dataset, scorers=[ failing_single_turn, always_pass, FailingSessionScorer(), WorkingSessionScorer(), ], ) assert "always_pass/mean" in result.metrics assert "working_session_scorer/mean" in result.metrics warning_calls = [call.args[0] for call in mock_warning.call_args_list] failure_warnings = [ msg for msg in warning_calls if "Some scorer invocations failed during evaluation" in msg ] assert len(failure_warnings) > 0 warning_message = failure_warnings[0] assert "'failing_single_turn': 4/4 failed" in warning_message assert "'failing_session_scorer': 2/2 failed" in warning_message assert "Check individual trace assessments for detailed error messages" in warning_message def test_no_rate_limit_backward_compat(monkeypatch): monkeypatch.setenv("MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT", "0") monkeypatch.setenv("MLFLOW_GENAI_EVAL_SCORER_RATE_LIMIT", "0") data = [ { "inputs": {"question": "What is MLflow?"}, "outputs": "MLflow is a tool for ML", "expectations": {"expected_response": "MLflow is a tool for ML", "max_length": 100}, }, ] result = mlflow.genai.evaluate(data=data, scorers=[exact_match, is_concise]) assert result.metrics["exact_match/mean"] == 1.0 assert result.metrics["is_concise/mean"] == 1.0 # ===================== Retry & Adaptive Rate Limiting Tests ===================== def test_predict_retries_on_429(monkeypatch): monkeypatch.setenv("MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT", "0") monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_RETRIES", "3") attempts = [] def flaky_predict(q): attempts.append(1) # First call is the pre-flight validation check — let it pass. # Fail on calls 2 and 3, succeed on call 4. if 1 < len(attempts) < 4: raise Exception("429 Too Many Requests") return "answer" data = [{"inputs": {"q": "Q1"}}] result = mlflow.genai.evaluate(data=data, predict_fn=flaky_predict, scorers=[always_pass]) # 1 validation + 1 fail + 1 fail + 1 success = 4 total assert len(attempts) == 4 assert result.metrics["always_pass/mean"] == 1.0 def test_scorer_retries_on_429(monkeypatch): monkeypatch.setenv("MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT", "0") monkeypatch.setenv("MLFLOW_GENAI_EVAL_SCORER_RATE_LIMIT", "0") monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_RETRIES", "3") attempts = [] @scorer def flaky_scorer(outputs): attempts.append(1) if len(attempts) < 3: raise Exception("429 Too Many Requests") return True data = [{"inputs": {"q": "Q1"}, "outputs": "a"}] result = mlflow.genai.evaluate(data=data, scorers=[flaky_scorer]) assert len(attempts) == 3 assert result.metrics["flaky_scorer/mean"] == 1.0 def test_adaptive_rate_reduces_on_429(monkeypatch): monkeypatch.setenv("MLFLOW_GENAI_EVAL_PREDICT_RATE_LIMIT", "auto") monkeypatch.setenv("MLFLOW_GENAI_EVAL_MAX_RETRIES", "3") rate_after_throttle = [] original_report_throttle = RPSRateLimiter.report_throttle def spy_report_throttle(self): original_report_throttle(self) rate_after_throttle.append(self._rps) attempts = [] def flaky_predict(q): attempts.append(1) # First call is the pre-flight validation check — let it pass. # Fail on call 2 to trigger a throttle event. if len(attempts) == 2: raise Exception("429 Too Many Requests") return "answer" data = [{"inputs": {"q": f"Q{i}"}} for i in range(3)] with mock.patch.object(RPSRateLimiter, "report_throttle", spy_report_throttle): result = mlflow.genai.evaluate(data=data, predict_fn=flaky_predict, scorers=[always_pass]) assert result.metrics["always_pass/mean"] == 1.0 # At least one throttle event observed, and the rate was reduced assert len(rate_after_throttle) >= 1 assert rate_after_throttle[0] < AUTO_INITIAL_RPS @pytest.mark.parametrize( ("trace_or_none", "run_id"), [ (None, "run-1"), ( Trace( info=create_test_trace_info_with_uc_table( trace_id="tr-uc", catalog_name="catalog", schema_name="schema" ), data=TraceData(spans=[]), ), "run-1", ), ], ids=["none_trace", "uc_schema_trace"], ) def test_should_clone_trace_returns_false_early(trace_or_none, run_id): assert _should_clone_trace(trace_or_none, run_id=run_id) is False @pytest.mark.parametrize( ("experiment_id", "expected"), [ ("exp-999", True), ("exp-123", False), ], ids=["different_experiment", "matching_experiment"], ) def test_should_clone_trace_with_explicit_experiment_id( mlflow_experiment_trace, experiment_id, expected ): with mock.patch( "mlflow.genai.evaluation.harness._does_store_support_trace_linking", return_value=True, ) as mock_store: result = _should_clone_trace( mlflow_experiment_trace, run_id="run-1", experiment_id=experiment_id ) assert result is expected if expected is False: mock_store.assert_called_once() def test_should_clone_trace_falls_back_to_get_experiment_id_when_none(mlflow_experiment_trace): with ( mock.patch( "mlflow.tracking.fluent._get_experiment_id", return_value="exp-999", ) as mock_get_exp, mock.patch( "mlflow.genai.evaluation.harness._does_store_support_trace_linking", return_value=True, ), ): result = _should_clone_trace(mlflow_experiment_trace, run_id="run-1", experiment_id=None) mock_get_exp.assert_called_once() assert result is True def test_should_clone_trace_does_not_call_get_experiment_id_when_provided(mlflow_experiment_trace): with ( mock.patch("mlflow.tracking.fluent._get_experiment_id") as mock_get_exp, mock.patch( "mlflow.genai.evaluation.harness._does_store_support_trace_linking", return_value=True, ), ): _should_clone_trace(mlflow_experiment_trace, run_id="run-1", experiment_id="exp-123") mock_get_exp.assert_not_called() def _make_eval_item(trace=None, expectations=None): return EvalItem( request_id="req-1", inputs={"question": "q"}, outputs="a", expectations=expectations or {}, trace=trace, ) def test_get_new_expectations_returns_empty_when_trace_is_none(): eval_item = _make_eval_item(trace=None, expectations={"expected": "value"}) assert _get_new_expectations(eval_item) == [] def test_get_new_expectations_filters_existing(mlflow_experiment_trace): eval_item = _make_eval_item(trace=mlflow_experiment_trace, expectations={"correctness": "yes"}) with mock.patch( "mlflow.genai.evaluation.entities.get_context", return_value=mock.Mock(**{"get_user_name.return_value": "tester"}), ): result = _get_new_expectations(eval_item) assert [e.name for e in result] == ["correctness"] def test_run_predict_clone_read_miss_records_error_and_nulls_trace(mlflow_experiment_trace): eval_item = _make_eval_item(trace=mlflow_experiment_trace) with ( mock.patch( "mlflow.genai.evaluation.harness._should_clone_trace", return_value=True, ) as mock_should_clone, mock.patch( "mlflow.genai.evaluation.harness.copy_trace_to_experiment", return_value="tr-cloned", ) as mock_copy, mock.patch("mlflow.get_trace", return_value=None) as mock_get_trace, ): _run_predict( eval_item, predict_fn=None, run_id=None, rate_limiter=NoOpRateLimiter(), experiment_id="exp-999", ) mock_should_clone.assert_called_once() mock_copy.assert_called_once() mock_get_trace.assert_called_once_with("tr-cloned", flush=True) assert eval_item.trace is None assert "could not be read back" in eval_item.error_message def test_run_predict_clone_read_hit_sets_trace_without_error(mlflow_experiment_trace): eval_item = _make_eval_item(trace=mlflow_experiment_trace) cloned = Trace( info=create_test_trace_info(trace_id="tr-cloned", experiment_id="exp-999"), data=TraceData(spans=[]), ) with ( mock.patch( "mlflow.genai.evaluation.harness._should_clone_trace", return_value=True, ) as mock_should_clone, mock.patch( "mlflow.genai.evaluation.harness.copy_trace_to_experiment", return_value="tr-cloned", ) as mock_copy, mock.patch("mlflow.get_trace", return_value=cloned) as mock_get_trace, ): _run_predict( eval_item, predict_fn=None, run_id=None, rate_limiter=NoOpRateLimiter(), experiment_id="exp-999", ) mock_should_clone.assert_called_once() mock_copy.assert_called_once() mock_get_trace.assert_called_once_with("tr-cloned", flush=True) assert eval_item.trace is cloned assert eval_item.error_message is None def test_evaluate_completes_when_cloned_trace_read_back_misses(): # Reproduce #24355: a real cross-experiment clone (searched traces live in a different # experiment than the eval run), where the cloned trace's async read-back misses and # returns None. evaluate() must complete without the AttributeError instead of crashing. exp_id = mlflow.set_experiment("traces exp").experiment_id with mlflow.start_span(name="qa") as span: span.set_inputs({"question": "What is MLflow?"}) span.set_outputs("MLflow is a tool for ML") mlflow.set_experiment("diff exp") trace_df = mlflow.search_traces(locations=[exp_id]) real_get_trace = mlflow.get_trace first_call = {"seen": False} def fake_get_trace(trace_id, *args, **kwargs): # The first get_trace in the pipeline is the clone read-back in _run_predict; null it # to simulate the read-after-write miss. Later calls pass through unchanged. if not first_call["seen"]: first_call["seen"] = True return None return real_get_trace(trace_id, *args, **kwargs) with mock.patch("mlflow.get_trace", side_effect=fake_get_trace) as mock_get_trace: result = mlflow.genai.evaluate(data=trace_df, scorers=[has_trace]) mock_get_trace.assert_any_call(mock.ANY, flush=True) assert result.metrics is not None def _make_score_submitter(session_groups): return _ScoreSubmitter( eval_items=[item for items in session_groups.values() for item in items], single_turn_scorers=[], multi_turn_scorers=[mock.Mock()], session_groups=session_groups, run_id=None, max_retries=0, rps=None, adaptive=False, max_rps_multiplier=1.0, pool_workers=1, ) def test_run_multi_turn_skips_session_with_only_none_traces(): submitter = _make_score_submitter({"session-1": [_make_eval_item(trace=None)]}) multi_turn_eval_results: dict[str, EvalResult] = {} with mock.patch( "mlflow.genai.evaluation.harness.evaluate_session_level_scorers", ) as mock_eval_session: submitter.run_multi_turn(multi_turn_eval_results, progress_bar=None) mock_eval_session.assert_not_called() assert multi_turn_eval_results == {} def test_run_multi_turn_filters_none_trace_items_from_session(mlflow_experiment_trace): valid_item = _make_eval_item(trace=mlflow_experiment_trace) none_item = _make_eval_item(trace=None) submitter = _make_score_submitter({"session-1": [none_item, valid_item]}) multi_turn_eval_results: dict[str, EvalResult] = {} with mock.patch( "mlflow.genai.evaluation.harness.evaluate_session_level_scorers", return_value=EvalResult(eval_item=valid_item), ) as mock_eval_session: submitter.run_multi_turn(multi_turn_eval_results, progress_bar=None) mock_eval_session.assert_called_once() assert mock_eval_session.call_args.kwargs["session_items"] == [valid_item] assert multi_turn_eval_results == {"tr-123": mock_eval_session.return_value}