134 lines
4.8 KiB
Python
134 lines
4.8 KiB
Python
import pytest
|
|
|
|
from mlflow.entities.trace_metrics import (
|
|
AggregationType,
|
|
MetricAggregation,
|
|
MetricDataPoint,
|
|
MetricViewType,
|
|
)
|
|
from mlflow.protos import service_pb2 as pb
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("view_type", "expected_proto"),
|
|
zip(MetricViewType, pb.MetricViewType.values(), strict=True),
|
|
)
|
|
def test_trace_metrics_view_type(view_type: MetricViewType, expected_proto: pb.MetricViewType):
|
|
assert view_type.to_proto() == expected_proto
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("aggregation_type", "expected_proto"),
|
|
zip(AggregationType, pb.AggregationType.values(), strict=True),
|
|
)
|
|
def test_trace_metrics_aggregation_type_to_proto(
|
|
aggregation_type: AggregationType, expected_proto: pb.AggregationType
|
|
):
|
|
assert aggregation_type.to_proto() == expected_proto
|
|
|
|
|
|
def test_metrics_aggregation_to_proto_without_percentile():
|
|
aggregation = MetricAggregation(aggregation_type=AggregationType.AVG)
|
|
proto = aggregation.to_proto()
|
|
assert proto.aggregation_type == pb.AggregationType.AVG
|
|
assert not proto.HasField("percentile_value")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("percentile_value"),
|
|
[50.0, 75.0, 90.0, 95.0, 99.0, 99.9],
|
|
)
|
|
def test_metrics_aggregation_percentile_values(percentile_value: float):
|
|
aggregation = MetricAggregation(
|
|
aggregation_type=AggregationType.PERCENTILE, percentile_value=percentile_value
|
|
)
|
|
proto = aggregation.to_proto()
|
|
assert proto.percentile_value == percentile_value
|
|
|
|
|
|
def test_metrics_aggregation_percentile_requires_value():
|
|
with pytest.raises(ValueError, match="Percentile value is required for PERCENTILE aggregation"):
|
|
MetricAggregation(aggregation_type=AggregationType.PERCENTILE)
|
|
|
|
|
|
@pytest.mark.parametrize("percentile_value", [-1.0, -0.1, 100.1, 101.0, 1000.0])
|
|
def test_metrics_aggregation_percentile_value_out_of_range(percentile_value: float):
|
|
with pytest.raises(ValueError, match="Percentile value must be between 0 and 100"):
|
|
MetricAggregation(
|
|
aggregation_type=AggregationType.PERCENTILE, percentile_value=percentile_value
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("percentile_value", [0.0, 0.1, 50.0, 99.9, 100.0])
|
|
def test_metrics_aggregation_percentile_value_valid_range(percentile_value: float):
|
|
aggregation = MetricAggregation(
|
|
aggregation_type=AggregationType.PERCENTILE, percentile_value=percentile_value
|
|
)
|
|
assert aggregation.percentile_value == percentile_value
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"agg_type",
|
|
[t for t in AggregationType if t is not AggregationType.PERCENTILE],
|
|
)
|
|
def test_metrics_aggregation_non_percentile_with_value_raises(agg_type: AggregationType):
|
|
with pytest.raises(
|
|
ValueError, match="Percentile value is only allowed for PERCENTILE aggregation"
|
|
):
|
|
MetricAggregation(aggregation_type=agg_type, percentile_value=50.0)
|
|
|
|
|
|
def test_trace_metrics_metric_data_point_from_proto():
|
|
metric_data_point_proto = pb.MetricDataPoint(
|
|
metric_name="latency",
|
|
dimensions={"status": "OK"},
|
|
values={"avg": 150.5, "p99": 200},
|
|
)
|
|
assert MetricDataPoint.from_proto(metric_data_point_proto) == MetricDataPoint(
|
|
metric_name="latency",
|
|
dimensions={"status": "OK"},
|
|
values={"avg": 150.5, "p99": 200},
|
|
)
|
|
|
|
|
|
def test_trace_metrics_metric_data_point_to_proto():
|
|
metric_data_point = MetricDataPoint(
|
|
metric_name="latency",
|
|
dimensions={"status": "OK", "model": "gpt-4"},
|
|
values={"avg": 150.5, "p99": 200.0},
|
|
)
|
|
proto = metric_data_point.to_proto()
|
|
assert proto.metric_name == "latency"
|
|
assert dict(proto.dimensions) == {"status": "OK", "model": "gpt-4"}
|
|
assert dict(proto.values) == {"avg": 150.5, "p99": 200.0}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("view_type", "expected_proto"),
|
|
zip(MetricViewType, pb.MetricViewType.values(), strict=True),
|
|
)
|
|
def test_trace_metrics_view_type_from_proto(view_type: MetricViewType, expected_proto: int):
|
|
assert MetricViewType.from_proto(expected_proto) == view_type
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"agg_type",
|
|
[t for t in AggregationType if t is not AggregationType.PERCENTILE],
|
|
)
|
|
def test_metrics_aggregation_from_proto_without_percentile(agg_type: AggregationType):
|
|
proto = pb.MetricAggregation(aggregation_type=agg_type.to_proto())
|
|
aggregation = MetricAggregation.from_proto(proto)
|
|
assert aggregation.aggregation_type == agg_type
|
|
assert aggregation.percentile_value is None
|
|
|
|
|
|
@pytest.mark.parametrize("percentile_value", [50.0, 75.0, 90.0, 95.0, 99.0, 99.9])
|
|
def test_metrics_aggregation_from_proto_with_percentile(percentile_value: float):
|
|
proto = pb.MetricAggregation(
|
|
aggregation_type=pb.AggregationType.PERCENTILE,
|
|
percentile_value=percentile_value,
|
|
)
|
|
aggregation = MetricAggregation.from_proto(proto)
|
|
assert aggregation.aggregation_type == AggregationType.PERCENTILE
|
|
assert aggregation.percentile_value == percentile_value
|