⚡️ Speed up function get_module_min_and_max_supported_ranges by 11%

[codeflash-ai]

📄 11% (0.11x) speedup for get_module_min_and_max_supported_ranges in mlflow/utils/docstring_utils.py

⏱️ Runtime : 22.1 microseconds → 20.0 microseconds (best of 54 runs)

📝 Explanation and details

The optimization achieves a 10% speedup by eliminating redundant dictionary lookups. The key improvement is storing _ML_PACKAGE_VERSIONS[flavor_name] in a variable pkg_info and reusing it, rather than performing the same dictionary lookup multiple times.

Specific optimizations:

1. Reduced dictionary lookups: The original code performed _ML_PACKAGE_VERSIONS[flavor_name] lookup twice (lines with 50% and 4.3% time in profiler), while the optimized version does it once and stores the result in pkg_info.
2. Cached nested dictionary access: Instead of repeatedly accessing pkg_info["package_info"] and pkg_info["models"], these are stored in variables and reused.
3. Streamlined return statement: Returns the values directly from the cached versions dictionary without intermediate variable assignments.

Why this works in Python:
Dictionary lookups in Python involve hash computation and collision resolution, which has measurable overhead even for small dictionaries. By reducing the number of hash lookups from multiple accesses to single cached accesses, we eliminate this repeated computational cost.

Performance impact:
The line profiler shows the optimization is most effective for the dictionary access operations - the time spent on the main lookup line decreased from 50% to 44.5% of total execution time. Test results confirm consistent 5-19% improvements across different flavors, with the best gains on cases like tensorflow and pyspark.ml that have longer processing paths.

This optimization is particularly valuable since this function appears to be used for version validation during ML model operations, where even small microsecond improvements can accumulate across many calls.

✅ Correctness verification report:

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

🌀 Generated Regression Tests and Runtime

import pytest
from mlflow.utils.docstring_utils import \
    get_module_min_and_max_supported_ranges

# --- Function to test (with a mock _ML_PACKAGE_VERSIONS for testability) ---

# Mocked _ML_PACKAGE_VERSIONS for testing purposes
_ML_PACKAGE_VERSIONS = {
    "sklearn": {
        "package_info": {"module_name": "sklearn"},
        "models": {"minimum": "0.18.0", "maximum": "1.3.0"}
    },
    "xgboost": {
        "package_info": {"module_name": "xgboost"},
        "models": {"minimum": "0.82", "maximum": "1.7.0"}
    },
    "tensorflow": {
        "package_info": {"module_name": "tensorflow"},
        "models": {"minimum": "1.15.0", "maximum": "2.13.0"}
    },
    "pytorch": {
        "package_info": {"module_name": "torch"},
        "models": {"minimum": "1.0.0", "maximum": "2.0.1"}
    },
    "spark": {
        "package_info": {"module_name": "pyspark"},
        "models": {"minimum": "2.3.0", "maximum": "3.4.1"}
    },
    # Edge case: missing module_name in package_info
    "catboost": {
        "package_info": {},
        "models": {"minimum": "0.20", "maximum": "1.2"}
    },
    # Edge case: minimum == maximum
    "single_version": {
        "package_info": {"module_name": "unique"},
        "models": {"minimum": "1.0.0", "maximum": "1.0.0"}
    },
    # Edge case: empty strings for min/max
    "empty_versions": {
        "package_info": {"module_name": "empty"},
        "models": {"minimum": "", "maximum": ""}
    }
}
from mlflow.utils.docstring_utils import \
    get_module_min_and_max_supported_ranges

# --- Unit tests ---

# 1. Basic Test Cases

def test_basic_sklearn():
    # Test standard flavor with explicit module_name
    module, min_v, max_v = get_module_min_and_max_supported_ranges("sklearn") # 1.54μs -> 1.34μs (14.8% faster)

def test_basic_xgboost():
    # Test another standard flavor
    module, min_v, max_v = get_module_min_and_max_supported_ranges("xgboost") # 1.58μs -> 1.33μs (18.9% faster)

def test_basic_tensorflow():
    # Test with longer version strings
    module, min_v, max_v = get_module_min_and_max_supported_ranges("tensorflow") # 1.70μs -> 1.43μs (19.1% faster)

def test_basic_pytorch():
    # Test with module_name different from flavor_name
    module, min_v, max_v = get_module_min_and_max_supported_ranges("pytorch") # 1.54μs -> 1.36μs (13.5% faster)

# 2. Edge Test Cases

def test_special_case_pyspark_ml():
    # pyspark.ml should map to spark
    module, min_v, max_v = get_module_min_and_max_supported_ranges("pyspark.ml") # 1.73μs -> 1.46μs (19.0% faster)

def test_missing_module_name_in_package_info():
    # If module_name is missing, should fallback to flavor_name
    module, min_v, max_v = get_module_min_and_max_supported_ranges("catboost") # 1.41μs -> 1.35μs (4.31% faster)


def test_flavor_not_found():
    # Should raise KeyError for unknown flavor
    with pytest.raises(KeyError):
        get_module_min_and_max_supported_ranges("nonexistent_flavor") # 1.20μs -> 1.25μs (4.17% slower)

def test_models_key_missing():
    # Edge: models key missing
    _ML_PACKAGE_VERSIONS["bad_flavor"] = {"package_info": {"module_name": "bad"}}
    try:
        with pytest.raises(KeyError):
            get_module_min_and_max_supported_ranges("bad_flavor")
    finally:
        del _ML_PACKAGE_VERSIONS["bad_flavor"]


def test_models_minimum_maximum_missing():
    # Edge: minimum/maximum missing in models
    _ML_PACKAGE_VERSIONS["missing_versions"] = {
        "package_info": {"module_name": "miss"},
        "models": {}
    }
    try:
        with pytest.raises(KeyError):
            get_module_min_and_max_supported_ranges("missing_versions")
    finally:
        del _ML_PACKAGE_VERSIONS["missing_versions"]

# 3. Large Scale Test Cases

import pytest
from mlflow.utils.docstring_utils import \
    get_module_min_and_max_supported_ranges

# Simulate the _ML_PACKAGE_VERSIONS global as used by the function.
# This would normally be imported from mlflow.ml_package_versions
_ML_PACKAGE_VERSIONS = {
    "sklearn": {
        "package_info": {"module_name": "sklearn"},
        "models": {"minimum": "0.20.0", "maximum": "1.3.2"}
    },
    "xgboost": {
        "package_info": {"module_name": "xgboost"},
        "models": {"minimum": "0.90", "maximum": "1.7.6"}
    },
    "spark": {
        "package_info": {"module_name": "pyspark"},
        "models": {"minimum": "2.3.0", "maximum": "3.4.1"}
    },
    "pytorch": {
        "package_info": {"module_name": "torch"},
        "models": {"minimum": "1.0.0", "maximum": "2.1.0"}
    },
    "custom_flavor": {
        "package_info": {"module_name": "custom.module"},
        "models": {"minimum": "1.2.3", "maximum": "4.5.6"}
    },
    "missing_module_name": {
        "package_info": {},  # No module_name key
        "models": {"minimum": "0.1", "maximum": "0.2"}
    },
    "identical_versions": {
        "package_info": {"module_name": "identical"},
        "models": {"minimum": "1.0.0", "maximum": "1.0.0"}
    },
    "min_greater_than_max": {
        "package_info": {"module_name": "badversions"},
        "models": {"minimum": "2.0.0", "maximum": "1.0.0"}
    },
    "empty_versions": {
        "package_info": {"module_name": "empty"},
        "models": {"minimum": "", "maximum": ""}
    },
}
from mlflow.utils.docstring_utils import \
    get_module_min_and_max_supported_ranges

# --------------------- UNIT TESTS ---------------------

# Basic Test Cases

def test_basic_sklearn():
    # Test normal case for sklearn flavor
    codeflash_output = get_module_min_and_max_supported_ranges("sklearn"); result = codeflash_output # 1.62μs -> 1.40μs (15.3% faster)

def test_basic_xgboost():
    # Test normal case for xgboost flavor
    codeflash_output = get_module_min_and_max_supported_ranges("xgboost"); result = codeflash_output # 1.46μs -> 1.38μs (5.71% faster)

def test_basic_pytorch():
    # Test normal case for pytorch flavor
    codeflash_output = get_module_min_and_max_supported_ranges("pytorch"); result = codeflash_output # 1.49μs -> 1.30μs (14.6% faster)


def test_pyspark_ml_special_case():
    # Test pyspark.ml maps to spark
    codeflash_output = get_module_min_and_max_supported_ranges("pyspark.ml"); result = codeflash_output # 1.73μs -> 1.49μs (15.6% faster)





def test_nonexistent_flavor_key():
    # Test case where flavor_name does not exist in _ML_PACKAGE_VERSIONS
    with pytest.raises(KeyError):
        get_module_min_and_max_supported_ranges("nonexistent_flavor") # 1.19μs -> 1.27μs (6.46% slower)



def test_flavor_name_with_spaces():
    # Test case where flavor_name has spaces, not present in dict (should raise KeyError)
    with pytest.raises(KeyError):
        get_module_min_and_max_supported_ranges("sk learn") # 1.37μs -> 1.23μs (11.8% faster)

# Large Scale Test Cases

To edit these changes git checkout codeflash/optimize-get_module_min_and_max_supported_ranges-mhx83lcj and push.