⚡️ Speed up method ModelManager.add_model by 25% in PR #1196 (revert-1175-depth-estimation-workflow-block)

[codeflash-ai]

⚡️ This pull request contains optimizations for PR #1196

If you approve this dependent PR, these changes will be merged into the original PR branch revert-1175-depth-estimation-workflow-block.

This PR will be automatically closed if the original PR is merged.

---

📄 25% (0.25x) speedup for ModelManager.add_model in inference/core/managers/base.py

⏱️ Runtime : 2.12 milliseconds → 1.70 millisecond (best of 42 runs)

📝 Explanation and details

To optimize the given Python program, we can make several improvements, mainly focusing on reducing the time spent in the add_model method. This includes reducing the repetitive calls to the logger.debug method and streamlining the model registration process. Below are the optimized versions of the classes.

Changes and Optimizations.

1. Improved Error Handling in get_model.

   • Replaced the if statement with a try-except block to catch KeyError directly, which is more efficient than checking and then accessing the dictionary.

2. Reduced Logger Calls in add_model.

   • Moved the logger debug statements to only occur at necessary points, avoiding redundant calls.
   • Combined the final log message into one call rather than multiple scattered calls during the process.

3. Streamlined Model Initialization in add_model.

   • Pulled the model class retrieval and initialization outside the logger debug call to minimize calculation duplication, focusing on reducing lookup overhead.

By reducing the number of logger debug calls, and optimizing model checks and retrievals, the overall performance of the add_model method is improved.

✅ Correctness verification report:

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

🌀 Generated Regression Tests Details

from concurrent.futures import ThreadPoolExecutor
from typing import Dict, Optional
from unittest.mock import MagicMock

# imports
import pytest  # used for our unit tests
# function to test
from inference.core.exceptions import ModelNotRecognisedError
from inference.core.logger import logger
from inference.core.managers.base import ModelManager
from inference.core.models.base import Model
from inference.core.registries.base import ModelRegistry

# unit tests

# Mock Model class for testing
class MockModel:
    def __init__(self, model_id, api_key):
        self.model_id = model_id
        self.api_key = api_key

# Mock ModelRegistry for testing
@pytest.fixture
def mock_model_registry():
    registry_dict = {
        "model_type_1": MockModel,
        "model_type_2": MockModel
    }
    return ModelRegistry(registry_dict)

@pytest.fixture
def model_manager(mock_model_registry):
    return ModelManager(mock_model_registry)

# Basic Functionality
def test_add_new_model_success(model_manager):
    model_manager.add_model("model_type_1", "api_key_1")


def test_prevent_duplicate_model_addition(model_manager):
    model_manager.add_model("model_type_1", "api_key_1")
    model_manager.add_model("model_type_1", "api_key_1")



def test_add_model_with_none_alias(model_manager):
    model_manager.add_model("model_type_1", "api_key_1", None)



def test_model_not_recognised_error(model_manager):
    with pytest.raises(ModelNotRecognisedError):
        model_manager.add_model("unknown_model_type", "api_key_1")

def test_model_not_recognised_error_with_alias(model_manager):
    with pytest.raises(ModelNotRecognisedError):
        model_manager.add_model("unknown_model_type", "api_key_1", "alias_1")

# Large Scale Test Cases


def test_model_initialization_and_retrieval(model_manager):
    model_manager.add_model("model_type_1", "api_key_1")
    model = model_manager._models["model_type_1"]

def test_multiple_model_initialization_and_retrieval(model_manager):
    model_manager.add_model("model_type_1", "api_key_1")
    model_manager.add_model("model_type_2", "api_key_2")

# Logging and Debugging
def test_logging_for_model_addition(model_manager, caplog):
    with caplog.at_level("DEBUG"):
        model_manager.add_model("model_type_1", "api_key_1")

def test_logging_for_duplicate_model_addition(model_manager, caplog):
    model_manager.add_model("model_type_1", "api_key_1")
    with caplog.at_level("DEBUG"):
        model_manager.add_model("model_type_1", "api_key_1")


def test_concurrent_model_addition(model_manager):
    def add_model_concurrently(i):
        model_manager.add_model(f"model_type_1_{i}", f"api_key_{i}")

    with ThreadPoolExecutor(max_workers=10) as executor:
        futures = [executor.submit(add_model_concurrently, i) for i in range(100)]
        for future in futures:
            future.result()

# Invalid API Keys





def test_model_retrieval_post_addition(model_manager):
    model_manager.add_model("model_type_1", "api_key_1")
    model = model_manager["model_type_1"]

def test_multiple_model_retrieval_post_addition(model_manager):
    model_manager.add_model("model_type_1", "api_key_1")
    model_manager.add_model("model_type_2", "api_key_2")
# 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, Optional
from unittest.mock import MagicMock

# imports
import pytest  # used for our unit tests
# function to test
from inference.core.exceptions import ModelNotRecognisedError
from inference.core.logger import logger
from inference.core.managers.base import ModelManager
from inference.core.models.base import Model
from inference.core.registries.base import ModelRegistry

# unit tests

# Mock Model class for testing
class MockModel:
    def __init__(self, model_id, api_key):
        self.model_id = model_id
        self.api_key = api_key

# Mock logger to avoid actual logging during tests
logger = MagicMock()

# Basic Functionality
def test_add_new_model():
    registry = ModelRegistry({"mock_model": MockModel})
    manager = ModelManager(registry)
    manager.add_model("mock_model", "api_key1")

def test_add_new_model_with_alias():
    registry = ModelRegistry({"alias_model": MockModel})
    manager = ModelManager(registry)
    manager.add_model("model_id", "api_key2", "alias_model")

# Duplicate Model Handling
def test_add_duplicate_model():
    registry = ModelRegistry({"mock_model": MockModel})
    manager = ModelManager(registry)
    manager.add_model("mock_model", "api_key1")
    manager.add_model("mock_model", "api_key1")

def test_add_duplicate_model_with_alias():
    registry = ModelRegistry({"alias_model": MockModel})
    manager = ModelManager(registry)
    manager.add_model("model_id", "api_key2", "alias_model")
    manager.add_model("model_id", "api_key2", "alias_model")

# Error Handling
def test_model_not_recognised():
    registry = ModelRegistry({})
    manager = ModelManager(registry)
    with pytest.raises(ModelNotRecognisedError):
        manager.add_model("unknown_model", "api_key1")

# Edge Cases
def test_add_model_with_empty_model_id():
    registry = ModelRegistry({"": MockModel})
    manager = ModelManager(registry)
    manager.add_model("", "api_key1")


def test_add_model_with_non_string_model_id():
    registry = ModelRegistry({123: MockModel})
    manager = ModelManager(registry)
    manager.add_model(123, "api_key1")

def test_add_model_with_non_string_alias():
    registry = ModelRegistry({456: MockModel})
    manager = ModelManager(registry)
    manager.add_model("model_id", "api_key2", 456)

# Large Scale Test Cases
def test_add_large_number_of_models():
    registry = ModelRegistry({f"model_{i}": MockModel for i in range(1000)})
    manager = ModelManager(registry)
    for i in range(1000):
        manager.add_model(f"model_{i}", f"api_key_{i}")

def test_add_model_with_long_identifiers():
    long_id = "a" * 1000
    registry = ModelRegistry({long_id: MockModel})
    manager = ModelManager(registry)
    manager.add_model(long_id, "api_key1")

# Special Characters in Identifiers

def test_case_sensitivity_in_identifiers():
    registry = ModelRegistry({"ModelID": MockModel, "modelid": MockModel})
    manager = ModelManager(registry)
    manager.add_model("ModelID", "api_key1")
    manager.add_model("modelid", "api_key2")

# Model Registry Variations
def test_different_model_registry_configurations():
    registry1 = ModelRegistry({"model1": MockModel})
    registry2 = ModelRegistry({"model2": MockModel})
    manager1 = ModelManager(registry1)
    manager2 = ModelManager(registry2)
    manager1.add_model("model1", "api_key1")
    manager2.add_model("model2", "api_key2")
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-pr1196-2025-04-17T09.36.29 and push.

[codeflash-ai]

This PR has been automatically closed because the original PR #1196 by PawelPeczek-Roboflow was closed.