wehub-resource-sync
2526 lines
78 KiB
Python
2526 lines
78 KiB
Python
import json
|
|
import time
|
|
from unittest import mock
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
import sklearn.neighbors as knn
|
|
from click.testing import CliRunner
|
|
from fastapi import Request
|
|
|
|
import mlflow
|
|
from mlflow import MlflowClient
|
|
from mlflow.entities import (
|
|
EvaluationDataset,
|
|
Expectation,
|
|
Feedback,
|
|
GatewayEndpointModelConfig,
|
|
IssueSeverity,
|
|
IssueStatus,
|
|
Metric,
|
|
Param,
|
|
RunTag,
|
|
)
|
|
from mlflow.entities.assessment_source import AssessmentSource, AssessmentSourceType
|
|
from mlflow.entities.gateway_budget_policy import (
|
|
BudgetAction,
|
|
BudgetDuration,
|
|
BudgetDurationUnit,
|
|
BudgetTargetScope,
|
|
BudgetUnit,
|
|
)
|
|
from mlflow.entities.gateway_endpoint import GatewayModelLinkageType
|
|
from mlflow.entities.gateway_guardrail import GuardrailAction, GuardrailStage
|
|
from mlflow.entities.trace import Trace
|
|
from mlflow.entities.webhook import WebhookAction, WebhookEntity, WebhookEvent
|
|
from mlflow.gateway.cli import start
|
|
from mlflow.gateway.constants import MLFLOW_GATEWAY_CALLER_HEADER
|
|
from mlflow.gateway.schemas import chat
|
|
from mlflow.genai.datasets import create_dataset
|
|
from mlflow.genai.discovery.entities import _TriageResult
|
|
from mlflow.genai.discovery.pipeline import discover_issues
|
|
from mlflow.genai.judges import make_judge
|
|
from mlflow.genai.judges.base import AlignmentOptimizer
|
|
from mlflow.genai.scorers import scorer
|
|
from mlflow.genai.scorers.base import Scorer
|
|
from mlflow.genai.scorers.builtin_scorers import (
|
|
Completeness,
|
|
Guidelines,
|
|
RelevanceToQuery,
|
|
Safety,
|
|
UserFrustration,
|
|
)
|
|
from mlflow.genai.simulators import ConversationSimulator
|
|
from mlflow.pyfunc.model import (
|
|
ResponsesAgent,
|
|
ResponsesAgentRequest,
|
|
ResponsesAgentResponse,
|
|
)
|
|
from mlflow.server.gateway_api import chat_completions, invocations
|
|
from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig, GatewayModelConfig
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
from mlflow.telemetry.client import TelemetryClient
|
|
from mlflow.telemetry.events import (
|
|
AiCommandRunEvent,
|
|
AlignJudgeEvent,
|
|
AutologgingEvent,
|
|
CreateDatasetEvent,
|
|
CreateExperimentEvent,
|
|
CreateLoggedModelEvent,
|
|
CreateModelVersionEvent,
|
|
CreatePromptEvent,
|
|
CreateRegisteredModelEvent,
|
|
CreateRunEvent,
|
|
CreateWebhookEvent,
|
|
DiscoverIssuesEvent,
|
|
EvaluateEvent,
|
|
GatewayCreateBudgetPolicyEvent,
|
|
GatewayCreateEndpointEvent,
|
|
GatewayCreateGuardrailEvent,
|
|
GatewayCreateModelDefinitionEvent,
|
|
GatewayCreateSecretEvent,
|
|
GatewayDeleteBudgetPolicyEvent,
|
|
GatewayDeleteEndpointEvent,
|
|
GatewayDeleteGuardrailEvent,
|
|
GatewayDeleteSecretEvent,
|
|
GatewayGetEndpointEvent,
|
|
GatewayInvocationEvent,
|
|
GatewayListBudgetPoliciesEvent,
|
|
GatewayListEndpointsEvent,
|
|
GatewayListSecretsEvent,
|
|
GatewayStartEvent,
|
|
GatewayUpdateBudgetPolicyEvent,
|
|
GatewayUpdateEndpointEvent,
|
|
GatewayUpdateGuardrailEvent,
|
|
GatewayUpdateSecretEvent,
|
|
GenAIEvaluateEvent,
|
|
GetLoggedModelEvent,
|
|
GitModelVersioningEvent,
|
|
InvokeCustomJudgeModelEvent,
|
|
LoadPromptEvent,
|
|
LogAssessmentEvent,
|
|
LogBatchEvent,
|
|
LogDatasetEvent,
|
|
LogMetricEvent,
|
|
LogParamEvent,
|
|
MakeJudgeEvent,
|
|
McpRunEvent,
|
|
MergeRecordsEvent,
|
|
PromptOptimizationEvent,
|
|
ScorerCallEvent,
|
|
SimulateConversationEvent,
|
|
StartTraceEvent,
|
|
TracingContextPropagation,
|
|
TrackingServerStartEvent,
|
|
UpdateIssueEvent,
|
|
)
|
|
from mlflow.tracing.distributed import (
|
|
get_tracing_context_headers_for_http_request,
|
|
set_tracing_context_from_http_request_headers,
|
|
)
|
|
from mlflow.tracking.fluent import (
|
|
_create_dataset_input,
|
|
_create_logged_model,
|
|
_initialize_logged_model,
|
|
)
|
|
from mlflow.utils.os import is_windows
|
|
|
|
from tests.telemetry.helper_functions import validate_telemetry_record
|
|
|
|
|
|
class TestModel(mlflow.pyfunc.PythonModel):
|
|
def predict(self, model_input: list[str]) -> str:
|
|
return "test"
|
|
|
|
|
|
@pytest.fixture
|
|
def mlflow_client():
|
|
return MlflowClient()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_get_telemetry_client(mock_telemetry_client: TelemetryClient):
|
|
with mock.patch(
|
|
"mlflow.telemetry.track.get_telemetry_client",
|
|
return_value=mock_telemetry_client,
|
|
):
|
|
yield
|
|
|
|
|
|
def test_create_logged_model(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
event_name = CreateLoggedModelEvent.name
|
|
mlflow.create_external_model(name="model")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, {"flavor": "external"}
|
|
)
|
|
|
|
mlflow.initialize_logged_model(name="model", tags={"key": "value"})
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, {"flavor": "initialize"}
|
|
)
|
|
|
|
_initialize_logged_model(name="model", flavor="keras")
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, event_name, {"flavor": "keras"})
|
|
|
|
mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=TestModel(),
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"flavor": "pyfunc.CustomPythonModel"},
|
|
)
|
|
|
|
mlflow.sklearn.log_model(
|
|
knn.KNeighborsClassifier(),
|
|
name="model",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"flavor": "sklearn", "serialization_format": "skops"},
|
|
)
|
|
|
|
mlflow.sklearn.log_model(
|
|
knn.KNeighborsClassifier(),
|
|
name="model",
|
|
serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE,
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"flavor": "sklearn", "serialization_format": "pickle"},
|
|
)
|
|
|
|
class SimpleResponsesAgent(ResponsesAgent):
|
|
def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
|
|
mock_response = {
|
|
"output": [
|
|
{
|
|
"type": "message",
|
|
"id": "1234",
|
|
"status": "completed",
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": request.input[0].content,
|
|
}
|
|
],
|
|
}
|
|
],
|
|
}
|
|
return ResponsesAgentResponse(**mock_response)
|
|
|
|
mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=SimpleResponsesAgent(),
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"flavor": "pyfunc.ResponsesAgent"},
|
|
)
|
|
|
|
_create_logged_model(name="model", flavor="pyfunc", uses_uv=True)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"flavor": "pyfunc", "uses_uv": True},
|
|
)
|
|
|
|
_create_logged_model(name="model", flavor="pyfunc", uses_uv=False)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"flavor": "pyfunc"},
|
|
)
|
|
|
|
|
|
def test_create_experiment(mock_requests, mlflow_client, mock_telemetry_client: TelemetryClient):
|
|
event_name = CreateExperimentEvent.name
|
|
exp_id = mlflow.create_experiment(name="test_experiment")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, {"experiment_id": exp_id}
|
|
)
|
|
|
|
exp_id = mlflow_client.create_experiment(name="test_experiment1")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, {"experiment_id": exp_id}
|
|
)
|
|
|
|
|
|
def test_create_run(mock_requests, mlflow_client, mock_telemetry_client: TelemetryClient):
|
|
event_name = CreateRunEvent.name
|
|
exp_id = mlflow.create_experiment(name="test_experiment")
|
|
with mlflow.start_run(experiment_id=exp_id):
|
|
record = validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, check_params=False
|
|
)
|
|
assert json.loads(record["params"])["experiment_id"] == exp_id
|
|
|
|
mlflow_client.create_run(experiment_id=exp_id)
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, event_name, check_params=False)
|
|
|
|
exp_id = mlflow.create_experiment(name="test_experiment2")
|
|
mlflow.set_experiment(experiment_id=exp_id)
|
|
with mlflow.start_run():
|
|
record = validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, check_params=False
|
|
)
|
|
params = json.loads(record["params"])
|
|
assert params["mlflow_experiment_id"] == exp_id
|
|
|
|
|
|
def test_create_run_with_imports(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
event_name = CreateRunEvent.name
|
|
import pyspark.ml # noqa: F401
|
|
|
|
with mlflow.start_run():
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, check_params=False
|
|
)
|
|
assert "pyspark.ml" in json.loads(data["params"])["imports"]
|
|
|
|
|
|
def test_create_registered_model(
|
|
mock_requests, mlflow_client, mock_telemetry_client: TelemetryClient
|
|
):
|
|
event_name = CreateRegisteredModelEvent.name
|
|
mlflow_client.create_registered_model(name="test_model1")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"is_prompt": False},
|
|
)
|
|
|
|
mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=TestModel(),
|
|
registered_model_name="test_model",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"is_prompt": False},
|
|
)
|
|
|
|
|
|
def test_create_model_version(mock_requests, mlflow_client, mock_telemetry_client: TelemetryClient):
|
|
event_name = CreateModelVersionEvent.name
|
|
mlflow_client.create_registered_model(name="test_model")
|
|
mlflow_client.create_model_version(
|
|
name="test_model", source="test_source", run_id="test_run_id"
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"is_prompt": False},
|
|
)
|
|
|
|
mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=TestModel(),
|
|
registered_model_name="test_model",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"is_prompt": False},
|
|
)
|
|
|
|
mlflow.genai.register_prompt(
|
|
name="ai_assistant_prompt",
|
|
template="Respond to the user's message as a {{style}} AI. {{greeting}}",
|
|
commit_message="Initial version of AI assistant",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
event_name,
|
|
{"is_prompt": True},
|
|
)
|
|
|
|
|
|
def test_start_trace(mock_requests, mlflow_client, mock_telemetry_client: TelemetryClient):
|
|
event_name = StartTraceEvent.name
|
|
with mlflow.start_span(name="test_span"):
|
|
pass
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, event_name, check_params=False)
|
|
|
|
@mlflow.trace
|
|
def test_func():
|
|
pass
|
|
|
|
test_func()
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, event_name, check_params=False)
|
|
|
|
trace_id = mlflow_client.start_trace(name="test_trace").trace_id
|
|
mlflow_client.end_trace(trace_id=trace_id)
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, event_name, check_params=False)
|
|
|
|
import openai # noqa: F401
|
|
|
|
test_func()
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, check_params=False
|
|
)
|
|
params = json.loads(data["params"])
|
|
assert "openai" in params["imports"]
|
|
assert params["format"] == "native"
|
|
|
|
|
|
def test_start_trace_genai_semconv(
|
|
mock_requests, monkeypatch, mock_telemetry_client: TelemetryClient
|
|
):
|
|
monkeypatch.setenv("MLFLOW_ENABLE_OTEL_GENAI_SEMCONV", "true")
|
|
event_name = StartTraceEvent.name
|
|
|
|
@mlflow.trace
|
|
def test_func():
|
|
pass
|
|
|
|
test_func()
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, event_name, check_params=False
|
|
)
|
|
assert json.loads(data["params"])["format"] == "genai_semconv"
|
|
|
|
|
|
def test_create_prompt(mock_requests, mlflow_client, mock_telemetry_client: TelemetryClient):
|
|
mlflow_client.create_prompt(name="test_prompt")
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, CreatePromptEvent.name)
|
|
|
|
# OSS prompt registry uses create_registered_model with a special tag
|
|
mlflow.genai.register_prompt(
|
|
name="greeting_prompt",
|
|
template="Respond to the user's message as a {{style}} AI. {{greeting}}",
|
|
)
|
|
expected_params = {"is_prompt": True}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
CreateRegisteredModelEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
|
|
def test_log_assessment(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
with mlflow.start_span(name="test_span") as span:
|
|
feedback = Feedback(
|
|
name="faithfulness",
|
|
value=0.9,
|
|
rationale="The model is faithful to the input.",
|
|
metadata={"model": "gpt-4o-mini"},
|
|
)
|
|
|
|
mlflow.log_assessment(trace_id=span.trace_id, assessment=feedback)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogAssessmentEvent.name,
|
|
{"type": "feedback", "source_type": "CODE"},
|
|
)
|
|
mlflow.log_feedback(trace_id=span.trace_id, value=0.9, name="faithfulness")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogAssessmentEvent.name,
|
|
{"type": "feedback", "source_type": "CODE"},
|
|
)
|
|
|
|
with mlflow.start_span(name="test_span2") as span:
|
|
expectation = Expectation(
|
|
name="expected_answer",
|
|
value="MLflow",
|
|
)
|
|
|
|
mlflow.log_assessment(trace_id=span.trace_id, assessment=expectation)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogAssessmentEvent.name,
|
|
{"type": "expectation", "source_type": "HUMAN"},
|
|
)
|
|
mlflow.log_expectation(trace_id=span.trace_id, value="MLflow", name="expected_answer")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogAssessmentEvent.name,
|
|
{"type": "expectation", "source_type": "HUMAN"},
|
|
)
|
|
|
|
|
|
def test_evaluate(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
mlflow.models.evaluate(
|
|
data=pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}),
|
|
model=lambda x: x["x"] * 2,
|
|
extra_metrics=[mlflow.metrics.latency()],
|
|
)
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, EvaluateEvent.name)
|
|
|
|
|
|
def test_create_webhook(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
client = MlflowClient()
|
|
client.create_webhook(
|
|
name="test_webhook",
|
|
url="https://example.com/webhook",
|
|
events=[WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED)],
|
|
)
|
|
expected_params = {"events": ["model_version.created"]}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, CreateWebhookEvent.name, expected_params
|
|
)
|
|
|
|
|
|
def test_genai_evaluate(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
@mlflow.genai.scorer
|
|
def decorator_scorer():
|
|
return 1.0
|
|
|
|
instructions_judge = make_judge(
|
|
name="quality_judge",
|
|
instructions="Evaluate if {{ outputs }} is high quality",
|
|
model="openai:/gpt-4",
|
|
)
|
|
|
|
session_level_instruction_judge = make_judge(
|
|
name="conversation_quality",
|
|
instructions="Evaluate if the {{ conversation }} is engaging and coherent",
|
|
model="openai:/gpt-4",
|
|
)
|
|
|
|
guidelines_scorer = Guidelines(
|
|
name="politeness",
|
|
guidelines=["Be polite", "Be respectful"],
|
|
)
|
|
|
|
builtin_scorer = RelevanceToQuery(name="relevance_check")
|
|
|
|
session_level_builtin_scorer = UserFrustration(name="frustration_check")
|
|
|
|
data = [
|
|
{
|
|
"inputs": {"model_input": ["What is MLflow?"]},
|
|
"outputs": "MLflow is an open source platform.",
|
|
}
|
|
]
|
|
|
|
model = TestModel()
|
|
|
|
with (
|
|
mock.patch("mlflow.genai.judges.utils.invocation_utils.invoke_judge_model"),
|
|
mock.patch("mlflow.genai.judges.builtin.invoke_judge_model"),
|
|
mock.patch("mlflow.genai.judges.instructions_judge.invoke_judge_model"),
|
|
):
|
|
# Test with all scorer kinds and scopes, without predict_fn
|
|
mlflow.genai.evaluate(
|
|
data=data,
|
|
scorers=[
|
|
decorator_scorer,
|
|
instructions_judge,
|
|
session_level_instruction_judge,
|
|
guidelines_scorer,
|
|
builtin_scorer,
|
|
session_level_builtin_scorer,
|
|
],
|
|
)
|
|
|
|
expected_params = {
|
|
"predict_fn_provided": False,
|
|
"scorer_info": [
|
|
{
|
|
"class": "UserDefinedScorer",
|
|
"kind": "decorator",
|
|
"scope": "trace",
|
|
},
|
|
{
|
|
"class": "UserDefinedScorer",
|
|
"kind": "instructions",
|
|
"scope": "trace",
|
|
},
|
|
{
|
|
"class": "UserDefinedScorer",
|
|
"kind": "instructions",
|
|
"scope": "session",
|
|
},
|
|
{"class": "Guidelines", "kind": "guidelines", "scope": "trace"},
|
|
{"class": "RelevanceToQuery", "kind": "builtin", "scope": "trace"},
|
|
{"class": "UserFrustration", "kind": "builtin", "scope": "session"},
|
|
],
|
|
"eval_data_type": "list[dict]",
|
|
"eval_data_size": 1,
|
|
"eval_data_provided_fields": ["inputs", "outputs"],
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GenAIEvaluateEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
# Test with predict_fn
|
|
mlflow.genai.evaluate(
|
|
data=data,
|
|
scorers=[builtin_scorer, guidelines_scorer],
|
|
predict_fn=model.predict,
|
|
)
|
|
expected_params = {
|
|
"predict_fn_provided": True,
|
|
"scorer_info": [
|
|
{"class": "RelevanceToQuery", "kind": "builtin", "scope": "trace"},
|
|
{"class": "Guidelines", "kind": "guidelines", "scope": "trace"},
|
|
],
|
|
"eval_data_type": "list[dict]",
|
|
"eval_data_size": 1,
|
|
"eval_data_provided_fields": ["inputs", "outputs"],
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GenAIEvaluateEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
|
|
def test_genai_evaluate_telemetry_data_fields(
|
|
mock_requests, mock_telemetry_client: TelemetryClient
|
|
):
|
|
@mlflow.genai.scorer
|
|
def sample_scorer():
|
|
return 1.0
|
|
|
|
with mock.patch("mlflow.genai.judges.utils.invocation_utils.invoke_judge_model"):
|
|
# Test with list of dicts
|
|
data_list = [
|
|
{
|
|
"inputs": {"question": "Q1"},
|
|
"outputs": "A1",
|
|
"expectations": {"answer": "Expected1"},
|
|
},
|
|
{
|
|
"inputs": {"question": "Q2"},
|
|
"outputs": "A2",
|
|
"expectations": {"answer": "Expected2"},
|
|
},
|
|
]
|
|
mlflow.genai.evaluate(data=data_list, scorers=[sample_scorer])
|
|
expected_params = {
|
|
"predict_fn_provided": False,
|
|
"scorer_info": [
|
|
{
|
|
"class": "UserDefinedScorer",
|
|
"kind": "decorator",
|
|
"scope": "trace",
|
|
},
|
|
],
|
|
"eval_data_type": "list[dict]",
|
|
"eval_data_size": 2,
|
|
"eval_data_provided_fields": ["expectations", "inputs", "outputs"],
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GenAIEvaluateEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
# Test with pandas DataFrame
|
|
df_data = pd.DataFrame([
|
|
{"inputs": {"question": "Q1"}, "outputs": "A1"},
|
|
{"inputs": {"question": "Q2"}, "outputs": "A2"},
|
|
{"inputs": {"question": "Q3"}, "outputs": "A3"},
|
|
])
|
|
mlflow.genai.evaluate(data=df_data, scorers=[sample_scorer])
|
|
expected_params = {
|
|
"predict_fn_provided": False,
|
|
"scorer_info": [
|
|
{
|
|
"class": "UserDefinedScorer",
|
|
"kind": "decorator",
|
|
"scope": "trace",
|
|
},
|
|
],
|
|
"eval_data_type": "pd.DataFrame",
|
|
"eval_data_size": 3,
|
|
"eval_data_provided_fields": ["inputs", "outputs"],
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GenAIEvaluateEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
# Test with list of Traces
|
|
trace_ids = []
|
|
for i in range(2):
|
|
with mlflow.start_span(name=f"test_span_{i}") as span:
|
|
span.set_inputs({"question": f"Q{i}"})
|
|
span.set_outputs({"answer": f"A{i}"})
|
|
trace_ids.append(span.trace_id)
|
|
|
|
traces = [mlflow.get_trace(trace_id) for trace_id in trace_ids]
|
|
mlflow.genai.evaluate(data=traces, scorers=[sample_scorer])
|
|
expected_params = {
|
|
"predict_fn_provided": False,
|
|
"scorer_info": [
|
|
{
|
|
"class": "UserDefinedScorer",
|
|
"kind": "decorator",
|
|
"scope": "trace",
|
|
},
|
|
],
|
|
"eval_data_type": "list[Trace]",
|
|
"eval_data_size": 2,
|
|
"eval_data_provided_fields": ["inputs", "outputs", "trace"],
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GenAIEvaluateEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
# Test with EvaluationDataset
|
|
from mlflow.genai.datasets import create_dataset
|
|
|
|
dataset = create_dataset("test_dataset")
|
|
dataset_data = [
|
|
{
|
|
"inputs": {"question": "Q1"},
|
|
"outputs": "A1",
|
|
"expectations": {"answer": "Expected1"},
|
|
},
|
|
{
|
|
"inputs": {"question": "Q2"},
|
|
"outputs": "A2",
|
|
"expectations": {"answer": "Expected2"},
|
|
},
|
|
]
|
|
dataset.merge_records(dataset_data)
|
|
mlflow.genai.evaluate(data=dataset, scorers=[sample_scorer])
|
|
expected_params = {
|
|
"predict_fn_provided": False,
|
|
"scorer_info": [
|
|
{
|
|
"class": "UserDefinedScorer",
|
|
"kind": "decorator",
|
|
"scope": "trace",
|
|
},
|
|
],
|
|
"eval_data_type": "EvaluationDataset",
|
|
"eval_data_size": 2,
|
|
"eval_data_provided_fields": ["expectations", "inputs", "outputs"],
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GenAIEvaluateEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
|
|
def test_simulate_conversation(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
simulator = ConversationSimulator(
|
|
test_cases=[
|
|
{"goal": "Learn about MLflow"},
|
|
{"goal": "Debug an issue"},
|
|
],
|
|
max_turns=2,
|
|
)
|
|
|
|
def mock_predict_fn(input, **kwargs):
|
|
return {"role": "assistant", "content": "Mock response"}
|
|
|
|
mock_trace = mock.Mock()
|
|
with (
|
|
mock.patch(
|
|
"mlflow.genai.simulators.simulator.invoke_model_without_tracing",
|
|
return_value="Mock user message",
|
|
),
|
|
mock.patch(
|
|
"mlflow.genai.simulators.simulator.ConversationSimulator._check_goal_achieved",
|
|
return_value=False,
|
|
),
|
|
mock.patch(
|
|
"mlflow.genai.simulators.simulator.mlflow.get_trace",
|
|
return_value=mock_trace,
|
|
),
|
|
):
|
|
result = simulator.simulate(predict_fn=mock_predict_fn)
|
|
|
|
assert len(result) == 2
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
SimulateConversationEvent.name,
|
|
{
|
|
"callsite": "conversation_simulator",
|
|
"simulated_conversation_info": [
|
|
{"turn_count": len(result[0])},
|
|
{"turn_count": len(result[1])},
|
|
],
|
|
},
|
|
)
|
|
|
|
|
|
def test_simulate_conversation_from_genai_evaluate(
|
|
mock_requests, mock_telemetry_client: TelemetryClient
|
|
):
|
|
simulator = ConversationSimulator(
|
|
test_cases=[
|
|
{"goal": "Learn about MLflow"},
|
|
],
|
|
max_turns=1,
|
|
)
|
|
|
|
def mock_predict_fn(input, **kwargs):
|
|
return {"role": "assistant", "content": "Mock response"}
|
|
|
|
@scorer
|
|
def simple_scorer(outputs) -> bool:
|
|
return len(outputs) > 0
|
|
|
|
with (
|
|
mock.patch(
|
|
"mlflow.genai.simulators.simulator.invoke_model_without_tracing",
|
|
return_value="Mock user message",
|
|
),
|
|
mock.patch(
|
|
"mlflow.genai.simulators.simulator.ConversationSimulator._check_goal_achieved",
|
|
return_value=True,
|
|
),
|
|
):
|
|
mlflow.genai.evaluate(data=simulator, predict_fn=mock_predict_fn, scorers=[simple_scorer])
|
|
|
|
mock_telemetry_client.flush()
|
|
|
|
simulate_events = [
|
|
record
|
|
for record in mock_requests
|
|
if record["data"]["event_name"] == SimulateConversationEvent.name
|
|
]
|
|
assert len(simulate_events) == 1
|
|
|
|
event_params = json.loads(simulate_events[0]["data"]["params"])
|
|
assert event_params == {
|
|
"callsite": "genai_evaluate",
|
|
"simulated_conversation_info": [{"turn_count": 1}],
|
|
}
|
|
|
|
|
|
def test_prompt_optimization(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
from mlflow.genai.optimize import optimize_prompts
|
|
from mlflow.genai.optimize.optimizers import BasePromptOptimizer
|
|
from mlflow.genai.optimize.types import PromptOptimizerOutput
|
|
|
|
class MockAdapter(BasePromptOptimizer):
|
|
def __init__(self):
|
|
self.model_name = "openai:/gpt-4o-mini"
|
|
|
|
def optimize(self, eval_fn, train_data, target_prompts, enable_tracking):
|
|
return PromptOptimizerOutput(optimized_prompts=target_prompts)
|
|
|
|
sample_prompt = mlflow.genai.register_prompt(
|
|
name="test_prompt_for_adaptation",
|
|
template="Translate {{input_text}} to {{language}}",
|
|
)
|
|
|
|
sample_data = [
|
|
{"inputs": {"input_text": "Hello", "language": "Spanish"}, "outputs": "Hola"},
|
|
{"inputs": {"input_text": "World", "language": "French"}, "outputs": "Monde"},
|
|
]
|
|
|
|
@mlflow.genai.scorers.scorer
|
|
def exact_match_scorer(outputs, expectations):
|
|
return 1.0 if outputs == expectations["expected_response"] else 0.0
|
|
|
|
def predict_fn(input_text, language):
|
|
mlflow.genai.load_prompt(f"prompts:/{sample_prompt.name}/{sample_prompt.version}")
|
|
return "translated"
|
|
|
|
optimize_prompts(
|
|
predict_fn=predict_fn,
|
|
train_data=sample_data,
|
|
prompt_uris=[f"prompts:/{sample_prompt.name}/{sample_prompt.version}"],
|
|
optimizer=MockAdapter(),
|
|
scorers=[exact_match_scorer],
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
PromptOptimizationEvent.name,
|
|
{
|
|
"optimizer_type": "MockAdapter",
|
|
"prompt_count": 1,
|
|
"scorer_count": 1,
|
|
"custom_aggregation": False,
|
|
},
|
|
)
|
|
|
|
|
|
def test_create_dataset(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
with mock.patch("mlflow.tracking._tracking_service.utils._get_store") as mock_store:
|
|
mock_store_instance = mock.MagicMock()
|
|
mock_store.return_value = mock_store_instance
|
|
mock_store_instance.create_dataset.return_value = mock.MagicMock(
|
|
dataset_id="test-dataset-id", name="test_dataset", tags={"test": "value"}
|
|
)
|
|
|
|
create_dataset(name="test_dataset", tags={"test": "value"})
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, CreateDatasetEvent.name)
|
|
|
|
|
|
def test_merge_records(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
with mock.patch("mlflow.tracking._tracking_service.utils._get_store") as mock_store:
|
|
mock_store_instance = mock.MagicMock()
|
|
mock_store.return_value = mock_store_instance
|
|
mock_store_instance.get_dataset.return_value = mock.MagicMock(dataset_id="test-id")
|
|
mock_store_instance.upsert_dataset_records.return_value = {
|
|
"inserted": 2,
|
|
"updated": 0,
|
|
}
|
|
|
|
evaluation_dataset = EvaluationDataset(
|
|
dataset_id="test-id",
|
|
name="test",
|
|
digest="digest",
|
|
created_time=123,
|
|
last_update_time=456,
|
|
)
|
|
|
|
records = [
|
|
{"inputs": {"q": "Q1"}, "expectations": {"a": "A1"}},
|
|
{"inputs": {"q": "Q2"}, "expectations": {"a": "A2"}},
|
|
]
|
|
evaluation_dataset.merge_records(records)
|
|
|
|
expected_params = {
|
|
"record_count": 2,
|
|
"input_type": "list[dict]",
|
|
"dataset_type": "trace",
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
MergeRecordsEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
|
|
def test_log_dataset(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
with mlflow.start_run() as run:
|
|
dataset = mlflow.data.from_pandas(pd.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]}))
|
|
mlflow.log_input(dataset)
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, LogDatasetEvent.name)
|
|
|
|
mlflow.log_inputs(datasets=[dataset], contexts=["training"], tags_list=[None])
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, LogDatasetEvent.name)
|
|
|
|
client = MlflowClient()
|
|
client.log_inputs(run_id=run.info.run_id, datasets=[_create_dataset_input(dataset)])
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, LogDatasetEvent.name)
|
|
|
|
|
|
def test_log_metric(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
with mlflow.start_run():
|
|
mlflow.log_metric("test_metric", 1.0)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogMetricEvent.name,
|
|
{"synchronous": True},
|
|
)
|
|
|
|
mlflow.log_metric("test_metric", 1.0, synchronous=False)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogMetricEvent.name,
|
|
{"synchronous": False},
|
|
)
|
|
|
|
client = MlflowClient()
|
|
client.log_metric(
|
|
run_id=mlflow.active_run().info.run_id,
|
|
key="test_metric",
|
|
value=1.0,
|
|
timestamp=int(time.time()),
|
|
step=0,
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogMetricEvent.name,
|
|
{"synchronous": True},
|
|
)
|
|
|
|
client.log_metric(
|
|
run_id=mlflow.active_run().info.run_id,
|
|
key="test_metric",
|
|
value=1.0,
|
|
timestamp=int(time.time()),
|
|
step=0,
|
|
synchronous=False,
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogMetricEvent.name,
|
|
{"synchronous": False},
|
|
)
|
|
|
|
|
|
def test_log_param(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
with mlflow.start_run():
|
|
mlflow.log_param("test_param", "test_value")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogParamEvent.name,
|
|
{"synchronous": True},
|
|
)
|
|
|
|
mlflow.log_param("test_param", "test_value", synchronous=False)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogParamEvent.name,
|
|
{"synchronous": False},
|
|
)
|
|
|
|
client = mlflow.MlflowClient()
|
|
client.log_param(
|
|
run_id=mlflow.active_run().info.run_id,
|
|
key="test_param",
|
|
value="test_value",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogParamEvent.name,
|
|
{"synchronous": True},
|
|
)
|
|
|
|
|
|
def test_log_batch(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
with mlflow.start_run():
|
|
mlflow.log_params(params={"test_param": "test_value"})
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogBatchEvent.name,
|
|
{"metrics": False, "params": True, "tags": False, "synchronous": True},
|
|
)
|
|
|
|
mlflow.log_params(params={"test_param": "test_value"}, synchronous=False)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogBatchEvent.name,
|
|
{"metrics": False, "params": True, "tags": False, "synchronous": False},
|
|
)
|
|
|
|
mlflow.log_metrics(metrics={"test_metric": 1.0})
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogBatchEvent.name,
|
|
{"metrics": True, "params": False, "tags": False, "synchronous": True},
|
|
)
|
|
|
|
mlflow.log_metrics(metrics={"test_metric": 1.0}, synchronous=False)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogBatchEvent.name,
|
|
{"metrics": True, "params": False, "tags": False, "synchronous": False},
|
|
)
|
|
|
|
mlflow.set_tags(tags={"test_tag": "test_value"})
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogBatchEvent.name,
|
|
{"metrics": False, "params": False, "tags": True, "synchronous": True},
|
|
)
|
|
|
|
mlflow.set_tags(tags={"test_tag": "test_value"}, synchronous=False)
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogBatchEvent.name,
|
|
{"metrics": False, "params": False, "tags": True, "synchronous": False},
|
|
)
|
|
|
|
client = mlflow.MlflowClient()
|
|
client.log_batch(
|
|
run_id=mlflow.active_run().info.run_id,
|
|
metrics=[Metric(key="test_metric", value=1.0, timestamp=int(time.time()), step=0)],
|
|
params=[Param(key="test_param", value="test_value")],
|
|
tags=[RunTag(key="test_tag", value="test_value")],
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LogBatchEvent.name,
|
|
{"metrics": True, "params": True, "tags": True, "synchronous": True},
|
|
)
|
|
|
|
|
|
def test_get_logged_model(mock_requests, mock_telemetry_client: TelemetryClient, tmp_path):
|
|
model_info = mlflow.sklearn.log_model(
|
|
knn.KNeighborsClassifier(),
|
|
name="model",
|
|
)
|
|
mock_telemetry_client.flush()
|
|
|
|
mlflow.sklearn.load_model(model_info.model_uri)
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GetLoggedModelEvent.name,
|
|
check_params=False,
|
|
)
|
|
assert "sklearn" in json.loads(data["params"])["imports"]
|
|
|
|
mlflow.pyfunc.load_model(model_info.model_uri)
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GetLoggedModelEvent.name,
|
|
check_params=False,
|
|
)
|
|
|
|
model_def = """
|
|
import mlflow
|
|
from mlflow.models import set_model
|
|
|
|
class TestModel(mlflow.pyfunc.PythonModel):
|
|
def predict(self, context, model_input: list[str], params=None) -> list[str]:
|
|
return model_input
|
|
|
|
set_model(TestModel())
|
|
"""
|
|
model_path = tmp_path / "model.py"
|
|
model_path.write_text(model_def)
|
|
model_info = mlflow.pyfunc.log_model(
|
|
name="model",
|
|
python_model=model_path,
|
|
)
|
|
mock_telemetry_client.flush()
|
|
|
|
mlflow.pyfunc.load_model(model_info.model_uri)
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GetLoggedModelEvent.name,
|
|
check_params=False,
|
|
)
|
|
|
|
# test load model after registry
|
|
mlflow.register_model(model_info.model_uri, name="test")
|
|
mock_telemetry_client.flush()
|
|
|
|
mlflow.pyfunc.load_model("models:/test/1")
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GetLoggedModelEvent.name,
|
|
check_params=False,
|
|
)
|
|
|
|
|
|
def test_mcp_run(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
from mlflow.mcp.cli import run
|
|
|
|
runner = CliRunner(catch_exceptions=False)
|
|
with mock.patch("mlflow.mcp.cli.run_server") as mock_run_server:
|
|
runner.invoke(run)
|
|
|
|
mock_run_server.assert_called_once()
|
|
mock_telemetry_client.flush()
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, McpRunEvent.name)
|
|
|
|
|
|
@pytest.mark.skipif(is_windows(), reason="Windows does not support gateway start")
|
|
def test_gateway_start(tmp_path, mock_requests, mock_telemetry_client: TelemetryClient):
|
|
config = tmp_path.joinpath("config.yml")
|
|
config.write_text(
|
|
"""
|
|
endpoints:
|
|
- name: test-endpoint
|
|
endpoint_type: llm/v1/completions
|
|
model:
|
|
provider: openai
|
|
name: gpt-3.5-turbo
|
|
config:
|
|
openai_api_key: test-key
|
|
"""
|
|
)
|
|
|
|
def assert_event_recorded_before_run_app(**kwargs):
|
|
mock_telemetry_client.flush()
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, GatewayStartEvent.name)
|
|
|
|
runner = CliRunner(catch_exceptions=False)
|
|
with mock.patch("mlflow.gateway.cli.run_app", side_effect=assert_event_recorded_before_run_app):
|
|
runner.invoke(start, ["--config-path", str(config)])
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("cli_args", "expected_params"),
|
|
[
|
|
(
|
|
["--backend-store-uri", "sqlite:///test.db"],
|
|
{
|
|
"auth_enabled": False,
|
|
"app_name": None,
|
|
"backend_store_type": "sqlite",
|
|
"serve_artifacts": True,
|
|
"artifacts_only": False,
|
|
"expose_prometheus": False,
|
|
"enable_workspaces": False,
|
|
"workers": None,
|
|
"dev": False,
|
|
},
|
|
),
|
|
(
|
|
["--backend-store-uri", "sqlite:///test.db", "--app-name", "basic-auth"],
|
|
{
|
|
"auth_enabled": True,
|
|
"app_name": "basic-auth",
|
|
"backend_store_type": "sqlite",
|
|
"serve_artifacts": True,
|
|
"artifacts_only": False,
|
|
"expose_prometheus": False,
|
|
"enable_workspaces": False,
|
|
"workers": None,
|
|
"dev": False,
|
|
},
|
|
),
|
|
(
|
|
[
|
|
"--backend-store-uri",
|
|
"sqlite:///test.db",
|
|
"--no-serve-artifacts",
|
|
"--expose-prometheus",
|
|
"/tmp/metrics",
|
|
"--enable-workspaces",
|
|
],
|
|
{
|
|
"auth_enabled": False,
|
|
"app_name": None,
|
|
"backend_store_type": "sqlite",
|
|
"serve_artifacts": False,
|
|
"artifacts_only": False,
|
|
"expose_prometheus": True,
|
|
"enable_workspaces": True,
|
|
"workers": None,
|
|
"dev": False,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_tracking_server_start(
|
|
tmp_path,
|
|
mock_requests,
|
|
mock_telemetry_client: TelemetryClient,
|
|
monkeypatch,
|
|
cli_args,
|
|
expected_params,
|
|
):
|
|
|
|
from mlflow.cli import server
|
|
|
|
# Isolate env vars that server() mutates so they don't leak into other tests
|
|
for key in (
|
|
"MLFLOW_ENABLE_WORKSPACES",
|
|
"MLFLOW_TRACE_ARCHIVAL_CONFIG",
|
|
"MLFLOW_WORKSPACE_STORE_URI",
|
|
"MLFLOW_SERVER_DISABLE_SECURITY_MIDDLEWARE",
|
|
"MLFLOW_SERVER_ALLOWED_HOSTS",
|
|
"MLFLOW_SERVER_CORS_ALLOWED_ORIGINS",
|
|
"MLFLOW_SERVER_X_FRAME_OPTIONS",
|
|
):
|
|
monkeypatch.delenv(key, raising=False)
|
|
|
|
def assert_event_recorded_before_run_server(**kwargs):
|
|
mock_telemetry_client.flush()
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
TrackingServerStartEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
runner = CliRunner(catch_exceptions=False)
|
|
with (
|
|
mock.patch(
|
|
"mlflow.server._run_server", side_effect=assert_event_recorded_before_run_server
|
|
),
|
|
mock.patch("mlflow.server.handlers.initialize_backend_stores"),
|
|
):
|
|
runner.invoke(server, cli_args)
|
|
|
|
|
|
def test_ai_command_run(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
from mlflow.ai_commands import commands
|
|
|
|
runner = CliRunner(catch_exceptions=False)
|
|
# Test CLI context
|
|
with mock.patch("mlflow.ai_commands.get_command", return_value="---\ntest\n---\nTest command"):
|
|
result = runner.invoke(commands, ["run", "test_command"])
|
|
assert result.exit_code == 0
|
|
|
|
mock_telemetry_client.flush()
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
AiCommandRunEvent.name,
|
|
{"command_key": "test_command", "context": "cli"},
|
|
)
|
|
|
|
|
|
def test_git_model_versioning(mock_requests, mock_telemetry_client):
|
|
from mlflow.genai import enable_git_model_versioning
|
|
|
|
with enable_git_model_versioning():
|
|
pass
|
|
|
|
mock_telemetry_client.flush()
|
|
validate_telemetry_record(mock_telemetry_client, mock_requests, GitModelVersioningEvent.name)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model_uri", "expected_provider"),
|
|
[
|
|
("databricks:/llama-3.1-70b", "databricks"),
|
|
("openai:/gpt-4o-mini", "openai"),
|
|
("endpoints:/my-endpoint", "endpoints"),
|
|
("anthropic:/claude-3-opus", "anthropic"),
|
|
],
|
|
)
|
|
def test_invoke_custom_judge_model(
|
|
mock_requests,
|
|
mock_telemetry_client: TelemetryClient,
|
|
model_uri,
|
|
expected_provider,
|
|
):
|
|
from mlflow.genai.judges.adapters.gateway_adapter import InvokeOutput
|
|
from mlflow.genai.judges.utils import invoke_judge_model
|
|
|
|
mock_response = json.dumps({"result": 0.8, "rationale": "Test rationale"})
|
|
|
|
mock_target = (
|
|
"mlflow.genai.judges.adapters.gateway_adapter._invoke_via_gateway"
|
|
if expected_provider == "endpoints"
|
|
else "mlflow.genai.judges.adapters.gateway_adapter.GatewayAdapter._invoke_and_handle_tools"
|
|
)
|
|
mock_return = (
|
|
mock_response
|
|
if expected_provider == "endpoints"
|
|
else InvokeOutput(
|
|
response=mock_response,
|
|
request_id=None,
|
|
num_prompt_tokens=None,
|
|
num_completion_tokens=None,
|
|
)
|
|
)
|
|
with mock.patch(mock_target, return_value=mock_return):
|
|
invoke_judge_model(
|
|
model_uri=model_uri,
|
|
prompt="Test prompt",
|
|
assessment_name="test_assessment",
|
|
)
|
|
|
|
expected_params = {"model_provider": expected_provider}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
InvokeCustomJudgeModelEvent.name,
|
|
expected_params,
|
|
)
|
|
|
|
|
|
def test_make_judge(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
make_judge(
|
|
name="test_judge",
|
|
instructions="Evaluate the {{ inputs }} and {{ outputs }}",
|
|
model="openai:/gpt-4",
|
|
feedback_value_type=str,
|
|
)
|
|
expected_params = {"model_provider": "openai"}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, MakeJudgeEvent.name, expected_params
|
|
)
|
|
|
|
make_judge(
|
|
name="test_judge",
|
|
instructions="Evaluate the {{ inputs }} and {{ outputs }}",
|
|
feedback_value_type=str,
|
|
)
|
|
expected_params = {"model_provider": None}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, MakeJudgeEvent.name, expected_params
|
|
)
|
|
|
|
|
|
def test_align_judge(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
judge = make_judge(
|
|
name="test_judge",
|
|
instructions="Evaluate the {{ inputs }} and {{ outputs }}",
|
|
model="openai:/gpt-4",
|
|
feedback_value_type=str,
|
|
)
|
|
|
|
traces = [
|
|
mock.MagicMock(spec=Trace),
|
|
mock.MagicMock(spec=Trace),
|
|
]
|
|
|
|
class MockOptimizer(AlignmentOptimizer):
|
|
def align(self, judge, traces):
|
|
return judge
|
|
|
|
custom_optimizer = MockOptimizer()
|
|
judge.align(traces, optimizer=custom_optimizer)
|
|
|
|
expected_params = {"trace_count": 2, "optimizer_type": "MockOptimizer"}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, AlignJudgeEvent.name, expected_params
|
|
)
|
|
|
|
|
|
def test_discover_issues(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
traces = [
|
|
mock.MagicMock(spec=Trace),
|
|
mock.MagicMock(spec=Trace),
|
|
mock.MagicMock(spec=Trace),
|
|
]
|
|
|
|
mock_triage_run_id = "abc123"
|
|
mock_eval_result = mock.MagicMock()
|
|
mock_eval_result.run_id = mock_triage_run_id
|
|
|
|
with (
|
|
patch("mlflow.genai.discovery.pipeline.get_session_id", return_value=None),
|
|
patch("mlflow.genai.discovery.pipeline.verify_scorer"),
|
|
patch(
|
|
"mlflow.genai.discovery.pipeline.mlflow.genai.evaluate",
|
|
return_value=mock_eval_result,
|
|
),
|
|
patch(
|
|
"mlflow.genai.discovery.pipeline.extract_failing_traces",
|
|
return_value=_TriageResult([], {}, {}),
|
|
),
|
|
patch("mlflow.genai.discovery.pipeline.mlflow.MlflowClient"),
|
|
patch("mlflow.genai.discovery.pipeline.mlflow.set_experiment"),
|
|
):
|
|
discover_issues(
|
|
traces=traces,
|
|
model="openai:/gpt-4",
|
|
categories=["hallucination", "accuracy"],
|
|
)
|
|
|
|
expected_params = {
|
|
"model": "openai:/gpt-4",
|
|
"trace_count": 3,
|
|
"categories": ["hallucination", "accuracy"],
|
|
"source_run_id": None,
|
|
"issue_count": 0,
|
|
"total_traces_analyzed": 3,
|
|
"total_cost_usd": None,
|
|
"triage_run_id": mock_triage_run_id,
|
|
}
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, DiscoverIssuesEvent.name, expected_params
|
|
)
|
|
|
|
|
|
def test_autologging(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
try:
|
|
mlflow.openai.autolog()
|
|
|
|
mlflow.autolog()
|
|
mock_telemetry_client.flush()
|
|
data = [record["data"] for record in mock_requests]
|
|
params = [event["params"] for event in data if event["event_name"] == AutologgingEvent.name]
|
|
assert (
|
|
json.dumps({
|
|
"flavor": mlflow.openai.FLAVOR_NAME,
|
|
"log_traces": True,
|
|
"disable": False,
|
|
})
|
|
in params
|
|
)
|
|
assert json.dumps({"flavor": "all", "log_traces": True, "disable": False}) in params
|
|
finally:
|
|
mlflow.autolog(disable=True)
|
|
|
|
|
|
def test_load_prompt(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
# Register a prompt first
|
|
prompt = mlflow.genai.register_prompt(
|
|
name="test_prompt",
|
|
template="Hello {{name}}",
|
|
)
|
|
mock_telemetry_client.flush()
|
|
|
|
# Set an alias for testing
|
|
mlflow.genai.set_prompt_alias(name="test_prompt", version=prompt.version, alias="production")
|
|
|
|
# Test load_prompt with version (no alias)
|
|
mlflow.genai.load_prompt(name_or_uri="test_prompt", version=prompt.version)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LoadPromptEvent.name,
|
|
{"uses_alias": False},
|
|
)
|
|
|
|
# Test load_prompt with URI and version (no alias)
|
|
mlflow.genai.load_prompt(name_or_uri=f"prompts:/test_prompt/{prompt.version}")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
LoadPromptEvent.name,
|
|
{"uses_alias": False},
|
|
)
|
|
|
|
# Test load_prompt with alias
|
|
mlflow.genai.load_prompt(name_or_uri="prompts:/test_prompt@production")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, LoadPromptEvent.name, {"uses_alias": True}
|
|
)
|
|
|
|
# Test load_prompt with @latest (special alias)
|
|
mlflow.genai.load_prompt(name_or_uri="prompts:/test_prompt@latest")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client, mock_requests, LoadPromptEvent.name, {"uses_alias": True}
|
|
)
|
|
|
|
|
|
def test_scorer_call_direct(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
@scorer
|
|
def custom_scorer(outputs) -> bool:
|
|
return len(outputs) > 0
|
|
|
|
result = custom_scorer(outputs="test output")
|
|
assert result is True
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "UserDefinedScorer",
|
|
"scorer_kind": "decorator",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
safety_scorer = Safety()
|
|
|
|
mock_feedback = Feedback(
|
|
name="test_feedback",
|
|
value="yes",
|
|
rationale="Test rationale",
|
|
)
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.builtin.invoke_judge_model",
|
|
return_value=mock_feedback,
|
|
):
|
|
safety_scorer(outputs="test output")
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "Safety",
|
|
"scorer_kind": "builtin",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
mock_requests.clear()
|
|
|
|
guidelines_scorer = Guidelines(guidelines="The response must be in English")
|
|
with mock.patch(
|
|
"mlflow.genai.judges.builtin.invoke_judge_model",
|
|
return_value=mock_feedback,
|
|
):
|
|
guidelines_scorer(
|
|
inputs={"question": "What is MLflow?"}, outputs="MLflow is an ML platform"
|
|
)
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "Guidelines",
|
|
"scorer_kind": "guidelines",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
mock_requests.clear()
|
|
|
|
class CustomClassScorer(Scorer):
|
|
name: str = "custom_class"
|
|
|
|
def __call__(self, *, outputs) -> bool:
|
|
return len(outputs) > 0
|
|
|
|
custom_class_scorer = CustomClassScorer()
|
|
result = custom_class_scorer(outputs="test output")
|
|
assert result is True
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "UserDefinedScorer",
|
|
"scorer_kind": "class",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
|
|
def test_scorer_call_from_genai_evaluate(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
@scorer
|
|
def simple_length_checker(outputs) -> bool:
|
|
return len(outputs) > 0
|
|
|
|
session_judge = make_judge(
|
|
name="conversation_quality",
|
|
instructions="Evaluate if the {{ conversation }} is engaging and coherent",
|
|
model="openai:/gpt-4",
|
|
)
|
|
|
|
# Create traces with session metadata for session-level scorer testing
|
|
@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("What is MLflow?", session_id="test_session")
|
|
trace_1 = mlflow.get_trace(mlflow.get_last_active_trace_id())
|
|
|
|
model("How does MLflow work?", session_id="test_session")
|
|
trace_2 = mlflow.get_trace(mlflow.get_last_active_trace_id())
|
|
|
|
test_data = pd.DataFrame([
|
|
{
|
|
"trace": trace_1,
|
|
},
|
|
{
|
|
"trace": trace_2,
|
|
},
|
|
])
|
|
|
|
mock_feedback = Feedback(
|
|
name="test_feedback",
|
|
value="yes",
|
|
rationale="Test",
|
|
)
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.instructions_judge.invoke_judge_model",
|
|
return_value=mock_feedback,
|
|
):
|
|
mlflow.genai.evaluate(data=test_data, scorers=[simple_length_checker, session_judge])
|
|
|
|
mock_telemetry_client.flush()
|
|
|
|
scorer_call_events = [
|
|
record for record in mock_requests if record["data"]["event_name"] == ScorerCallEvent.name
|
|
]
|
|
|
|
# Should have 3 events: 2 response-level calls (one per trace)
|
|
# + 1 session-level call (one per session)
|
|
assert len(scorer_call_events) == 3
|
|
|
|
event_params = [json.loads(event["data"]["params"]) for event in scorer_call_events]
|
|
|
|
# Validate response-level scorer was called twice (once per trace)
|
|
response_level_events = [
|
|
params
|
|
for params in event_params
|
|
if params["scorer_class"] == "UserDefinedScorer"
|
|
and params["scorer_kind"] == "decorator"
|
|
and params["scope"] == "trace"
|
|
and params["callsite"] == "genai_evaluate"
|
|
and params["has_feedback_error"] is False
|
|
]
|
|
assert len(response_level_events) == 2
|
|
|
|
# Validate session-level scorer was called once (once per session)
|
|
session_level_events = [
|
|
params
|
|
for params in event_params
|
|
if params["scorer_class"] == "UserDefinedScorer"
|
|
and params["scorer_kind"] == "instructions"
|
|
and params["scope"] == "session"
|
|
and params["callsite"] == "genai_evaluate"
|
|
and params["has_feedback_error"] is False
|
|
]
|
|
assert len(session_level_events) == 1
|
|
|
|
mock_requests.clear()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("job_name", "expected_callsite"),
|
|
[
|
|
("run_online_trace_scorer", "online_scoring"),
|
|
("run_online_session_scorer", "online_scoring"),
|
|
# Counterexample: non-online-scoring job should be treated as direct call
|
|
("invoke_scorer", "direct_scorer_call"),
|
|
],
|
|
)
|
|
def test_scorer_call_online_scoring_callsite(
|
|
mock_requests, mock_telemetry_client: TelemetryClient, monkeypatch, job_name, expected_callsite
|
|
):
|
|
# Import here to avoid circular imports
|
|
from mlflow.server.jobs.utils import MLFLOW_SERVER_JOB_NAME_ENV_VAR
|
|
|
|
monkeypatch.setenv(MLFLOW_SERVER_JOB_NAME_ENV_VAR, job_name)
|
|
|
|
@scorer
|
|
def custom_scorer(outputs: str) -> bool:
|
|
return True
|
|
|
|
custom_scorer(outputs="test output")
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "UserDefinedScorer",
|
|
"scorer_kind": "decorator",
|
|
"scope": "trace",
|
|
"callsite": expected_callsite,
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
|
|
def test_scorer_call_tracks_feedback_errors(mock_requests, mock_telemetry_client: TelemetryClient):
|
|
error_judge = make_judge(
|
|
name="quality_judge",
|
|
instructions="Evaluate if {{ outputs }} is high quality",
|
|
model="openai:/gpt-4",
|
|
)
|
|
|
|
error_feedback = Feedback(
|
|
name="quality_judge",
|
|
error="Model invocation failed",
|
|
source=AssessmentSource(
|
|
source_type=AssessmentSourceType.LLM_JUDGE, source_id="openai:/gpt-4"
|
|
),
|
|
)
|
|
with mock.patch(
|
|
"mlflow.genai.judges.instructions_judge.invoke_judge_model",
|
|
return_value=error_feedback,
|
|
):
|
|
result = error_judge(outputs="test output")
|
|
assert result.error is not None
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "UserDefinedScorer",
|
|
"scorer_kind": "instructions",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": True,
|
|
},
|
|
)
|
|
|
|
mock_requests.clear()
|
|
|
|
# Test Scorer returns list of Feedback with mixed errors
|
|
@scorer
|
|
def multi_feedback_scorer(outputs) -> list[Feedback]:
|
|
return [
|
|
Feedback(name="feedback1", value=1.0),
|
|
Feedback(name="feedback2", error=ValueError("Error in feedback 2")),
|
|
Feedback(name="feedback3", value=0.5),
|
|
]
|
|
|
|
multi_feedback_scorer(outputs="test")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "UserDefinedScorer",
|
|
"scorer_kind": "decorator",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": True,
|
|
},
|
|
)
|
|
|
|
mock_requests.clear()
|
|
|
|
# Test Scorer returns primitive type (no Feedback error possible)
|
|
@scorer
|
|
def primitive_scorer(outputs) -> bool:
|
|
return True
|
|
|
|
primitive_scorer(outputs="test")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "UserDefinedScorer",
|
|
"scorer_kind": "decorator",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
|
|
def test_scorer_call_wrapped_builtin_scorer_direct(
|
|
mock_requests, mock_telemetry_client: TelemetryClient
|
|
):
|
|
completeness_scorer = Completeness()
|
|
|
|
mock_feedback = Feedback(
|
|
name="completeness",
|
|
value="yes",
|
|
rationale="Test rationale",
|
|
)
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.instructions_judge.invoke_judge_model",
|
|
return_value=mock_feedback,
|
|
):
|
|
completeness_scorer(inputs={"question": "What is MLflow?"}, outputs="MLflow is a platform")
|
|
|
|
mock_telemetry_client.flush()
|
|
|
|
# Verify exactly 1 scorer_call event was created
|
|
# (only top-level Completeness, not nested InstructionsJudge)
|
|
scorer_call_events = [
|
|
record for record in mock_requests if record["data"]["event_name"] == ScorerCallEvent.name
|
|
]
|
|
assert len(scorer_call_events) == 1, (
|
|
f"Expected 1 scorer call event for Completeness scorer (nested calls should be skipped), "
|
|
f"got {len(scorer_call_events)}"
|
|
)
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "Completeness",
|
|
"scorer_kind": "builtin",
|
|
"scope": "trace",
|
|
"callsite": "direct_scorer_call",
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
|
|
def test_scorer_call_wrapped_builtin_scorer_from_genai_evaluate(
|
|
mock_requests, mock_telemetry_client: TelemetryClient
|
|
):
|
|
user_frustration_scorer = UserFrustration()
|
|
|
|
@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("What is MLflow?", session_id="test_session")
|
|
trace_1 = mlflow.get_trace(mlflow.get_last_active_trace_id())
|
|
|
|
model("How does MLflow work?", session_id="test_session")
|
|
trace_2 = mlflow.get_trace(mlflow.get_last_active_trace_id())
|
|
|
|
test_data = pd.DataFrame([
|
|
{"trace": trace_1},
|
|
{"trace": trace_2},
|
|
])
|
|
|
|
mock_feedback = Feedback(
|
|
name="user_frustration",
|
|
value="no",
|
|
rationale="Test rationale",
|
|
)
|
|
|
|
with mock.patch(
|
|
"mlflow.genai.judges.instructions_judge.invoke_judge_model",
|
|
return_value=mock_feedback,
|
|
):
|
|
mlflow.genai.evaluate(data=test_data, scorers=[user_frustration_scorer])
|
|
|
|
mock_telemetry_client.flush()
|
|
|
|
# Verify exactly 1 scorer_call event was created for the session-level scorer
|
|
# (one call at the session level and no nested InstructionsJudge event)
|
|
scorer_call_events = [
|
|
record for record in mock_requests if record["data"]["event_name"] == ScorerCallEvent.name
|
|
]
|
|
assert len(scorer_call_events) == 1, (
|
|
f"Expected 1 scorer call event for UserFrustration scorer "
|
|
f"(nested calls should be skipped), got {len(scorer_call_events)}"
|
|
)
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
ScorerCallEvent.name,
|
|
{
|
|
"scorer_class": "UserFrustration",
|
|
"scorer_kind": "builtin",
|
|
"scope": "session",
|
|
"callsite": "genai_evaluate",
|
|
"has_feedback_error": False,
|
|
},
|
|
)
|
|
|
|
|
|
def test_gateway_crud_telemetry(mock_requests, mock_telemetry_client: TelemetryClient, tmp_path):
|
|
db_path = tmp_path / "mlflow.db"
|
|
store = SqlAlchemyStore(f"sqlite:///{db_path}", tmp_path.as_posix())
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-secret",
|
|
secret_value={"api_key": "test-api-key"},
|
|
provider="openai",
|
|
created_by="test-user",
|
|
)
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_id=secret.secret_id,
|
|
created_by="test-user",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayCreateModelDefinitionEvent.name,
|
|
{"model_name": "gpt-4", "provider": "openai"},
|
|
)
|
|
|
|
model_config = GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=100,
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="test-endpoint",
|
|
model_configs=[model_config],
|
|
created_by="test-user",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayCreateEndpointEvent.name,
|
|
{
|
|
"has_fallback_config": False,
|
|
"routing_strategy": None,
|
|
"num_model_configs": 1,
|
|
"usage_tracking": True,
|
|
},
|
|
)
|
|
|
|
store.get_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayGetEndpointEvent.name,
|
|
)
|
|
|
|
store.list_gateway_endpoints()
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayListEndpointsEvent.name,
|
|
{"filter_by_provider": False},
|
|
)
|
|
|
|
store.list_gateway_endpoints(provider="openai")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayListEndpointsEvent.name,
|
|
{"filter_by_provider": True},
|
|
)
|
|
|
|
store.update_gateway_endpoint(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
name="updated-endpoint",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayUpdateEndpointEvent.name,
|
|
{
|
|
"has_fallback_config": False,
|
|
"routing_strategy": None,
|
|
"num_model_configs": None,
|
|
"usage_tracking": None,
|
|
},
|
|
)
|
|
|
|
store.delete_gateway_endpoint(endpoint_id=endpoint.endpoint_id)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayDeleteEndpointEvent.name,
|
|
)
|
|
|
|
|
|
def test_gateway_secret_crud_telemetry(
|
|
mock_requests, mock_telemetry_client: TelemetryClient, tmp_path
|
|
):
|
|
db_path = tmp_path / "mlflow.db"
|
|
store = SqlAlchemyStore(f"sqlite:///{db_path}", tmp_path.as_posix())
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-secret",
|
|
secret_value={"api_key": "test-api-key"},
|
|
provider="openai",
|
|
created_by="test-user",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayCreateSecretEvent.name,
|
|
{"provider": "openai"},
|
|
)
|
|
|
|
secret2 = store.create_gateway_secret(
|
|
secret_name="test-secret-2",
|
|
secret_value={"api_key": "test-api-key-2"},
|
|
created_by="test-user",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayCreateSecretEvent.name,
|
|
{"provider": None},
|
|
)
|
|
|
|
store.list_secret_infos()
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayListSecretsEvent.name,
|
|
{"filter_by_provider": False},
|
|
)
|
|
|
|
store.list_secret_infos(provider="openai")
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayListSecretsEvent.name,
|
|
{"filter_by_provider": True},
|
|
)
|
|
|
|
store.update_gateway_secret(
|
|
secret_id=secret.secret_id,
|
|
secret_value={"api_key": "updated-api-key"},
|
|
updated_by="test-user",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayUpdateSecretEvent.name,
|
|
)
|
|
|
|
store.delete_gateway_secret(secret_id=secret.secret_id)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayDeleteSecretEvent.name,
|
|
)
|
|
|
|
store.delete_gateway_secret(secret_id=secret2.secret_id)
|
|
|
|
|
|
def test_gateway_budget_policy_crud_telemetry(
|
|
mock_requests, mock_telemetry_client: TelemetryClient, tmp_path
|
|
):
|
|
db_path = tmp_path / "mlflow.db"
|
|
store = SqlAlchemyStore(f"sqlite:///{db_path}", tmp_path.as_posix())
|
|
|
|
policy = store.create_budget_policy(
|
|
budget_unit=BudgetUnit.USD,
|
|
budget_amount=100.0,
|
|
duration=BudgetDuration(unit=BudgetDurationUnit.DAYS, value=30),
|
|
target_scope=BudgetTargetScope.GLOBAL,
|
|
budget_action=BudgetAction.ALERT,
|
|
created_by="test-user",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayCreateBudgetPolicyEvent.name,
|
|
{
|
|
"budget_unit": "USD",
|
|
"duration_unit": "DAYS",
|
|
"target_scope": "GLOBAL",
|
|
"budget_action": "ALERT",
|
|
},
|
|
)
|
|
|
|
store.list_budget_policies()
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayListBudgetPoliciesEvent.name,
|
|
)
|
|
|
|
store.update_budget_policy(
|
|
budget_policy_id=policy.budget_policy_id,
|
|
budget_amount=200.0,
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayUpdateBudgetPolicyEvent.name,
|
|
)
|
|
|
|
store.delete_budget_policy(budget_policy_id=policy.budget_policy_id)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayDeleteBudgetPolicyEvent.name,
|
|
)
|
|
|
|
|
|
def test_gateway_guardrail_crud_telemetry(
|
|
mock_requests, mock_telemetry_client: TelemetryClient, tmp_path
|
|
):
|
|
db_path = tmp_path / "mlflow.db"
|
|
store = SqlAlchemyStore(f"sqlite:///{db_path}", tmp_path.as_posix())
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-secret",
|
|
secret_value={"api_key": "test-api-key"},
|
|
provider="openai",
|
|
created_by="test-user",
|
|
)
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_id=secret.secret_id,
|
|
created_by="test-user",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="test-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=100,
|
|
)
|
|
],
|
|
created_by="test-user",
|
|
usage_tracking=False,
|
|
)
|
|
scorer_experiment_id = store.create_experiment("guardrail-scorer-exp")
|
|
serialized_scorer = json.dumps({
|
|
"instructions_judge_pydantic_data": {
|
|
"model": "gateway:/test-endpoint",
|
|
"instructions": "Is this input safe?",
|
|
}
|
|
})
|
|
scorer = store.register_scorer(
|
|
experiment_id=scorer_experiment_id,
|
|
name="safety-judge",
|
|
serialized_scorer=serialized_scorer,
|
|
)
|
|
guardrail = store.create_gateway_guardrail(
|
|
name="guardrail-1",
|
|
scorer_id=scorer.scorer_id,
|
|
scorer_version=scorer.scorer_version,
|
|
stage=GuardrailStage.BEFORE,
|
|
action=GuardrailAction.VALIDATION,
|
|
created_by="test-user",
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayCreateGuardrailEvent.name,
|
|
{
|
|
"stage": "BEFORE",
|
|
"action": "VALIDATION",
|
|
},
|
|
)
|
|
|
|
# Guardrail update telemetry is emitted by endpoint guardrail config updates.
|
|
store.add_guardrail_to_endpoint(endpoint.endpoint_id, guardrail.guardrail_id, execution_order=1)
|
|
store.update_endpoint_guardrail_config(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
guardrail_id=guardrail.guardrail_id,
|
|
execution_order=2,
|
|
)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayUpdateGuardrailEvent.name,
|
|
{"stage": None, "action": None},
|
|
)
|
|
|
|
store.delete_gateway_guardrail(guardrail.guardrail_id)
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayDeleteGuardrailEvent.name,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_gateway_invocation_telemetry(
|
|
mock_requests, mock_telemetry_client: TelemetryClient, tmp_path
|
|
):
|
|
db_path = tmp_path / "mlflow.db"
|
|
store = SqlAlchemyStore(f"sqlite:///{db_path}", tmp_path.as_posix())
|
|
|
|
secret = store.create_gateway_secret(
|
|
secret_name="test-secret",
|
|
secret_value={"api_key": "test-api-key"},
|
|
provider="openai",
|
|
created_by="test-user",
|
|
)
|
|
mock_telemetry_client.flush()
|
|
mock_requests.clear()
|
|
|
|
model_def = store.create_gateway_model_definition(
|
|
name="test-model",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_id=secret.secret_id,
|
|
created_by="test-user",
|
|
)
|
|
endpoint = store.create_gateway_endpoint(
|
|
name="test-endpoint",
|
|
model_configs=[
|
|
GatewayEndpointModelConfig(
|
|
model_definition_id=model_def.model_definition_id,
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
weight=100,
|
|
)
|
|
],
|
|
created_by="test-user",
|
|
)
|
|
mock_telemetry_client.flush()
|
|
mock_requests.clear()
|
|
|
|
mock_response = chat.ResponsePayload(
|
|
id="test-id",
|
|
object="chat.completion",
|
|
created=1234567890,
|
|
model="gpt-4",
|
|
choices=[
|
|
chat.Choice(
|
|
index=0,
|
|
message=chat.ResponseMessage(role="assistant", content="Hello!"),
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
usage=chat.ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
|
)
|
|
|
|
mock_model = GatewayModelConfig(
|
|
model_definition_id="test-model-def",
|
|
provider="openai",
|
|
model_name="gpt-4",
|
|
secret_value={"api_key": "test"},
|
|
linkage_type=GatewayModelLinkageType.PRIMARY,
|
|
)
|
|
|
|
# Test invocations endpoint (chat)
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_request.headers = {}
|
|
mock_request.json = AsyncMock(
|
|
return_value={
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
"temperature": 0.7,
|
|
"stream": False,
|
|
}
|
|
)
|
|
|
|
with (
|
|
patch("mlflow.server.gateway_api._get_store", return_value=store),
|
|
patch(
|
|
"mlflow.server.gateway_api._create_provider_from_endpoint_name"
|
|
) as mock_create_provider,
|
|
):
|
|
mock_provider = MagicMock()
|
|
mock_provider.chat = AsyncMock(return_value=mock_response)
|
|
mock_endpoint_config = GatewayEndpointConfig(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
endpoint_name=endpoint.name,
|
|
models=[mock_model],
|
|
)
|
|
mock_create_provider.return_value = (mock_provider, mock_endpoint_config)
|
|
|
|
await invocations(endpoint.name, mock_request)
|
|
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayInvocationEvent.name,
|
|
check_params=False,
|
|
)
|
|
params = json.loads(data["params"])
|
|
assert params["is_streaming"] is False
|
|
assert params["invocation_type"] == "mlflow_invocations"
|
|
assert params["has_traceparent"] is False
|
|
assert params["auth_enabled"] is False
|
|
assert params["endpoint_id"] == endpoint.endpoint_id
|
|
assert params["provider"] == "openai"
|
|
# Non-streaming includes timing fields
|
|
assert "provider_duration_ms" in params
|
|
assert "gateway_overhead_ms" in params
|
|
|
|
# Test chat_completions endpoint
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_request.headers = {}
|
|
mock_request.json = AsyncMock(
|
|
return_value={
|
|
"model": endpoint.name,
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
"temperature": 0.7,
|
|
"stream": False,
|
|
}
|
|
)
|
|
|
|
with (
|
|
patch("mlflow.server.gateway_api._get_store", return_value=store),
|
|
patch(
|
|
"mlflow.server.gateway_api._create_provider_from_endpoint_name"
|
|
) as mock_create_provider,
|
|
):
|
|
mock_provider = MagicMock()
|
|
mock_provider.chat = AsyncMock(return_value=mock_response)
|
|
mock_endpoint_config = GatewayEndpointConfig(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
endpoint_name=endpoint.name,
|
|
models=[mock_model],
|
|
)
|
|
mock_create_provider.return_value = (mock_provider, mock_endpoint_config)
|
|
|
|
await chat_completions(mock_request)
|
|
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayInvocationEvent.name,
|
|
check_params=False,
|
|
)
|
|
params = json.loads(data["params"])
|
|
assert params["is_streaming"] is False
|
|
assert params["invocation_type"] == "mlflow_chat_completions"
|
|
assert params["endpoint_id"] == endpoint.endpoint_id
|
|
assert params["provider"] == "openai"
|
|
|
|
# Test streaming invocation — timing fields should be absent
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_request.headers = {}
|
|
mock_request.json = AsyncMock(
|
|
return_value={
|
|
"model": endpoint.name,
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
"stream": True,
|
|
}
|
|
)
|
|
|
|
async def mock_stream():
|
|
yield chat.StreamResponsePayload(
|
|
id="test-id",
|
|
object="chat.completion.chunk",
|
|
created=1234567890,
|
|
model="gpt-4",
|
|
choices=[
|
|
chat.StreamChoice(
|
|
index=0,
|
|
delta=chat.StreamDelta(role="assistant", content="Hello"),
|
|
finish_reason=None,
|
|
)
|
|
],
|
|
)
|
|
|
|
with (
|
|
patch("mlflow.server.gateway_api._get_store", return_value=store),
|
|
patch(
|
|
"mlflow.server.gateway_api._create_provider_from_endpoint_name"
|
|
) as mock_create_provider,
|
|
):
|
|
mock_provider = MagicMock()
|
|
mock_provider.chat_stream = MagicMock(return_value=mock_stream())
|
|
mock_endpoint_config = GatewayEndpointConfig(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
endpoint_name=endpoint.name,
|
|
models=[mock_model],
|
|
)
|
|
mock_create_provider.return_value = (mock_provider, mock_endpoint_config)
|
|
|
|
await chat_completions(mock_request)
|
|
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayInvocationEvent.name,
|
|
check_params=False,
|
|
)
|
|
params = json.loads(data["params"])
|
|
assert params["is_streaming"] is True
|
|
assert params["invocation_type"] == "mlflow_chat_completions"
|
|
# Streaming responses should NOT include timing fields
|
|
assert "provider_duration_ms" not in params
|
|
assert "gateway_overhead_ms" not in params
|
|
|
|
# Test that caller header and traceparent are included in telemetry when present
|
|
mock_request = MagicMock(spec=Request)
|
|
mock_request.json = AsyncMock(
|
|
return_value={
|
|
"model": endpoint.name,
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
"stream": False,
|
|
}
|
|
)
|
|
mock_request.headers = {MLFLOW_GATEWAY_CALLER_HEADER: "judge", "traceparent": "00-abc-def-01"}
|
|
|
|
mock_auth_module = MagicMock()
|
|
mock_auth_module.is_auth_enabled = MagicMock(return_value=True)
|
|
|
|
with (
|
|
patch("mlflow.server.gateway_api._get_store", return_value=store),
|
|
patch(
|
|
"mlflow.server.gateway_api._create_provider_from_endpoint_name"
|
|
) as mock_create_provider,
|
|
patch.dict("sys.modules", {"mlflow.server.auth": mock_auth_module}),
|
|
):
|
|
mock_provider = MagicMock()
|
|
mock_provider.chat = AsyncMock(return_value=mock_response)
|
|
mock_endpoint_config = GatewayEndpointConfig(
|
|
endpoint_id=endpoint.endpoint_id,
|
|
endpoint_name=endpoint.name,
|
|
models=[mock_model],
|
|
)
|
|
mock_create_provider.return_value = (mock_provider, mock_endpoint_config)
|
|
|
|
await chat_completions(mock_request)
|
|
|
|
data = validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
GatewayInvocationEvent.name,
|
|
check_params=False,
|
|
)
|
|
params = json.loads(data["params"])
|
|
assert params["is_streaming"] is False
|
|
assert params["invocation_type"] == "mlflow_chat_completions"
|
|
assert params["caller"] == "judge"
|
|
assert params["has_traceparent"] is True
|
|
assert params["auth_enabled"] is True
|
|
|
|
|
|
def test_tracing_context_propagation_get_and_set_success(
|
|
mock_requests, mock_telemetry_client: TelemetryClient
|
|
):
|
|
with mock.patch(
|
|
"mlflow.telemetry.track.get_telemetry_client", return_value=mock_telemetry_client
|
|
):
|
|
with mlflow.start_span("client span"):
|
|
headers = get_tracing_context_headers_for_http_request()
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
TracingContextPropagation.name,
|
|
)
|
|
|
|
with mock.patch(
|
|
"mlflow.telemetry.track.get_telemetry_client", return_value=mock_telemetry_client
|
|
):
|
|
with set_tracing_context_from_http_request_headers(headers):
|
|
with mlflow.start_span("server span"):
|
|
pass
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
TracingContextPropagation.name,
|
|
)
|
|
|
|
|
|
def test_update_issue_telemetry(mock_requests, mock_telemetry_client: TelemetryClient, db_uri):
|
|
store = SqlAlchemyStore(db_uri, "/tmp")
|
|
|
|
exp_id = store.create_experiment("test-exp")
|
|
issue = store.create_issue(
|
|
experiment_id=exp_id,
|
|
name="Original name",
|
|
description="Original description",
|
|
status=IssueStatus.PENDING,
|
|
)
|
|
mock_telemetry_client.flush()
|
|
mock_requests.clear()
|
|
|
|
store.update_issue(
|
|
issue_id=issue.issue_id,
|
|
status=IssueStatus.RESOLVED,
|
|
name="Updated name",
|
|
description="Updated description",
|
|
severity=IssueSeverity.HIGH,
|
|
)
|
|
|
|
validate_telemetry_record(
|
|
mock_telemetry_client,
|
|
mock_requests,
|
|
UpdateIssueEvent.name,
|
|
{
|
|
"status": "resolved",
|
|
"has_name": True,
|
|
"has_description": True,
|
|
"severity": "high",
|
|
"source_run_id": None,
|
|
},
|
|
)
|