genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracing/processor/base_mlflow.py

@@ -0,0 +1,199 @@
+import json
+import logging
+from typing import Any, Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.tracing.constant import (
+    MAX_CHARS_IN_TRACE_INFO_METADATA,
+    TRUNCATION_SUFFIX,
+    SpanAttributeKey,
+    TraceMetadataKey,
+    TraceTagKey,
+)
+from mlflow.tracing.trace_manager import InMemoryTraceManager, _Trace
+from mlflow.tracing.utils import (
+    aggregate_usage_from_spans,
+    deduplicate_span_names_in_place,
+    get_otel_attribute,
+    maybe_get_dependencies_schemas,
+    maybe_get_logged_model_id,
+    maybe_get_request_id,
+    update_trace_state_from_span_conditionally,
+)
+from mlflow.tracing.utils.environment import resolve_env_metadata
+from mlflow.tracking.fluent import (
+    _get_active_model_id_global,
+    _get_experiment_id,
+    _get_latest_active_run,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+class BaseMlflowSpanProcessor(SimpleSpanProcessor):
+    """
+    Defines custom hooks to be executed when a span is started or ended (before exporting).
+
+    """
+
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        experiment_id: Optional[str] = None,
+    ):
+        self.span_exporter = span_exporter
+        self._experiment_id = experiment_id
+        self._trace_manager = InMemoryTraceManager.get_instance()
+        self._env_metadata = resolve_env_metadata()
+
+    def on_start(self, span: OTelSpan, parent_context: Optional[Context] = None):
+        """
+        Handle the start of a span. This method is called when an OpenTelemetry span is started.
+
+        Args:
+            span: An OpenTelemetry Span object that is started.
+            parent_context: The context of the span. Note that this is only passed when the context
+                object is explicitly specified to OpenTelemetry start_span call. If the parent span
+                is obtained from the global context, it won't be passed here so we should not rely
+                on it.
+        """
+        trace_id = self._trace_manager.get_mlflow_trace_id_from_otel_id(span.context.trace_id)
+
+        if not trace_id and span.parent is not None:
+            _logger.debug(
+                "Received a non-root span but the trace ID is not found."
+                "The trace has likely been halted due to a timeout expiration."
+            )
+            return
+
+        if span.parent is None:
+            trace_info = self._start_trace(span)
+            trace_id = trace_info.trace_id
+
+        span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(trace_id))
+
+    def _start_trace(self, root_span: OTelSpan) -> TraceInfo:
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    def on_end(self, span: OTelReadableSpan) -> None:
+        """
+        Handle the end of a span. This method is called when an OpenTelemetry span is ended.
+
+        Args:
+            span: An OpenTelemetry ReadableSpan object that is ended.
+        """
+        # Processing the trace only when the root span is found.
+        if span._parent is not None:
+            return
+
+        trace_id = get_otel_attribute(span, SpanAttributeKey.REQUEST_ID)
+        with self._trace_manager.get_trace(trace_id) as trace:
+            if trace is None:
+                _logger.debug(f"Trace data with request ID {trace_id} not found.")
+                return
+
+            self._update_trace_info(trace, span)
+            deduplicate_span_names_in_place(list(trace.span_dict.values()))
+
+        super().on_end(span)
+
+    def _get_experiment_id_for_trace(self, span: OTelReadableSpan) -> str:
+        """
+        Determine the experiment ID to associate with the trace.
+
+        The experiment ID can be configured in multiple ways, in order of precedence:
+          1. An experiment ID specified via the span creation API i.e. MlflowClient().start_trace()
+          2. An experiment ID specified via the processor constructor
+          3. An experiment ID of an active run.
+          4. The default experiment ID
+        """
+        from mlflow.tracking.fluent import _get_latest_active_run
+
+        if experiment_id := get_otel_attribute(span, SpanAttributeKey.EXPERIMENT_ID):
+            return experiment_id
+
+        if self._experiment_id:
+            return self._experiment_id
+
+        if run := _get_latest_active_run():
+            return run.info.experiment_id
+
+        return _get_experiment_id()
+
+    def _get_basic_trace_metadata(self) -> dict[str, Any]:
+        metadata = self._env_metadata.copy()
+
+        # If the span is started within an active MLflow run, we should record it as a trace tag
+        # Note `mlflow.active_run()` can only get thread-local active run,
+        # but tracing routine might be applied to model inference worker threads
+        # in the following cases:
+        #  - langchain model `chain.batch` which uses thread pool to spawn workers.
+        #  - MLflow langchain pyfunc model `predict` which calls `api_request_parallel_processor`.
+        # Therefore, we use `_get_global_active_run()` instead to get the active run from
+        # all threads and set it as the tracing source run.
+        if run := _get_latest_active_run():
+            metadata[TraceMetadataKey.SOURCE_RUN] = run.info.run_id
+
+        # The order is:
+        # 1. model_id of the current active model set by `set_active_model`
+        # 2. model_id from the current prediction context
+        #   (set by mlflow pyfunc predict, or explicitly using set_prediction_context)
+        if active_model_id := _get_active_model_id_global():
+            metadata[TraceMetadataKey.MODEL_ID] = active_model_id
+        elif model_id := maybe_get_logged_model_id():
+            metadata[TraceMetadataKey.MODEL_ID] = model_id
+
+        return metadata
+
+    def _get_basic_trace_tags(self, span: OTelReadableSpan) -> dict[str, Any]:
+        # If the trace is created in the context of MLflow model evaluation, we extract the request
+        # ID from the prediction context. Otherwise, we create a new trace info by calling the
+        # backend API.
+        tags = {}
+        if request_id := maybe_get_request_id(is_evaluate=True):
+            tags.update({TraceTagKey.EVAL_REQUEST_ID: request_id})
+        if dependencies_schema := maybe_get_dependencies_schemas():
+            tags.update(dependencies_schema)
+        tags.update({TraceTagKey.TRACE_NAME: span.name})
+        return tags
+
+    def _update_trace_info(self, trace: _Trace, root_span: OTelReadableSpan):
+        """Update the trace info with the final values from the root span."""
+        # The trace/span start time needs adjustment to exclude the latency of
+        # the backend API call. We already adjusted the span start time in the
+        # on_start method, so we reflect the same to the trace start time here.
+        trace.info.request_time = root_span.start_time // 1_000_000  # nanosecond to millisecond
+        trace.info.execution_duration = (root_span.end_time - root_span.start_time) // 1_000_000
+
+        # Update trace state from span status, but only if the user hasn't explicitly set
+        # a different trace status
+        update_trace_state_from_span_conditionally(trace, root_span)
+        trace.info.trace_metadata.update(
+            {
+                TraceMetadataKey.INPUTS: self._truncate_metadata(
+                    root_span.attributes.get(SpanAttributeKey.INPUTS)
+                ),
+                TraceMetadataKey.OUTPUTS: self._truncate_metadata(
+                    root_span.attributes.get(SpanAttributeKey.OUTPUTS)
+                ),
+            }
+        )
+
+        # Aggregate token usage information from all spans
+        if usage := aggregate_usage_from_spans(trace.span_dict.values()):
+            trace.info.request_metadata[TraceMetadataKey.TOKEN_USAGE] = json.dumps(usage)
+
+    def _truncate_metadata(self, value: Optional[str]) -> str:
+        """Get truncated value of the attribute if it exceeds the maximum length."""
+        if not value:
+            return ""
+
+        if len(value) > MAX_CHARS_IN_TRACE_INFO_METADATA:
+            trunc_length = MAX_CHARS_IN_TRACE_INFO_METADATA - len(TRUNCATION_SUFFIX)
+            value = value[:trunc_length] + TRUNCATION_SUFFIX
+        return value

mlflow/tracing/processor/inference_table.py

@@ -0,0 +1,175 @@
+import json
+import logging
+from typing import Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_location import TraceLocation
+from mlflow.entities.trace_state import TraceState
+from mlflow.environment_variables import MLFLOW_EXPERIMENT_ID
+from mlflow.tracing.constant import (
+    TRACE_SCHEMA_VERSION,
+    TRACE_SCHEMA_VERSION_KEY,
+    SpanAttributeKey,
+    TraceMetadataKey,
+)
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import (
+    _try_get_prediction_context,
+    aggregate_usage_from_spans,
+    deduplicate_span_names_in_place,
+    generate_trace_id_v3,
+    get_otel_attribute,
+    maybe_get_dependencies_schemas,
+    maybe_get_request_id,
+    update_trace_state_from_span_conditionally,
+)
+from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_MODEL_SERVING_ENDPOINT_NAME
+
+_logger = logging.getLogger(__name__)
+
+
+_HEADER_REQUEST_ID_KEY = "X-Request-Id"
+
+
+# Extracting for testing purposes
+def _get_flask_request():
+    import flask
+
+    if flask.has_request_context():
+        return flask.request
+
+
+class InferenceTableSpanProcessor(SimpleSpanProcessor):
+    """
+    Defines custom hooks to be executed when a span is started or ended (before exporting).
+
+    This processor is used when the tracing destination is Databricks Inference Table.
+    """
+
+    def __init__(self, span_exporter: SpanExporter):
+        self.span_exporter = span_exporter
+        self._trace_manager = InMemoryTraceManager.get_instance()
+
+    def on_start(self, span: OTelSpan, parent_context: Optional[Context] = None):
+        """
+        Handle the start of a span. This method is called when an OpenTelemetry span is started.
+
+        Args:
+            span: An OpenTelemetry Span object that is started.
+            parent_context: The context of the span. Note that this is only passed when the context
+                object is explicitly specified to OpenTelemetry start_span call. If the parent
+                span is obtained from the global context, it won't be passed here so we should not
+                rely on it.
+        """
+        databricks_request_id = maybe_get_request_id()
+        if databricks_request_id is None:
+            # NB: This is currently used for streaming inference in Databricks Model Serving.
+            # In normal prediction, serving logic pass the request ID using the
+            # `with set_prediction_context` context manager that wraps `model.predict`
+            # call. However, in streaming case, the context manager is not applicable
+            # so we still need to rely on Flask request context (which is set to the
+            # stream response via flask.stream_with_context()
+            if flask_request := _get_flask_request():
+                databricks_request_id = flask_request.headers.get(_HEADER_REQUEST_ID_KEY)
+                if not databricks_request_id:
+                    _logger.warning(
+                        "Databricks request ID not found in the request headers. "
+                        "Skipping trace processing."
+                    )
+                    return
+            else:
+                _logger.warning(
+                    "Failed to get Databricks request ID from the request headers because "
+                    "request context is not available. Skipping trace processing."
+                )
+                return
+
+        trace_id = generate_trace_id_v3(span)
+        span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(trace_id))
+        tags = {}
+        if dependencies_schema := maybe_get_dependencies_schemas():
+            tags.update(dependencies_schema)
+
+        if span._parent is None:
+            trace_info = TraceInfo(
+                trace_id=trace_id,
+                client_request_id=databricks_request_id,
+                # NB: Agent framework populate the MLFLOW_EXPERIMENT_ID environment variable
+                #   with the experiment ID to which the model is logged. We don't use the
+                #   _get_experiment_id() method because it will fallback to the default
+                #   experiment if the MLFLOW_EXPERIMENT_ID is not set.
+                trace_location=TraceLocation.from_experiment_id(MLFLOW_EXPERIMENT_ID.get()),
+                request_time=span.start_time // 1_000_000,  # nanosecond to millisecond
+                execution_duration=None,
+                state=TraceState.IN_PROGRESS,
+                trace_metadata=self._get_trace_metadata(),
+                tags=tags,
+            )
+            self._trace_manager.register_trace(span.context.trace_id, trace_info)
+
+    def on_end(self, span: OTelReadableSpan) -> None:
+        """
+        Handle the end of a span. This method is called when an OpenTelemetry span is ended.
+
+        Args:
+            span: An OpenTelemetry ReadableSpan object that is ended.
+        """
+        # Processing the trace only when the root span is found.
+        if span._parent is not None:
+            return
+
+        trace_id = get_otel_attribute(span, SpanAttributeKey.REQUEST_ID)
+        with self._trace_manager.get_trace(trace_id) as trace:
+            if trace is None:
+                _logger.debug(f"Trace data with trace ID {trace_id} not found.")
+                return
+
+            trace.info.execution_duration = (span.end_time - span.start_time) // 1_000_000
+
+            # Update trace state from span status, but only if the user hasn't explicitly set
+            # a different trace status
+            update_trace_state_from_span_conditionally(trace, span)
+
+            spans = list(trace.span_dict.values())
+            deduplicate_span_names_in_place(spans)
+
+            # Aggregate token usage information from all spans
+            if usage := aggregate_usage_from_spans(spans):
+                trace.info.request_metadata[TraceMetadataKey.TOKEN_USAGE] = json.dumps(usage)
+
+        super().on_end(span)
+
+    def _get_trace_metadata(self) -> dict[str, str]:
+        metadata = {TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)}
+
+        context = _try_get_prediction_context()
+        if context:
+            metadata[MLFLOW_DATABRICKS_MODEL_SERVING_ENDPOINT_NAME] = context.endpoint_name or ""
+
+        # The model ID fetch order is
+        # 1. from the context, this could be set by databricks serving
+        # 2. from the active model, this could be set by model loading or with environment variable
+        if context and context.model_id:
+            metadata[TraceMetadataKey.MODEL_ID] = context.model_id
+            _logger.debug(f"Model id fetched from the context: {context.model_id}")
+        else:
+            try:
+                from mlflow.tracking.fluent import _get_active_model_id_global
+
+                if active_model_id := _get_active_model_id_global():
+                    metadata[TraceMetadataKey.MODEL_ID] = active_model_id
+                    _logger.debug(
+                        f"Adding active model ID {active_model_id} to the trace metadata."
+                    )
+            except Exception as e:
+                _logger.debug(
+                    f"Failed to get model ID from the active model: {e}. "
+                    "Skipping adding model ID to trace metadata.",
+                    exc_info=True,
+                )
+        return metadata

mlflow/tracing/processor/mlflow_v3.py

@@ -0,0 +1,47 @@
+from typing import Optional
+
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_location import TraceLocation
+from mlflow.entities.trace_state import TraceState
+from mlflow.tracing.processor.base_mlflow import BaseMlflowSpanProcessor
+from mlflow.tracing.utils import generate_trace_id_v3
+
+
+class MlflowV3SpanProcessor(BaseMlflowSpanProcessor):
+    """
+    Defines custom hooks to be executed when a span is started or ended (before exporting).
+
+    This processor is used for exporting traces to MLflow Tracking Server
+    using the V3 trace schema and API.
+    """
+
+    def __init__(
+        self,
+        span_exporter: SpanExporter,
+        experiment_id: Optional[str] = None,
+    ):
+        super().__init__(span_exporter, experiment_id)
+
+    def _start_trace(self, root_span: OTelSpan) -> TraceInfo:
+        """
+        Create a new TraceInfo object and register it with the trace manager.
+
+        This method is called in the on_start method of the base class.
+        """
+        trace_info = TraceInfo(
+            trace_id=generate_trace_id_v3(root_span),
+            trace_location=TraceLocation.from_experiment_id(
+                self._get_experiment_id_for_trace(root_span)
+            ),
+            request_time=root_span.start_time // 1_000_000,  # nanosecond to millisecond
+            execution_duration=None,
+            state=TraceState.IN_PROGRESS,
+            trace_metadata=self._get_basic_trace_metadata(),
+            tags=self._get_basic_trace_tags(root_span),
+        )
+        self._trace_manager.register_trace(root_span.context.trace_id, trace_info)
+
+        return trace_info

mlflow/tracing/processor/otel.py

@@ -0,0 +1,73 @@
+import json
+from typing import Optional
+
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan as OTelReadableSpan
+from opentelemetry.sdk.trace import Span as OTelSpan
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+
+from mlflow.entities.trace_info import TraceInfo
+from mlflow.entities.trace_location import TraceLocation
+from mlflow.entities.trace_state import TraceState
+from mlflow.tracing.constant import TRACE_SCHEMA_VERSION, TRACE_SCHEMA_VERSION_KEY, SpanAttributeKey
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+from mlflow.tracing.utils import generate_trace_id_v3
+
+
+class OtelSpanProcessor(BatchSpanProcessor):
+    """
+    SpanProcessor implementation to export MLflow traces to a OpenTelemetry collector.
+
+    Extending OpenTelemetry BatchSpanProcessor to add some custom hooks to be executed when a span
+    is started or ended (before exporting).
+    """
+
+    def __init__(self, span_exporter: SpanExporter):
+        super().__init__(span_exporter)
+        # In opentelemetry-sdk 1.34.0, the `span_exporter` field was removed from the
+        # `BatchSpanProcessor` class.
+        # https://github.com/open-telemetry/opentelemetry-python/issues/4616
+        #
+        # The `span_exporter` field was restored as a property in 1.34.1
+        # https://github.com/open-telemetry/opentelemetry-python/pull/4621
+        #
+        # We use a try-except block to maintain compatibility with both versions.
+        try:
+            self.span_exporter = span_exporter
+        except AttributeError:
+            pass
+        self._trace_manager = InMemoryTraceManager.get_instance()
+
+    def on_start(self, span: OTelSpan, parent_context: Optional[Context] = None):
+        """
+        Handle the start of a span. This method is called when an OpenTelemetry span is started.
+
+        Args:
+            span: An OpenTelemetry Span object that is started.
+            parent_context: The context of the span. Note that this is only passed when the context
+                object is explicitly specified to OpenTelemetry start_span call. If the parent
+                span is obtained from the global context, it won't be passed here so we should not
+                rely on it.
+        """
+        trace_id = generate_trace_id_v3(span)
+        trace_info = TraceInfo(
+            trace_id=trace_id,
+            trace_location=TraceLocation.from_experiment_id(None),
+            request_time=span.start_time // 1_000_000,  # nanosecond to millisecond
+            execution_duration=None,
+            state=TraceState.IN_PROGRESS,
+            trace_metadata={TRACE_SCHEMA_VERSION_KEY: str(TRACE_SCHEMA_VERSION)},
+            tags={},
+        )
+        span.set_attribute(SpanAttributeKey.REQUEST_ID, json.dumps(trace_id))
+
+        self._trace_manager.register_trace(span.context.trace_id, trace_info)
+
+        super().on_start(span, parent_context)
+
+    def on_end(self, span: OTelReadableSpan):
+        # Pops the trace entry from the in-memory trace manager to avoid memory leak
+        if span._parent is None:
+            self._trace_manager.pop_trace(span.context.trace_id)
+
+        super().on_end(span)