⚡️ Speed up function _get_models_from_access_groups by 49%

[codeflash-ai]

📄 49% (0.49x) speedup for _get_models_from_access_groups in litellm/proxy/auth/model_checks.py

⏱️ Runtime : 603 microseconds → 404 microseconds (best of 314 runs)

📝 Explanation and details

The optimized code achieves a 49% speedup by eliminating expensive list operations and improving lookup efficiency:

Key optimizations:

1. Eliminated costly pop() operations: The original code collected indices to remove, sorted them in reverse, then called pop() on each index. Each pop() operation on a list triggers O(n) element shifts, making this quadratic for multiple removals. The optimization builds the result list directly, avoiding all pop() operations.

2. Replaced dictionary lookup with set lookup: For the non-include case, the optimization creates access_set = set(model_access_groups) and uses model in access_set instead of model in model_access_groups. Set membership testing is O(1) vs O(n) for dictionary key iteration.

3. Added early return optimization: When model_access_groups is empty, the function immediately returns all_models, avoiding unnecessary processing. This provides dramatic speedups (9000%+ faster) for empty access group cases as seen in the test results.

4. Streamlined control flow: The optimization separates the include/exclude logic into distinct branches, eliminating redundant condition checks inside the main loop.

Performance benefits by test case type:

• Empty access groups: Massive speedups (9000%+ faster) due to early return
• Small to medium datasets: 20-70% faster due to eliminated pop() operations
• Large datasets with many access groups: 40-100% faster, as the quadratic pop() penalty becomes more severe
• Include flag scenarios: 37-68% faster due to simplified logic flow

The optimization maintains the original's in-place mutation behavior using all_models[:] = models_out, ensuring identical functionality while dramatically improving performance.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

from typing import Dict, List, Optional

# imports
import pytest  # used for our unit tests
from litellm.proxy.auth.model_checks import _get_models_from_access_groups

# unit tests

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

def test_basic_no_access_groups():
    # No access groups, should return all_models unchanged
    models = ["a", "b", "c"]
    access_groups = {}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.74μs -> 378ns (360% faster)

def test_basic_one_access_group_removal():
    # One model in access group, should be replaced by its mapped models
    models = ["a", "b", "c"]
    access_groups = {"b": ["x", "y"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.23μs -> 1.87μs (19.4% faster)

def test_basic_one_access_group_include_flag():
    # One model in access group, but include_model_access_groups=True, so don't remove
    models = ["a", "b", "c"]
    access_groups = {"b": ["x", "y"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy(), include_model_access_groups=True); result = codeflash_output # 2.20μs -> 1.33μs (65.4% faster)

def test_basic_multiple_access_groups():
    # Multiple models in access groups, all replaced
    models = ["a", "b", "c", "d"]
    access_groups = {"b": ["x"], "d": ["y", "z"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.83μs -> 2.10μs (34.6% faster)

def test_basic_access_group_with_empty_list():
    # Access group maps to empty list, model is removed
    models = ["a", "b", "c"]
    access_groups = {"b": []}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.00μs -> 1.77μs (12.8% faster)

def test_basic_access_group_with_duplicate_models():
    # Access group maps to models already present
    models = ["a", "b", "c"]
    access_groups = {"b": ["a", "d"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.08μs -> 1.76μs (17.9% faster)

# ------------------ EDGE TEST CASES ------------------

def test_edge_empty_all_models():
    # No models to process
    models = []
    access_groups = {"a": ["x"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.39μs -> 1.26μs (10.8% faster)

def test_edge_empty_access_groups_and_models():
    # Both empty
    models = []
    access_groups = {}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.32μs -> 392ns (238% faster)

def test_edge_access_group_model_not_in_all_models():
    # Access group for model not in all_models
    models = ["a", "b"]
    access_groups = {"c": ["x", "y"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.61μs -> 1.65μs (2.84% slower)

def test_edge_access_group_maps_to_self():
    # Access group maps model to itself
    models = ["a", "b"]
    access_groups = {"b": ["b"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.07μs -> 1.64μs (26.0% faster)

def test_edge_access_group_multiple_models_same_target():
    # Multiple models map to same target model
    models = ["a", "b", "c"]
    access_groups = {"b": ["x"], "c": ["x"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.55μs -> 1.79μs (42.0% faster)

def test_edge_access_group_model_appears_multiple_times():
    # Model in all_models appears multiple times, only those entries replaced
    models = ["a", "b", "b", "c"]
    access_groups = {"b": ["x"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.47μs -> 1.82μs (35.6% faster)

def test_edge_access_group_removal_order():
    # Ensure removal does not affect index of later removals
    models = ["a", "b", "c", "b"]
    access_groups = {"b": ["x"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.42μs -> 1.73μs (39.8% faster)

def test_edge_include_model_access_groups_true_with_duplicates():
    # include_model_access_groups True, access group maps to existing model
    models = ["a", "b"]
    access_groups = {"b": ["a"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy(), include_model_access_groups=True); result = codeflash_output # 2.07μs -> 1.30μs (59.0% faster)

# ------------------ LARGE SCALE TEST CASES ------------------

def test_large_scale_many_models_some_in_access_groups():
    # Large number of models, some in access groups
    models = [f"model_{i}" for i in range(1000)]
    access_groups = {f"model_{i}": [f"special_{i}"] for i in range(0, 1000, 100)}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 60.1μs -> 62.5μs (3.82% slower)
    # Every 100th model removed, replaced by special_{i}
    expected = [m for i, m in enumerate(models) if i % 100 != 0]
    expected += [f"special_{i}" for i in range(0, 1000, 100)]

def test_large_scale_all_models_in_access_groups():
    # All models are access groups
    models = [f"model_{i}" for i in range(500)]
    access_groups = {f"model_{i}": [f"special_{i}"] for i in range(500)}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 73.2μs -> 51.9μs (41.1% faster)
    # All removed, all special_{i} appended
    expected = [f"special_{i}" for i in range(500)]

def test_large_scale_access_groups_with_large_lists():
    # Access groups map to large lists
    models = ["a", "b"]
    access_groups = {"a": [f"x{i}" for i in range(500)], "b": [f"y{i}" for i in range(500)]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 6.73μs -> 8.15μs (17.4% slower)
    # "a","b" removed, 500 x's and 500 y's appended
    expected = [f"x{i}" for i in range(500)] + [f"y{i}" for i in range(500)]

def test_large_scale_include_model_access_groups_true():
    # Large number of models, include_model_access_groups True
    models = [f"model_{i}" for i in range(1000)]
    access_groups = {f"model_{i}": [f"special_{i}"] for i in range(0, 1000, 200)}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy(), include_model_access_groups=True); result = codeflash_output # 58.1μs -> 48.5μs (19.8% faster)
    # All models kept, special_{i} appended for every 200th model
    expected = models + [f"special_{i}" for i in range(0, 1000, 200)]

def test_large_scale_no_access_groups():
    # Large number of models, no access groups
    models = [f"model_{i}" for i in range(1000)]
    access_groups = {}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 39.0μs -> 396ns (9751% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from typing import Dict, List, Optional

# imports
import pytest  # used for our unit tests
from litellm.proxy.auth.model_checks import _get_models_from_access_groups

# unit tests

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

def test_no_access_groups_returns_all_models():
    # No access groups, should return models unchanged
    models = ["modelA", "modelB"]
    access_groups = {}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.81μs -> 420ns (332% faster)

def test_single_access_group_removal():
    # One model is an access group, should be replaced
    models = ["modelA", "group1", "modelB"]
    access_groups = {"group1": ["modelX", "modelY"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.35μs -> 1.94μs (21.0% faster)

def test_single_access_group_include_flag():
    # include_model_access_groups=True should keep the group
    models = ["modelA", "group1", "modelB"]
    access_groups = {"group1": ["modelX", "modelY"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy(), include_model_access_groups=True); result = codeflash_output # 2.21μs -> 1.31μs (68.1% faster)

def test_multiple_access_groups():
    # Multiple access groups in list
    models = ["group1", "modelA", "group2", "modelB"]
    access_groups = {"group1": ["modelX"], "group2": ["modelY"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.92μs -> 2.01μs (45.3% faster)

def test_access_group_with_empty_models():
    # Access group maps to empty list
    models = ["group1", "modelA"]
    access_groups = {"group1": []}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.10μs -> 1.71μs (22.9% faster)

def test_access_group_with_duplicate_models():
    # Access group maps to models already in all_models
    models = ["group1", "modelA"]
    access_groups = {"group1": ["modelA", "modelB"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.16μs -> 1.71μs (26.7% faster)

def test_access_group_not_in_models():
    # Access group key not present in all_models
    models = ["modelA", "modelB"]
    access_groups = {"groupX": ["modelY"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.66μs -> 1.56μs (6.88% faster)

def test_access_group_multiple_occurrences():
    # Access group appears multiple times in all_models
    models = ["group1", "modelA", "group1", "modelB"]
    access_groups = {"group1": ["modelX", "modelY"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.85μs -> 1.93μs (47.8% faster)

def test_access_group_with_non_string_keys_and_values():
    # Models and groups are not strictly strings
    models = [1, "modelA", 2]
    access_groups = {1: ["modelX"], 2: ["modelY"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.68μs -> 1.83μs (46.9% faster)

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

def test_empty_all_models():
    # all_models is empty
    models = []
    access_groups = {"group1": ["modelX"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.34μs -> 1.26μs (6.26% faster)

def test_empty_access_groups():
    # access_groups is empty
    models = ["modelA", "modelB"]
    access_groups = {}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.62μs -> 375ns (333% faster)

def test_all_models_are_access_groups():
    # Every model is an access group
    models = ["group1", "group2"]
    access_groups = {"group1": ["modelA"], "group2": ["modelB"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.73μs -> 1.61μs (69.8% faster)

def test_access_group_with_self_reference():
    # Access group includes itself (should not infinite loop)
    models = ["group1"]
    access_groups = {"group1": ["group1", "modelA"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.02μs -> 1.56μs (28.9% faster)

def test_access_group_with_circular_reference():
    # Two groups referencing each other
    models = ["group1", "group2"]
    access_groups = {"group1": ["group2"], "group2": ["group1"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.67μs -> 1.60μs (67.2% faster)

def test_access_group_with_none_and_empty_strings():
    # None and empty string as model names
    models = ["", None, "group1"]
    access_groups = {"group1": ["modelA"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.01μs -> 1.71μs (17.7% faster)


def test_access_group_with_nested_lists():
    # Access group maps to nested list (should flatten only one level)
    models = ["group1"]
    access_groups = {"group1": [["modelA", "modelB"]]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.28μs -> 1.80μs (26.5% faster)

def test_access_group_with_empty_string_key():
    # Empty string as key
    models = ["", "modelA"]
    access_groups = {"": ["modelX"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 2.11μs -> 1.73μs (22.4% faster)

def test_access_group_with_nonexistent_models():
    # Access group maps to models not in all_models
    models = ["group1"]
    access_groups = {"group1": ["modelX", "modelY"]}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 1.99μs -> 1.62μs (22.5% faster)

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

def test_large_number_of_models_and_groups():
    # Large number of models and access groups
    n = 500
    models = [f"group{i}" for i in range(n)] + [f"model{i}" for i in range(n)]
    access_groups = {f"group{i}": [f"modelX{i}", f"modelY{i}"] for i in range(n)}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 115μs -> 90.2μs (27.6% faster)
    # All groups removed, replaced by their models
    expected = [f"model{i}" for i in range(n)]
    for i in range(n):
        expected.extend([f"modelX{i}", f"modelY{i}"])

def test_large_models_no_access_groups():
    # Large models, no access groups
    models = [f"model{i}" for i in range(1000)]
    access_groups = {}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 38.3μs -> 406ns (9338% faster)

def test_large_access_groups_some_models():
    # Large access group dict, only some present in all_models
    models = [f"group{i}" for i in range(10)]
    access_groups = {f"group{i}": [f"modelX{i}"] for i in range(1000)}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 5.00μs -> 22.8μs (78.1% slower)

def test_large_models_with_duplicates():
    # Large models with duplicates and access groups
    models = [f"group{i%10}" for i in range(1000)]
    access_groups = {f"group{i}": [f"modelX{i}"] for i in range(10)}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy()); result = codeflash_output # 132μs -> 66.2μs (99.8% faster)
    # Each group appears 100 times, so modelX0...modelX9 appended 100 times each
    expected = []
    for i in range(10):
        expected.extend([f"modelX{i}"] * 100)

def test_large_models_include_flag():
    # Large models, include_model_access_groups=True
    models = [f"group{i}" for i in range(10)] + [f"model{i}" for i in range(10)]
    access_groups = {f"group{i}": [f"modelX{i}"] for i in range(10)}
    codeflash_output = _get_models_from_access_groups(access_groups, models.copy(), include_model_access_groups=True); result = codeflash_output # 4.04μs -> 2.94μs (37.2% faster)
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.
#------------------------------------------------
from litellm.proxy.auth.model_checks import _get_models_from_access_groups

def test__get_models_from_access_groups():
    _get_models_from_access_groups({'': []}, [''], include_model_access_groups=None)

🔎 Concolic Coverage Tests and Runtime

Test File::Test Function	Original ⏱️	Optimized ⏱️	Speedup
codeflash_concolic_zbim32de/tmp_crc7tzh/test_concolic_coverage.py::test__get_models_from_access_groups	2.24μs	1.81μs	24.0%✅

To edit these changes git checkout codeflash/optimize-_get_models_from_access_groups-mhdzjfwd and push.