wehub-resource-sync
397 lines
14 KiB
Python
397 lines
14 KiB
Python
from abc import ABC, abstractmethod
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Any
|
|
|
|
from mlflow.entities._mlflow_object import _MlflowObject
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos import service_pb2 as pb
|
|
from mlflow.utils.annotations import deprecated, experimental
|
|
|
|
_UC_SCHEMA_DEFAULT_SPANS_TABLE_NAME = "mlflow_experiment_trace_otel_spans"
|
|
_UC_SCHEMA_DEFAULT_LOGS_TABLE_NAME = "mlflow_experiment_trace_otel_logs"
|
|
|
|
|
|
@dataclass
|
|
class TraceLocationBase(_MlflowObject, ABC):
|
|
"""
|
|
Base class for trace location classes.
|
|
"""
|
|
|
|
@abstractmethod
|
|
def to_dict(self) -> dict[str, Any]: ...
|
|
|
|
@classmethod
|
|
@abstractmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "TraceLocationBase": ...
|
|
|
|
|
|
@dataclass
|
|
class MlflowExperimentLocation(TraceLocationBase):
|
|
"""
|
|
Represents the location of an MLflow experiment.
|
|
|
|
Args:
|
|
experiment_id: The ID of the MLflow experiment where the trace is stored.
|
|
"""
|
|
|
|
experiment_id: str
|
|
|
|
def to_proto(self):
|
|
return pb.TraceLocation.MlflowExperimentLocation(experiment_id=self.experiment_id)
|
|
|
|
@classmethod
|
|
def from_proto(cls, proto) -> "MlflowExperimentLocation":
|
|
return cls(experiment_id=proto.experiment_id)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"experiment_id": self.experiment_id}
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "MlflowExperimentLocation":
|
|
return cls(experiment_id=d["experiment_id"])
|
|
|
|
|
|
@deprecated(since="3.7.0")
|
|
@dataclass
|
|
class InferenceTableLocation(TraceLocationBase):
|
|
"""
|
|
Represents the location of a Databricks inference table.
|
|
|
|
Args:
|
|
full_table_name: The fully qualified name of the inference table where
|
|
the trace is stored, in the format of `<catalog>.<schema>.<table>`.
|
|
"""
|
|
|
|
full_table_name: str
|
|
|
|
def to_proto(self):
|
|
return pb.TraceLocation.InferenceTableLocation(full_table_name=self.full_table_name)
|
|
|
|
@classmethod
|
|
def from_proto(cls, proto) -> "InferenceTableLocation":
|
|
return cls(full_table_name=proto.full_table_name)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
return {"full_table_name": self.full_table_name}
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "InferenceTableLocation":
|
|
return cls(full_table_name=d["full_table_name"])
|
|
|
|
|
|
@dataclass
|
|
class UCSchemaLocation(TraceLocationBase):
|
|
"""
|
|
Represents the location of a Databricks Unity Catalog (UC) schema.
|
|
|
|
Args:
|
|
catalog_name: The name of the Unity Catalog catalog name.
|
|
schema_name: The name of the Unity Catalog schema.
|
|
"""
|
|
|
|
catalog_name: str
|
|
schema_name: str
|
|
|
|
# These table names are set by the backend
|
|
_otel_spans_table_name: str | None = _UC_SCHEMA_DEFAULT_SPANS_TABLE_NAME
|
|
_otel_logs_table_name: str | None = _UC_SCHEMA_DEFAULT_LOGS_TABLE_NAME
|
|
|
|
@property
|
|
def schema_location(self) -> str:
|
|
return f"{self.catalog_name}.{self.schema_name}"
|
|
|
|
@property
|
|
def full_otel_spans_table_name(self) -> str | None:
|
|
if self._otel_spans_table_name:
|
|
return f"{self.catalog_name}.{self.schema_name}.{self._otel_spans_table_name}"
|
|
|
|
@property
|
|
def full_otel_logs_table_name(self) -> str | None:
|
|
if self._otel_logs_table_name:
|
|
return f"{self.catalog_name}.{self.schema_name}.{self._otel_logs_table_name}"
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
d = {
|
|
"catalog_name": self.catalog_name,
|
|
"schema_name": self.schema_name,
|
|
}
|
|
if self._otel_spans_table_name:
|
|
d["otel_spans_table_name"] = self._otel_spans_table_name
|
|
if self._otel_logs_table_name:
|
|
d["otel_logs_table_name"] = self._otel_logs_table_name
|
|
return d
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "UCSchemaLocation":
|
|
location = cls(catalog_name=d["catalog_name"], schema_name=d["schema_name"])
|
|
if otel_spans_table_name := d.get("otel_spans_table_name"):
|
|
location._otel_spans_table_name = otel_spans_table_name
|
|
if otel_logs_table_name := d.get("otel_logs_table_name"):
|
|
location._otel_logs_table_name = otel_logs_table_name
|
|
return location
|
|
|
|
@classmethod
|
|
def from_proto(cls, proto) -> "UCSchemaLocation":
|
|
from mlflow.utils.databricks_tracing_utils import uc_schema_location_from_proto
|
|
|
|
return uc_schema_location_from_proto(proto)
|
|
|
|
|
|
@experimental(version="3.11.0")
|
|
@dataclass
|
|
class UnityCatalog(TraceLocationBase):
|
|
"""
|
|
Represents a Databricks Unity Catalog location with a table prefix.
|
|
|
|
Note: Arclight catalogs are not supported.
|
|
|
|
Args:
|
|
catalog_name: The name of the Unity Catalog catalog.
|
|
schema_name: The name of the Unity Catalog schema.
|
|
table_prefix: The prefix for tables in this location.
|
|
"""
|
|
|
|
catalog_name: str
|
|
schema_name: str
|
|
table_prefix: str | None = None
|
|
|
|
# These are fully qualified table names (catalog.schema.table) set by the backend.
|
|
_otel_spans_table_name: str | None = None
|
|
_otel_logs_table_name: str | None = None
|
|
_annotations_table_name: str | None = None
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
if not isinstance(other, UnityCatalog):
|
|
return NotImplemented
|
|
return (
|
|
self.catalog_name == other.catalog_name
|
|
and self.schema_name == other.schema_name
|
|
and self.table_prefix == other.table_prefix
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return (
|
|
f"UnityCatalog(catalog_name={self.catalog_name!r}, "
|
|
f"schema_name={self.schema_name!r}, "
|
|
f"table_prefix={self.table_prefix!r})"
|
|
)
|
|
|
|
@property
|
|
def schema_location(self) -> str:
|
|
return f"{self.catalog_name}.{self.schema_name}"
|
|
|
|
@property
|
|
def full_table_prefix(self) -> str:
|
|
if self.table_prefix is None:
|
|
raise MlflowException.invalid_parameter_value(
|
|
"table_prefix is required but was not set."
|
|
)
|
|
return f"{self.catalog_name}.{self.schema_name}.{self.table_prefix}"
|
|
|
|
@property
|
|
def full_otel_spans_table_name(self) -> str | None:
|
|
return self._otel_spans_table_name
|
|
|
|
@property
|
|
def full_otel_logs_table_name(self) -> str | None:
|
|
return self._otel_logs_table_name
|
|
|
|
@property
|
|
def full_annotations_table_name(self) -> str | None:
|
|
return self._annotations_table_name
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
d = {
|
|
"catalog_name": self.catalog_name,
|
|
"schema_name": self.schema_name,
|
|
}
|
|
if self.table_prefix is not None:
|
|
d["table_prefix"] = self.table_prefix
|
|
if self._otel_spans_table_name:
|
|
d["otel_spans_table_name"] = self._otel_spans_table_name
|
|
if self._otel_logs_table_name:
|
|
d["otel_logs_table_name"] = self._otel_logs_table_name
|
|
if self._annotations_table_name:
|
|
d["annotations_table_name"] = self._annotations_table_name
|
|
return d
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "UnityCatalog":
|
|
location = cls(
|
|
catalog_name=d["catalog_name"],
|
|
schema_name=d["schema_name"],
|
|
table_prefix=d.get("table_prefix"),
|
|
)
|
|
if otel_spans_table_name := d.get("otel_spans_table_name"):
|
|
location._otel_spans_table_name = otel_spans_table_name
|
|
if otel_logs_table_name := d.get("otel_logs_table_name"):
|
|
location._otel_logs_table_name = otel_logs_table_name
|
|
if annotations_table_name := d.get("annotations_table_name"):
|
|
location._annotations_table_name = annotations_table_name
|
|
return location
|
|
|
|
@classmethod
|
|
def from_proto(cls, proto) -> "UnityCatalog":
|
|
from mlflow.utils.databricks_tracing_utils import uc_table_prefix_location_from_proto
|
|
|
|
return uc_table_prefix_location_from_proto(proto)
|
|
|
|
|
|
class TraceLocationType(str, Enum):
|
|
TRACE_LOCATION_TYPE_UNSPECIFIED = "TRACE_LOCATION_TYPE_UNSPECIFIED"
|
|
MLFLOW_EXPERIMENT = "MLFLOW_EXPERIMENT"
|
|
INFERENCE_TABLE = "INFERENCE_TABLE"
|
|
UC_SCHEMA = "UC_SCHEMA"
|
|
UC_TABLE_PREFIX = "UC_TABLE_PREFIX"
|
|
|
|
def to_proto(self):
|
|
return pb.TraceLocation.TraceLocationType.Value(self)
|
|
|
|
@classmethod
|
|
def from_proto(cls, proto: int) -> "TraceLocationType":
|
|
return TraceLocationType(pb.TraceLocation.TraceLocationType.Name(proto))
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "TraceLocationType":
|
|
return cls(d["type"])
|
|
|
|
|
|
@dataclass
|
|
class TraceLocation(_MlflowObject):
|
|
"""
|
|
Represents the location where the trace is stored.
|
|
|
|
Currently, MLflow supports two types of trace locations:
|
|
|
|
- MLflow experiment: The trace is stored in an MLflow experiment.
|
|
- Inference table: The trace is stored in a Databricks inference table.
|
|
|
|
Args:
|
|
type: The type of the trace location, should be one of the
|
|
:py:class:`TraceLocationType` enum values.
|
|
mlflow_experiment: The MLflow experiment location. Set this when the
|
|
location type is MLflow experiment.
|
|
inference_table: The inference table location. Set this when the
|
|
location type is Databricks Inference table.
|
|
"""
|
|
|
|
type: TraceLocationType
|
|
mlflow_experiment: MlflowExperimentLocation | None = None
|
|
inference_table: InferenceTableLocation | None = None
|
|
uc_schema: UCSchemaLocation | None = None
|
|
uc_table_prefix: UnityCatalog | None = None
|
|
|
|
def __post_init__(self) -> None:
|
|
if (
|
|
sum([
|
|
self.mlflow_experiment is not None,
|
|
self.inference_table is not None,
|
|
self.uc_schema is not None,
|
|
self.uc_table_prefix is not None,
|
|
])
|
|
> 1
|
|
):
|
|
raise MlflowException.invalid_parameter_value(
|
|
"Only one of mlflow_experiment, inference_table, uc_schema, "
|
|
"or uc_table_prefix can be provided."
|
|
)
|
|
|
|
if (
|
|
(self.mlflow_experiment and self.type != TraceLocationType.MLFLOW_EXPERIMENT)
|
|
or (self.inference_table and self.type != TraceLocationType.INFERENCE_TABLE)
|
|
or (self.uc_schema and self.type != TraceLocationType.UC_SCHEMA)
|
|
or (self.uc_table_prefix and self.type != TraceLocationType.UC_TABLE_PREFIX)
|
|
):
|
|
location = (
|
|
self.mlflow_experiment
|
|
or self.inference_table
|
|
or self.uc_schema
|
|
or self.uc_table_prefix
|
|
)
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Trace location type {self.type} does not match the provided location {location}."
|
|
)
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
d = {"type": self.type.value}
|
|
if self.mlflow_experiment:
|
|
d["mlflow_experiment"] = self.mlflow_experiment.to_dict()
|
|
elif self.inference_table:
|
|
d["inference_table"] = self.inference_table.to_dict()
|
|
elif self.uc_schema:
|
|
d["uc_schema"] = self.uc_schema.to_dict()
|
|
elif self.uc_table_prefix:
|
|
d["uc_table_prefix"] = self.uc_table_prefix.to_dict()
|
|
return d
|
|
|
|
@classmethod
|
|
def from_dict(cls, d: dict[str, Any]) -> "TraceLocation":
|
|
return cls(
|
|
type=TraceLocationType(d["type"]),
|
|
mlflow_experiment=(
|
|
MlflowExperimentLocation.from_dict(v) if (v := d.get("mlflow_experiment")) else None
|
|
),
|
|
inference_table=(
|
|
InferenceTableLocation.from_dict(v) if (v := d.get("inference_table")) else None
|
|
),
|
|
uc_schema=(UCSchemaLocation.from_dict(v) if (v := d.get("uc_schema")) else None),
|
|
uc_table_prefix=(
|
|
UnityCatalog.from_dict(v) if (v := d.get("uc_table_prefix")) else None
|
|
),
|
|
)
|
|
|
|
def to_proto(self) -> pb.TraceLocation:
|
|
if self.mlflow_experiment:
|
|
return pb.TraceLocation(
|
|
type=self.type.to_proto(),
|
|
mlflow_experiment=self.mlflow_experiment.to_proto(),
|
|
)
|
|
elif self.inference_table:
|
|
return pb.TraceLocation(
|
|
type=self.type.to_proto(),
|
|
inference_table=self.inference_table.to_proto(),
|
|
)
|
|
elif self.uc_table_prefix:
|
|
return pb.TraceLocation(type=self.type.to_proto())
|
|
# uc schema is not supported in to_proto since it's databricks specific, should use
|
|
# databricks_service_utils to convert to proto
|
|
else:
|
|
return pb.TraceLocation(type=self.type.to_proto())
|
|
|
|
@classmethod
|
|
def from_proto(cls, proto) -> "TraceLocation":
|
|
from mlflow.utils.databricks_tracing_utils import trace_location_from_proto
|
|
|
|
return trace_location_from_proto(proto)
|
|
|
|
@classmethod
|
|
def from_experiment_id(cls, experiment_id: str) -> "TraceLocation":
|
|
return cls(
|
|
type=TraceLocationType.MLFLOW_EXPERIMENT,
|
|
mlflow_experiment=MlflowExperimentLocation(experiment_id=experiment_id),
|
|
)
|
|
|
|
@classmethod
|
|
def from_databricks_uc_schema(cls, catalog_name: str, schema_name: str) -> "TraceLocation":
|
|
return cls(
|
|
type=TraceLocationType.UC_SCHEMA,
|
|
uc_schema=UCSchemaLocation(
|
|
catalog_name=catalog_name,
|
|
schema_name=schema_name,
|
|
),
|
|
)
|
|
|
|
@classmethod
|
|
def from_databricks_uc_table_prefix(
|
|
cls, catalog_name: str, schema_name: str, table_prefix: str
|
|
) -> "TraceLocation":
|
|
return cls(
|
|
type=TraceLocationType.UC_TABLE_PREFIX,
|
|
uc_table_prefix=UnityCatalog(
|
|
catalog_name=catalog_name,
|
|
schema_name=schema_name,
|
|
table_prefix=table_prefix,
|
|
),
|
|
)
|