wehub-resource-sync
286 lines
10 KiB
Python
286 lines
10 KiB
Python
import base64
|
|
import gzip
|
|
import os
|
|
import zlib
|
|
from typing import Any
|
|
|
|
from opentelemetry.proto.common.v1.common_pb2 import AnyValue, ArrayValue, KeyValueList
|
|
from opentelemetry.proto.resource.v1.resource_pb2 import Resource as OTelProtoResource
|
|
from opentelemetry.sdk.resources import Resource as OTelResource
|
|
from opentelemetry.sdk.trace.export import SpanExporter
|
|
|
|
from mlflow.environment_variables import MLFLOW_ENABLE_OTLP_EXPORTER, MLFLOW_TRACKING_TOKEN
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
|
|
from mlflow.utils.credentials import read_mlflow_creds
|
|
|
|
# Constants for OpenTelemetry integration
|
|
MLFLOW_EXPERIMENT_ID_HEADER = "x-mlflow-experiment-id"
|
|
MLFLOW_RUN_ID_HEADER = "x-mlflow-run-id"
|
|
OTLP_TRACES_PATH = "/v1/traces"
|
|
OTLP_METRICS_PATH = "/v1/metrics"
|
|
|
|
|
|
def build_otlp_headers(experiment_id: str) -> dict[str, str]:
|
|
"""Build OTLP exporter headers with experiment ID and auth credentials."""
|
|
headers: dict[str, str] = {MLFLOW_EXPERIMENT_ID_HEADER: experiment_id}
|
|
|
|
creds = read_mlflow_creds()
|
|
if creds.username and creds.password:
|
|
basic_auth_str = f"{creds.username}:{creds.password}".encode()
|
|
headers["Authorization"] = f"Basic {base64.standard_b64encode(basic_auth_str).decode()}"
|
|
elif token := MLFLOW_TRACKING_TOKEN.get():
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
|
|
return headers
|
|
|
|
|
|
def should_use_otlp_exporter() -> bool:
|
|
"""
|
|
Determine if OTLP traces should be exported based on environment configuration.
|
|
"""
|
|
return _get_otlp_traces_endpoint() is not None and MLFLOW_ENABLE_OTLP_EXPORTER.get()
|
|
|
|
|
|
def should_export_otlp_metrics() -> bool:
|
|
"""
|
|
Determine if OTLP metrics should be exported based on environment configuration.
|
|
|
|
Returns True if metrics endpoint is configured.
|
|
"""
|
|
return _get_otlp_metrics_endpoint() is not None
|
|
|
|
|
|
def get_otlp_exporter() -> SpanExporter:
|
|
"""
|
|
Get the OTLP exporter based on the configured protocol.
|
|
"""
|
|
endpoint = _get_otlp_traces_endpoint()
|
|
protocol = _get_otlp_protocol()
|
|
if protocol == "grpc":
|
|
try:
|
|
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
|
|
except ImportError:
|
|
raise MlflowException(
|
|
"gRPC OTLP exporter is not available. Please install the required dependency by "
|
|
"running `pip install opentelemetry-exporter-otlp-proto-grpc`.",
|
|
error_code=RESOURCE_DOES_NOT_EXIST,
|
|
)
|
|
|
|
return OTLPSpanExporter(endpoint=endpoint)
|
|
elif protocol == "http/protobuf":
|
|
try:
|
|
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
|
|
except ImportError as e:
|
|
raise MlflowException(
|
|
"HTTP OTLP exporter is not available. Please install the required dependency by "
|
|
"running `pip install opentelemetry-exporter-otlp-proto-http`.",
|
|
error_code=RESOURCE_DOES_NOT_EXIST,
|
|
) from e
|
|
|
|
return OTLPSpanExporter(endpoint=endpoint)
|
|
else:
|
|
raise MlflowException.invalid_parameter_value(
|
|
f"Unsupported OTLP protocol '{protocol}' is configured. Please set "
|
|
"the protocol to either 'grpc' or 'http/protobuf'."
|
|
)
|
|
|
|
|
|
def _get_otlp_traces_endpoint() -> str | None:
|
|
"""
|
|
Get the OTLP endpoint from the environment variables.
|
|
Ref: https://opentelemetry.io/docs/languages/sdk-configuration/otlp-exporter/#endpoint-configuration
|
|
|
|
Per the OTel spec:
|
|
- OTEL_EXPORTER_OTLP_TRACES_ENDPOINT: Full URL used as-is
|
|
- OTEL_EXPORTER_OTLP_ENDPOINT: Base URL, requires appending signal path
|
|
"""
|
|
if traces_endpoint := os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"):
|
|
return traces_endpoint
|
|
|
|
if base_endpoint := os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"):
|
|
return base_endpoint.rstrip("/") + OTLP_TRACES_PATH
|
|
|
|
return None
|
|
|
|
|
|
def _get_otlp_metrics_endpoint() -> str | None:
|
|
"""
|
|
Get the OTLP metrics endpoint from the environment variables.
|
|
|
|
Per the OTel spec:
|
|
- OTEL_EXPORTER_OTLP_METRICS_ENDPOINT: Full URL used as-is
|
|
- OTEL_EXPORTER_OTLP_ENDPOINT: Base URL, requires appending signal path
|
|
"""
|
|
if metrics_endpoint := os.environ.get("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT"):
|
|
return metrics_endpoint
|
|
|
|
if base_endpoint := os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"):
|
|
return base_endpoint.rstrip("/") + OTLP_METRICS_PATH
|
|
|
|
return None
|
|
|
|
|
|
def _get_otlp_protocol(default_value: str = "grpc") -> str:
|
|
"""
|
|
Get the OTLP traces protocol from environment variables.
|
|
|
|
Returns the value of OTEL_EXPORTER_OTLP_TRACES_PROTOCOL if set,
|
|
otherwise falls back to OTEL_EXPORTER_OTLP_PROTOCOL, then to default_value.
|
|
|
|
Args:
|
|
default_value: The default protocol to use if no environment variables are set.
|
|
"""
|
|
return os.environ.get("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL") or os.environ.get(
|
|
"OTEL_EXPORTER_OTLP_PROTOCOL", default_value
|
|
)
|
|
|
|
|
|
def _get_otlp_metrics_protocol(default_value: str = "grpc") -> str:
|
|
"""
|
|
Get the OTLP metrics protocol from environment variables.
|
|
|
|
Returns the value of OTEL_EXPORTER_OTLP_METRICS_PROTOCOL if set,
|
|
otherwise falls back to OTEL_EXPORTER_OTLP_PROTOCOL, then to default_value.
|
|
|
|
Args:
|
|
default_value: The default protocol to use if no environment variables are set.
|
|
"""
|
|
return os.environ.get("OTEL_EXPORTER_OTLP_METRICS_PROTOCOL") or os.environ.get(
|
|
"OTEL_EXPORTER_OTLP_PROTOCOL", default_value
|
|
)
|
|
|
|
|
|
def _otel_proto_bytes_to_id(id_bytes: bytes) -> int:
|
|
"""Convert OTel protobuf bytes to integer ID."""
|
|
return int.from_bytes(id_bytes, byteorder="big", signed=False)
|
|
|
|
|
|
def _sanitize_otel_string(value: str) -> str:
|
|
"""Replace characters that cannot be encoded as UTF-8 for protobuf string assignment."""
|
|
return value.encode("utf-8", errors="replace").decode("utf-8")
|
|
|
|
|
|
def _set_otel_proto_anyvalue(pb_any_value: AnyValue, value: Any) -> None:
|
|
"""Set a value on an OTel protobuf AnyValue message.
|
|
|
|
Args:
|
|
pb_any_value: The OTel protobuf AnyValue message to populate.
|
|
value: The value to set.
|
|
"""
|
|
if value is None:
|
|
# Leave the value unset for None
|
|
pass
|
|
elif isinstance(value, bool):
|
|
pb_any_value.bool_value = value
|
|
elif isinstance(value, str):
|
|
pb_any_value.string_value = _sanitize_otel_string(value)
|
|
elif isinstance(value, int):
|
|
pb_any_value.int_value = value
|
|
elif isinstance(value, float):
|
|
pb_any_value.double_value = value
|
|
elif isinstance(value, bytes):
|
|
pb_any_value.bytes_value = value
|
|
elif isinstance(value, (list, tuple)):
|
|
# Explicitly set array_value using CopyFrom to ensure the field is set even for empty lists
|
|
array_value = ArrayValue()
|
|
for item in value:
|
|
_set_otel_proto_anyvalue(array_value.values.add(), item)
|
|
pb_any_value.array_value.CopyFrom(array_value)
|
|
elif isinstance(value, dict):
|
|
# Explicitly set kvlist_value using CopyFrom to ensure the field is set even for empty dicts
|
|
kvlist_value = KeyValueList()
|
|
for k, v in value.items():
|
|
kv = kvlist_value.values.add()
|
|
kv.key = _sanitize_otel_string(str(k))
|
|
_set_otel_proto_anyvalue(kv.value, v)
|
|
pb_any_value.kvlist_value.CopyFrom(kvlist_value)
|
|
else:
|
|
# For unknown types, convert to string
|
|
pb_any_value.string_value = _sanitize_otel_string(str(value))
|
|
|
|
|
|
def _decode_otel_proto_anyvalue(pb_any_value: AnyValue) -> Any:
|
|
"""Decode an OTel protobuf AnyValue.
|
|
|
|
Args:
|
|
pb_any_value: The OTel protobuf AnyValue message to decode.
|
|
|
|
Returns:
|
|
The decoded value.
|
|
"""
|
|
value_type = pb_any_value.WhichOneof("value")
|
|
if not value_type:
|
|
return None
|
|
|
|
# Handle complex types that need recursion
|
|
if value_type == "array_value":
|
|
return [_decode_otel_proto_anyvalue(v) for v in pb_any_value.array_value.values]
|
|
elif value_type == "kvlist_value":
|
|
return {
|
|
kv.key: _decode_otel_proto_anyvalue(kv.value) for kv in pb_any_value.kvlist_value.values
|
|
}
|
|
else:
|
|
# For simple types, just get the attribute directly
|
|
return getattr(pb_any_value, value_type)
|
|
|
|
|
|
def decompress_otlp_body(raw_body: bytes, content_encoding: str) -> bytes:
|
|
"""
|
|
Decompress OTLP request body according to Content-Encoding.
|
|
|
|
Supported encodings:
|
|
- gzip
|
|
- deflate (RFC-compliant and raw deflate)
|
|
|
|
Raises HTTPException if the payload cannot be decompressed.
|
|
"""
|
|
from fastapi import HTTPException, status
|
|
|
|
match content_encoding:
|
|
case "gzip":
|
|
try:
|
|
return gzip.decompress(raw_body)
|
|
except Exception:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="Failed to decompress gzip payload",
|
|
)
|
|
|
|
case "deflate":
|
|
try:
|
|
return zlib.decompress(raw_body)
|
|
except Exception:
|
|
# Try raw DEFLATE stream (some clients send this)
|
|
try:
|
|
return zlib.decompress(raw_body, -zlib.MAX_WBITS)
|
|
except Exception:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail="Failed to decompress deflate payload",
|
|
)
|
|
case _:
|
|
raise HTTPException(
|
|
status_code=status.HTTP_400_BAD_REQUEST,
|
|
detail=f"Unsupported Content-Encoding: {content_encoding}",
|
|
)
|
|
|
|
|
|
def resource_to_otel_proto(resource: OTelResource | None) -> OTelProtoResource:
|
|
"""
|
|
Convert an OpenTelemetry SDK Resource to protobuf Resource format.
|
|
|
|
Args:
|
|
resource: The OpenTelemetry SDK Resource object, or None.
|
|
|
|
Returns:
|
|
An OpenTelemetry protobuf Resource message.
|
|
"""
|
|
otel_resource = OTelProtoResource()
|
|
if resource is not None:
|
|
for key, value in resource.attributes.items():
|
|
attr = otel_resource.attributes.add()
|
|
attr.key = key
|
|
_set_otel_proto_anyvalue(attr.value, value)
|
|
return otel_resource
|