Skip to content

⚡️ Speed up function _calculate_percentile by 10% #166

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
codeflash-ai wants to merge 1 commit into
base: master
Choose a base branch
from
Open

⚡️ Speed up function _calculate_percentile by 10% #166

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 11 additions & 5 deletions mlflow/tracing/utils/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,8 @@
from mlflow.pyfunc.context import Context
from mlflow.types.chat import ChatTool

_IMMUTABLE_TAGS_SET = set(IMMUTABLE_TAGS)


def capture_function_input_args(func, args, kwargs) -> dict[str, Any] | None:
try:
Expand Down Expand Up @@ -289,7 +291,10 @@ def maybe_get_logged_model_id() -> str | None:

def exclude_immutable_tags(tags: dict[str, str]) -> dict[str, str]:
"""Exclude immutable tags e.g. "mlflow.user" from the given tags."""
return {k: v for k, v in tags.items() if k not in IMMUTABLE_TAGS}
# Use set subtraction for performance if tags is large
if not tags:
return {}
return {k: v for k, v in tags.items() if k not in _IMMUTABLE_TAGS_SET}


def generate_mlflow_trace_id_from_otel_trace_id(otel_trace_id: int) -> str:
Expand Down Expand Up @@ -461,10 +466,9 @@ def _calculate_percentile(sorted_data: list[float], percentile: float) -> float:
Returns:
The percentile value
"""
if not sorted_data:
return 0.0

n = len(sorted_data)
if n == 0:
return 0.0
index = percentile * (n - 1)
lower = int(index)
upper = lower + 1
Expand All @@ -474,7 +478,9 @@ def _calculate_percentile(sorted_data: list[float], percentile: float) -> float:

# Linear interpolation between two nearest values
weight = index - lower
return sorted_data[lower] * (1 - weight) + sorted_data[upper] * weight
lower_value = sorted_data[lower]
upper_value = sorted_data[upper]
return lower_value + weight * (upper_value - lower_value)


def add_size_stats_to_trace_metadata(trace: Trace):
Expand Down