⚡️ Speed up method GrpcRunner.run by 818%

[codeflash-ai]

📄 818% (8.18x) speedup for GrpcRunner.run in pinecone/grpc/grpc_runner.py

⏱️ Runtime : 1.26 milliseconds → 137 microseconds (best of 182 runs)

📝 Explanation and details

The optimization achieves an 818% speedup by eliminating the unnecessary function wrapping overhead and optimizing metadata handling in the hot execution path.

Key optimizations:

1. Removed unnecessary function wrapper: The original code used @wraps(func) to create an inner wrapped() function that was immediately called. This added significant overhead - the line profiler shows 86% of time was spent in return wrapped(). The optimized version executes the logic directly, eliminating this call stack bloat.

2. Optimized metadata merging: Added _prepare_metadata() helper that handles the common case where user metadata is empty (returns fixed_metadata directly) and efficiently merges when needed. This avoids unnecessary dict operations in the hot path.

Performance characteristics from tests:

• Small/empty metadata cases see the biggest gains (800-1100% faster) - these benefit most from avoiding the wrapper overhead and efficient empty metadata handling
• Large metadata cases still see substantial improvements (259-625% faster) - the optimized merging strategy reduces allocation overhead
• Repeated calls maintain consistent performance without state leakage

The optimization is particularly effective because most gRPC calls in typical usage have minimal or no user-provided metadata, making the empty metadata fast-path highly valuable. The wrapper removal provides consistent benefits across all call patterns by reducing Python function call overhead in the critical execution path.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from functools import wraps
from typing import Dict, Optional

# imports
import pytest  # used for our unit tests
from pinecone.grpc.grpc_runner import GrpcRunner

# --- Dummy Classes and Setup for Testing ---

# Dummy Message class to simulate google.protobuf.message.Message
class DummyMessage:
    def __init__(self, value=None):
        self.value = value

# Dummy CallCredentials and Compression classes to simulate grpc types
class DummyCallCredentials:
    pass

class DummyCompression:
    pass

# Dummy Config and GRPCClientConfig classes to simulate pinecone config objects
class DummyConfig:
    def __init__(self, api_key):
        self.api_key = api_key

class DummyGRPCClientConfig:
    def __init__(self, additional_metadata=None):
        self.additional_metadata = additional_metadata or {}

# Dummy constants for CLIENT_VERSION and API_VERSION
CLIENT_VERSION = "1.2.3"
API_VERSION = "2024-06-01"

# Dummy PineconeException
class PineconeException(Exception):
    pass

# Dummy _InactiveRpcError to simulate gRPC error
class DummyInactiveRpcError(Exception):
    def __init__(self, debug_error_string):
        super().__init__(debug_error_string)
        self._state = type("State", (), {"debug_error_string": debug_error_string})()
from pinecone.grpc.grpc_runner import GrpcRunner

# --- Unit Tests ---

# Basic Test Cases

def test_run_basic_success():
    """Test that run calls the function with correct arguments and returns its result."""
    config = DummyConfig(api_key="testkey")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("test-index", config, grpc_config)
    msg = DummyMessage("data")
    # Dummy function that just returns the value
    def dummy_func(request, **kwargs):
        return "success"
    codeflash_output = runner.run(dummy_func, msg); result = codeflash_output # 17.3μs -> 1.43μs (1107% faster)

def test_run_metadata_merging():
    """Test that user-provided metadata overrides fixed metadata."""
    config = DummyConfig(api_key="fixedkey")
    grpc_config = DummyGRPCClientConfig({"extra": "fixed"})
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    def dummy_func(request, metadata, **kwargs):
        # Metadata should be a list of tuples
        mdict = dict(metadata)
        return "ok"
    user_metadata = {"api-key": "userkey", "extra": "user"}
    codeflash_output = runner.run(dummy_func, msg, metadata=user_metadata); result = codeflash_output # 16.9μs -> 2.42μs (600% faster)

def test_run_all_optional_args():
    """Test that run passes all optional arguments through."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    credentials = DummyCallCredentials()
    compression = DummyCompression()
    def dummy_func(request, timeout, metadata, credentials, wait_for_ready, compression):
        return "done"
    codeflash_output = runner.run(
        dummy_func, msg, timeout=10, credentials=credentials, wait_for_ready=True, compression=compression
    ); result = codeflash_output # 15.3μs -> 1.41μs (985% faster)

# Edge Test Cases

def test_run_empty_metadata():
    """Test that run works when no metadata is provided."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    def dummy_func(request, metadata, **kwargs):
        mdict = dict(metadata)
        return True
    codeflash_output = runner.run(dummy_func, msg); result = codeflash_output # 15.7μs -> 1.73μs (808% faster)

def test_run_empty_additional_metadata():
    """Test that run works when additional_metadata is empty."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig({})
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    def dummy_func(request, metadata, **kwargs):
        mdict = dict(metadata)
        return True
    codeflash_output = runner.run(dummy_func, msg); result = codeflash_output # 15.8μs -> 1.65μs (858% faster)

def test_run_exception_translation():
    """Test that _InactiveRpcError is translated to PineconeException."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    def dummy_func(request, **kwargs):
        raise DummyInactiveRpcError("debug string")
    with pytest.raises(PineconeException) as excinfo:
        runner.run(dummy_func, msg)

def test_run_metadata_collision():
    """Test that user metadata overrides both fixed and additional metadata."""
    config = DummyConfig(api_key="fixedkey")
    grpc_config = DummyGRPCClientConfig({"api-key": "additional", "foo": "bar"})
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    def dummy_func(request, metadata, **kwargs):
        mdict = dict(metadata)
        return "ok"
    user_metadata = {"api-key": "userkey", "foo": "baz"}
    codeflash_output = runner.run(dummy_func, msg, metadata=user_metadata); result = codeflash_output # 21.9μs -> 2.54μs (763% faster)

def test_run_none_request():
    """Test that run passes None as request if given."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    def dummy_func(request, **kwargs):
        return "none"
    codeflash_output = runner.run(dummy_func, None); result = codeflash_output # 16.6μs -> 1.38μs (1103% faster)

def test_run_large_metadata_keys():
    """Test that run handles large number of metadata keys."""
    config = DummyConfig(api_key="A")
    # 500 additional metadata keys
    extra_md = {f"k{i}": f"v{i}" for i in range(500)}
    grpc_config = DummyGRPCClientConfig(extra_md)
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    def dummy_func(request, metadata, **kwargs):
        mdict = dict(metadata)
        for i in range(500):
            pass
        return "ok"
    codeflash_output = runner.run(dummy_func, msg); result = codeflash_output # 65.8μs -> 9.08μs (625% faster)

# Large Scale Test Cases

def test_run_large_request_object():
    """Test that run handles large request objects."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    # Large DummyMessage (simulate a large payload)
    big_data = "x" * 1000
    msg = DummyMessage(big_data)
    def dummy_func(request, **kwargs):
        return len(request.value)
    codeflash_output = runner.run(dummy_func, msg); result = codeflash_output # 16.5μs -> 1.50μs (998% faster)

def test_run_large_user_metadata():
    """Test that run handles large user metadata."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    user_metadata = {f"user{i}": f"val{i}" for i in range(900)}
    def dummy_func(request, metadata, **kwargs):
        mdict = dict(metadata)
        for i in range(900):
            pass
        return "ok"
    codeflash_output = runner.run(dummy_func, msg, metadata=user_metadata); result = codeflash_output # 112μs -> 21.2μs (429% faster)

def test_run_many_runs():
    """Test that run can be called many times in a row without leaking state."""
    config = DummyConfig(api_key="A")
    grpc_config = DummyGRPCClientConfig()
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    def dummy_func(request, **kwargs):
        return "ok"
    for _ in range(100):
        codeflash_output = runner.run(dummy_func, msg) # 473μs -> 40.0μs (1082% faster)

def test_run_large_additional_metadata_and_user_metadata():
    """Test that run merges large additional and user metadata correctly."""
    config = DummyConfig(api_key="A")
    additional_md = {f"add{i}": f"addv{i}" for i in range(500)}
    grpc_config = DummyGRPCClientConfig(additional_md)
    runner = GrpcRunner("idx", config, grpc_config)
    msg = DummyMessage()
    user_md = {f"add{i}": f"userv{i}" for i in range(500)}
    def dummy_func(request, metadata, **kwargs):
        mdict = dict(metadata)
        for i in range(500):
            pass
        return "ok"
    codeflash_output = runner.run(dummy_func, msg, metadata=user_md); result = codeflash_output # 73.3μs -> 20.4μs (259% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from unittest.mock import Mock

# imports
import pytest
from google.protobuf.message import Message
from grpc import CallCredentials, Compression
from grpc._channel import _InactiveRpcError
from pinecone.grpc.grpc_runner import GrpcRunner


# Dummy classes and constants to allow tests to run without pinecone dependencies
class DummyConfig:
    def __init__(self, api_key):
        self.api_key = api_key

class DummyGRPCClientConfig:
    def __init__(self, additional_metadata=None):
        self.additional_metadata = additional_metadata or {}

CLIENT_VERSION = "1.2.3"
API_VERSION = "2024-06"

class PineconeException(Exception):
    pass
from pinecone.grpc.grpc_runner import GrpcRunner


# Dummy Message for test purposes
class DummyMessage(Message):
    pass

# unit tests

# 1. Basic Test Cases











def test_run_with_non_dict_metadata():
    """Test run with invalid metadata type (should raise TypeError)."""
    runner = GrpcRunner("index", DummyConfig("key"), DummyGRPCClientConfig())
    dummy_request = DummyMessage()
    def dummy_func(request, timeout, metadata, credentials, wait_for_ready, compression):
        return metadata
    with pytest.raises(Exception):
        # Passing a list instead of dict should fail when updating metadata
        runner.run(dummy_func, dummy_request, metadata=["not", "a", "dict"]) # 16.5μs -> 4.09μs (303% faster)

# 3. Large Scale Test Cases

To edit these changes git checkout codeflash/optimize-GrpcRunner.run-mh9r8xlu and push.