wehub-resource-sync
744 lines
23 KiB
Python
744 lines
23 KiB
Python
import json
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
|
|
|
|
from mlflow.entities.dataset_record import DatasetRecord
|
|
from mlflow.entities.dataset_record_source import DatasetRecordSourceType
|
|
from mlflow.entities.evaluation_dataset import EvaluationDataset
|
|
from mlflow.entities.span import Span, SpanType
|
|
from mlflow.entities.trace import Trace
|
|
from mlflow.entities.trace_data import TraceData
|
|
from mlflow.entities.trace_info import TraceInfo
|
|
from mlflow.entities.trace_location import TraceLocation
|
|
from mlflow.entities.trace_state import TraceState
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.tracing.utils import build_otel_context
|
|
|
|
|
|
def test_evaluation_dataset_creation():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="abc123",
|
|
created_time=123456789,
|
|
last_update_time=987654321,
|
|
tags={"source": "manual", "type": "HUMAN"},
|
|
schema='{"fields": ["input", "output"]}',
|
|
profile='{"count": 100}',
|
|
created_by="user1",
|
|
last_updated_by="user2",
|
|
)
|
|
|
|
assert dataset.dataset_id == "dataset123"
|
|
assert dataset.name == "test_dataset"
|
|
assert dataset.tags == {"source": "manual", "type": "HUMAN"}
|
|
assert dataset.schema == '{"fields": ["input", "output"]}'
|
|
assert dataset.profile == '{"count": 100}'
|
|
assert dataset.digest == "abc123"
|
|
assert dataset.created_by == "user1"
|
|
assert dataset.last_updated_by == "user2"
|
|
assert dataset.created_time == 123456789
|
|
assert dataset.last_update_time == 987654321
|
|
|
|
dataset.experiment_ids = ["exp1", "exp2"]
|
|
assert dataset.experiment_ids == ["exp1", "exp2"]
|
|
|
|
|
|
def test_evaluation_dataset_timestamps_required():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=987654321,
|
|
)
|
|
|
|
assert dataset.created_time == 123456789
|
|
assert dataset.last_update_time == 987654321
|
|
|
|
|
|
def test_evaluation_dataset_experiment_ids_setter():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
new_experiment_ids = ["exp1", "exp2"]
|
|
dataset.experiment_ids = new_experiment_ids
|
|
assert dataset._experiment_ids == new_experiment_ids
|
|
assert dataset.experiment_ids == new_experiment_ids
|
|
|
|
dataset.experiment_ids = []
|
|
assert dataset._experiment_ids == []
|
|
assert dataset.experiment_ids == []
|
|
|
|
dataset.experiment_ids = None
|
|
assert dataset._experiment_ids == []
|
|
assert dataset.experiment_ids == []
|
|
|
|
|
|
def test_evaluation_dataset_to_from_proto():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
tags={"source": "manual", "type": "HUMAN"},
|
|
schema='{"fields": ["input", "output"]}',
|
|
profile='{"count": 100}',
|
|
digest="abc123",
|
|
created_time=123456789,
|
|
last_update_time=987654321,
|
|
created_by="user1",
|
|
last_updated_by="user2",
|
|
)
|
|
dataset.experiment_ids = ["exp1", "exp2"]
|
|
|
|
proto = dataset.to_proto()
|
|
assert proto.name == "test_dataset"
|
|
assert proto.tags == '{"source": "manual", "type": "HUMAN"}'
|
|
assert proto.schema == '{"fields": ["input", "output"]}'
|
|
assert proto.profile == '{"count": 100}'
|
|
assert proto.digest == "abc123"
|
|
assert proto.created_time == 123456789
|
|
assert proto.last_update_time == 987654321
|
|
assert proto.created_by == "user1"
|
|
assert proto.last_updated_by == "user2"
|
|
assert list(proto.experiment_ids) == ["exp1", "exp2"]
|
|
|
|
dataset2 = EvaluationDataset.from_proto(proto)
|
|
assert dataset2.dataset_id == dataset.dataset_id
|
|
assert dataset2.name == dataset.name
|
|
assert dataset2.tags == dataset.tags
|
|
assert dataset2.schema == dataset.schema
|
|
assert dataset2.profile == dataset.profile
|
|
assert dataset2.digest == dataset.digest
|
|
assert dataset2.created_time == dataset.created_time
|
|
assert dataset2.last_update_time == dataset.last_update_time
|
|
assert dataset2.created_by == dataset.created_by
|
|
assert dataset2.last_updated_by == dataset.last_updated_by
|
|
assert dataset2._experiment_ids == ["exp1", "exp2"]
|
|
assert dataset2.experiment_ids == ["exp1", "exp2"]
|
|
|
|
|
|
def test_evaluation_dataset_to_from_proto_minimal():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
proto = dataset.to_proto()
|
|
dataset2 = EvaluationDataset.from_proto(proto)
|
|
|
|
assert dataset2.dataset_id == "dataset123"
|
|
assert dataset2.name == "test_dataset"
|
|
assert dataset2.tags is None
|
|
assert dataset2.schema is None
|
|
assert dataset2.profile is None
|
|
assert dataset2.digest == "digest123"
|
|
assert dataset2.created_by is None
|
|
assert dataset2.last_updated_by is None
|
|
assert dataset2._experiment_ids is None
|
|
|
|
|
|
def test_evaluation_dataset_to_from_dict():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
tags={"source": "manual", "type": "HUMAN"},
|
|
schema='{"fields": ["input", "output"]}',
|
|
profile='{"count": 100}',
|
|
digest="abc123",
|
|
created_time=123456789,
|
|
last_update_time=987654321,
|
|
created_by="user1",
|
|
last_updated_by="user2",
|
|
)
|
|
dataset.experiment_ids = ["exp1", "exp2"]
|
|
|
|
dataset._records = [
|
|
DatasetRecord(
|
|
dataset_record_id="rec789",
|
|
dataset_id="dataset123",
|
|
inputs={"question": "What is MLflow?"},
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
]
|
|
|
|
data = dataset.to_dict()
|
|
assert data["dataset_id"] == "dataset123"
|
|
assert data["name"] == "test_dataset"
|
|
assert data["tags"] == {"source": "manual", "type": "HUMAN"}
|
|
assert data["schema"] == '{"fields": ["input", "output"]}'
|
|
assert data["profile"] == '{"count": 100}'
|
|
assert data["digest"] == "abc123"
|
|
assert data["created_time"] == 123456789
|
|
assert data["last_update_time"] == 987654321
|
|
assert data["created_by"] == "user1"
|
|
assert data["last_updated_by"] == "user2"
|
|
assert data["experiment_ids"] == ["exp1", "exp2"]
|
|
assert len(data["records"]) == 1
|
|
assert data["records"][0]["inputs"]["question"] == "What is MLflow?"
|
|
|
|
dataset2 = EvaluationDataset.from_dict(data)
|
|
assert dataset2.dataset_id == dataset.dataset_id
|
|
assert dataset2.name == dataset.name
|
|
assert dataset2.tags == dataset.tags
|
|
assert dataset2.schema == dataset.schema
|
|
assert dataset2.profile == dataset.profile
|
|
assert dataset2.digest == dataset.digest
|
|
assert dataset2.created_time == dataset.created_time
|
|
assert dataset2.last_update_time == dataset.last_update_time
|
|
assert dataset2.created_by == dataset.created_by
|
|
assert dataset2.last_updated_by == dataset.last_updated_by
|
|
assert dataset2._experiment_ids == ["exp1", "exp2"]
|
|
assert dataset2.experiment_ids == ["exp1", "exp2"]
|
|
assert len(dataset2._records) == 1
|
|
assert dataset2._records[0].inputs["question"] == "What is MLflow?"
|
|
|
|
|
|
def test_evaluation_dataset_to_from_dict_minimal():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
dataset._experiment_ids = []
|
|
dataset._records = []
|
|
|
|
data = dataset.to_dict()
|
|
dataset2 = EvaluationDataset.from_dict(data)
|
|
|
|
assert dataset2.dataset_id == "dataset123"
|
|
assert dataset2.name == "test_dataset"
|
|
assert dataset2.tags is None
|
|
assert dataset2.schema is None
|
|
assert dataset2.profile is None
|
|
assert dataset2.digest == "digest123"
|
|
assert dataset2.created_by is None
|
|
assert dataset2.last_updated_by is None
|
|
assert dataset2._experiment_ids == []
|
|
assert dataset2._records == []
|
|
|
|
|
|
def test_evaluation_dataset_has_records():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
assert dataset.has_records() is False
|
|
|
|
dataset._records = [
|
|
DatasetRecord(
|
|
dataset_record_id="rec123",
|
|
dataset_id="dataset123",
|
|
inputs={"test": "data"},
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
]
|
|
assert dataset.has_records() is True
|
|
|
|
dataset._records = []
|
|
assert dataset.has_records() is True
|
|
|
|
|
|
def test_evaluation_dataset_proto_with_unloaded_experiment_ids():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
assert dataset._experiment_ids is None
|
|
|
|
proto = dataset.to_proto()
|
|
assert len(proto.experiment_ids) == 0
|
|
assert dataset._experiment_ids is None
|
|
|
|
|
|
def test_evaluation_dataset_complex_tags():
|
|
complex_tags = {
|
|
"source": "automated",
|
|
"metadata": {"version": "1.0", "config": {"temperature": 0.7, "max_tokens": 100}},
|
|
"labels": ["production", "evaluated"],
|
|
}
|
|
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
tags=complex_tags,
|
|
)
|
|
|
|
proto = dataset.to_proto()
|
|
dataset2 = EvaluationDataset.from_proto(proto)
|
|
assert dataset2.tags == complex_tags
|
|
|
|
dataset._experiment_ids = []
|
|
dataset._records = []
|
|
|
|
data = dataset.to_dict()
|
|
dataset3 = EvaluationDataset.from_dict(data)
|
|
assert dataset3.tags == complex_tags
|
|
|
|
|
|
def test_evaluation_dataset_to_df():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
# Test empty dataset
|
|
df_empty = dataset.to_df()
|
|
assert isinstance(df_empty, pd.DataFrame)
|
|
expected_columns = [
|
|
"inputs",
|
|
"outputs",
|
|
"expectations",
|
|
"tags",
|
|
"source_type",
|
|
"source_id",
|
|
"source",
|
|
"created_time",
|
|
"dataset_record_id",
|
|
]
|
|
assert list(df_empty.columns) == expected_columns
|
|
assert len(df_empty) == 0
|
|
|
|
# Test dataset with records
|
|
dataset._records = [
|
|
DatasetRecord(
|
|
dataset_record_id="rec123",
|
|
dataset_id="dataset123",
|
|
inputs={"question": "What is MLflow?"},
|
|
outputs={
|
|
"answer": "MLflow is an ML platform for managing machine learning lifecycle",
|
|
"key1": "value1",
|
|
},
|
|
expectations={"answer": "MLflow is an ML platform"},
|
|
tags={"source": "manual"},
|
|
source_type="HUMAN",
|
|
source_id="user123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
),
|
|
DatasetRecord(
|
|
dataset_record_id="rec456",
|
|
dataset_id="dataset123",
|
|
inputs={"question": "What is Spark?"},
|
|
outputs={"answer": "Apache Spark is a unified analytics engine for data processing"},
|
|
expectations={"answer": "Spark is a data engine"},
|
|
tags={"source": "automated"},
|
|
source_type="CODE",
|
|
source_id="script456",
|
|
created_time=123456790,
|
|
last_update_time=123456790,
|
|
),
|
|
]
|
|
|
|
df = dataset.to_df()
|
|
assert isinstance(df, pd.DataFrame)
|
|
assert list(df.columns) == expected_columns
|
|
assert len(df) == 2
|
|
|
|
# Check that outputs column exists and contains actual values
|
|
assert "outputs" in df.columns
|
|
assert df["outputs"].iloc[0] == {
|
|
"answer": "MLflow is an ML platform for managing machine learning lifecycle",
|
|
"key1": "value1",
|
|
}
|
|
assert df["outputs"].iloc[1] == {
|
|
"answer": "Apache Spark is a unified analytics engine for data processing"
|
|
}
|
|
|
|
# Check other columns have expected values
|
|
assert df["inputs"].iloc[0] == {"question": "What is MLflow?"}
|
|
assert df["inputs"].iloc[1] == {"question": "What is Spark?"}
|
|
assert df["expectations"].iloc[0] == {"answer": "MLflow is an ML platform"}
|
|
assert df["expectations"].iloc[1] == {"answer": "Spark is a data engine"}
|
|
assert df["tags"].iloc[0] == {"source": "manual"}
|
|
assert df["tags"].iloc[1] == {"source": "automated"}
|
|
assert df["source_type"].iloc[0] == "HUMAN"
|
|
assert df["source_type"].iloc[1] == "CODE"
|
|
assert df["source_id"].iloc[0] == "user123"
|
|
assert df["source_id"].iloc[1] == "script456"
|
|
assert df["dataset_record_id"].iloc[0] == "rec123"
|
|
assert df["dataset_record_id"].iloc[1] == "rec456"
|
|
|
|
|
|
def create_test_span(
|
|
span_id=1,
|
|
parent_id=None,
|
|
name="test_span",
|
|
inputs=None,
|
|
outputs=None,
|
|
span_type=SpanType.UNKNOWN,
|
|
):
|
|
attributes = {
|
|
"mlflow.spanType": json.dumps(span_type),
|
|
}
|
|
|
|
if inputs is not None:
|
|
attributes["mlflow.spanInputs"] = json.dumps(inputs)
|
|
|
|
if outputs is not None:
|
|
attributes["mlflow.spanOutputs"] = json.dumps(outputs)
|
|
|
|
otel_span = OTelReadableSpan(
|
|
name=name,
|
|
context=build_otel_context(trace_id=123456789, span_id=span_id),
|
|
parent=build_otel_context(trace_id=123456789, span_id=parent_id) if parent_id else None,
|
|
start_time=100000000,
|
|
end_time=200000000,
|
|
attributes=attributes,
|
|
)
|
|
return Span(otel_span)
|
|
|
|
|
|
def create_test_trace(
|
|
trace_id="test-trace-123",
|
|
inputs=None,
|
|
outputs=None,
|
|
expectations=None,
|
|
trace_metadata=None,
|
|
_no_defaults=False,
|
|
):
|
|
assessments = []
|
|
if expectations:
|
|
from mlflow.entities.assessment import AssessmentSource, AssessmentSourceType, Expectation
|
|
|
|
for name, value in expectations.items():
|
|
expectation = Expectation(
|
|
name=name,
|
|
value=value,
|
|
source=AssessmentSource(
|
|
source_type=AssessmentSourceType.HUMAN, source_id="test_user"
|
|
),
|
|
)
|
|
assessments.append(expectation)
|
|
|
|
trace_info = TraceInfo(
|
|
trace_id=trace_id,
|
|
trace_location=TraceLocation.from_experiment_id("0"),
|
|
request_time=1234567890,
|
|
execution_duration=1000,
|
|
state=TraceState.OK,
|
|
assessments=assessments,
|
|
trace_metadata=trace_metadata or {},
|
|
)
|
|
|
|
default_inputs = {"question": "What is MLflow?"}
|
|
default_outputs = {"answer": "MLflow is a platform"}
|
|
|
|
if _no_defaults:
|
|
span_inputs = inputs
|
|
span_outputs = outputs
|
|
else:
|
|
span_inputs = inputs if inputs is not None else default_inputs
|
|
span_outputs = outputs if outputs is not None else default_outputs
|
|
|
|
spans = [
|
|
create_test_span(
|
|
span_id=1,
|
|
parent_id=None,
|
|
name="root_span",
|
|
inputs=span_inputs,
|
|
outputs=span_outputs,
|
|
span_type=SpanType.CHAIN,
|
|
)
|
|
]
|
|
|
|
trace_data = TraceData(spans=spans)
|
|
return Trace(info=trace_info, data=trace_data)
|
|
|
|
|
|
def test_process_trace_records_with_dict_outputs():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
trace = create_test_trace(
|
|
trace_id="trace1",
|
|
inputs={"question": "What is MLflow?"},
|
|
outputs={"answer": "MLflow is a platform", "confidence": 0.95},
|
|
)
|
|
|
|
record_dicts = dataset._process_trace_records([trace])
|
|
|
|
assert len(record_dicts) == 1
|
|
assert record_dicts[0]["inputs"] == {"question": "What is MLflow?"}
|
|
assert record_dicts[0]["outputs"] == {"answer": "MLflow is a platform", "confidence": 0.95}
|
|
assert record_dicts[0]["expectations"] == {}
|
|
assert record_dicts[0]["source"]["source_type"] == DatasetRecordSourceType.TRACE.value
|
|
assert record_dicts[0]["source"]["source_data"]["trace_id"] == "trace1"
|
|
|
|
|
|
def test_process_trace_records_with_string_outputs():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
trace = create_test_trace(
|
|
trace_id="trace2",
|
|
inputs={"query": "Tell me about Python"},
|
|
outputs="Python is a programming language",
|
|
)
|
|
|
|
record_dicts = dataset._process_trace_records([trace])
|
|
|
|
assert len(record_dicts) == 1
|
|
assert record_dicts[0]["inputs"] == {"query": "Tell me about Python"}
|
|
assert record_dicts[0]["outputs"] == "Python is a programming language"
|
|
assert record_dicts[0]["expectations"] == {}
|
|
assert record_dicts[0]["source"]["source_type"] == DatasetRecordSourceType.TRACE.value
|
|
|
|
|
|
def test_process_trace_records_with_non_dict_non_string_outputs():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
trace = create_test_trace(
|
|
trace_id="trace3", inputs={"x": 1, "y": 2}, outputs=["result1", "result2", "result3"]
|
|
)
|
|
|
|
record_dicts = dataset._process_trace_records([trace])
|
|
|
|
assert len(record_dicts) == 1
|
|
assert record_dicts[0]["inputs"] == {"x": 1, "y": 2}
|
|
assert record_dicts[0]["outputs"] == ["result1", "result2", "result3"]
|
|
assert record_dicts[0]["source"]["source_type"] == DatasetRecordSourceType.TRACE.value
|
|
|
|
|
|
def test_process_trace_records_with_numeric_outputs():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
trace = create_test_trace(trace_id="trace4", inputs={"number": 42}, outputs=42)
|
|
|
|
record_dicts = dataset._process_trace_records([trace])
|
|
|
|
assert len(record_dicts) == 1
|
|
assert record_dicts[0]["outputs"] == 42
|
|
|
|
|
|
def test_process_trace_records_with_none_outputs():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
trace = create_test_trace(
|
|
trace_id="trace5", inputs={"input": "test"}, outputs=None, _no_defaults=True
|
|
)
|
|
|
|
record_dicts = dataset._process_trace_records([trace])
|
|
|
|
assert len(record_dicts) == 1
|
|
assert record_dicts[0]["outputs"] is None
|
|
|
|
|
|
def test_process_trace_records_with_expectations():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
trace = create_test_trace(
|
|
trace_id="trace6",
|
|
inputs={"question": "What is 2+2?"},
|
|
outputs={"answer": "4"},
|
|
expectations={"correctness": True, "tone": "neutral"},
|
|
)
|
|
|
|
record_dicts = dataset._process_trace_records([trace])
|
|
|
|
assert len(record_dicts) == 1
|
|
assert record_dicts[0]["expectations"] == {"correctness": True, "tone": "neutral"}
|
|
|
|
|
|
def test_process_trace_records_multiple_traces():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
traces = [
|
|
create_test_trace(trace_id="trace1", outputs={"result": "answer1"}),
|
|
create_test_trace(trace_id="trace2", outputs="string answer"),
|
|
create_test_trace(trace_id="trace3", outputs=[1, 2, 3]),
|
|
]
|
|
|
|
record_dicts = dataset._process_trace_records(traces)
|
|
|
|
assert len(record_dicts) == 3
|
|
assert record_dicts[0]["outputs"] == {"result": "answer1"}
|
|
assert record_dicts[1]["outputs"] == "string answer"
|
|
assert record_dicts[2]["outputs"] == [1, 2, 3]
|
|
|
|
|
|
def test_process_trace_records_mixed_types_error():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
trace = create_test_trace(trace_id="trace1")
|
|
not_a_trace = {"not": "a trace"}
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
"Mixed types in trace list.*Expected all elements to be Trace objects.*"
|
|
"element at index 1 is dict"
|
|
),
|
|
):
|
|
dataset._process_trace_records([trace, not_a_trace])
|
|
|
|
|
|
def test_process_trace_records_preserves_session_metadata():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
# Create trace with session metadata
|
|
trace_with_session = create_test_trace(
|
|
trace_id="tr-123",
|
|
trace_metadata={"mlflow.trace.session": "session_1"},
|
|
)
|
|
|
|
# Create trace without session metadata
|
|
trace_without_session = create_test_trace(
|
|
trace_id="tr-456",
|
|
trace_metadata={},
|
|
)
|
|
|
|
records = dataset._process_trace_records([trace_with_session, trace_without_session])
|
|
|
|
# Trace with session should have session_id in source_data
|
|
assert records[0]["source"]["source_data"]["trace_id"] == "tr-123"
|
|
assert records[0]["source"]["source_data"]["session_id"] == "session_1"
|
|
|
|
# Trace without session should only have trace_id
|
|
assert records[1]["source"]["source_data"]["trace_id"] == "tr-456"
|
|
assert "session_id" not in records[1]["source"]["source_data"]
|
|
|
|
|
|
def test_to_df_includes_source_column():
|
|
from mlflow.entities.dataset_record import DatasetRecord
|
|
from mlflow.entities.dataset_record_source import DatasetRecordSource
|
|
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
# Manually add a record with source to the dataset
|
|
source = DatasetRecordSource(
|
|
source_type=DatasetRecordSourceType.TRACE,
|
|
source_data={"trace_id": "tr-123"},
|
|
)
|
|
record = DatasetRecord(
|
|
dataset_id="dataset123",
|
|
dataset_record_id="record123",
|
|
inputs={"question": "test"},
|
|
outputs={"answer": "test answer"},
|
|
expectations={},
|
|
tags={},
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
source=source,
|
|
)
|
|
dataset._records = [record]
|
|
|
|
df = dataset.to_df()
|
|
|
|
assert "source" in df.columns
|
|
assert df["source"].notna().all()
|
|
assert df["source"].iloc[0] == source
|
|
|
|
|
|
def test_delete_records():
|
|
dataset = EvaluationDataset(
|
|
dataset_id="dataset123",
|
|
name="test_dataset",
|
|
digest="digest123",
|
|
created_time=123456789,
|
|
last_update_time=123456789,
|
|
)
|
|
|
|
# Add some records to cache
|
|
dataset._records = [Mock(), Mock()]
|
|
|
|
mock_store = Mock()
|
|
mock_store.delete_dataset_records.return_value = 2
|
|
|
|
with patch("mlflow.tracking._tracking_service.utils._get_store", return_value=mock_store):
|
|
deleted_count = dataset.delete_records(["record1", "record2"])
|
|
|
|
assert deleted_count == 2
|
|
mock_store.delete_dataset_records.assert_called_once_with(
|
|
dataset_id="dataset123",
|
|
dataset_record_ids=["record1", "record2"],
|
|
)
|
|
# Verify cache was cleared
|
|
assert dataset._records is None
|