⚡️ Speed up method AWSKeyManagementService_V2.load_aws_kms by 47%

[codeflash-ai]

📄 47% (0.47x) speedup for AWSKeyManagementService_V2.load_aws_kms in litellm/secret_managers/aws_secret_manager.py

⏱️ Runtime : 271 microseconds → 184 microseconds (best of 108 runs)

📝 Explanation and details

The optimized code achieves a 46% speedup through three key optimizations that reduce redundant environment variable lookups:

1. Eliminated duplicate region validation calls:

• Original: Called validate_environment() inside load_aws_kms(), then separately called os.getenv("AWS_REGION_NAME") - resulting in two environment lookups
• Optimized: Performs validation inline with a single os.environ.get("AWS_REGION_NAME") call, eliminating the redundant function call and second lookup

2. More efficient environment variable access:

• Original: "AWS_REGION_NAME" not in os.environ creates a membership test across all environment variables
• Optimized: os.environ.get("AWS_REGION_NAME") directly accesses the specific key with None check, which is faster for individual key lookups

3. Simplified boolean condition:

• Original: if use_aws_kms is None or use_aws_kms is False: performs two comparisons
• Optimized: if not use_aws_kms: uses Python's truthiness evaluation, which is more direct

The line profiler shows the biggest performance gain comes from eliminating the validate_environment() call (originally 51.6% of execution time), which was performing redundant work since the region name was looked up again immediately after. The optimization particularly benefits scenarios with frequent AWS client creation, as shown in the test cases where region switching improved by 44.9% and basic client loading improved by 35.2%.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import os
# function to test
import sys

# imports
import pytest
from litellm.secret_managers.aws_secret_manager import \
    AWSKeyManagementService_V2


def load_aws_kms(use_aws_kms):
    """
    Loads AWS KMS client if use_aws_kms is True.
    Returns None if use_aws_kms is False or None.
    Raises ValueError if required environment variables are missing.
    """
    if use_aws_kms is None or use_aws_kms is False:
        return
    try:
        import boto3

        validate_environment()

        # Create a Secrets Manager client
        kms_client = boto3.client("kms", region_name=os.getenv("AWS_REGION_NAME"))

        return kms_client
    except Exception as e:
        raise e

# ---- BASIC TEST CASES ----















#------------------------------------------------
import os
# function to test
import types

# imports
import pytest
from litellm.secret_managers.aws_secret_manager import \
    AWSKeyManagementService_V2

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




def test_missing_aws_region(monkeypatch):
    """Test that load_aws_kms raises ValueError if AWS_REGION_NAME is missing."""
    monkeypatch.delenv("AWS_REGION_NAME", raising=False)
    monkeypatch.setenv("LITELLM_LICENSE", "valid_license")
    kms = AWSKeyManagementService_V2
    # Patch boto3.client to avoid real AWS calls
    monkeypatch.setitem(__import__("sys").modules, "boto3", types.SimpleNamespace(client=lambda *_, **__: None))
    instance = kms.__new__(kms)
    with pytest.raises(ValueError) as excinfo:
        instance.load_aws_kms(use_aws_kms=True) # 4.75μs -> 4.24μs (11.9% faster)

def test_missing_license(monkeypatch):
    """Test that __init__ raises ValueError if neither LITELLM_LICENSE nor LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE is set."""
    monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
    monkeypatch.delenv("LITELLM_LICENSE", raising=False)
    monkeypatch.delenv("LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE", raising=False)
    with pytest.raises(ValueError) as excinfo:
        AWSKeyManagementService_V2()

def test_license_from_secret(monkeypatch):
    """Test that __init__ works if LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE is set instead of LITELLM_LICENSE."""
    monkeypatch.setenv("AWS_REGION_NAME", "us-east-1")
    monkeypatch.setenv("LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE", "valid_license")
    # Patch boto3.client to dummy
    monkeypatch.setitem(__import__("sys").modules, "boto3", types.SimpleNamespace(client=lambda *_, **__: object()))
    kms = AWSKeyManagementService_V2()



def test_load_aws_kms_region_env_var(monkeypatch):
    """Test that load_aws_kms uses the value of AWS_REGION_NAME from the environment."""
    monkeypatch.setenv("AWS_REGION_NAME", "eu-west-2")
    monkeypatch.setenv("LITELLM_LICENSE", "valid_license")
    def dummy_client(service, region_name=None):
        return object()
    monkeypatch.setitem(__import__("sys").modules, "boto3", types.SimpleNamespace(client=dummy_client))
    kms = AWSKeyManagementService_V2()
    codeflash_output = kms.load_aws_kms(use_aws_kms=True); client = codeflash_output # 2.46μs -> 1.82μs (35.2% faster)

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


def test_load_aws_kms_many_regions(monkeypatch):
    """Test load_aws_kms with 50 different region names."""
    monkeypatch.setenv("LITELLM_LICENSE", "valid_license")
    class DummyClient:
        def __init__(self, region): self.region = region
    def dummy_client(service, region_name=None):
        return DummyClient(region_name)
    monkeypatch.setitem(__import__("sys").modules, "boto3", types.SimpleNamespace(client=dummy_client))
    kms = AWSKeyManagementService_V2()
    for i in range(50):
        region = f"region-{i}"
        monkeypatch.setenv("AWS_REGION_NAME", region)
        codeflash_output = kms.load_aws_kms(use_aws_kms=True); client = codeflash_output # 91.8μs -> 63.3μs (44.9% faster)

def test_load_aws_kms_concurrent(monkeypatch):
    """Test load_aws_kms can be called in parallel (simulate with threads)."""
    import threading
    monkeypatch.setenv("AWS_REGION_NAME", "ap-southeast-1")
    monkeypatch.setenv("LITELLM_LICENSE", "valid_license")
    class DummyClient: pass
    def dummy_client(service, region_name=None):
        return DummyClient()
    monkeypatch.setitem(__import__("sys").modules, "boto3", types.SimpleNamespace(client=dummy_client))
    kms = AWSKeyManagementService_V2()
    results = []
    def worker():
        codeflash_output = kms.load_aws_kms(use_aws_kms=True); client = codeflash_output
        results.append(isinstance(client, DummyClient))
    threads = [threading.Thread(target=worker) for _ in range(20)]
    for t in threads: t.start()
    for t in threads: t.join()
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from litellm.secret_managers.aws_secret_manager import AWSKeyManagementService_V2

To edit these changes git checkout codeflash/optimize-AWSKeyManagementService_V2.load_aws_kms-mhddj05r and push.