⚡️ Speed up function build_otel_context by 16%

[codeflash-ai]

📄 16% (0.16x) speedup for build_otel_context in mlflow/tracing/utils/__init__.py

⏱️ Runtime : 1.81 milliseconds → 1.56 milliseconds (best of 66 runs)

📝 Explanation and details

The optimization replaces keyword arguments with positional arguments when calling trace_api.SpanContext(), achieving a 15% speedup.

Key Change:

• Before: trace_api.SpanContext(trace_id=trace_id, span_id=span_id, is_remote=False)
• After: trace_api.SpanContext(trace_id, span_id, False)

Why This is Faster:
In Python, keyword arguments require dictionary lookups and additional processing overhead. Positional arguments are passed directly to the function, avoiding the parameter name resolution step. The line profiler shows the original version spent significant time on individual parameter assignments (8%, 6.6%, and 7% of total time), while the optimized version eliminates this overhead entirely.

Performance Impact:
Based on the function references, build_otel_context is called frequently in MLflow's span deserialization pipeline - specifically in Span.from_dict(), Span.from_dict_v2(), and Span.from_otel_proto() methods. These methods are used when reconstructing span objects from stored data or network protocols, making this a hot path optimization.

Test Case Performance:
The annotated tests show consistent 15-25% improvements across all scenarios:

• Basic cases with valid IDs: 18-25% faster
• Edge cases with invalid inputs: 7-11% faster
• Large-scale batch processing: 15-16% faster

This optimization is particularly effective for high-throughput tracing workloads where spans are frequently deserialized, as the cumulative effect of the 15% improvement per call becomes significant.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 4028 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import pytest  # used for our unit tests
from mlflow.tracing.utils.__init__ import build_otel_context
# function to test
from opentelemetry import trace as trace_api

# unit tests

# ----------- Basic Test Cases -----------

def test_basic_valid_ids():
    """Test with typical valid trace_id and span_id values."""
    trace_id = 12345678901234567890123456789012
    span_id = 9876543210987654
    codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 2.39μs -> 1.93μs (23.7% faster)

def test_basic_zero_ids():
    """Test with zero trace_id and span_id."""
    trace_id = 0
    span_id = 0
    codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 2.26μs -> 1.80μs (25.8% faster)

def test_basic_maximum_ids():
    """Test with maximum allowed trace_id and span_id values (per OpenTelemetry spec)."""
    max_trace_id = (1 << 128) - 1  # 128-bit unsigned int max
    max_span_id = (1 << 64) - 1    # 64-bit unsigned int max
    codeflash_output = build_otel_context(max_trace_id, max_span_id); ctx = codeflash_output # 2.34μs -> 1.92μs (21.6% faster)

# ----------- Edge Test Cases -----------





def test_non_integer_trace_id():
    """Test with non-integer trace_id (should raise TypeError)."""
    with pytest.raises(TypeError):
        build_otel_context("not_an_int", 123456) # 2.94μs -> 2.74μs (7.22% faster)

def test_non_integer_span_id():
    """Test with non-integer span_id (should raise TypeError)."""
    with pytest.raises(TypeError):
        build_otel_context(123456, "not_an_int") # 3.28μs -> 2.94μs (11.5% faster)





def test_none_trace_id():
    """Test with None as trace_id (should raise TypeError)."""
    with pytest.raises(TypeError):
        build_otel_context(None, 123456) # 2.99μs -> 2.69μs (11.3% faster)

def test_none_span_id():
    """Test with None as span_id (should raise TypeError)."""
    with pytest.raises(TypeError):
        build_otel_context(123456, None) # 3.26μs -> 2.95μs (10.3% faster)

def test_is_remote_always_false():
    """Test that is_remote is always False for any valid input."""
    for trace_id, span_id in [
        (1, 1),
        ((1 << 128) - 1, (1 << 64) - 1),
        (0, 0),
        (42, 99),
    ]:
        codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 4.52μs -> 3.75μs (20.7% faster)

# ----------- Large Scale Test Cases -----------

def test_many_unique_contexts():
    """Test creating a large number of unique SpanContext objects."""
    num_contexts = 1000
    trace_ids = [i for i in range(num_contexts)]
    span_ids = [num_contexts - i for i in range(num_contexts)]
    contexts = []
    for t_id, s_id in zip(trace_ids, span_ids):
        codeflash_output = build_otel_context(t_id, s_id); ctx = codeflash_output # 436μs -> 378μs (15.3% faster)
        contexts.append(ctx)
    # Ensure all contexts are unique and correct
    for idx, ctx in enumerate(contexts):
        pass

def test_large_ids_near_limits():
    """Test with trace_id and span_id values near their maximum limits."""
    for offset in range(1000):
        trace_id = (1 << 128) - 1 - offset
        span_id = (1 << 64) - 1 - offset
        codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 438μs -> 379μs (15.6% faster)

def test_performance_large_batch():
    """Test that building many contexts does not raise and is reasonably fast."""
    import time
    num_contexts = 1000
    trace_ids = [i for i in range(num_contexts)]
    span_ids = [i for i in range(num_contexts)]
    start = time.time()
    for t_id, s_id in zip(trace_ids, span_ids):
        codeflash_output = build_otel_context(t_id, s_id); ctx = codeflash_output # 436μs -> 376μs (16.0% faster)
    duration = time.time() - start
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import pytest  # used for our unit tests
from mlflow.tracing.utils.__init__ import build_otel_context
# function to test
# (copied from mlflow/tracing/utils/__init__.py)
from opentelemetry import trace as trace_api
from opentelemetry.trace import SpanContext

# unit tests

# --- Basic Test Cases ---

def test_basic_valid_ids():
    """Test with typical positive trace_id and span_id."""
    trace_id = 12345678901234567890123456789012
    span_id = 9876543210987654
    codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 2.30μs -> 1.94μs (18.3% faster)

def test_zero_ids():
    """Test with zero trace_id and span_id."""
    trace_id = 0
    span_id = 0
    codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 2.16μs -> 1.89μs (14.1% faster)

def test_maximum_valid_ids():
    """Test with maximum valid trace_id (2^128-1) and span_id (2^64-1)."""
    max_trace_id = (1 << 128) - 1
    max_span_id = (1 << 64) - 1
    codeflash_output = build_otel_context(max_trace_id, max_span_id); ctx = codeflash_output # 2.38μs -> 1.92μs (24.1% faster)

def test_minimum_valid_ids():
    """Test with minimum valid trace_id and span_id (should be 0)."""
    codeflash_output = build_otel_context(0, 0); ctx = codeflash_output # 2.29μs -> 1.86μs (22.8% faster)

# --- Edge Test Cases ---




def test_span_id_too_large():
    """Test with span_id greater than 2^64-1, should raise ValueError."""
    too_large_span_id = (1 << 64)
    with pytest.raises(ValueError):
        build_otel_context(1, too_large_span_id)

@pytest.mark.parametrize(
    "trace_id,span_id",
    [
        (None, 1),
        (1, None),
        (None, None),
        ("abc", 1),
        (1, "xyz"),
        (1.5, 2),
        (2, 1.5),
        ([], 1),
        (1, {}),
    ]
)
def test_invalid_types(trace_id, span_id):
    """Test with invalid types for trace_id and span_id."""
    with pytest.raises((TypeError, ValueError)):
        build_otel_context(trace_id, span_id) # 22.0μs -> 20.6μs (7.21% faster)

def test_is_remote_always_false():
    """Ensure is_remote is always False regardless of input."""
    codeflash_output = build_otel_context(123, 456); ctx = codeflash_output # 2.40μs -> 1.99μs (20.6% faster)

# --- Large Scale Test Cases ---

def test_many_unique_contexts():
    """Test creating a large number of unique contexts for uniqueness and correctness."""
    # 1000 unique trace_id and span_id pairs
    for i in range(1000):
        trace_id = i + 1
        span_id = 1000 - i
        codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 436μs -> 373μs (16.6% faster)

def test_large_ids_near_upper_bound():
    """Test IDs just below the maximum allowed values."""
    trace_id = (1 << 128) - 2
    span_id = (1 << 64) - 2
    codeflash_output = build_otel_context(trace_id, span_id); ctx = codeflash_output # 2.32μs -> 1.99μs (16.8% faster)


def test_trace_id_one_span_id_zero():
    """Test with trace_id=1 and span_id=0."""
    codeflash_output = build_otel_context(1, 0); ctx = codeflash_output # 2.36μs -> 2.00μs (18.5% faster)

def test_trace_id_zero_span_id_one():
    """Test with trace_id=0 and span_id=1."""
    codeflash_output = build_otel_context(0, 1); ctx = codeflash_output # 2.28μs -> 1.85μs (23.4% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

from mlflow.tracing.utils.__init__ import build_otel_context

def test_build_otel_context():
    build_otel_context(0, 0)

To edit these changes git checkout codeflash/optimize-build_otel_context-mhwxwtvq and push.