1  from enum import Enum
 2  
 3  from opentelemetry import trace as trace_api
 4  
 5  from mlflow.protos import service_pb2 as pb
 6  
 7  
 8  class TraceState(str, Enum):
 9      """Enum representing the state of a trace.
10  
11      - ``STATE_UNSPECIFIED``: Unspecified trace state.
12      - ``OK``: Trace successfully completed.
13      - ``ERROR``: Trace encountered an error.
14      - ``IN_PROGRESS``: Trace is currently in progress.
15      """
16  
17      STATE_UNSPECIFIED = "STATE_UNSPECIFIED"
18      OK = "OK"
19      ERROR = "ERROR"
20      IN_PROGRESS = "IN_PROGRESS"
21  
22      def __str__(self):
23          return self.value
24  
25      def to_proto(self):
26          return pb.TraceInfoV3.State.Value(self)
27  
28      @classmethod
29      def from_proto(cls, proto: int) -> "TraceState":
30          return TraceState(pb.TraceInfoV3.State.Name(proto))
31  
32      @staticmethod
33      def from_otel_status(otel_status: trace_api.Status):
34          """Convert OpenTelemetry status code to MLflow TraceState."""
35          return _OTEL_STATUS_CODE_TO_MLFLOW[otel_status.status_code]
36  
37  
38  _OTEL_STATUS_CODE_TO_MLFLOW = {
39      trace_api.StatusCode.OK: TraceState.OK,
40      trace_api.StatusCode.ERROR: TraceState.ERROR,
41      trace_api.StatusCode.UNSET: TraceState.STATE_UNSPECIFIED,
42  }