wehub-resource-sync
198 lines
7.0 KiB
Python
198 lines
7.0 KiB
Python
import pytest
|
|
|
|
from mlflow.entities.trace_location import (
|
|
InferenceTableLocation,
|
|
MlflowExperimentLocation,
|
|
TraceLocation,
|
|
TraceLocationType,
|
|
UCSchemaLocation,
|
|
UnityCatalog,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos import databricks_tracing_pb2 as tracing_pb
|
|
from mlflow.protos import service_pb2 as pb
|
|
|
|
|
|
def test_trace_location():
|
|
trace_location = TraceLocation(
|
|
type=TraceLocationType.MLFLOW_EXPERIMENT,
|
|
mlflow_experiment=MlflowExperimentLocation(experiment_id="123"),
|
|
)
|
|
assert trace_location.type == TraceLocationType.MLFLOW_EXPERIMENT
|
|
assert trace_location.mlflow_experiment.experiment_id == "123"
|
|
|
|
trace_location = TraceLocation(
|
|
type=TraceLocationType.INFERENCE_TABLE,
|
|
inference_table=InferenceTableLocation(full_table_name="a.b.c"),
|
|
)
|
|
assert trace_location.type == TraceLocationType.INFERENCE_TABLE
|
|
assert trace_location.inference_table.full_table_name == "a.b.c"
|
|
|
|
from_proto = TraceLocation.from_proto(trace_location.to_proto())
|
|
assert from_proto == trace_location
|
|
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
"Only one of mlflow_experiment, inference_table, uc_schema, "
|
|
"or uc_table_prefix can be provided"
|
|
),
|
|
):
|
|
TraceLocation(
|
|
type=TraceLocationType.TRACE_LOCATION_TYPE_UNSPECIFIED,
|
|
mlflow_experiment=MlflowExperimentLocation(experiment_id="123"),
|
|
inference_table=InferenceTableLocation(full_table_name="a.b.c"),
|
|
uc_schema=UCSchemaLocation(catalog_name="a", schema_name="b"),
|
|
)
|
|
|
|
|
|
def test_trace_location_mismatch():
|
|
with pytest.raises(
|
|
MlflowException, match="Trace location .+ does not match the provided location"
|
|
):
|
|
TraceLocation(
|
|
type=TraceLocationType.INFERENCE_TABLE,
|
|
mlflow_experiment=MlflowExperimentLocation(experiment_id="123"),
|
|
)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match="Trace location .+ does not match the provided location"
|
|
):
|
|
TraceLocation(
|
|
type=TraceLocationType.MLFLOW_EXPERIMENT,
|
|
inference_table=InferenceTableLocation(full_table_name="a.b.c"),
|
|
)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match="Trace location .+ does not match the provided location"
|
|
):
|
|
TraceLocation(
|
|
type=TraceLocationType.INFERENCE_TABLE,
|
|
uc_schema=UCSchemaLocation(catalog_name="a", schema_name="b"),
|
|
)
|
|
|
|
with pytest.raises(
|
|
MlflowException, match="Trace location .+ does not match the provided location"
|
|
):
|
|
TraceLocation(
|
|
type=TraceLocationType.INFERENCE_TABLE,
|
|
uc_table_prefix=UnityCatalog(catalog_name="a", schema_name="b", table_prefix="p"),
|
|
)
|
|
|
|
|
|
def test_trace_location_from_v4_proto_mlflow_experiment():
|
|
proto = pb.TraceLocation(
|
|
type=pb.TraceLocation.TraceLocationType.MLFLOW_EXPERIMENT,
|
|
mlflow_experiment=pb.TraceLocation.MlflowExperimentLocation(experiment_id="1234"),
|
|
)
|
|
trace_location = TraceLocation.from_proto(proto)
|
|
assert trace_location.type == TraceLocationType.MLFLOW_EXPERIMENT
|
|
assert trace_location.mlflow_experiment.experiment_id == "1234"
|
|
|
|
|
|
def test_trace_location_from_v4_proto_inference_table():
|
|
proto = pb.TraceLocation(
|
|
type=pb.TraceLocation.TraceLocationType.INFERENCE_TABLE,
|
|
inference_table=pb.TraceLocation.InferenceTableLocation(
|
|
full_table_name="test_catalog.test_schema.test_table"
|
|
),
|
|
)
|
|
trace_location = TraceLocation.from_proto(proto)
|
|
assert trace_location.type == TraceLocationType.INFERENCE_TABLE
|
|
assert trace_location.inference_table.full_table_name == "test_catalog.test_schema.test_table"
|
|
|
|
|
|
def test_uc_schema_location_full_otel_spans_table_name():
|
|
uc_schema = UCSchemaLocation(
|
|
catalog_name="test_catalog",
|
|
schema_name="test_schema",
|
|
)
|
|
uc_schema._otel_spans_table_name = "otel_spans"
|
|
assert uc_schema.full_otel_spans_table_name == "test_catalog.test_schema.otel_spans"
|
|
|
|
|
|
def test_uc_schema_location_round_trip():
|
|
uc_schema = UCSchemaLocation(
|
|
catalog_name="test_catalog",
|
|
schema_name="test_schema",
|
|
)
|
|
assert uc_schema.schema_location == "test_catalog.test_schema"
|
|
assert UCSchemaLocation.from_dict(uc_schema.to_dict()) == uc_schema
|
|
|
|
|
|
def test_unity_catalog_default_table_prefix():
|
|
location = UnityCatalog(catalog_name="catalog", schema_name="schema")
|
|
assert location.table_prefix is None
|
|
with pytest.raises(MlflowException, match="table_prefix is required"):
|
|
location.full_table_prefix
|
|
|
|
|
|
def test_unity_catalog_from_dict_without_table_prefix():
|
|
d = {"catalog_name": "catalog", "schema_name": "schema"}
|
|
location = UnityCatalog.from_dict(d)
|
|
assert location.table_prefix is None
|
|
assert location.catalog_name == "catalog"
|
|
assert location.schema_name == "schema"
|
|
|
|
|
|
def test_unity_catalog_factory_for_table_prefix():
|
|
location = UnityCatalog(catalog_name="catalog", schema_name="schema", table_prefix="pref")
|
|
assert isinstance(location, UnityCatalog)
|
|
assert location.full_table_prefix == "catalog.schema.pref"
|
|
|
|
|
|
def test_uc_table_prefix_location_round_trip():
|
|
location = UnityCatalog(
|
|
catalog_name="catalog",
|
|
schema_name="schema",
|
|
table_prefix="prefix",
|
|
)
|
|
assert location.full_table_prefix == "catalog.schema.prefix"
|
|
assert UnityCatalog.from_dict(location.to_dict()) == location
|
|
|
|
|
|
def test_unity_catalog_equality_ignores_private_fields():
|
|
a = UnityCatalog(catalog_name="cat", schema_name="sch", table_prefix="pfx")
|
|
b = UnityCatalog(catalog_name="cat", schema_name="sch", table_prefix="pfx")
|
|
b._otel_spans_table_name = "cat.sch.pfx_otel_spans"
|
|
b._otel_logs_table_name = "cat.sch.pfx_otel_logs"
|
|
assert a == b
|
|
|
|
|
|
def test_unity_catalog_inequality_on_different_prefix():
|
|
a = UnityCatalog(catalog_name="cat", schema_name="sch", table_prefix="pfx1")
|
|
b = UnityCatalog(catalog_name="cat", schema_name="sch", table_prefix="pfx2")
|
|
assert a != b
|
|
|
|
|
|
def test_uc_schema_location_from_proto():
|
|
proto = tracing_pb.UCSchemaLocation(
|
|
catalog_name="c",
|
|
schema_name="s",
|
|
otel_spans_table_name="spans",
|
|
otel_logs_table_name="logs",
|
|
)
|
|
location = UCSchemaLocation.from_proto(proto)
|
|
assert location.catalog_name == "c"
|
|
assert location.schema_name == "s"
|
|
assert location._otel_spans_table_name == "spans"
|
|
assert location._otel_logs_table_name == "logs"
|
|
|
|
|
|
def test_unity_catalog_from_proto():
|
|
proto = tracing_pb.UcTablePrefixLocation(
|
|
catalog_name="c",
|
|
schema_name="s",
|
|
table_prefix="p",
|
|
spans_table_name="spans",
|
|
logs_table_name="logs",
|
|
annotations_table_name="anns",
|
|
)
|
|
location = UnityCatalog.from_proto(proto)
|
|
assert location.catalog_name == "c"
|
|
assert location.schema_name == "s"
|
|
assert location.table_prefix == "p"
|
|
assert location._otel_spans_table_name == "spans"
|
|
assert location._otel_logs_table_name == "logs"
|
|
assert location._annotations_table_name == "anns"
|