genesis-flow 1.0.0__py3-none-any.whl

mlflow/tracing/utils/environment.py

@@ -0,0 +1,55 @@
+import logging
+import os
+from functools import lru_cache
+
+from mlflow.tracking.context.git_context import GitRunContext
+from mlflow.tracking.context.registry import resolve_tags
+from mlflow.utils.databricks_utils import is_in_databricks_notebook
+from mlflow.utils.git_utils import get_git_branch, get_git_commit, get_git_repo_url
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_GIT_BRANCH,
+    MLFLOW_GIT_COMMIT,
+    MLFLOW_GIT_REPO_URL,
+    TRACE_RESOLVE_TAGS_ALLOWLIST,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+@lru_cache(maxsize=1)
+def resolve_env_metadata():
+    """
+    Resolve common environment metadata to be saved in the trace info. These should not
+    # change over time, so we resolve them only once. These will be stored in trace
+    # metadata rather than tags, because they are immutable.
+    """
+    # GitRunContext does not property work in notebook because _get_main_file()
+    # points to the kernel launcher file, not the actual notebook file.
+    metadata = resolve_tags(ignore=[GitRunContext])
+    if not is_in_databricks_notebook():
+        # Get Git metadata for the script or notebook. If the notebook is in a
+        # Databricks managed Git repo, DatabricksRepoRunContext the metadata
+        # so we don't need to run this logic.
+        metadata.update(_resolve_git_metadata())
+
+    return {key: value for key, value in metadata.items() if key in TRACE_RESOLVE_TAGS_ALLOWLIST}
+
+
+def _resolve_git_metadata():
+    try:
+        import git  # noqa: F401
+    except ImportError:
+        _logger.debug("Git python package is not installed. Skipping git metadata resolution.")
+        return {}
+
+    try:
+        repo = os.getcwd()
+        return {
+            MLFLOW_GIT_COMMIT: get_git_commit(repo) or "",
+            MLFLOW_GIT_REPO_URL: get_git_repo_url(repo) or "",
+            MLFLOW_GIT_BRANCH: get_git_branch(repo) or "",
+        }
+    except Exception:
+        _logger.debug("Failed to resolve git metadata", exc_info=True)
+
+    return {}

mlflow/tracing/utils/exception.py

@@ -0,0 +1,21 @@
+import functools
+
+from mlflow.exceptions import MlflowTracingException
+
+
+def raise_as_trace_exception(f):
+    """
+    A decorator to make sure that the decorated function only raises MlflowTracingException.
+
+    Any exceptions are caught and translated to MlflowTracingException before exiting the function.
+    This is helpful for upstream functions to handle tracing related exceptions properly.
+    """
+
+    @functools.wraps(f)
+    def wrapper(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except Exception as e:
+            raise MlflowTracingException(e) from e
+
+    return wrapper

mlflow/tracing/utils/once.py

@@ -0,0 +1,35 @@
+# Customized from https://github.com/open-telemetry/opentelemetry-python/blob/754fc36a408dd45e86d4a0f820f84e692f14b4c1/opentelemetry-api/src/opentelemetry/util/_once.py
+from threading import Lock
+from typing import Callable
+
+
+class Once:
+    """Execute a function exactly once and block all callers until the function returns"""
+
+    def __init__(self) -> None:
+        self.__lock = Lock()
+        self.__done = False
+
+    @property
+    def done(self):
+        with self.__lock:
+            return self.__done
+
+    @done.setter
+    def done(self, value):
+        with self.__lock:
+            self.__done = value
+
+    def do_once(self, func: Callable[[], None]):
+        """
+        Execute ``func`` if it hasn't been executed or return.
+        Will block until ``func`` has been called by one thread.
+        """
+        if self.__done:
+            return
+
+        with self.__lock:
+            if not self.__done:
+                func()
+                self.__done = True
+                return

mlflow/tracing/utils/otlp.py

@@ -0,0 +1,63 @@
+import os
+from typing import Optional
+
+from opentelemetry.sdk.trace.export import SpanExporter
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import RESOURCE_DOES_NOT_EXIST
+
+
+def should_use_otlp_exporter() -> bool:
+    return _get_otlp_endpoint() is not None
+
+
+def get_otlp_exporter() -> SpanExporter:
+    """
+    Get the OTLP exporter based on the configured protocol.
+    """
+    endpoint = _get_otlp_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_endpoint() -> Optional[str]:
+    """
+    Get the OTLP endpoint from the environment variables.
+    Ref: https://opentelemetry.io/docs/languages/sdk-configuration/otlp-exporter/#endpoint-configuration
+    """
+    # Use `or` instead of default value to do lazy eval
+    return os.environ.get("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT") or os.environ.get(
+        "OTEL_EXPORTER_OTLP_ENDPOINT"
+    )
+
+
+def _get_otlp_protocol() -> str:
+    return os.environ.get("OTEL_EXPORTER_OTLP_TRACES_PROTOCOL") or os.environ.get(
+        "OTEL_EXPORTER_OTLP_PROTOCOL", "grpc"
+    )

mlflow/tracing/utils/processor.py

@@ -0,0 +1,54 @@
+import logging
+
+from mlflow.exceptions import MlflowException
+
+_logger = logging.getLogger(__name__)
+
+
+def apply_span_processors(span):
+    """Apply configured span processors sequentially to the span."""
+    from mlflow.tracing.config import get_config
+
+    config = get_config()
+    if not config.span_processors:
+        return
+
+    non_null_return_processors = []
+    for processor in config.span_processors:
+        try:
+            result = processor(span)
+            if result is not None:
+                non_null_return_processors.append(processor.__name__)
+        except Exception as e:
+            _logger.warning(
+                f"Span processor {processor.__name__} failed: {e}",
+                exc_info=_logger.isEnabledFor(logging.DEBUG),
+            )
+
+    if non_null_return_processors:
+        _logger.warning(
+            f"Span processors {non_null_return_processors} returned a non-null value, "
+            "but it will be ignored. Span processors should not return a value."
+        )
+
+
+def validate_span_processors(span_processors):
+    """Validate that the span processor is a valid function."""
+    span_processors = span_processors or []
+
+    for span_processor in span_processors:
+        if not callable(span_processor):
+            raise MlflowException.invalid_parameter_value(
+                "Span processor must be a callable function."
+            )
+
+        # Skip validation for builtin functions and partial functions that don't have __code__
+        if not hasattr(span_processor, "__code__"):
+            continue
+
+        if span_processor.__code__.co_argcount != 1:
+            raise MlflowException.invalid_parameter_value(
+                "Span processor must take exactly one argument that accepts a LiveSpan object."
+            )
+
+    return span_processors

mlflow/tracing/utils/search.py

@@ -0,0 +1,292 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, Optional, Union
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+SPANS_COLUMN_NAME = "spans"
+
+if TYPE_CHECKING:
+    import pandas
+
+    import mlflow.entities
+    from mlflow.entities import Trace
+
+
+def traces_to_df(traces: list[Trace]) -> "pandas.DataFrame":
+    """
+    Convert a list of MLflow Traces to a pandas DataFrame with one column called "traces"
+    containing string representations of each Trace.
+    """
+    import pandas as pd
+
+    from mlflow.entities.trace import Trace  # import here to avoid circular import
+
+    rows = [trace.to_pandas_dataframe_row() for trace in traces]
+    return pd.DataFrame.from_records(data=rows, columns=Trace.pandas_dataframe_columns())
+
+
+def extract_span_inputs_outputs(
+    traces: Union[list["mlflow.entities.Trace"], "pandas.DataFrame"],
+    fields: list[str],
+    col_name: Optional[str] = None,
+) -> "pandas.DataFrame":
+    """
+    Extracts the specified input and output fields from the spans contained in the specified traces.
+
+    Args:
+        traces: A list of :py:class:`mlflow.entities.Trace` or a pandas DataFrame containing traces.
+        fields: A list of field strings of the form 'span_name.[inputs|outputs]' or
+            'span_name.[inputs|outputs].field_name'.
+        col_name: The name of the column in the traces DataFrame containing the spans. If `traces`
+            is a list of MLflow Traces, this argument should not be provided.
+    """
+    try:
+        import pandas as pd
+    except ImportError as e:
+        raise MlflowException(
+            message=(
+                "The `pandas` library is not installed. Please install `pandas` to use the"
+                f"`mlflow.tracing.extract` function. Error: {e}"
+            ),
+        )
+
+    parsed_fields = _parse_fields(fields)
+
+    if isinstance(traces, list):
+        if col_name is not None:
+            raise MlflowException(
+                message=(
+                    "If `traces` is a list of MLflow Traces, `col_name` should not be provided."
+                ),
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+        traces = traces_to_df(traces)
+        col_name = SPANS_COLUMN_NAME
+
+    if isinstance(traces, pd.DataFrame):
+        return _extract_from_traces_pandas_df(df=traces, col_name=col_name, fields=parsed_fields)
+
+    raise MlflowException(
+        message=(
+            "`traces` must be a list of MLflow Traces or a pandas DataFrame. Got: {type(traces)}"
+        ),
+        error_code=INVALID_PARAMETER_VALUE,
+    )
+
+
+class _PeekableIterator:
+    """
+    Wraps an iterator and allows peeking at the next element without consuming it.
+    """
+
+    def __init__(self, it):
+        self.it = iter(it)
+        self._next = None
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._next is not None:
+            next_value = self._next
+            self._next = None
+            return next_value
+        return next(self.it)
+
+    def peek(self):
+        if self._next is None:
+            try:
+                self._next = next(self.it)
+            except StopIteration:
+                return None
+        return self._next
+
+
+class _ParsedField(NamedTuple):
+    """
+    Represents a parsed field from a string of the form 'span_name.[inputs|outputs]' or
+    'span_name.[inputs|outputs].field_name'.
+    """
+
+    span_name: str
+    field_type: Literal["inputs", "outputs"]
+    field_name: Optional[str]
+
+    def __str__(self) -> str:
+        return (
+            f"{self.span_name}.{self.field_type}.{self.field_name}"
+            if self.field_name is not None
+            else f"{self.span_name}.{self.field_type}"
+        )
+
+
+_BACKTICK = "`"
+
+
+class _FieldParser:
+    def __init__(self, field: str) -> None:
+        self.field = field
+        self.chars = _PeekableIterator(field)
+
+    def peek(self) -> str:
+        return self.chars.peek()
+
+    def next(self) -> str:
+        return next(self.chars)
+
+    def has_next(self) -> bool:
+        return self.peek() is not None
+
+    def consume_until_char_or_end(self, stop_char: Optional[str] = None) -> str:
+        """
+        Consume characters until the specified character is encountered or the end of the
+        string. If char is None, consume until the end of the string.
+        """
+        consumed = ""
+        while (c := self.peek()) and c != stop_char:
+            consumed += self.next()
+        return consumed
+
+    def _parse_span_name(self) -> str:
+        if self.peek() == _BACKTICK:
+            self.next()
+            span_name = self.consume_until_char_or_end(_BACKTICK)
+            if self.peek() != _BACKTICK:
+                raise MlflowException.invalid_parameter_value(
+                    f"Expected closing backtick: {self.field!r}"
+                )
+            self.next()
+        else:
+            span_name = self.consume_until_char_or_end(".")
+
+        if self.peek() != ".":
+            raise MlflowException.invalid_parameter_value(
+                f"Expected dot after span name: {self.field!r}"
+            )
+        self.next()
+        return span_name
+
+    def _parse_field_type(self) -> str:
+        field_type = self.consume_until_char_or_end(".")
+        if field_type not in ("inputs", "outputs"):
+            raise MlflowException.invalid_parameter_value(
+                f"Invalid field type: {field_type!r}. Expected 'inputs' or 'outputs'."
+            )
+
+        if self.has_next():
+            self.next()  # Consume the dot
+        return field_type
+
+    def _parse_field_name(self) -> str:
+        if self.peek() == _BACKTICK:
+            self.next()
+            field_name = self.consume_until_char_or_end(_BACKTICK)
+            if self.peek() != _BACKTICK:
+                raise MlflowException.invalid_parameter_value(
+                    f"Expected closing backtick: {self.field!r}"
+                )
+            self.next()
+
+            # There should be no more characters after the closing backtick
+            if self.has_next():
+                raise MlflowException.invalid_parameter_value(
+                    f"Unexpected characters after closing backtick: {self.field!r}"
+                )
+
+        else:
+            field_name = self.consume_until_char_or_end()
+
+        return field_name
+
+    def parse(self) -> _ParsedField:
+        span_name = self._parse_span_name()
+        field_type = self._parse_field_type()
+        field_name = self._parse_field_name() if self.has_next() else None
+        return _ParsedField(span_name=span_name, field_type=field_type, field_name=field_name)
+
+
+def _parse_fields(fields: list[str]) -> list[_ParsedField]:
+    """
+    Parses the specified field strings of the form 'span_name.[inputs|outputs]' or
+    'span_name.[inputs|outputs].field_name' into _ParsedField objects.
+    """
+    return [_FieldParser(field).parse() for field in fields]
+
+
+def _extract_from_traces_pandas_df(
+    df: "pandas.DataFrame", col_name: str, fields: list[_ParsedField]
+) -> "pandas.DataFrame":
+    """
+    Extracts the specified fields from the spans contained in the specified column of the
+    specified traces DataFrame.
+    """
+
+    from mlflow.entities import Span
+
+    if col_name not in df.columns:
+        raise MlflowException(
+            message=(
+                f"Column '{col_name}' not found in traces DataFrame."
+                f" Available columns: {df.columns}"
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    new_columns: dict[str, list[Any]] = defaultdict(list)
+    for _, row in df.iterrows():
+        spans_dict: dict[str, list[Span]] = defaultdict(list)
+        for span in _extract_spans_from_row(row[col_name]):
+            spans_dict[span.name].append(span)
+
+        for field in fields:
+            matching_spans = spans_dict.get(field.span_name, [])
+            matching_value = _find_matching_value(field, matching_spans)
+            new_columns[str(field)].append(matching_value)
+
+    df_with_new_fields = df.copy()
+    for field in fields:
+        df_with_new_fields[str(field)] = new_columns[str(field)]
+
+    return df_with_new_fields
+
+
+def _find_matching_value(field: _ParsedField, spans: list["mlflow.entities.Span"]) -> Optional[Any]:
+    """
+    Find the value of the field in the list of spans. If the field is not found, return None.
+    """
+    for span in spans:
+        span_inputs_or_outputs = getattr(span, field.field_type)
+        if (
+            isinstance(span_inputs_or_outputs, dict)
+            and field.field_name is not None
+            and field.field_name in span_inputs_or_outputs
+        ):
+            return span_inputs_or_outputs.get(field.field_name)
+        elif field.field_name is None:
+            return span_inputs_or_outputs
+
+
+def _extract_spans_from_row(
+    row_content: Optional[list[dict[str, Any]]],
+) -> list["mlflow.entities.Span"]:
+    """
+    Parses and extracts MLflow Spans from the row content of a traces pandas DataFrame.
+    """
+    from mlflow.entities import Span
+
+    if row_content is None:
+        return []
+
+    try:
+        return [Span.from_dict(span_dict) for span_dict in row_content]
+    except Exception as e:
+        raise MlflowException(
+            message=(
+                f"Failed to extract spans from traces DataFrame row content: {row_content}."
+                f" Error: {e}"
+            ),
+            error_code=INVALID_PARAMETER_VALUE,
+        ) from e