From 52ecc8dae3cc3093ff7ee89d65b72b9423c96a04 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 15 Sep 2026 15:53:40 -0230 Subject: [PATCH 01/21] Cache explicit package selections independently and validate serving models --- inference/core/exceptions.py | 4 + inference/core/managers/active_learning.py | 12 +- inference/core/managers/base.py | 128 +++++++++++- inference/core/managers/decorators/base.py | 28 ++- .../managers/decorators/fixed_size_cache.py | 25 ++- .../core/managers/decorators/locked_load.py | 8 + inference/core/managers/decorators/logger.py | 9 +- inference/core/managers/stub_loader.py | 11 + .../core/managers/test_model_package_cache.py | 190 ++++++++++++++++++ 9 files changed, 403 insertions(+), 12 deletions(-) create mode 100644 tests/inference/unit_tests/core/managers/test_model_package_cache.py diff --git a/inference/core/exceptions.py b/inference/core/exceptions.py index 2327d4c2c1..1861f0f5c0 100644 --- a/inference/core/exceptions.py +++ b/inference/core/exceptions.py @@ -38,6 +38,10 @@ class InferenceModelNotFound(Exception): """ +class ModelPackageSelectionError(Exception): + pass + + class InvalidEnvironmentVariableError(Exception): """Raised when an environment variable is invalid. diff --git a/inference/core/managers/active_learning.py b/inference/core/managers/active_learning.py index ba5c5b24c4..c4737138a3 100644 --- a/inference/core/managers/active_learning.py +++ b/inference/core/managers/active_learning.py @@ -7,6 +7,7 @@ from inference.core.active_learning.middlewares import ActiveLearningMiddleware from inference.core.cache.base import BaseCache from inference.core.entities.requests.inference import InferenceRequest +from inference.core.entities.requests.model_selection import model_selection_kwargs from inference.core.entities.responses.inference import InferenceResponse from inference.core.env import DISABLE_PREPROC_AUTO_ORIENT from inference.core.managers.base import ModelManager @@ -82,7 +83,12 @@ def register( self, prediction: InferenceResponse, model_id: str, request: InferenceRequest ) -> None: try: - resolved_model_id = resolve_roboflow_model_alias(model_id=model_id) + public_model_id = ( + (request.model_id or model_id) + if model_selection_kwargs(request) + else model_id + ) + resolved_model_id = resolve_roboflow_model_alias(model_id=public_model_id) if not hasattr(request, "active_learning_target_dataset"): return None target_dataset = ( @@ -98,7 +104,9 @@ def register( ) self.register_datapoint( prediction=prediction, - model_id=resolved_model_id, + model_id=( + model_id if model_selection_kwargs(request) else resolved_model_id + ), request=request, middleware_key=middleware_key, ) diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index c51c785c1b..3f2f02ff26 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -10,6 +10,10 @@ from inference.core.cache.serializers import to_cachable_inference_item from inference.core.devices.utils import GLOBAL_INFERENCE_SERVER_ID from inference.core.entities.requests.inference import InferenceRequest +from inference.core.entities.requests.model_selection import ( + ModelSelectionRequest, + model_selection_kwargs, +) from inference.core.entities.responses.inference import InferenceResponse from inference.core.env import ( DISABLE_INFERENCE_CACHE, @@ -25,6 +29,7 @@ from inference.core.exceptions import ( InferenceModelNotFound, ModelManagerLockAcquisitionError, + ModelPackageSelectionError, RoboflowAPINotAuthorizedError, ) from inference.core.logger import logger @@ -128,6 +133,10 @@ def add_model( endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, service_secret: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + model_cache_key: Optional[str] = None, ) -> None: """Adds a new model to the manager. @@ -136,6 +145,26 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ + validate_public_model_id(model_id, model_id_alias) + selectors = { + name: value + for name, value in ( + ("model_package_id", model_package_id), + ("backend", backend), + ("quantization", quantization), + ) + if value is not None + } + if selectors and not USE_INFERENCE_MODELS: + raise ModelPackageSelectionError( + "Model package selection requires USE_INFERENCE_MODELS=true." + ) + if model_package_id is not None and ( + backend is not None or quantization is not None + ): + raise ModelPackageSelectionError( + "model_package_id cannot be combined with backend or quantization." + ) if MODELS_CACHE_AUTH_ENABLED and not OFFLINE_MODE: if not _check_if_api_key_has_access_to_model( api_key=api_key, @@ -151,7 +180,8 @@ def add_model( logger.debug( f"ModelManager - Adding model with model_id={model_id}, model_id_alias={model_id_alias}" ) - resolved_identifier = model_id if model_id_alias is None else model_id_alias + lookup_identifier = model_id if model_id_alias is None else model_id_alias + resolved_identifier = model_cache_key or lookup_identifier self.record_request_metadata( model_id=resolved_identifier, original_model_id=model_id, @@ -168,6 +198,7 @@ def add_model( f"Could not acquire lock for model with id={resolved_identifier}." ) if resolved_identifier in self._models: + self.validate_model_selection(resolved_identifier, **selectors) logger.debug( f"ModelManager - model with model_id={resolved_identifier} is already loaded." ) @@ -178,13 +209,14 @@ def add_model( t_load_start = time.perf_counter() vram_before = _get_cuda_memory_allocated() model_class = self.model_registry.get_model( - resolved_identifier, + lookup_identifier, api_key, countinference=countinference, service_secret=service_secret, ) extra_init_kwargs = {} + extra_init_kwargs.update(selectors) if USE_INFERENCE_MODELS: extra_init_kwargs["torchscript_state_global_lock"] = ( self.torchscript_state_global_lock @@ -199,8 +231,9 @@ def add_model( service_secret=service_secret, **extra_init_kwargs, ) + self._validate_loaded_model_selection(model, selectors) bind_usage_model_descriptor( - model, model_id, resolved_identifier, model_id_alias + model, model_id, lookup_identifier, model_id_alias ) vram_after = _get_cuda_memory_allocated() if vram_before is not None and vram_after is not None: @@ -249,6 +282,59 @@ def add_model( self._dispose_model_lock(model_id=resolved_identifier) raise error + def validate_model_selection( + self, + model_id: str, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + ) -> None: + selectors = { + name: value + for name, value in ( + ("model_package_id", model_package_id), + ("backend", backend), + ("quantization", quantization), + ) + if value is not None + } + if selectors: + self._validate_loaded_model_selection( + self._get_model_reference(model_id), selectors + ) + + @staticmethod + def _validate_loaded_model_selection( + model: Model, selectors: Dict[str, str] + ) -> None: + if not selectors: + return + if not USE_INFERENCE_MODELS: + raise ModelPackageSelectionError( + "Model package selection requires USE_INFERENCE_MODELS=true." + ) + resolved = getattr(model, "resolved_model", None) + if resolved is None: + raise ModelPackageSelectionError( + "This model does not support explicit package selection." + ) + from inference.core.env import ( + DISABLED_INFERENCE_MODELS_BACKENDS, + VALID_INFERENCE_MODELS_BACKENDS, + ) + + if resolved.backend not in ( + VALID_INFERENCE_MODELS_BACKENDS - DISABLED_INFERENCE_MODELS_BACKENDS + ): + raise ModelPackageSelectionError( + "The resolved model package uses a backend disabled on this server." + ) + for name, expected in selectors.items(): + if getattr(resolved, name, None) != expected: + raise ModelPackageSelectionError( + f"The loaded model does not satisfy the requested {name}." + ) + def load_action_recognition_model( self, model_id: str, api_key: Optional[str] = None, **kwargs ): @@ -518,12 +604,20 @@ def infer_from_request_sync( async def model_infer(self, model_id: str, request: InferenceRequest, **kwargs): model = self._get_model_reference(model_id=model_id) + if isinstance(request, ModelSelectionRequest): + self._validate_loaded_model_selection( + model, model_selection_kwargs(request) + ) return model.infer_from_request(request) def model_infer_sync( self, model_id: str, request: InferenceRequest, **kwargs ) -> Union[List[InferenceResponse], InferenceResponse]: model = self._get_model_reference(model_id=model_id) + if isinstance(request, ModelSelectionRequest): + self._validate_loaded_model_selection( + model, model_selection_kwargs(request) + ) return model.infer_from_request(request) def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: @@ -851,6 +945,34 @@ def _get_cuda_memory_allocated() -> Optional[int]: return None +def validate_public_model_id( + model_id: str, model_id_alias: Optional[str] = None +) -> None: + if any(":package:" in value for value in (model_id, model_id_alias) if value): + raise ModelPackageSelectionError( + "Cache handles cannot be used for inference or model loading. " + "Use the public model_id with package selection parameters." + ) + + +def model_load_options( + model_package_id: Optional[str], + backend: Optional[str], + quantization: Optional[str], + model_cache_key: Optional[str], +) -> Dict[str, str]: + return { + name: value + for name, value in ( + ("model_package_id", model_package_id), + ("backend", backend), + ("quantization", quantization), + ("model_cache_key", model_cache_key), + ) + if value is not None + } + + def try_releasing_cuda_memory() -> None: if not USE_INFERENCE_MODELS: # For now, we only want to have purge CUDA memory when USE_INFERENCE_MODELS is True diff --git a/inference/core/managers/decorators/base.py b/inference/core/managers/decorators/base.py index d17384a400..0774379799 100644 --- a/inference/core/managers/decorators/base.py +++ b/inference/core/managers/decorators/base.py @@ -6,7 +6,7 @@ from inference.core.entities.requests.inference import InferenceRequest from inference.core.entities.responses.inference import InferenceResponse from inference.core.env import API_KEY -from inference.core.managers.base import Model, ModelManager +from inference.core.managers.base import Model, ModelManager, model_load_options from inference.core.managers.model_load_collector import request_model_ids from inference.core.models.types import PreprocessReturnMetadata from inference.core.roboflow_api import ModelEndpointType @@ -63,6 +63,10 @@ def add_model( endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, service_secret: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + model_cache_key: Optional[str] = None, ): """Adds a model to the manager. @@ -71,15 +75,25 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ - if model_id in self: + from inference.core.managers.base import validate_public_model_id + + validate_public_model_id(model_id, model_id_alias) + cache_key = model_cache_key or model_id + if cache_key in self: + self.validate_model_selection( + cache_key, + model_package_id=model_package_id, + backend=backend, + quantization=quantization, + ) self.model_manager.record_request_metadata( - model_id=model_id, + model_id=cache_key, original_model_id=model_id, model_id_alias=model_id_alias, ) ids_collector = request_model_ids.get(None) if ids_collector is not None: - ids_collector.add(model_id) + ids_collector.add(cache_key) return self.model_manager.add_model( model_id, @@ -88,8 +102,14 @@ def add_model( endpoint_type=endpoint_type, countinference=countinference, service_secret=service_secret, + **model_load_options( + model_package_id, backend, quantization, model_cache_key + ), ) + def validate_model_selection(self, model_id: str, **selectors) -> None: + self.model_manager.validate_model_selection(model_id, **selectors) + def load_action_recognition_model( self, model_id: str, api_key: Optional[str] = None, **kwargs ): diff --git a/inference/core/managers/decorators/fixed_size_cache.py b/inference/core/managers/decorators/fixed_size_cache.py index d4e8c66347..73dbd66876 100644 --- a/inference/core/managers/decorators/fixed_size_cache.py +++ b/inference/core/managers/decorators/fixed_size_cache.py @@ -18,7 +18,12 @@ ModelManagerLockAcquisitionError, RoboflowAPINotAuthorizedError, ) -from inference.core.managers.base import Model, ModelManager, acquire_with_timeout +from inference.core.managers.base import ( + Model, + ModelManager, + acquire_with_timeout, + model_load_options, +) from inference.core.managers.decorators.base import ModelManagerDecorator from inference.core.managers.entities import ModelDescription from inference.core.managers.model_load_collector import request_model_ids @@ -59,6 +64,10 @@ def add_model( endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, service_secret: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + model_cache_key: Optional[str] = None, ) -> None: """Adds a model to the manager and evicts the least recently used if the cache is full. @@ -67,6 +76,9 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ + from inference.core.managers.base import validate_public_model_id + + validate_public_model_id(model_id, model_id_alias) if MODELS_CACHE_AUTH_ENABLED and not OFFLINE_MODE: if not _check_if_api_key_has_access_to_model( api_key=api_key, @@ -79,13 +91,19 @@ def add_model( f"API key {api_key} does not have access to model {model_id}" ) - queue_id = self._resolve_queue_id( + queue_id = model_cache_key or self._resolve_queue_id( model_id=model_id, model_id_alias=model_id_alias ) ids_collector = request_model_ids.get(None) if ids_collector is not None: ids_collector.add(queue_id) if queue_id in self: + self.validate_model_selection( + queue_id, + model_package_id=model_package_id, + backend=backend, + quantization=quantization, + ) logger.debug( f"Detected {queue_id} in WithFixedSizeCache models queue -> marking as most recently used." ) @@ -165,6 +183,9 @@ def add_model( endpoint_type=endpoint_type, countinference=countinference, service_secret=service_secret, + **model_load_options( + model_package_id, backend, quantization, model_cache_key + ), ) except Exception as error: logger.debug( diff --git a/inference/core/managers/decorators/locked_load.py b/inference/core/managers/decorators/locked_load.py index ead6a97e82..3eeafa859d 100644 --- a/inference/core/managers/decorators/locked_load.py +++ b/inference/core/managers/decorators/locked_load.py @@ -1,6 +1,7 @@ from typing import Optional from inference.core.cache import cache +from inference.core.managers.base import model_load_options from inference.core.managers.decorators.base import ModelManagerDecorator from inference.core.roboflow_api import ModelEndpointType @@ -18,6 +19,10 @@ def add_model( endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, service_secret: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + model_cache_key: Optional[str] = None, ): with cache.lock(lock_str(model_id), expire=180.0): return super().add_model( @@ -27,4 +32,7 @@ def add_model( endpoint_type=endpoint_type, countinference=countinference, service_secret=service_secret, + **model_load_options( + model_package_id, backend, quantization, model_cache_key + ), ) diff --git a/inference/core/managers/decorators/logger.py b/inference/core/managers/decorators/logger.py index cc00b775f8..8d6733d5e5 100644 --- a/inference/core/managers/decorators/logger.py +++ b/inference/core/managers/decorators/logger.py @@ -3,7 +3,7 @@ from inference.core.entities.requests.inference import InferenceRequest from inference.core.entities.responses.inference import InferenceResponse from inference.core.logger import logger -from inference.core.managers.base import Model +from inference.core.managers.base import Model, model_load_options from inference.core.managers.decorators.base import ModelManagerDecorator from inference.core.roboflow_api import ModelEndpointType @@ -19,6 +19,10 @@ def add_model( endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, service_secret: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + model_cache_key: Optional[str] = None, ): """Adds a model to the manager and logs the action. @@ -37,6 +41,9 @@ def add_model( endpoint_type=endpoint_type, countinference=countinference, service_secret=service_secret, + **model_load_options( + model_package_id, backend, quantization, model_cache_key + ), ) async def infer_from_request( diff --git a/inference/core/managers/stub_loader.py b/inference/core/managers/stub_loader.py index 85dfcddac4..dbf59cffa1 100644 --- a/inference/core/managers/stub_loader.py +++ b/inference/core/managers/stub_loader.py @@ -1,5 +1,6 @@ from typing import Optional +from inference.core.exceptions import ModelPackageSelectionError from inference.core.managers.base import ModelManager from inference.core.registries.base import ModelRegistry from inference.core.roboflow_api import ModelEndpointType @@ -30,6 +31,10 @@ def add_model( endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, service_secret: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + model_cache_key: Optional[str] = None, ) -> None: """Adds a new model to the manager. @@ -38,6 +43,12 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ + if any( + value is not None for value in (model_package_id, backend, quantization) + ): + raise ModelPackageSelectionError( + "Model package selection is unavailable in parallel stub mode." + ) if model_id in self._models: return model_class = self.model_registry.get_model( diff --git a/tests/inference/unit_tests/core/managers/test_model_package_cache.py b/tests/inference/unit_tests/core/managers/test_model_package_cache.py new file mode 100644 index 0000000000..3140037864 --- /dev/null +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -0,0 +1,190 @@ +from concurrent.futures import ThreadPoolExecutor +from threading import Event +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import pytest + +from inference.core import env +from inference.core.exceptions import ModelPackageSelectionError +from inference.core.managers import base as base_module +from inference.core.managers.active_learning import ActiveLearningManager +from inference.core.managers.base import ModelManager +from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache +from inference_models.entities import ResolvedModelMetadata + + +class PackageModel: + task_type = "object-detection" + batch_size = 1 + img_size_h = 32 + img_size_w = 32 + + def __init__(self, model_id, api_key, backend=None, **kwargs): + self.model_id = model_id + self.resolved_model = ResolvedModelMetadata( + model_id=model_id, + model_package_id=f"{backend or 'onnx'}-package", + backend=backend or "onnx", + quantization="fp16" if backend == "trt" else "fp32", + ) + + def infer_from_request(self, request): + return {"backend": self.resolved_model.backend} + + def clear_cache(self, delete_from_disk=True): + pass + + +def test_package_variants_coexist_without_changing_automatic_model(): + registry = Mock() + registry.get_model.return_value = PackageModel + with patch.object(base_module, "USE_INFERENCE_MODELS", True): + manager = WithFixedSizeCache( + ModelManager( + model_registry=registry, content_addressed_artifact_cache=Mock() + ), + max_size=3, + ) + manager.add_model("project/1", "key") + manager.add_model( + "project/1", "key", backend="trt", model_cache_key="project/1:trt" + ) + + assert set(manager.keys()) == {"project/1", "project/1:trt"} + assert manager.infer_from_request_sync("project/1", SimpleNamespace()) == { + "backend": "onnx" + } + assert manager.infer_from_request_sync("project/1:trt", SimpleNamespace()) == { + "backend": "trt" + } + assert all( + call.args[0] == "project/1" for call in registry.get_model.call_args_list + ) + + +@pytest.fixture +def package_manager(monkeypatch): + monkeypatch.setattr(base_module, "USE_INFERENCE_MODELS", True) + registry = Mock() + registry.get_model.return_value = PackageModel + return WithFixedSizeCache( + ModelManager(model_registry=registry, content_addressed_artifact_cache=Mock()), + max_size=3, + ) + + +def test_explicit_first_load_does_not_bind_automatic_selection(package_manager): + package_manager.add_model("project/1", "key", backend="trt", model_cache_key="trt") + package_manager.add_model("project/1", "key") + assert package_manager.infer_from_request_sync("project/1", SimpleNamespace()) == { + "backend": "onnx" + } + + +def test_cached_variant_is_validated_without_reloading(package_manager): + package_manager.add_model("project/1", "key", backend="trt", model_cache_key="trt") + package_manager.add_model("project/1", "key", backend="trt", model_cache_key="trt") + with pytest.raises(ModelPackageSelectionError, match="requested backend"): + package_manager.add_model( + "project/1", "key", backend="onnx", model_cache_key="trt" + ) + package_manager.model_manager.model_registry.get_model.assert_called_once() + + +def test_exact_package_id_cannot_bypass_disabled_backend(package_manager, monkeypatch): + monkeypatch.setattr(env, "DISABLED_INFERENCE_MODELS_BACKENDS", {"onnx"}) + with pytest.raises(ModelPackageSelectionError, match="disabled"): + package_manager.add_model( + "project/1", "key", model_package_id="onnx-package", model_cache_key="exact" + ) + assert not package_manager.keys() + + +def test_feature_flag_cannot_be_bypassed_by_warm_variant(package_manager, monkeypatch): + package_manager.add_model("project/1", "key", backend="trt", model_cache_key="trt") + monkeypatch.setattr(base_module, "USE_INFERENCE_MODELS", False) + with pytest.raises(ModelPackageSelectionError, match="USE_INFERENCE_MODELS"): + package_manager.add_model( + "project/1", "key", backend="trt", model_cache_key="trt" + ) + + +def test_removing_variant_leaves_automatic_model(package_manager): + package_manager.add_model("project/1", "key") + package_manager.add_model("project/1", "key", backend="trt", model_cache_key="trt") + package_manager.remove("trt") + assert list(package_manager.keys()) == ["project/1"] + + +@pytest.mark.parametrize("as_alias", [False, True]) +def test_cache_handle_cannot_bypass_selection_validation(package_manager, as_alias): + handle = "project/1:package:secret" + package_manager.add_model("project/1", "key", backend="trt", model_cache_key=handle) + with pytest.raises(ModelPackageSelectionError, match="Cache handles"): + package_manager.add_model( + "project/1" if as_alias else handle, + "different-key", + model_id_alias=handle if as_alias else None, + ) + + +def test_active_learning_uses_public_identity_and_selected_model_for_task_type(): + manager = ActiveLearningManager( + model_registry=Mock(), cache=Mock(), content_addressed_artifact_cache=Mock() + ) + manager._models["project/1:package:variant"] = PackageModel( + "project/1", "key", backend="trt" + ) + request = SimpleNamespace( + model_id="project/1", + backend="trt", + api_key="key", + image=[], + active_learning_target_dataset=None, + id="request-id", + ) + with patch( + "inference.core.managers.active_learning.ActiveLearningMiddleware.init" + ) as initialize: + manager.register( + prediction=Mock(), model_id="project/1:package:variant", request=request + ) + assert initialize.call_args.kwargs["model_id"] == "project/1" + assert ( + initialize.return_value.register_batch.call_args.kwargs["prediction_type"] + == "object-detection" + ) + + +def test_active_request_finishes_on_its_instance_after_lru_eviction(package_manager): + started, finish = Event(), Event() + + class BlockingPackageModel(PackageModel): + def infer_from_request(self, request): + if self.resolved_model.backend == "onnx": + started.set() + assert finish.wait(timeout=10) + return super().infer_from_request(request) + + package_manager.model_manager.model_registry.get_model.return_value = ( + BlockingPackageModel + ) + package_manager.max_size = 1 + package_manager.add_model("project/1", "key") + with ThreadPoolExecutor(max_workers=1) as executor: + active = executor.submit( + package_manager.infer_from_request_sync, "project/1", SimpleNamespace() + ) + try: + assert started.wait(timeout=10) + package_manager.add_model( + "project/1", "key", backend="trt", model_cache_key="trt" + ) + assert set(package_manager.keys()) == {"trt"} + assert package_manager.infer_from_request_sync( + "trt", SimpleNamespace() + ) == {"backend": "trt"} + finally: + finish.set() + assert active.result(timeout=10) == {"backend": "onnx"} From 896b061794b066c60b3ab4796665a18897571169 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 15 Sep 2026 15:55:01 -0230 Subject: [PATCH 02/21] Add explicit HTTP model package selection --- docs/http-model-package-selection.md | 41 ++++ inference/core/entities/requests/inference.py | 3 +- .../core/entities/requests/model_selection.py | 48 +++++ .../core/entities/requests/server_state.py | 8 +- .../core/interfaces/http/error_handlers.py | 5 + inference/core/interfaces/http/http_api.py | 90 +++++++- .../core/models/inference_models_adapters.py | 26 ++- inference_sdk/http/client.py | 44 +++- inference_sdk/http/entities.py | 45 ++++ inference_sdk/http/utils/executors.py | 11 +- inference_sdk/http/utils/model_selection.py | 17 ++ .../interfaces/http/test_model_selection.py | 202 ++++++++++++++++++ .../unit_tests/http/test_entities.py | 22 ++ .../unit_tests/http/test_model_selection.py | 119 +++++++++++ 14 files changed, 656 insertions(+), 25 deletions(-) create mode 100644 docs/http-model-package-selection.md create mode 100644 inference/core/entities/requests/model_selection.py create mode 100644 inference_sdk/http/utils/model_selection.py create mode 100644 tests/inference/unit_tests/core/interfaces/http/test_model_selection.py create mode 100644 tests/inference_sdk/unit_tests/http/test_model_selection.py diff --git a/docs/http-model-package-selection.md b/docs/http-model-package-selection.md new file mode 100644 index 0000000000..6792b2a363 --- /dev/null +++ b/docs/http-model-package-selection.md @@ -0,0 +1,41 @@ +# Select a model package over HTTP + +Model package selection requires a server with `USE_INFERENCE_MODELS=true`. +A request can specify `model_package_id`, or `backend`, `quantization`, or both. +The server rejects a package ID combined with either other selector. +A successful request uses a package that satisfies all explicit selectors. +An unavailable, incompatible, or disabled package produces an error. + +Standard computer vision endpoints accept these fields in the v1 JSON body. +The legacy `/{project}/{version}` endpoint accepts them as query parameters. +`/model/add` and `/model/remove` also accept them in the JSON body. + +```python +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +client = InferenceHTTPClient(api_url="http://localhost:9001", api_key="YOUR_API_KEY") +with client.use_configuration(InferenceConfiguration(backend="trt", quantization="fp16")): + result = client.infer("image.jpg", model_id="project/1") + client.unload_model("project/1") + +with client.use_configuration(InferenceConfiguration(model_package_id="PACKAGE_ID")): + client.load_model("project/1") + result = client.infer("image.jpg", model_id="project/1") +``` + +The configuration applies to synchronous and asynchronous SDK calls. +The SDK requires a server acknowledgment for explicit selection. +It raises an error if an older server silently ignores the selectors. +HTTP success responses include `X-Roboflow-Model-Selection: applied` for selected inference and model loading. + +Each distinct selector combination and client credential uses a separate cache entry. +Each entry counts toward the existing LRU limit. +The registry exposes its cache handle as `model_id`. +Two selector combinations can load separate instances of the same package. +Requests without selectors retain their existing automatic selection and cache entry. +Registration of a preferred package does not replace a resident instance. + +`unload_model()` removes the entry for the active configuration. +`/model/remove` accepts the original model ID with the same selectors, or the registry handle without selectors. +An active request retains its model reference until inference completes. +Concurrent package loads still require sufficient device memory. diff --git a/inference/core/entities/requests/inference.py b/inference/core/entities/requests/inference.py index f195e3adf0..a0e7b59b17 100644 --- a/inference/core/entities/requests/inference.py +++ b/inference/core/entities/requests/inference.py @@ -6,6 +6,7 @@ from inference.core import logger from inference.core.entities.common import ApiKey, ModelID, ModelType +from inference.core.entities.requests.model_selection import ModelSelectionRequest from inference_sdk.http.entities import Confidence @@ -74,7 +75,7 @@ class InferenceRequestImage(BaseModel): ) -class CVInferenceRequest(InferenceRequest): +class CVInferenceRequest(InferenceRequest, ModelSelectionRequest): """Computer Vision inference request. Attributes: diff --git a/inference/core/entities/requests/model_selection.py b/inference/core/entities/requests/model_selection.py new file mode 100644 index 0000000000..ce5edbd33c --- /dev/null +++ b/inference/core/entities/requests/model_selection.py @@ -0,0 +1,48 @@ +from typing import Optional + +from pydantic import BaseModel, Field, model_validator + +from inference.core import env +from inference_sdk.http.entities import ( + model_selection_cache_key as model_selection_cache_key, +) + + +class ModelSelectionRequest(BaseModel): + model_package_id: Optional[str] = Field( + default=None, + min_length=1, + description="Exact model package ID. Cannot be combined with backend or quantization.", + ) + backend: Optional[str] = Field( + default=None, + min_length=1, + description="Required package backend, such as trt or onnx.", + ) + quantization: Optional[str] = Field( + default=None, + min_length=1, + description="Required package quantization, such as fp16 or fp32.", + ) + + @model_validator(mode="after") + def validate_model_selection(self): + if self.model_package_id is not None and ( + self.backend is not None or self.quantization is not None + ): + raise ValueError( + "model_package_id cannot be combined with backend or quantization." + ) + if model_selection_kwargs(self) and not env.USE_INFERENCE_MODELS: + raise ValueError( + "Model package selection requires USE_INFERENCE_MODELS=true." + ) + return self + + +def model_selection_kwargs(request) -> dict: + return { + name: value + for name in ("model_package_id", "backend", "quantization") + if (value := getattr(request, name, None)) is not None + } diff --git a/inference/core/entities/requests/server_state.py b/inference/core/entities/requests/server_state.py index 917c665d66..547ecaf1d2 100644 --- a/inference/core/entities/requests/server_state.py +++ b/inference/core/entities/requests/server_state.py @@ -1,11 +1,12 @@ from typing import Optional -from pydantic import BaseModel, ConfigDict +from pydantic import ConfigDict from inference.core.entities.common import ApiKey, ModelID, ModelType +from inference.core.entities.requests.model_selection import ModelSelectionRequest -class AddModelRequest(BaseModel): +class AddModelRequest(ModelSelectionRequest): """Request to add a model to the inference server. Attributes: @@ -20,7 +21,7 @@ class AddModelRequest(BaseModel): api_key: Optional[str] = ApiKey -class ClearModelRequest(BaseModel): +class ClearModelRequest(ModelSelectionRequest): """Request to clear a model from the inference server. Attributes: @@ -29,3 +30,4 @@ class ClearModelRequest(BaseModel): model_config = ConfigDict(protected_namespaces=()) model_id: str = ModelID + api_key: Optional[str] = ApiKey diff --git a/inference/core/interfaces/http/error_handlers.py b/inference/core/interfaces/http/error_handlers.py index 5faa2436a7..d3f3831013 100644 --- a/inference/core/interfaces/http/error_handlers.py +++ b/inference/core/interfaces/http/error_handlers.py @@ -23,6 +23,7 @@ ModelArtefactError, ModelDeploymentNotSupportedError, ModelManagerLockAcquisitionError, + ModelPackageSelectionError, OnnxProviderNotAvailable, PayloadTooLargeError, PaymentRequiredError, @@ -499,6 +500,8 @@ def wrapped_route(*args, **kwargs): status_code=502, content={"message": "Internal error. Request to Roboflow API failed."}, ) + except ModelPackageSelectionError as error: + resp = JSONResponse(status_code=400, content={"message": str(error)}) except InferenceModelNotFound as error: logger.exception("%s: %s", type(error).__name__, error) resp = JSONResponse( @@ -979,6 +982,8 @@ async def wrapped_route(*args, **kwargs): status_code=502, content={"message": "Internal error. Request to Roboflow API failed."}, ) + except ModelPackageSelectionError as error: + resp = JSONResponse(status_code=400, content={"message": str(error)}) except InferenceModelNotFound as error: logger.exception("%s: %s", type(error).__name__, error) resp = JSONResponse( diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index 07d2a371cd..170f5d4f88 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -68,6 +68,11 @@ SemanticSegmentationInferenceRequest, ensure_wire_safe_mask_format, ) +from inference.core.entities.requests.model_selection import ( + ModelSelectionRequest, + model_selection_cache_key, + model_selection_kwargs, +) from inference.core.entities.requests.owlv2 import OwlV2InferenceRequest from inference.core.entities.requests.perception_encoder import ( PerceptionEncoderCompareRequest, @@ -391,6 +396,7 @@ from inference.usage_tracking.decorator_helpers import ( non_billable_intent_is_authenticated, ) +from inference_sdk.http.utils.model_selection import MODEL_SELECTION_HEADER if LAMBDA and not OFFLINE_MODE: from inference.core.usage import trackUsage @@ -1557,7 +1563,7 @@ def process_inference_request( countinference: Optional[bool] = None, service_secret: Optional[str] = None, **kwargs, - ) -> InferenceResponse: + ) -> Response: """Processes an inference request by calling the appropriate model. Args: @@ -1573,6 +1579,8 @@ def process_inference_request( inference_request.api_key = api_key ensure_wire_safe_mask_format(inference_request) requested_model_id = inference_request.model_id + if requested_model_id is None: + raise HTTPException(status_code=422, detail="model_id is required.") de_aliased_model_id = resolve_roboflow_model_alias( model_id=requested_model_id ) @@ -1581,12 +1589,25 @@ def process_inference_request( if de_aliased_model_id != requested_model_id else None ) + selectors = model_selection_kwargs(inference_request) + if selectors and isinstance(inference_request, LMMInferenceRequest): + raise HTTPException( + status_code=422, + detail="Model package selection is not supported for LMM requests.", + ) + cache_key = model_selection_cache_key( + de_aliased_model_id, selectors, inference_request.api_key + ) + selection_args = ( + {**selectors, "model_cache_key": cache_key} if selectors else {} + ) self.model_manager.add_model( de_aliased_model_id, inference_request.api_key, model_id_alias=model_id_alias, countinference=countinference, service_secret=service_secret, + **selection_args, ) inference_model_id = ( requested_model_id @@ -1594,11 +1615,14 @@ def process_inference_request( else de_aliased_model_id ) resp = self.model_manager.infer_from_request_sync( - inference_model_id, + cache_key if selectors else inference_model_id, inference_request, **kwargs, ) - return orjson_response(resp) + response = orjson_response(resp) + if selectors: + response.headers[MODEL_SELECTION_HEADER] = "applied" + return response def process_workflow_inference_request( workflow_request: WorkflowInferenceRequest, @@ -1943,6 +1967,7 @@ def registry(): @with_route_exceptions def model_add( request: AddModelRequest, + response: Response, countinference: Optional[bool] = None, service_secret: Optional[str] = None, ): @@ -1962,12 +1987,26 @@ def model_add( model_id=request.model_id ) logger.info(f"Loading model: {de_aliased_model_id}") + selectors = model_selection_kwargs(request) + selection_args = ( + { + **selectors, + "model_cache_key": model_selection_cache_key( + de_aliased_model_id, selectors, request.api_key + ), + } + if selectors + else {} + ) self.model_manager.add_model( de_aliased_model_id, request.api_key, countinference=countinference, service_secret=service_secret, + **selection_args, ) + if selectors: + response.headers[MODEL_SELECTION_HEADER] = "applied" models_descriptions = self.model_manager.describe_models() return ModelsDescriptions.from_models_descriptions( models_descriptions=models_descriptions @@ -1993,7 +2032,13 @@ def model_remove(request: ClearModelRequest): de_aliased_model_id = resolve_roboflow_model_alias( model_id=request.model_id ) - self.model_manager.remove(de_aliased_model_id) + self.model_manager.remove( + model_selection_cache_key( + de_aliased_model_id, + model_selection_kwargs(request), + api_key_override(request.api_key), + ) + ) models_descriptions = self.model_manager.describe_models() return ModelsDescriptions.from_models_descriptions( models_descriptions=models_descriptions @@ -4621,6 +4666,9 @@ def legacy_infer_from_request( None, description="Roboflow API Key that will be passed to the model during initialization for artifact retrieval", ), + model_package_id: Optional[str] = Query(None, min_length=1), + backend: Optional[str] = Query(None, min_length=1), + quantization: Optional[str] = Query(None, min_length=1), confidence: Confidence = Query( 0.4, description=( @@ -4815,15 +4863,34 @@ def legacy_infer_from_request( logger.debug( f"State of model registry: {self.model_manager.describe_models()}" ) + try: + selection = ModelSelectionRequest( + model_package_id=model_package_id, + backend=backend, + quantization=quantization, + ) + except ValidationError as error: + raise HTTPException(status_code=422, detail=str(error)) from error + selectors = model_selection_kwargs(selection) + cache_key = model_selection_cache_key(model_id, selectors, api_key) + selection_args = ( + {**selectors, "model_cache_key": cache_key} if selectors else {} + ) self.model_manager.add_model( request_model_id, api_key, model_id_alias=model_id, + **selection_args, countinference=countinference, service_secret=service_secret, ) - task_type = self.model_manager.get_task_type(model_id, api_key=api_key) + task_type = self.model_manager.get_task_type(cache_key, api_key=api_key) + if selectors and task_type == "action-recognition": + raise HTTPException( + status_code=422, + detail="Model package selection is not supported for action recognition requests.", + ) if task_type == "action-recognition": # The payload is a clip, so none of the image-shaped # arguments below apply to it. The `image` query parameter @@ -4850,7 +4917,7 @@ def legacy_infer_from_request( logger.debug("Response ready.") return orjson_response(inference_response) inference_request_type = ObjectDetectionInferenceRequest - args = dict() + args: Dict[str, Any] = {} if task_type == "instance-segmentation": inference_request_type = InstanceSegmentationInferenceRequest args = { @@ -4889,10 +4956,11 @@ def legacy_infer_from_request( source_info=source_info, usage_billable=countinference, disable_model_monitoring=disable_model_monitoring, + **selectors, **args, ) inference_response = self.model_manager.infer_from_request_sync( - inference_request.model_id, + cache_key, inference_request, active_learning_eligible=True, background_tasks=background_tasks, @@ -4902,9 +4970,15 @@ def legacy_infer_from_request( return Response( content=inference_response.visualization, media_type="image/jpeg", + headers=( + {MODEL_SELECTION_HEADER: "applied"} if selectors else None + ), ) else: - return orjson_response(inference_response) + response = orjson_response(inference_response) + if selectors: + response.headers[MODEL_SELECTION_HEADER] = "applied" + return response if not (LAMBDA or GCP_SERVERLESS): # Legacy clear cache endpoint for backwards compatibility diff --git a/inference/core/models/inference_models_adapters.py b/inference/core/models/inference_models_adapters.py index 85814074f8..46ab6f0f3d 100644 --- a/inference/core/models/inference_models_adapters.py +++ b/inference/core/models/inference_models_adapters.py @@ -214,6 +214,20 @@ def _get_enabled_inference_models_backends() -> List[str]: ) +def _get_requested_inference_models_backends(kwargs: dict) -> Union[str, List[str]]: + from inference.core.exceptions import ModelPackageSelectionError + + requested_backend = kwargs.pop("backend", None) + enabled_backends = _get_enabled_inference_models_backends() + if requested_backend is None: + return enabled_backends + if requested_backend not in enabled_backends: + raise ModelPackageSelectionError( + f"Backend {requested_backend!r} is not enabled on this server." + ) + return requested_backend + + def _supports_independent_stage_execution(pre_process) -> bool: """Return whether preprocessing declares the composed-execution control.""" try: @@ -271,7 +285,7 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): countinference=kwargs.get("countinference"), service_secret=kwargs.get("service_secret"), ) - backend = _get_enabled_inference_models_backends() + backend = _get_requested_inference_models_backends(kwargs) self._model: ObjectDetectionModel = AutoModel.from_pretrained( model_id_or_path=model_id, api_key=self.api_key, @@ -437,7 +451,7 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): countinference=kwargs.get("countinference"), service_secret=kwargs.get("service_secret"), ) - backend = _get_enabled_inference_models_backends() + backend = _get_requested_inference_models_backends(kwargs) self._model: InstanceSegmentationModel = AutoModel.from_pretrained( model_id_or_path=model_id, api_key=self.api_key, @@ -1150,7 +1164,7 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): countinference=kwargs.get("countinference"), service_secret=kwargs.get("service_secret"), ) - backend = _get_enabled_inference_models_backends() + backend = _get_requested_inference_models_backends(kwargs) self._model: KeyPointsDetectionModel = AutoModel.from_pretrained( model_id_or_path=model_id, api_key=self.api_key, @@ -1370,7 +1384,7 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): countinference=kwargs.get("countinference"), service_secret=kwargs.get("service_secret"), ) - backend = _get_enabled_inference_models_backends() + backend = _get_requested_inference_models_backends(kwargs) self._model: Union[ClassificationModel, MultiLabelClassificationModel] = ( AutoModel.from_pretrained( model_id_or_path=model_id, @@ -1774,7 +1788,7 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): countinference=kwargs.get("countinference"), service_secret=kwargs.get("service_secret"), ) - backend = _get_enabled_inference_models_backends() + backend = _get_requested_inference_models_backends(kwargs) self._model: SemanticSegmentationModel = AutoModel.from_pretrained( model_id_or_path=model_id, api_key=self.api_key, @@ -2172,7 +2186,7 @@ def load_action_recognition_model( countinference=kwargs.get("countinference"), service_secret=kwargs.get("service_secret"), ), - backend=_get_enabled_inference_models_backends(), + backend=_get_requested_inference_models_backends(kwargs), **kwargs, ) return _as_action_recognition_model(model=loaded_model, model_id=model_id) diff --git a/inference_sdk/http/client.py b/inference_sdk/http/client.py index 343289c30a..252b3129d4 100644 --- a/inference_sdk/http/client.py +++ b/inference_sdk/http/client.py @@ -41,6 +41,7 @@ RegisteredModels, ServerInfo, VideoReference, + model_selection_cache_key, ) from inference_sdk.http.errors import ( APIKeyNotProvided, @@ -81,6 +82,7 @@ load_stream_inference_input, uri_is_http_link, ) +from inference_sdk.http.utils.model_selection import ensure_model_selection_applied from inference_sdk.http.utils.post_processing import ( adjust_prediction_to_client_scaling_factor, combine_clip_embeddings, @@ -908,15 +910,20 @@ def get_model_description( self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) registered_models = self.list_loaded_models() + cache_key = model_selection_cache_key( + de_aliased_model_id, + self.__inference_configuration.to_model_selection_parameters(), + self.__api_key, + ) matching_model = filter_model_descriptions( descriptions=registered_models.models, - model_id=de_aliased_model_id, + model_id=cache_key, ) if matching_model is None and allow_loading is True: registered_models = self.load_model(model_id=de_aliased_model_id) matching_model = filter_model_descriptions( descriptions=registered_models.models, - model_id=de_aliased_model_id, + model_id=cache_key, ) if matching_model is not None: return matching_model @@ -944,9 +951,14 @@ async def get_model_description_async( self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) registered_models = await self.list_loaded_models_async() + cache_key = model_selection_cache_key( + de_aliased_model_id, + self.__inference_configuration.to_model_selection_parameters(), + self.__api_key, + ) matching_model = filter_model_descriptions( descriptions=registered_models.models, - model_id=de_aliased_model_id, + model_id=cache_key, ) if matching_model is None and allow_loading is True: registered_models = await self.load_model_async( @@ -954,7 +966,7 @@ async def get_model_description_async( ) matching_model = filter_model_descriptions( descriptions=registered_models.models, - model_id=de_aliased_model_id, + model_id=cache_key, ) if matching_model is not None: return matching_model @@ -1032,11 +1044,16 @@ def load_model( f"{self.__api_url}/model/add", json={ "model_id": de_aliased_model_id, + **self.__inference_configuration.to_model_selection_parameters(), **self.__legacy_api_key_payload(), }, headers=self.__headers_with_auth(DEFAULT_HEADERS), ) response.raise_for_status() + ensure_model_selection_applied( + response.headers, + self.__inference_configuration.to_model_selection_parameters(), + ) response_payload = response.json() if set_as_default: self.__selected_model = de_aliased_model_id @@ -1064,6 +1081,7 @@ async def load_model_async( de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) payload = { "model_id": de_aliased_model_id, + **self.__inference_configuration.to_model_selection_parameters(), **self.__legacy_api_key_payload(), } async with aiohttp.ClientSession() as session: @@ -1073,6 +1091,10 @@ async def load_model_async( headers=self.__headers_with_auth(DEFAULT_HEADERS), ) as response: response.raise_for_status() + ensure_model_selection_applied( + response.headers, + self.__inference_configuration.to_model_selection_parameters(), + ) response_payload = await response.json() if set_as_default: self.__selected_model = de_aliased_model_id @@ -1095,10 +1117,15 @@ def unload_model(self, model_id: str) -> RegisteredModels: """ self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) + cache_key = model_selection_cache_key( + de_aliased_model_id, + self.__inference_configuration.to_model_selection_parameters(), + self.__api_key, + ) response = requests.post( f"{self.__api_url}/model/remove", json={ - "model_id": de_aliased_model_id, + "model_id": cache_key, }, headers=self.__headers_with_auth(DEFAULT_HEADERS), ) @@ -1115,11 +1142,16 @@ def unload_model(self, model_id: str) -> RegisteredModels: async def unload_model_async(self, model_id: str) -> RegisteredModels: self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) + cache_key = model_selection_cache_key( + de_aliased_model_id, + self.__inference_configuration.to_model_selection_parameters(), + self.__api_key, + ) async with aiohttp.ClientSession() as session: async with session.post( f"{self.__api_url}/model/remove", json={ - "model_id": de_aliased_model_id, + "model_id": cache_key, }, headers=self.__headers_with_auth(DEFAULT_HEADERS), ) as response: diff --git a/inference_sdk/http/entities.py b/inference_sdk/http/entities.py index 0e82e17783..0117a82bc8 100644 --- a/inference_sdk/http/entities.py +++ b/inference_sdk/http/entities.py @@ -1,3 +1,5 @@ +import hashlib +import json from dataclasses import dataclass, field from enum import Enum from typing import Any, Dict, List, Literal, Optional, Tuple, Union @@ -141,6 +143,9 @@ class InferenceConfiguration: stroke_width: The stroke width for the inference. """ + model_package_id: Optional[str] = None + backend: Optional[str] = None + quantization: Optional[str] = None confidence_threshold: Optional[Confidence] = None keypoint_confidence_threshold: Optional[float] = None format: Optional[str] = None @@ -185,6 +190,12 @@ class InferenceConfiguration: api_key_transport: Optional[Union[str, ApiKeyTransport]] = None def __post_init__(self) -> None: + if self.model_package_id is not None and ( + self.backend is not None or self.quantization is not None + ): + raise InvalidParameterError( + "model_package_id cannot be combined with backend or quantization." + ) # Normalise the transport to the enum so the client can rely on # identity checks. NOTE: this field configures the credential CHANNEL # only - it is deliberately absent from every to_*_parameters() @@ -201,6 +212,13 @@ def __post_init__(self) -> None: f"one of: {[transport.value for transport in ApiKeyTransport]}." ) + def to_model_selection_parameters(self) -> Dict[str, str]: + return { + name: value + for name in ("model_package_id", "backend", "quantization") + if (value := getattr(self, name)) is not None + } + @classmethod def init_default(cls) -> "InferenceConfiguration": return cls() @@ -275,6 +293,9 @@ def to_object_detection_parameters(self) -> Dict[str, Any]: Dict[str, Any]: The object detection parameters. """ parameters_specs = [ + ("model_package_id", "model_package_id"), + ("backend", "backend"), + ("quantization", "quantization"), ("disable_preproc_auto_orientation", "disable_preproc_auto_orient"), ("disable_preproc_contrast", "disable_preproc_contrast"), ("disable_preproc_grayscale", "disable_preproc_grayscale"), @@ -317,6 +338,9 @@ def to_instance_segmentation_parameters(self) -> Dict[str, Any]: """ parameters = self.to_object_detection_parameters() parameters_specs = [ + ("model_package_id", "model_package_id"), + ("backend", "backend"), + ("quantization", "quantization"), ("mask_decode_mode", "mask_decode_mode"), ("tradeoff_factor", "tradeoff_factor"), ("response_mask_format", "response_mask_format"), @@ -333,6 +357,9 @@ def to_classification_parameters(self) -> Dict[str, Any]: """ parameters_specs = [ ("include_anomaly_map", "include_anomaly_map"), + ("model_package_id", "model_package_id"), + ("backend", "backend"), + ("quantization", "quantization"), ("disable_preproc_auto_orientation", "disable_preproc_auto_orient"), ("disable_preproc_contrast", "disable_preproc_contrast"), ("disable_preproc_grayscale", "disable_preproc_grayscale"), @@ -361,6 +388,9 @@ def to_legacy_call_parameters(self) -> Dict[str, Any]: Dict[str, Any]: The legacy call parameters. """ parameters_specs = [ + ("model_package_id", "model_package_id"), + ("backend", "backend"), + ("quantization", "quantization"), ("confidence_threshold", "confidence"), ("keypoint_confidence_threshold", "keypoint_confidence"), ("format", "format"), @@ -407,3 +437,18 @@ def get_non_empty_attributes( for internal_name, external_name in specification } return remove_empty_values(dictionary=attributes) + + +def model_selection_cache_key( + model_id: str, selectors: dict, api_key: Optional[str] = None +) -> str: + if not selectors: + return model_id + digest = hashlib.sha256( + json.dumps( + {"selectors": selectors, "api_key": api_key}, + sort_keys=True, + separators=(",", ":"), + ).encode() + ).hexdigest() + return f"{model_id}:package:{digest}" diff --git a/inference_sdk/http/utils/executors.py b/inference_sdk/http/utils/executors.py index 723121ff8e..9533247d4e 100644 --- a/inference_sdk/http/utils/executors.py +++ b/inference_sdk/http/utils/executors.py @@ -27,6 +27,7 @@ ) from inference_sdk.http.errors import RetryError from inference_sdk.http.utils.iterables import make_batches +from inference_sdk.http.utils.model_selection import ensure_model_selection_applied from inference_sdk.http.utils.request_building import RequestData from inference_sdk.http.utils.requests import api_key_safe_raise_for_status @@ -81,8 +82,12 @@ def execute_requests_packages( results.extend(responses) all_request_data.extend(requests_data_package) _collect_remote_processing_times(results, all_request_data) - for response in results: + for response, request_data in zip(results, all_request_data): api_key_safe_raise_for_status(response=response) + ensure_model_selection_applied( + response.headers, + {**(request_data.parameters or {}), **(request_data.payload or {})}, + ) return results @@ -499,6 +504,10 @@ async def make_request_async( if response_is_not_retryable_error(response=response): response.raise_for_status() if response.status == 200: + ensure_model_selection_applied( + response.headers, + {**(request_data.parameters or {}), **(request_data.payload or {})}, + ) collect_remote_processing_metadata_from_headers( headers=response.headers, request_data=request_data, diff --git a/inference_sdk/http/utils/model_selection.py b/inference_sdk/http/utils/model_selection.py new file mode 100644 index 0000000000..1f0f0dd183 --- /dev/null +++ b/inference_sdk/http/utils/model_selection.py @@ -0,0 +1,17 @@ +from typing import Mapping + +from inference_sdk.http.errors import InvalidParameterError + +MODEL_SELECTION_HEADER = "X-Roboflow-Model-Selection" + + +def ensure_model_selection_applied(headers: Mapping[str, str], selectors: dict) -> None: + if any( + selectors.get(name) is not None + for name in ("model_package_id", "backend", "quantization") + ): + if headers.get(MODEL_SELECTION_HEADER) != "applied": + raise InvalidParameterError( + "The server did not confirm model package selection. Upgrade the server " + "and enable USE_INFERENCE_MODELS before requesting a specific package." + ) diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py new file mode 100644 index 0000000000..817d40bcec --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -0,0 +1,202 @@ +from unittest.mock import MagicMock + +import pytest +from pydantic import BaseModel +from starlette.testclient import TestClient + +from inference.core.interfaces.http import http_api + + +@pytest.fixture +def client(monkeypatch): + monkeypatch.setattr(http_api, "InferenceInstrumentator", MagicMock()) + monkeypatch.setattr(http_api, "OFFLINE_MODE", True) + monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) + monkeypatch.setattr(http_api, "LAMBDA", False) + monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) + interface = http_api.HttpInterface(model_manager=MagicMock()) + with TestClient(interface.app) as client: + yield client + + +@pytest.mark.parametrize("selector", [{"backend": "trt"}, {"quantization": "fp16"}]) +def test_http_rejects_package_id_combined_with_other_selectors(client, selector): + response = client.post( + "/infer/object_detection", + json={ + "model_id": "project/1", + "model_package_id": "package-1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + **selector, + }, + ) + assert response.status_code == 422 + assert "model_package_id" in response.text + + +@pytest.mark.parametrize("path", ["/infer/object_detection", "/model/add"]) +def test_http_rejects_selection_when_flag_is_disabled(client, monkeypatch, path): + from inference.core import env + + monkeypatch.setattr(env, "USE_INFERENCE_MODELS", False) + response = client.post( + path, + json={ + "model_id": "project/1", + "backend": "trt", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + }, + ) + assert response.status_code == 422 + assert "USE_INFERENCE_MODELS" in response.text + + +@pytest.fixture +def package_client(monkeypatch): + from types import SimpleNamespace + + from inference.core import env + from inference.core.managers import base + from inference.core.managers.base import ModelManager + from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache + + class PackageResponse(BaseModel): + served_package: str + + class PackageModel: + task_type = "object-detection" + batch_size = 1 + img_size_h = 32 + img_size_w = 32 + + def __init__( + self, model_id, api_key, backend=None, model_package_id=None, **kwargs + ): + if api_key == "other-workspace" and model_package_id == "engine-1": + from inference.core.exceptions import RoboflowAPINotAuthorizedError + + raise RoboflowAPINotAuthorizedError( + "Package is not accessible to this workspace." + ) + selected_backend = backend or ( + "trt" if model_package_id == "engine-1" else "onnx" + ) + self.resolved_model = SimpleNamespace( + model_id=model_id, + model_package_id="engine-1" if selected_backend == "trt" else "onnx-1", + backend=selected_backend, + quantization="fp16" if selected_backend == "trt" else "fp32", + ) + + def infer_from_request(self, request): + return PackageResponse(served_package=self.resolved_model.model_package_id) + + def clear_cache(self, delete_from_disk=True): + pass + + monkeypatch.setattr(env, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr(base, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr(base, "OFFLINE_MODE", True) + monkeypatch.setattr(http_api, "InferenceInstrumentator", MagicMock()) + monkeypatch.setattr(http_api, "OFFLINE_MODE", True) + monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) + monkeypatch.setattr(http_api, "LAMBDA", False) + monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) + registry = MagicMock() + registry.get_model.return_value = PackageModel + manager = WithFixedSizeCache( + ModelManager(registry, content_addressed_artifact_cache=MagicMock()), max_size=3 + ) + interface = http_api.HttpInterface(model_manager=manager) + with TestClient(interface.app) as client: + yield client + + +@pytest.mark.parametrize( + "selectors", + [{"backend": "trt", "quantization": "fp16"}, {"model_package_id": "engine-1"}], +) +def test_http_selected_package_coexists_with_automatic_package( + package_client, selectors +): + payload = { + "model_id": "project/1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + "disable_model_monitoring": True, + } + automatic = package_client.post("/infer/object_detection", json=payload) + selected = package_client.post( + "/infer/object_detection", json={**payload, **selectors} + ) + again = package_client.post("/infer/object_detection", json=payload) + assert automatic.status_code == selected.status_code == again.status_code == 200 + assert ( + automatic.json()["served_package"] == again.json()["served_package"] == "onnx-1" + ) + assert selected.json()["served_package"] == "engine-1" + registered = package_client.get("/model/registry").json()["models"] + assert len(registered) == 2 + removed = package_client.post( + "/model/remove", json={"model_id": "project/1", **selectors} + ) + assert removed.status_code == 200 + assert [model["model_id"] for model in removed.json()["models"]] == ["project/1"] + + +def test_legacy_http_selection_is_enforced_and_acknowledged(package_client): + response = package_client.post( + "/project/1", + params={ + "backend": "trt", + "quantization": "fp16", + "image": "https://example.com/image.jpg", + }, + ) + assert response.status_code == 200 + assert response.json()["served_package"] == "engine-1" + assert response.headers["X-Roboflow-Model-Selection"] == "applied" + + +def test_legacy_http_rejects_conflicting_selectors(package_client): + response = package_client.post( + "/project/1", + params={ + "backend": "trt", + "model_package_id": "engine-1", + "image": "https://example.com/image.jpg", + }, + ) + assert response.status_code == 422 + assert "model_package_id" in response.text + + +def test_http_cannot_reuse_another_credentials_restricted_package(package_client): + payload = { + "model_id": "project/1", + "model_package_id": "engine-1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + "disable_model_monitoring": True, + } + allowed = package_client.post( + "/infer/object_detection", json={**payload, "api_key": "owner"} + ) + denied = package_client.post( + "/infer/object_detection", json={**payload, "api_key": "other-workspace"} + ) + assert allowed.status_code == 200 + assert denied.status_code in (401, 403) + assert "engine-1" not in denied.json().get("served_package", "") + + +def test_lmm_rejects_model_package_selection(package_client): + response = package_client.post( + "/infer/lmm", + json={ + "model_id": "project/1", + "backend": "trt", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + "prompt": "Describe the image.", + }, + ) + assert response.status_code == 422 + assert "not supported for LMM" in response.text diff --git a/tests/inference_sdk/unit_tests/http/test_entities.py b/tests/inference_sdk/unit_tests/http/test_entities.py index 01f1e99186..30861c9f55 100644 --- a/tests/inference_sdk/unit_tests/http/test_entities.py +++ b/tests/inference_sdk/unit_tests/http/test_entities.py @@ -1,3 +1,5 @@ +from typing import Any + from inference_sdk.config import outbound_service_secret from inference_sdk.http.entities import ( CLASSIFICATION_TASK, @@ -282,3 +284,23 @@ def test_to_legacy_call_parameters_context_overrides_explicit_count_inference_tr # then assert parameters["countinference"] is False assert parameters["service_secret"] == "ctx-secret" + + +def test_model_selection_is_sent_in_both_api_versions() -> None: + configuration = InferenceConfiguration(backend="trt", quantization="fp16") + for mode in (HTTPClientMode.V0, HTTPClientMode.V1): + for task in (OBJECT_DETECTION_TASK, CLASSIFICATION_TASK): + parameters = configuration.to_api_call_parameters(mode, task) + assert parameters["backend"] == "trt" + assert parameters["quantization"] == "fp16" + + +def test_exact_package_cannot_be_combined_with_other_selectors() -> None: + import pytest + + from inference_sdk.http.errors import InvalidParameterError + + selectors: list[dict[str, Any]] = [{"backend": "trt"}, {"quantization": "fp16"}] + for selector in selectors: + with pytest.raises(InvalidParameterError, match="model_package_id"): + InferenceConfiguration(model_package_id="package-1", **selector) diff --git a/tests/inference_sdk/unit_tests/http/test_model_selection.py b/tests/inference_sdk/unit_tests/http/test_model_selection.py new file mode 100644 index 0000000000..42056c3da8 --- /dev/null +++ b/tests/inference_sdk/unit_tests/http/test_model_selection.py @@ -0,0 +1,119 @@ +from typing import Any + +import numpy as np +import pytest +from aioresponses import aioresponses + +from inference_sdk import InferenceConfiguration, InferenceHTTPClient +from inference_sdk.http.entities import model_selection_cache_key +from inference_sdk.http.errors import InvalidParameterError + + +@pytest.mark.parametrize("acknowledged", [True, False]) +def test_v0_selection_requires_server_acknowledgment(requests_mock, acknowledged): + headers = {"X-Roboflow-Model-Selection": "applied"} if acknowledged else {} + requests_mock.post( + "http://server/project/1", json={"predictions": []}, headers=headers + ) + client = InferenceHTTPClient(api_url="http://server", api_key="key").select_api_v0() + client.configure(InferenceConfiguration(backend="trt", quantization="fp16")) + if acknowledged: + assert client.infer( + np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1" + ) == {"predictions": []} + else: + with pytest.raises(InvalidParameterError, match="did not confirm"): + client.infer(np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1") + assert requests_mock.last_request.qs["backend"] == ["trt"] + assert requests_mock.last_request.qs["quantization"] == ["fp16"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("acknowledged", [True, False]) +async def test_async_v0_selection_requires_server_acknowledgment(acknowledged): + import re + + headers = {"X-Roboflow-Model-Selection": "applied"} if acknowledged else {} + client = InferenceHTTPClient(api_url="http://server", api_key="key").select_api_v0() + client.configure(InferenceConfiguration(model_package_id="engine-1")) + with aioresponses() as responses: + responses.post( + re.compile(r"http://server/project/1.*"), + payload={"predictions": []}, + headers=headers, + ) + if acknowledged: + assert await client.infer_async( + np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1" + ) == {"predictions": []} + else: + with pytest.raises(InvalidParameterError, match="did not confirm"): + await client.infer_async( + np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1" + ) + + +def test_v1_description_loads_selected_variant_without_loading_default(requests_mock): + selectors: dict[str, Any] = {"backend": "trt", "quantization": "fp16"} + cache_key = model_selection_cache_key("project/1", selectors, "key") + requests_mock.get( + "http://server/model/registry", + json={"models": [{"model_id": "project/1", "task_type": "object-detection"}]}, + ) + requests_mock.post( + "http://server/model/add", + json={"models": [{"model_id": cache_key, "task_type": "object-detection"}]}, + headers={"X-Roboflow-Model-Selection": "applied"}, + ) + client = InferenceHTTPClient(api_url="http://server", api_key="key") + client.configure(InferenceConfiguration(**selectors)) + assert client.get_model_description("project/1").model_id == cache_key + assert requests_mock.last_request.json() == { + "model_id": "project/1", + "api_key": "key", + **selectors, + } + + +@pytest.mark.asyncio +async def test_async_v1_description_loads_selected_variant_without_loading_default(): + selectors: dict[str, Any] = {"model_package_id": "engine-1"} + cache_key = model_selection_cache_key("project/1", selectors, "key") + client = InferenceHTTPClient(api_url="http://server", api_key="key") + client.configure(InferenceConfiguration(**selectors)) + with aioresponses() as responses: + responses.get( + "http://server/model/registry?api_key=key", + payload={ + "models": [{"model_id": "project/1", "task_type": "object-detection"}] + }, + ) + responses.post( + "http://server/model/add", + payload={ + "models": [{"model_id": cache_key, "task_type": "object-detection"}] + }, + headers={"X-Roboflow-Model-Selection": "applied"}, + ) + assert ( + await client.get_model_description_async("project/1") + ).model_id == cache_key + post = [ + calls[0] for key, calls in responses.requests.items() if key[0] == "POST" + ][0] + assert post.kwargs["json"] == { + "model_id": "project/1", + "api_key": "key", + **selectors, + } + + +def test_unload_selected_variant_never_sends_default_model_id(requests_mock): + selectors: dict[str, Any] = {"model_package_id": "engine-1"} + requests_mock.post("http://server/model/remove", json={"models": []}) + client = InferenceHTTPClient(api_url="http://server", api_key="key") + client.configure(InferenceConfiguration(**selectors)) + client.unload_model("project/1") + assert requests_mock.last_request.json() == { + "model_id": model_selection_cache_key("project/1", selectors, "key") + } From 337bf4407ac965b3b925a05bd9b94ed16c1f1dab Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 15 Sep 2026 15:56:21 -0230 Subject: [PATCH 03/21] Preserve positional SDK configuration arguments --- docs/http-model-package-selection.md | 4 ++-- inference_sdk/http/entities.py | 9 +++------ tests/inference_sdk/unit_tests/http/test_entities.py | 6 ++++++ 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/docs/http-model-package-selection.md b/docs/http-model-package-selection.md index 6792b2a363..3462a329fc 100644 --- a/docs/http-model-package-selection.md +++ b/docs/http-model-package-selection.md @@ -28,7 +28,7 @@ The SDK requires a server acknowledgment for explicit selection. It raises an error if an older server silently ignores the selectors. HTTP success responses include `X-Roboflow-Model-Selection: applied` for selected inference and model loading. -Each distinct selector combination and client credential uses a separate cache entry. +Each model ID, selector combination, and client credential identifies a separate cache entry. Each entry counts toward the existing LRU limit. The registry exposes its cache handle as `model_id`. Two selector combinations can load separate instances of the same package. @@ -36,6 +36,6 @@ Requests without selectors retain their existing automatic selection and cache e Registration of a preferred package does not replace a resident instance. `unload_model()` removes the entry for the active configuration. -`/model/remove` accepts the original model ID with the same selectors, or the registry handle without selectors. +`/model/remove` accepts the original model ID with the same selectors and credential, or the registry handle without selectors. An active request retains its model reference until inference completes. Concurrent package loads still require sufficient device memory. diff --git a/inference_sdk/http/entities.py b/inference_sdk/http/entities.py index 0117a82bc8..5cf9940033 100644 --- a/inference_sdk/http/entities.py +++ b/inference_sdk/http/entities.py @@ -143,9 +143,6 @@ class InferenceConfiguration: stroke_width: The stroke width for the inference. """ - model_package_id: Optional[str] = None - backend: Optional[str] = None - quantization: Optional[str] = None confidence_threshold: Optional[Confidence] = None keypoint_confidence_threshold: Optional[float] = None format: Optional[str] = None @@ -188,6 +185,9 @@ class InferenceConfiguration: # and emits a one-time recommendation to move to the header transport. # Pass "legacy" explicitly to keep the old behaviour silently. api_key_transport: Optional[Union[str, ApiKeyTransport]] = None + model_package_id: Optional[str] = None + backend: Optional[str] = None + quantization: Optional[str] = None def __post_init__(self) -> None: if self.model_package_id is not None and ( @@ -338,9 +338,6 @@ def to_instance_segmentation_parameters(self) -> Dict[str, Any]: """ parameters = self.to_object_detection_parameters() parameters_specs = [ - ("model_package_id", "model_package_id"), - ("backend", "backend"), - ("quantization", "quantization"), ("mask_decode_mode", "mask_decode_mode"), ("tradeoff_factor", "tradeoff_factor"), ("response_mask_format", "response_mask_format"), diff --git a/tests/inference_sdk/unit_tests/http/test_entities.py b/tests/inference_sdk/unit_tests/http/test_entities.py index 30861c9f55..9a73502dc7 100644 --- a/tests/inference_sdk/unit_tests/http/test_entities.py +++ b/tests/inference_sdk/unit_tests/http/test_entities.py @@ -304,3 +304,9 @@ def test_exact_package_cannot_be_combined_with_other_selectors() -> None: for selector in selectors: with pytest.raises(InvalidParameterError, match="model_package_id"): InferenceConfiguration(model_package_id="package-1", **selector) + + +def test_positional_configuration_preserves_confidence_threshold() -> None: + configuration = InferenceConfiguration(0.5) + assert configuration.confidence_threshold == 0.5 + assert configuration.to_model_selection_parameters() == {} From 30ec16dddf5d87a689b9d1dde07166af9438f669 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 15 Sep 2026 15:58:01 -0230 Subject: [PATCH 04/21] Document library rollout requirement for package selection --- docs/http-model-package-selection.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/http-model-package-selection.md b/docs/http-model-package-selection.md index 3462a329fc..1c092398a7 100644 --- a/docs/http-model-package-selection.md +++ b/docs/http-model-package-selection.md @@ -1,6 +1,8 @@ # Select a model package over HTTP Model package selection requires a server with `USE_INFERENCE_MODELS=true`. +The server also needs an `inference-models` build that exposes the resolved package descriptor. +Older library builds remain importable, but explicit selection fails without that descriptor. A request can specify `model_package_id`, or `backend`, `quantization`, or both. The server rejects a package ID combined with either other selector. A successful request uses a package that satisfies all explicit selectors. @@ -24,7 +26,7 @@ with client.use_configuration(InferenceConfiguration(model_package_id="PACKAGE_I ``` The configuration applies to synchronous and asynchronous SDK calls. -The SDK requires a server acknowledgment for explicit selection. +For inference and model loading, the SDK requires a server acknowledgment for explicit selection. It raises an error if an older server silently ignores the selectors. HTTP success responses include `X-Roboflow-Model-Selection: applied` for selected inference and model loading. From eba5f49419eb0c7fd12a95d295ad99e28af2a831 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 15 Sep 2026 15:59:16 -0230 Subject: [PATCH 05/21] Exercise package cache with typed inference requests --- .../core/managers/test_model_package_cache.py | 44 ++++++++++++------- 1 file changed, 27 insertions(+), 17 deletions(-) diff --git a/tests/inference/unit_tests/core/managers/test_model_package_cache.py b/tests/inference/unit_tests/core/managers/test_model_package_cache.py index 3140037864..8dfa3efbb2 100644 --- a/tests/inference/unit_tests/core/managers/test_model_package_cache.py +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -1,20 +1,25 @@ from concurrent.futures import ThreadPoolExecutor from threading import Event -from types import SimpleNamespace +from typing import Any from unittest.mock import Mock, patch import pytest from inference.core import env +from inference.core.entities.requests.inference import ( + InferenceRequest, + ObjectDetectionInferenceRequest, +) from inference.core.exceptions import ModelPackageSelectionError from inference.core.managers import base as base_module from inference.core.managers.active_learning import ActiveLearningManager from inference.core.managers.base import ModelManager from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache +from inference.core.models.base import Model from inference_models.entities import ResolvedModelMetadata -class PackageModel: +class PackageModel(Model): task_type = "object-detection" batch_size = 1 img_size_h = 32 @@ -29,7 +34,7 @@ def __init__(self, model_id, api_key, backend=None, **kwargs): quantization="fp16" if backend == "trt" else "fp32", ) - def infer_from_request(self, request): + def infer_from_request(self, request) -> Any: return {"backend": self.resolved_model.backend} def clear_cache(self, delete_from_disk=True): @@ -52,12 +57,12 @@ def test_package_variants_coexist_without_changing_automatic_model(): ) assert set(manager.keys()) == {"project/1", "project/1:trt"} - assert manager.infer_from_request_sync("project/1", SimpleNamespace()) == { - "backend": "onnx" - } - assert manager.infer_from_request_sync("project/1:trt", SimpleNamespace()) == { - "backend": "trt" - } + assert manager.infer_from_request_sync( + "project/1", InferenceRequest(id="request-id", model_id="project/1") + ) == {"backend": "onnx"} + assert manager.infer_from_request_sync( + "project/1:trt", InferenceRequest(id="request-id", model_id="project/1") + ) == {"backend": "trt"} assert all( call.args[0] == "project/1" for call in registry.get_model.call_args_list ) @@ -77,9 +82,9 @@ def package_manager(monkeypatch): def test_explicit_first_load_does_not_bind_automatic_selection(package_manager): package_manager.add_model("project/1", "key", backend="trt", model_cache_key="trt") package_manager.add_model("project/1", "key") - assert package_manager.infer_from_request_sync("project/1", SimpleNamespace()) == { - "backend": "onnx" - } + assert package_manager.infer_from_request_sync( + "project/1", InferenceRequest(id="request-id", model_id="project/1") + ) == {"backend": "onnx"} def test_cached_variant_is_validated_without_reloading(package_manager): @@ -129,14 +134,17 @@ def test_cache_handle_cannot_bypass_selection_validation(package_manager, as_ali ) -def test_active_learning_uses_public_identity_and_selected_model_for_task_type(): +def test_active_learning_uses_public_identity_and_selected_model_for_task_type( + monkeypatch, +): + monkeypatch.setattr(env, "USE_INFERENCE_MODELS", True) manager = ActiveLearningManager( model_registry=Mock(), cache=Mock(), content_addressed_artifact_cache=Mock() ) manager._models["project/1:package:variant"] = PackageModel( "project/1", "key", backend="trt" ) - request = SimpleNamespace( + request = ObjectDetectionInferenceRequest( model_id="project/1", backend="trt", api_key="key", @@ -161,7 +169,7 @@ def test_active_request_finishes_on_its_instance_after_lru_eviction(package_mana started, finish = Event(), Event() class BlockingPackageModel(PackageModel): - def infer_from_request(self, request): + def infer_from_request(self, request) -> Any: if self.resolved_model.backend == "onnx": started.set() assert finish.wait(timeout=10) @@ -174,7 +182,9 @@ def infer_from_request(self, request): package_manager.add_model("project/1", "key") with ThreadPoolExecutor(max_workers=1) as executor: active = executor.submit( - package_manager.infer_from_request_sync, "project/1", SimpleNamespace() + package_manager.infer_from_request_sync, + "project/1", + InferenceRequest(id="request-id", model_id="project/1"), ) try: assert started.wait(timeout=10) @@ -183,7 +193,7 @@ def infer_from_request(self, request): ) assert set(package_manager.keys()) == {"trt"} assert package_manager.infer_from_request_sync( - "trt", SimpleNamespace() + "trt", InferenceRequest(id="request-id", model_id="project/1") ) == {"backend": "trt"} finally: finish.set() From f6babbb19e909fa28ecd5ea878654840015d6ef7 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 15 Sep 2026 16:35:22 -0230 Subject: [PATCH 06/21] Require the updated inference-models release for package selection --- docs/http-model-package-selection.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/http-model-package-selection.md b/docs/http-model-package-selection.md index 1c092398a7..b82e2ac08b 100644 --- a/docs/http-model-package-selection.md +++ b/docs/http-model-package-selection.md @@ -1,8 +1,8 @@ # Select a model package over HTTP Model package selection requires a server with `USE_INFERENCE_MODELS=true`. -The server also needs an `inference-models` build that exposes the resolved package descriptor. -Older library builds remain importable, but explicit selection fails without that descriptor. +The server requires the `inference-models` release that exposes the resolved package descriptor and an updated dependency pin. +Before that release, a Docker build with the updated library source can validate the change. A request can specify `model_package_id`, or `backend`, `quantization`, or both. The server rejects a package ID combined with either other selector. A successful request uses a package that satisfies all explicit selectors. From bf21ee7a1b6c22b5d7663114f5c5648f8161bf1e Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Wed, 16 Sep 2026 12:02:57 -0230 Subject: [PATCH 07/21] Run package selection tests with the pinned inference-models wheel --- docs/http-model-package-selection.md | 3 ++- .../unit_tests/core/managers/test_model_package_cache.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/http-model-package-selection.md b/docs/http-model-package-selection.md index b82e2ac08b..12a9113de2 100644 --- a/docs/http-model-package-selection.md +++ b/docs/http-model-package-selection.md @@ -1,7 +1,8 @@ # Select a model package over HTTP Model package selection requires a server with `USE_INFERENCE_MODELS=true`. -The server requires the `inference-models` release that exposes the resolved package descriptor and an updated dependency pin. +Explicit selection requires the `inference-models` release that exposes the resolved package descriptor and an updated dependency pin. +With an older library, ordinary inference works, but explicit selection fails because the server cannot verify the loaded package. Before that release, a Docker build with the updated library source can validate the change. A request can specify `model_package_id`, or `backend`, `quantization`, or both. The server rejects a package ID combined with either other selector. diff --git a/tests/inference/unit_tests/core/managers/test_model_package_cache.py b/tests/inference/unit_tests/core/managers/test_model_package_cache.py index 8dfa3efbb2..7446ef405f 100644 --- a/tests/inference/unit_tests/core/managers/test_model_package_cache.py +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -1,5 +1,6 @@ from concurrent.futures import ThreadPoolExecutor from threading import Event +from types import SimpleNamespace from typing import Any from unittest.mock import Mock, patch @@ -16,7 +17,6 @@ from inference.core.managers.base import ModelManager from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache from inference.core.models.base import Model -from inference_models.entities import ResolvedModelMetadata class PackageModel(Model): @@ -27,7 +27,7 @@ class PackageModel(Model): def __init__(self, model_id, api_key, backend=None, **kwargs): self.model_id = model_id - self.resolved_model = ResolvedModelMetadata( + self.resolved_model = SimpleNamespace( model_id=model_id, model_package_id=f"{backend or 'onnx'}-package", backend=backend or "onnx", From a374ec87ac615df103b0f726ebe40100331b00ac Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Wed, 16 Sep 2026 12:48:06 -0230 Subject: [PATCH 08/21] Keep package handles server-owned and correct selection errors --- docs/http-model-package-selection.md | 14 ++ inference/core/constants.py | 1 + .../core/entities/requests/model_selection.py | 28 ++- .../core/entities/responses/server_state.py | 7 + inference/core/exceptions.py | 4 + .../core/interfaces/http/error_handlers.py | 11 + inference/core/interfaces/http/http_api.py | 20 +- inference/core/managers/base.py | 23 +- inference_sdk/http/client.py | 140 ++++++++---- inference_sdk/http/entities.py | 18 +- .../interfaces/http/test_model_selection.py | 64 ++++++ .../http/test_model_selection_errors.py | 93 ++++++++ .../core/managers/test_model_package_cache.py | 7 +- .../unit_tests/http/test_model_selection.py | 211 +++++++++++++++++- 14 files changed, 560 insertions(+), 81 deletions(-) create mode 100644 tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py diff --git a/docs/http-model-package-selection.md b/docs/http-model-package-selection.md index 12a9113de2..d41e0e3718 100644 --- a/docs/http-model-package-selection.md +++ b/docs/http-model-package-selection.md @@ -9,6 +9,16 @@ The server rejects a package ID combined with either other selector. A successful request uses a package that satisfies all explicit selectors. An unavailable, incompatible, or disabled package produces an error. +| SDK configuration field | Selection | +| --- | --- | +| `model_package_id` | Exact package ID; cannot accompany either other field. | +| `backend` | Required backend, such as `onnx` or `trt`. | +| `quantization` | Required precision, such as `fp32` or `fp16`; can accompany `backend`. | + +Malformed or conflicting selectors return HTTP 422. +Unsupported selector values and unavailable backend/precision combinations return HTTP 400. +Unknown package IDs return HTTP 404 with `Model or model package not found.` + Standard computer vision endpoints accept these fields in the v1 JSON body. The legacy `/{project}/{version}` endpoint accepts them as query parameters. `/model/add` and `/model/remove` also accept them in the JSON body. @@ -34,11 +44,15 @@ HTTP success responses include `X-Roboflow-Model-Selection: applied` for selecte Each model ID, selector combination, and client credential identifies a separate cache entry. Each entry counts toward the existing LRU limit. The registry exposes its cache handle as `model_id`. +The server generates opaque handles with a private random key. +`/model/add` returns `selected_model_id` for explicit selections; the SDK uses this identifier to find the loaded model. +Handles can change after a server restart. They cannot be used for inference or model loading. Two selector combinations can load separate instances of the same package. Requests without selectors retain their existing automatic selection and cache entry. Registration of a preferred package does not replace a resident instance. `unload_model()` removes the entry for the active configuration. +The SDK sends the public model ID, selectors, and credential so the server can find the current entry. `/model/remove` accepts the original model ID with the same selectors and credential, or the registry handle without selectors. An active request retains its model reference until inference completes. Concurrent package loads still require sufficient device memory. diff --git a/inference/core/constants.py b/inference/core/constants.py index d9d51c2f1d..18ca464d68 100644 --- a/inference/core/constants.py +++ b/inference/core/constants.py @@ -8,6 +8,7 @@ MODEL_LOAD_TIME_HEADER = "X-Model-Load-Time" MODEL_LOAD_DETAILS_HEADER = "X-Model-Load-Details" MODEL_ID_HEADER = "X-Model-Id" +MODEL_SELECTION_HEADER = "X-Roboflow-Model-Selection" WORKFLOW_ID_HEADER = "X-Workflow-Id" WORKSPACE_ID_HEADER = "X-Workspace-Id" TRACE_ID_HEADER = "X-Trace-Id" diff --git a/inference/core/entities/requests/model_selection.py b/inference/core/entities/requests/model_selection.py index ce5edbd33c..8d180d0e85 100644 --- a/inference/core/entities/requests/model_selection.py +++ b/inference/core/entities/requests/model_selection.py @@ -1,14 +1,19 @@ +import hmac +import json +import secrets from typing import Optional -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, ConfigDict, Field, model_validator from inference.core import env -from inference_sdk.http.entities import ( - model_selection_cache_key as model_selection_cache_key, -) +from inference.models.aliases import resolve_roboflow_model_alias + +_MODEL_SELECTION_SECRET = secrets.token_bytes(32) class ModelSelectionRequest(BaseModel): + model_config = ConfigDict(protected_namespaces=()) + model_package_id: Optional[str] = Field( default=None, min_length=1, @@ -46,3 +51,18 @@ def model_selection_kwargs(request) -> dict: for name in ("model_package_id", "backend", "quantization") if (value := getattr(request, name, None)) is not None } + + +def model_selection_cache_key( + model_id: str, selectors: dict, api_key: Optional[str] = None +) -> str: + if not selectors: + return model_id + model_id = resolve_roboflow_model_alias(model_id) + payload = json.dumps( + {"model_id": model_id, "selectors": selectors, "api_key": api_key}, + sort_keys=True, + separators=(",", ":"), + ).encode() + digest = hmac.digest(_MODEL_SELECTION_SECRET, payload, "sha256").hex() + return f"{model_id}:package:{digest}" diff --git a/inference/core/entities/responses/server_state.py b/inference/core/entities/responses/server_state.py index 8571061005..f29b2035de 100644 --- a/inference/core/entities/responses/server_state.py +++ b/inference/core/entities/responses/server_state.py @@ -143,6 +143,13 @@ def from_models_descriptions( ) +class ModelLoadResponse(ModelsDescriptions): + selected_model_id: Optional[str] = Field( + default=None, + description="Opaque identifier of the loaded entry for an explicit package selection.", + ) + + def _get_gpu_memory_stats() -> tuple: try: import torch diff --git a/inference/core/exceptions.py b/inference/core/exceptions.py index 1861f0f5c0..c737ffa9d0 100644 --- a/inference/core/exceptions.py +++ b/inference/core/exceptions.py @@ -42,6 +42,10 @@ class ModelPackageSelectionError(Exception): pass +class ModelPackageNotFoundError(Exception): + pass + + class InvalidEnvironmentVariableError(Exception): """Raised when an environment variable is invalid. diff --git a/inference/core/interfaces/http/error_handlers.py b/inference/core/interfaces/http/error_handlers.py index d3f3831013..346a34a743 100644 --- a/inference/core/interfaces/http/error_handlers.py +++ b/inference/core/interfaces/http/error_handlers.py @@ -23,6 +23,7 @@ ModelArtefactError, ModelDeploymentNotSupportedError, ModelManagerLockAcquisitionError, + ModelPackageNotFoundError, ModelPackageSelectionError, OnnxProviderNotAvailable, PayloadTooLargeError, @@ -500,6 +501,11 @@ def wrapped_route(*args, **kwargs): status_code=502, content={"message": "Internal error. Request to Roboflow API failed."}, ) + except ModelPackageNotFoundError: + resp = JSONResponse( + status_code=404, + content={"message": "Model or model package not found."}, + ) except ModelPackageSelectionError as error: resp = JSONResponse(status_code=400, content={"message": str(error)}) except InferenceModelNotFound as error: @@ -982,6 +988,11 @@ async def wrapped_route(*args, **kwargs): status_code=502, content={"message": "Internal error. Request to Roboflow API failed."}, ) + except ModelPackageNotFoundError: + resp = JSONResponse( + status_code=404, + content={"message": "Model or model package not found."}, + ) except ModelPackageSelectionError as error: resp = JSONResponse(status_code=400, content={"message": str(error)}) except InferenceModelNotFound as error: diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index 170f5d4f88..fc976cf897 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -38,6 +38,7 @@ MODEL_ID_HEADER, MODEL_LOAD_DETAILS_HEADER, MODEL_LOAD_TIME_HEADER, + MODEL_SELECTION_HEADER, PROCESSING_TIME_HEADER, TRACE_ID_HEADER, WORKFLOW_ID_HEADER, @@ -144,6 +145,7 @@ ) from inference.core.entities.responses.secure_gateway import SecureGatewayHealthResponse from inference.core.entities.responses.server_state import ( + ModelLoadResponse, ModelsDescriptions, ServerVersionInfo, ) @@ -396,7 +398,6 @@ from inference.usage_tracking.decorator_helpers import ( non_billable_intent_is_authenticated, ) -from inference_sdk.http.utils.model_selection import MODEL_SELECTION_HEADER if LAMBDA and not OFFLINE_MODE: from inference.core.usage import trackUsage @@ -759,6 +760,7 @@ async def on_shutdown(): MODEL_LOAD_TIME_HEADER, MODEL_LOAD_DETAILS_HEADER, MODEL_ID_HEADER, + MODEL_SELECTION_HEADER, WORKFLOW_ID_HEADER, WORKSPACE_ID_HEADER, TRACE_ID_HEADER, @@ -1960,7 +1962,7 @@ def registry(): @app.post( "/model/add", - response_model=ModelsDescriptions, + response_model=ModelLoadResponse, summary="Load a model", description="Load the model with the given model ID", ) @@ -1988,12 +1990,13 @@ def model_add( ) logger.info(f"Loading model: {de_aliased_model_id}") selectors = model_selection_kwargs(request) + cache_key = model_selection_cache_key( + de_aliased_model_id, selectors, request.api_key + ) selection_args = ( { **selectors, - "model_cache_key": model_selection_cache_key( - de_aliased_model_id, selectors, request.api_key - ), + "model_cache_key": cache_key, } if selectors else {} @@ -2008,10 +2011,15 @@ def model_add( if selectors: response.headers[MODEL_SELECTION_HEADER] = "applied" models_descriptions = self.model_manager.describe_models() - return ModelsDescriptions.from_models_descriptions( + descriptions = ModelsDescriptions.from_models_descriptions( models_descriptions=models_descriptions ) + return ModelLoadResponse( + **descriptions.model_dump(), + selected_model_id=cache_key if selectors else None, + ) + @app.post( "/model/remove", response_model=ModelsDescriptions, diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index 3f2f02ff26..93c7c6504a 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -29,6 +29,7 @@ from inference.core.exceptions import ( InferenceModelNotFound, ModelManagerLockAcquisitionError, + ModelPackageNotFoundError, ModelPackageSelectionError, RoboflowAPINotAuthorizedError, ) @@ -55,6 +56,11 @@ start_span, ) from inference.usage_tracking.model_types import bind_usage_model_descriptor +from inference_models.errors import ( + NoModelPackagesAvailableError, + UnknownBackendTypeError, + UnknownQuantizationError, +) if TYPE_CHECKING: from inference_models.utils.content_addressed_artifact_cache import ( @@ -280,6 +286,18 @@ def add_model( except Exception as error: record_error(error) self._dispose_model_lock(model_id=resolved_identifier) + if selectors and isinstance(error, NoModelPackagesAvailableError): + if model_package_id is not None: + raise ModelPackageNotFoundError() from error + raise ModelPackageSelectionError( + "No model package satisfies the requested selection on this server." + ) from error + if backend is not None and isinstance(error, UnknownBackendTypeError): + raise ModelPackageSelectionError("Unknown backend.") from error + if quantization is not None and isinstance( + error, UnknownQuantizationError + ): + raise ModelPackageSelectionError("Unknown quantization.") from error raise error def validate_model_selection( @@ -949,10 +967,7 @@ def validate_public_model_id( model_id: str, model_id_alias: Optional[str] = None ) -> None: if any(":package:" in value for value in (model_id, model_id_alias) if value): - raise ModelPackageSelectionError( - "Cache handles cannot be used for inference or model loading. " - "Use the public model_id with package selection parameters." - ) + raise ModelPackageNotFoundError() def model_load_options( diff --git a/inference_sdk/http/client.py b/inference_sdk/http/client.py index 252b3129d4..0b282e8220 100644 --- a/inference_sdk/http/client.py +++ b/inference_sdk/http/client.py @@ -41,7 +41,6 @@ RegisteredModels, ServerInfo, VideoReference, - model_selection_cache_key, ) from inference_sdk.http.errors import ( APIKeyNotProvided, @@ -295,6 +294,9 @@ def __init__( self.__inference_configuration = InferenceConfiguration.init_default() self.__client_mode = _determine_client_mode(api_url=api_url) self.__selected_model: Optional[str] = None + self.__model_selection_ids: Dict[ + Tuple[str, Tuple[Tuple[str, str], ...]], str + ] = {} self.__webrtc_client: Optional["WebRTCClient"] = None self.__webrtc_client_transport: Optional[ApiKeyTransport] = None self.__webrtc_transport_stickiness_warned = False @@ -910,20 +912,35 @@ def get_model_description( self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) registered_models = self.list_loaded_models() - cache_key = model_selection_cache_key( - de_aliased_model_id, - self.__inference_configuration.to_model_selection_parameters(), - self.__api_key, - ) - matching_model = filter_model_descriptions( - descriptions=registered_models.models, - model_id=cache_key, + selectors = self.__inference_configuration.to_model_selection_parameters() + selection_key = (de_aliased_model_id, tuple(sorted(selectors.items()))) + cache_key = ( + self.__model_selection_ids.get(selection_key) + if selectors + else de_aliased_model_id + ) + matching_model = ( + filter_model_descriptions( + descriptions=registered_models.models, + model_id=cache_key, + ) + if cache_key is not None + else None ) if matching_model is None and allow_loading is True: registered_models = self.load_model(model_id=de_aliased_model_id) - matching_model = filter_model_descriptions( - descriptions=registered_models.models, - model_id=cache_key, + cache_key = ( + registered_models.selected_model_id + if selectors + else de_aliased_model_id + ) + matching_model = ( + filter_model_descriptions( + descriptions=registered_models.models, + model_id=cache_key, + ) + if cache_key is not None + else None ) if matching_model is not None: return matching_model @@ -951,22 +968,37 @@ async def get_model_description_async( self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) registered_models = await self.list_loaded_models_async() - cache_key = model_selection_cache_key( - de_aliased_model_id, - self.__inference_configuration.to_model_selection_parameters(), - self.__api_key, - ) - matching_model = filter_model_descriptions( - descriptions=registered_models.models, - model_id=cache_key, + selectors = self.__inference_configuration.to_model_selection_parameters() + selection_key = (de_aliased_model_id, tuple(sorted(selectors.items()))) + cache_key = ( + self.__model_selection_ids.get(selection_key) + if selectors + else de_aliased_model_id + ) + matching_model = ( + filter_model_descriptions( + descriptions=registered_models.models, + model_id=cache_key, + ) + if cache_key is not None + else None ) if matching_model is None and allow_loading is True: registered_models = await self.load_model_async( model_id=de_aliased_model_id ) - matching_model = filter_model_descriptions( - descriptions=registered_models.models, - model_id=cache_key, + cache_key = ( + registered_models.selected_model_id + if selectors + else de_aliased_model_id + ) + matching_model = ( + filter_model_descriptions( + descriptions=registered_models.models, + model_id=cache_key, + ) + if cache_key is not None + else None ) if matching_model is not None: return matching_model @@ -1040,11 +1072,12 @@ def load_model( """ self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) + selectors = self.__inference_configuration.to_model_selection_parameters() response = requests.post( f"{self.__api_url}/model/add", json={ "model_id": de_aliased_model_id, - **self.__inference_configuration.to_model_selection_parameters(), + **selectors, **self.__legacy_api_key_payload(), }, headers=self.__headers_with_auth(DEFAULT_HEADERS), @@ -1052,12 +1085,15 @@ def load_model( response.raise_for_status() ensure_model_selection_applied( response.headers, - self.__inference_configuration.to_model_selection_parameters(), + selectors, ) response_payload = response.json() + registered_models = self.__remember_model_selection( + de_aliased_model_id, selectors, response_payload + ) if set_as_default: self.__selected_model = de_aliased_model_id - return RegisteredModels.from_dict(response_payload) + return registered_models @wrap_errors_async async def load_model_async( @@ -1079,9 +1115,10 @@ async def load_model_async( """ self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) + selectors = self.__inference_configuration.to_model_selection_parameters() payload = { "model_id": de_aliased_model_id, - **self.__inference_configuration.to_model_selection_parameters(), + **selectors, **self.__legacy_api_key_payload(), } async with aiohttp.ClientSession() as session: @@ -1093,12 +1130,29 @@ async def load_model_async( response.raise_for_status() ensure_model_selection_applied( response.headers, - self.__inference_configuration.to_model_selection_parameters(), + selectors, ) response_payload = await response.json() + registered_models = self.__remember_model_selection( + de_aliased_model_id, selectors, response_payload + ) if set_as_default: self.__selected_model = de_aliased_model_id - return RegisteredModels.from_dict(response_payload) + return registered_models + + def __remember_model_selection( + self, model_id: str, selectors: Dict[str, str], response_payload: dict + ) -> RegisteredModels: + registered_models = RegisteredModels.from_dict(response_payload) + if selectors: + if not registered_models.selected_model_id: + raise InvalidParameterError( + "The server did not return the selected model ID. Upgrade the inference server to use model selection." + ) + self.__model_selection_ids[(model_id, tuple(sorted(selectors.items())))] = ( + registered_models.selected_model_id + ) + return registered_models @wrap_errors def unload_model(self, model_id: str) -> RegisteredModels: @@ -1117,20 +1171,21 @@ def unload_model(self, model_id: str) -> RegisteredModels: """ self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) - cache_key = model_selection_cache_key( - de_aliased_model_id, - self.__inference_configuration.to_model_selection_parameters(), - self.__api_key, - ) + selectors = self.__inference_configuration.to_model_selection_parameters() response = requests.post( f"{self.__api_url}/model/remove", json={ - "model_id": cache_key, + "model_id": de_aliased_model_id, + **selectors, + **(self.__legacy_api_key_payload() if selectors else {}), }, headers=self.__headers_with_auth(DEFAULT_HEADERS), ) response.raise_for_status() response_payload = response.json() + self.__model_selection_ids.pop( + (de_aliased_model_id, tuple(sorted(selectors.items()))), None + ) if ( de_aliased_model_id == self.__selected_model or model_id == self.__selected_model @@ -1142,21 +1197,22 @@ def unload_model(self, model_id: str) -> RegisteredModels: async def unload_model_async(self, model_id: str) -> RegisteredModels: self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) - cache_key = model_selection_cache_key( - de_aliased_model_id, - self.__inference_configuration.to_model_selection_parameters(), - self.__api_key, - ) + selectors = self.__inference_configuration.to_model_selection_parameters() async with aiohttp.ClientSession() as session: async with session.post( f"{self.__api_url}/model/remove", json={ - "model_id": cache_key, + "model_id": de_aliased_model_id, + **selectors, + **(self.__legacy_api_key_payload() if selectors else {}), }, headers=self.__headers_with_auth(DEFAULT_HEADERS), ) as response: response.raise_for_status() response_payload = await response.json() + self.__model_selection_ids.pop( + (de_aliased_model_id, tuple(sorted(selectors.items()))), None + ) if ( de_aliased_model_id == self.__selected_model or model_id == self.__selected_model @@ -1174,6 +1230,7 @@ def unload_all_models(self) -> RegisteredModels: response.raise_for_status() response_payload = response.json() self.__selected_model = None + self.__model_selection_ids.clear() return RegisteredModels.from_dict(response_payload) @wrap_errors_async @@ -1187,6 +1244,7 @@ async def unload_all_models_async(self) -> RegisteredModels: response.raise_for_status() response_payload = await response.json() self.__selected_model = None + self.__model_selection_ids.clear() return RegisteredModels.from_dict(response_payload) @wrap_errors diff --git a/inference_sdk/http/entities.py b/inference_sdk/http/entities.py index 5cf9940033..bea1729260 100644 --- a/inference_sdk/http/entities.py +++ b/inference_sdk/http/entities.py @@ -1,5 +1,3 @@ -import hashlib -import json from dataclasses import dataclass, field from enum import Enum from typing import Any, Dict, List, Literal, Optional, Tuple, Union @@ -77,6 +75,7 @@ class RegisteredModels(DataClassJsonMixin): """ models: List[ModelDescription] + selected_model_id: Optional[str] = None class HTTPClientMode(str, Enum): @@ -434,18 +433,3 @@ def get_non_empty_attributes( for internal_name, external_name in specification } return remove_empty_values(dictionary=attributes) - - -def model_selection_cache_key( - model_id: str, selectors: dict, api_key: Optional[str] = None -) -> str: - if not selectors: - return model_id - digest = hashlib.sha256( - json.dumps( - {"selectors": selectors, "api_key": api_key}, - sort_keys=True, - separators=(",", ":"), - ).encode() - ).hexdigest() - return f"{model_id}:package:{digest}" diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py index 817d40bcec..2ad6bcf2f5 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -103,6 +103,7 @@ def clear_cache(self, delete_from_disk=True): monkeypatch.setattr(http_api, "LAMBDA", False) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) registry = MagicMock() + monkeypatch.setattr(http_api, "ALLOW_ORIGINS", ["https://client.example"]) registry.get_model.return_value = PackageModel manager = WithFixedSizeCache( ModelManager(registry, content_addressed_artifact_cache=MagicMock()), max_size=3 @@ -170,6 +171,69 @@ def test_legacy_http_rejects_conflicting_selectors(package_client): assert "model_package_id" in response.text +def test_aliased_legacy_selection_shares_v1_entry_and_can_be_removed( + package_client, monkeypatch +): + from inference.models.aliases import REGISTERED_ALIASES + + monkeypatch.setitem(REGISTERED_ALIASES, "alias/1", "project/1") + selectors = {"backend": "trt", "quantization": "fp16"} + legacy = package_client.post( + "/alias/1", params={**selectors, "image": "https://example.com/image.jpg"} + ) + modern = package_client.post( + "/infer/object_detection", + json={ + "model_id": "project/1", + **selectors, + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + }, + ) + assert legacy.status_code == modern.status_code == 200 + registered = package_client.get("/model/registry").json()["models"] + assert len(registered) == 1 + loaded = package_client.post( + "/model/add", json={"model_id": "alias/1", **selectors} + ) + assert loaded.status_code == 200 + assert loaded.json()["selected_model_id"] == registered[0]["model_id"] + removed = package_client.post( + "/model/remove", json={"model_id": "alias/1", **selectors} + ) + assert removed.status_code == 200 + assert removed.json()["models"] == [] + + +def test_selected_model_handle_depends_on_server_secret(package_client, monkeypatch): + from inference.core.entities.requests import model_selection + + payload = {"model_id": "project/1", "backend": "trt", "api_key": "owner"} + first = package_client.post("/model/add", json=payload) + assert first.status_code == 200 + handle = first.json()["selected_model_id"] + assert ( + package_client.post("/model/add", json=payload).json()["selected_model_id"] + == handle + ) + monkeypatch.setattr(model_selection, "_MODEL_SELECTION_SECRET", b"different-server") + second = package_client.post("/model/add", json=payload) + assert second.status_code == 200 + assert second.json()["selected_model_id"] != handle + + +def test_browser_can_read_package_selection_acknowledgment(package_client): + response = package_client.post( + "/model/add", + json={"model_id": "project/1", "backend": "trt"}, + headers={"Origin": "https://client.example"}, + ) + assert response.status_code == 200 + assert response.headers["X-Roboflow-Model-Selection"] == "applied" + assert "X-Roboflow-Model-Selection" in response.headers[ + "Access-Control-Expose-Headers" + ].split(", ") + + def test_http_cannot_reuse_another_credentials_restricted_package(package_client): payload = { "model_id": "project/1", diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py new file mode 100644 index 0000000000..cd6cb9c525 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py @@ -0,0 +1,93 @@ +from unittest.mock import MagicMock + +import pytest +from starlette.testclient import TestClient + +from inference.core import env +from inference.core.interfaces.http import http_api +from inference.core.managers import base +from inference_models.errors import ( + ModelLoadingError, + ModelPackageNegotiationError, + NoModelPackagesAvailableError, + RuntimeIntrospectionError, + UnknownBackendTypeError, + UnknownQuantizationError, +) + + +@pytest.fixture +def failing_package_client(monkeypatch): + monkeypatch.setattr(env, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr(base, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr(base, "OFFLINE_MODE", True) + monkeypatch.setattr(http_api, "InferenceInstrumentator", MagicMock()) + monkeypatch.setattr(http_api, "OFFLINE_MODE", True) + monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) + monkeypatch.setattr(http_api, "LAMBDA", False) + monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) + constructor = MagicMock() + registry = MagicMock() + registry.get_model.return_value = constructor + manager = base.ModelManager(registry, content_addressed_artifact_cache=MagicMock()) + interface = http_api.HttpInterface(model_manager=manager) + with TestClient(interface.app) as client: + yield client, constructor + + +@pytest.mark.parametrize( + "path", ["/infer/object_detection", "/model/add", "/project/1"] +) +@pytest.mark.parametrize("selector", [{"model_package_id": "missing"}, {}]) +def test_missing_package_and_cache_handle_return_same_public_error( + failing_package_client, path, selector +): + client, constructor = failing_package_client + constructor.side_effect = NoModelPackagesAvailableError("private package details") + model_id = "project/1" if selector else "project/1:package:private" + if path == "/project/1": + response = client.post( + f"/{model_id}", + params={"image": "https://example.com/image.jpg", **selector}, + ) + else: + response = client.post( + path, + json={ + "model_id": model_id, + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + **selector, + }, + ) + assert response.status_code == 404 + assert response.json() == {"message": "Model or model package not found."} + + +@pytest.mark.parametrize("path", ["/infer/object_detection", "/model/add"]) +@pytest.mark.parametrize( + "selectors,error,status", + [ + ({"backend": "unknown"}, UnknownBackendTypeError("details"), 400), + ({"quantization": "unknown"}, UnknownQuantizationError("details"), 400), + ({"backend": "trt"}, NoModelPackagesAvailableError("details"), 400), + ({}, NoModelPackagesAvailableError("details"), 500), + ({"backend": "trt"}, RuntimeIntrospectionError("details"), 500), + ({"backend": "trt"}, ModelPackageNegotiationError("details"), 500), + ({"model_package_id": "broken"}, ModelLoadingError("details"), 500), + ({"backend": "trt"}, UnknownQuantizationError("server configuration"), 500), + ], +) +def test_selector_errors_do_not_mask_server_failures( + failing_package_client, path, selectors, error, status +): + client, constructor = failing_package_client + constructor.side_effect = error + response = client.post( + path, + json={ + "model_id": "project/1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + **selectors, + }, + ) + assert response.status_code == status diff --git a/tests/inference/unit_tests/core/managers/test_model_package_cache.py b/tests/inference/unit_tests/core/managers/test_model_package_cache.py index 7446ef405f..62352b290e 100644 --- a/tests/inference/unit_tests/core/managers/test_model_package_cache.py +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -11,7 +11,10 @@ InferenceRequest, ObjectDetectionInferenceRequest, ) -from inference.core.exceptions import ModelPackageSelectionError +from inference.core.exceptions import ( + ModelPackageNotFoundError, + ModelPackageSelectionError, +) from inference.core.managers import base as base_module from inference.core.managers.active_learning import ActiveLearningManager from inference.core.managers.base import ModelManager @@ -126,7 +129,7 @@ def test_removing_variant_leaves_automatic_model(package_manager): def test_cache_handle_cannot_bypass_selection_validation(package_manager, as_alias): handle = "project/1:package:secret" package_manager.add_model("project/1", "key", backend="trt", model_cache_key=handle) - with pytest.raises(ModelPackageSelectionError, match="Cache handles"): + with pytest.raises(ModelPackageNotFoundError): package_manager.add_model( "project/1" if as_alias else handle, "different-key", diff --git a/tests/inference_sdk/unit_tests/http/test_model_selection.py b/tests/inference_sdk/unit_tests/http/test_model_selection.py index 42056c3da8..99b8811caf 100644 --- a/tests/inference_sdk/unit_tests/http/test_model_selection.py +++ b/tests/inference_sdk/unit_tests/http/test_model_selection.py @@ -5,10 +5,36 @@ from aioresponses import aioresponses from inference_sdk import InferenceConfiguration, InferenceHTTPClient -from inference_sdk.http.entities import model_selection_cache_key +from inference_sdk.http.entities import ApiKeyTransport from inference_sdk.http.errors import InvalidParameterError +def register_response( + requests_mock, async_responses, method, path, payload, headers=None +): + url = f"http://server{path}" + getattr(requests_mock, method)(url, json=payload, headers=headers or {}) + getattr(async_responses, method)(url, payload=payload, headers=headers or {}) + + +async def call_client(client, method, asynchronous, *args, **kwargs): + if asynchronous: + return await getattr(client, f"{method}_async")(*args, **kwargs) + return getattr(client, method)(*args, **kwargs) + + +def registry_payload(*model_ids, selected_model_id=None): + payload = { + "models": [ + {"model_id": model_id, "task_type": "object-detection"} + for model_id in model_ids + ] + } + if selected_model_id is not None: + payload["selected_model_id"] = selected_model_id + return payload + + @pytest.mark.parametrize("acknowledged", [True, False]) def test_v0_selection_requires_server_acknowledgment(requests_mock, acknowledged): headers = {"X-Roboflow-Model-Selection": "applied"} if acknowledged else {} @@ -55,14 +81,17 @@ async def test_async_v0_selection_requires_server_acknowledgment(acknowledged): def test_v1_description_loads_selected_variant_without_loading_default(requests_mock): selectors: dict[str, Any] = {"backend": "trt", "quantization": "fp16"} - cache_key = model_selection_cache_key("project/1", selectors, "key") + cache_key = "opaque-server-handle" requests_mock.get( "http://server/model/registry", json={"models": [{"model_id": "project/1", "task_type": "object-detection"}]}, ) requests_mock.post( "http://server/model/add", - json={"models": [{"model_id": cache_key, "task_type": "object-detection"}]}, + json={ + "models": [{"model_id": cache_key, "task_type": "object-detection"}], + "selected_model_id": cache_key, + }, headers={"X-Roboflow-Model-Selection": "applied"}, ) client = InferenceHTTPClient(api_url="http://server", api_key="key") @@ -78,7 +107,7 @@ def test_v1_description_loads_selected_variant_without_loading_default(requests_ @pytest.mark.asyncio async def test_async_v1_description_loads_selected_variant_without_loading_default(): selectors: dict[str, Any] = {"model_package_id": "engine-1"} - cache_key = model_selection_cache_key("project/1", selectors, "key") + cache_key = "another-server-handle" client = InferenceHTTPClient(api_url="http://server", api_key="key") client.configure(InferenceConfiguration(**selectors)) with aioresponses() as responses: @@ -91,7 +120,8 @@ async def test_async_v1_description_loads_selected_variant_without_loading_defau responses.post( "http://server/model/add", payload={ - "models": [{"model_id": cache_key, "task_type": "object-detection"}] + "models": [{"model_id": cache_key, "task_type": "object-detection"}], + "selected_model_id": cache_key, }, headers={"X-Roboflow-Model-Selection": "applied"}, ) @@ -108,12 +138,179 @@ async def test_async_v1_description_loads_selected_variant_without_loading_defau } -def test_unload_selected_variant_never_sends_default_model_id(requests_mock): +def test_unload_selected_variant_sends_public_id_and_selectors(requests_mock): selectors: dict[str, Any] = {"model_package_id": "engine-1"} requests_mock.post("http://server/model/remove", json={"models": []}) client = InferenceHTTPClient(api_url="http://server", api_key="key") client.configure(InferenceConfiguration(**selectors)) client.unload_model("project/1") assert requests_mock.last_request.json() == { - "model_id": model_selection_cache_key("project/1", selectors, "key") + "model_id": "project/1", + "api_key": "key", + **selectors, } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_selection_remembers_server_handles_across_configurations_and_reload( + requests_mock, asynchronous +): + client = InferenceHTTPClient(api_url="http://server", api_key="key") + with aioresponses() as responses: + for backend, handle in [("onnx", "server-onnx"), ("trt", "server-trt")]: + client.configure(InferenceConfiguration(backend=backend)) + register_response( + requests_mock, + responses, + "post", + "/model/add", + registry_payload(handle, selected_model_id=handle), + {"X-Roboflow-Model-Selection": "applied"}, + ) + await call_client(client, "load_model", asynchronous, "project/1") + + # Existing handles are selected by the active configuration, without a load. + for backend, handle in [("onnx", "server-onnx"), ("trt", "server-trt")]: + client.configure(InferenceConfiguration(backend=backend)) + register_response( + requests_mock, + responses, + "get", + "/model/registry?api_key=key", + registry_payload("project/1", "server-onnx", "server-trt"), + ) + description = await call_client( + client, + "get_model_description", + asynchronous, + "project/1", + allow_loading=False, + ) + assert description.model_id == handle + + # A restarted worker supplies a new handle for the same selection. + register_response( + requests_mock, + responses, + "get", + "/model/registry?api_key=key", + registry_payload("project/1", "unrecognized-handle"), + ) + register_response( + requests_mock, + responses, + "post", + "/model/add", + registry_payload( + "replacement-handle", selected_model_id="replacement-handle" + ), + {"X-Roboflow-Model-Selection": "applied"}, + ) + description = await call_client( + client, "get_model_description", asynchronous, "project/1" + ) + assert description.model_id == "replacement-handle" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("method", ["load_model", "get_model_description"]) +async def test_selected_load_rejects_missing_server_handle( + requests_mock, asynchronous, method +): + client = InferenceHTTPClient(api_url="http://server", api_key="key") + client.configure(InferenceConfiguration(backend="trt")) + with aioresponses() as responses: + register_response( + requests_mock, + responses, + "get", + "/model/registry?api_key=key", + registry_payload("project/1"), + ) + register_response( + requests_mock, + responses, + "post", + "/model/add", + registry_payload("project/1"), + {"X-Roboflow-Model-Selection": "applied"}, + ) + with pytest.raises(InvalidParameterError, match="selected model ID"): + await call_client(client, method, asynchronous, "project/1") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("unload_all", [False, True]) +async def test_unload_forgets_selected_handle(requests_mock, asynchronous, unload_all): + from inference_sdk.http.errors import ModelNotInitializedError + + client = InferenceHTTPClient(api_url="http://server", api_key="key") + client.configure(InferenceConfiguration(model_package_id="engine-1")) + with aioresponses() as responses: + register_response( + requests_mock, + responses, + "post", + "/model/add", + registry_payload("server-handle", selected_model_id="server-handle"), + {"X-Roboflow-Model-Selection": "applied"}, + ) + await call_client(client, "load_model", asynchronous, "project/1") + register_response( + requests_mock, + responses, + "post", + "/model/clear" if unload_all else "/model/remove", + registry_payload(), + ) + if unload_all: + await call_client(client, "unload_all_models", asynchronous) + else: + await call_client(client, "unload_model", asynchronous, "project/1") + register_response( + requests_mock, + responses, + "get", + "/model/registry?api_key=key", + registry_payload("server-handle"), + ) + with pytest.raises(ModelNotInitializedError): + await call_client( + client, + "get_model_description", + asynchronous, + "project/1", + allow_loading=False, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("transport", [ApiKeyTransport.LEGACY, ApiKeyTransport.HEADER]) +async def test_unload_selected_model_uses_public_alias_and_auth( + requests_mock, asynchronous, transport +): + client = InferenceHTTPClient(api_url="http://server", api_key="key") + client.configure(InferenceConfiguration(backend="trt", api_key_transport=transport)) + with aioresponses() as responses: + register_response( + requests_mock, responses, "post", "/model/remove", registry_payload() + ) + await call_client(client, "unload_model", asynchronous, "yolov8n-640") + if asynchronous: + request = next(iter(responses.requests.values()))[0].kwargs + payload, headers = request["json"], request["headers"] + else: + payload, headers = ( + requests_mock.last_request.json(), + requests_mock.last_request.headers, + ) + expected = {"model_id": "coco/3", "backend": "trt"} + if transport is ApiKeyTransport.LEGACY: + expected["api_key"] = "key" + else: + assert headers["Authorization"] == "Bearer key" + assert payload == expected From 25589d09076e6087191ded4a92da944963351b2c Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Wed, 16 Sep 2026 16:04:37 -0230 Subject: [PATCH 09/21] Simplify HTTP model selection request and response types --- inference/core/entities/requests/inference.py | 52 +++++++++++++++++-- .../core/entities/requests/model_selection.py | 37 ------------- .../core/entities/requests/server_state.py | 6 +-- .../core/entities/responses/server_state.py | 11 ++-- inference/core/interfaces/http/http_api.py | 16 +++--- inference/core/managers/base.py | 5 +- .../interfaces/http/test_model_selection.py | 24 +++++++++ 7 files changed, 88 insertions(+), 63 deletions(-) diff --git a/inference/core/entities/requests/inference.py b/inference/core/entities/requests/inference.py index a0e7b59b17..d96dc332a3 100644 --- a/inference/core/entities/requests/inference.py +++ b/inference/core/entities/requests/inference.py @@ -1,12 +1,19 @@ from typing import Any, ClassVar, List, Literal, Optional, Union from uuid import uuid4 -from pydantic import BaseModel, ConfigDict, Field, field_validator, validator +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_validator, + model_validator, + validator, +) from pydantic.json_schema import SkipJsonSchema -from inference.core import logger +from inference.core import env, logger from inference.core.entities.common import ApiKey, ModelID, ModelType -from inference.core.entities.requests.model_selection import ModelSelectionRequest +from inference.core.entities.requests.model_selection import model_selection_kwargs from inference_sdk.http.entities import Confidence @@ -54,6 +61,35 @@ class InferenceRequest(BaseRequest): model_id: Optional[str] = ModelID model_type: Optional[str] = ModelType + model_package_id: Optional[str] = Field( + default=None, + min_length=1, + description="Exact model package ID. Cannot be combined with backend or quantization.", + ) + backend: Optional[str] = Field( + default=None, + min_length=1, + description="Required package backend, such as trt or onnx.", + ) + quantization: Optional[str] = Field( + default=None, + min_length=1, + description="Required package quantization, such as fp16 or fp32.", + ) + + @model_validator(mode="after") + def validate_model_selection(self): + if self.model_package_id is not None and ( + self.backend is not None or self.quantization is not None + ): + raise ValueError( + "model_package_id cannot be combined with backend or quantization." + ) + if model_selection_kwargs(self) and not env.USE_INFERENCE_MODELS: + raise ValueError( + "Model package selection requires USE_INFERENCE_MODELS=true." + ) + return self class InferenceRequestImage(BaseModel): @@ -75,7 +111,7 @@ class InferenceRequestImage(BaseModel): ) -class CVInferenceRequest(InferenceRequest, ModelSelectionRequest): +class CVInferenceRequest(InferenceRequest): """Computer Vision inference request. Attributes: @@ -135,6 +171,14 @@ class DepthEstimationRequest(InferenceRequest): "lossy for geometric use).", ) + @model_validator(mode="after") + def reject_model_selection(self): + if model_selection_kwargs(self): + raise ValueError( + "Model package selection is not supported for depth estimation requests." + ) + return self + @validator("model_id", always=True) def validate_model_id(cls, value, values): if value is not None: diff --git a/inference/core/entities/requests/model_selection.py b/inference/core/entities/requests/model_selection.py index 8d180d0e85..cd716e96eb 100644 --- a/inference/core/entities/requests/model_selection.py +++ b/inference/core/entities/requests/model_selection.py @@ -3,48 +3,11 @@ import secrets from typing import Optional -from pydantic import BaseModel, ConfigDict, Field, model_validator - -from inference.core import env from inference.models.aliases import resolve_roboflow_model_alias _MODEL_SELECTION_SECRET = secrets.token_bytes(32) -class ModelSelectionRequest(BaseModel): - model_config = ConfigDict(protected_namespaces=()) - - model_package_id: Optional[str] = Field( - default=None, - min_length=1, - description="Exact model package ID. Cannot be combined with backend or quantization.", - ) - backend: Optional[str] = Field( - default=None, - min_length=1, - description="Required package backend, such as trt or onnx.", - ) - quantization: Optional[str] = Field( - default=None, - min_length=1, - description="Required package quantization, such as fp16 or fp32.", - ) - - @model_validator(mode="after") - def validate_model_selection(self): - if self.model_package_id is not None and ( - self.backend is not None or self.quantization is not None - ): - raise ValueError( - "model_package_id cannot be combined with backend or quantization." - ) - if model_selection_kwargs(self) and not env.USE_INFERENCE_MODELS: - raise ValueError( - "Model package selection requires USE_INFERENCE_MODELS=true." - ) - return self - - def model_selection_kwargs(request) -> dict: return { name: value diff --git a/inference/core/entities/requests/server_state.py b/inference/core/entities/requests/server_state.py index 547ecaf1d2..b38b2a4499 100644 --- a/inference/core/entities/requests/server_state.py +++ b/inference/core/entities/requests/server_state.py @@ -3,10 +3,10 @@ from pydantic import ConfigDict from inference.core.entities.common import ApiKey, ModelID, ModelType -from inference.core.entities.requests.model_selection import ModelSelectionRequest +from inference.core.entities.requests.inference import InferenceRequest -class AddModelRequest(ModelSelectionRequest): +class AddModelRequest(InferenceRequest): """Request to add a model to the inference server. Attributes: @@ -21,7 +21,7 @@ class AddModelRequest(ModelSelectionRequest): api_key: Optional[str] = ApiKey -class ClearModelRequest(ModelSelectionRequest): +class ClearModelRequest(InferenceRequest): """Request to clear a model from the inference server. Attributes: diff --git a/inference/core/entities/responses/server_state.py b/inference/core/entities/responses/server_state.py index f29b2035de..d89c6cb493 100644 --- a/inference/core/entities/responses/server_state.py +++ b/inference/core/entities/responses/server_state.py @@ -76,6 +76,10 @@ def from_model_description( class ModelsDescriptions(BaseModel): + selected_model_id: Optional[str] = Field( + default=None, + description="Opaque identifier of the loaded entry for an explicit package selection.", + ) models: List[ModelDescriptionEntity] = Field( description="List of models that are loaded by model manager.", ) @@ -143,13 +147,6 @@ def from_models_descriptions( ) -class ModelLoadResponse(ModelsDescriptions): - selected_model_id: Optional[str] = Field( - default=None, - description="Opaque identifier of the loaded entry for an explicit package selection.", - ) - - def _get_gpu_memory_stats() -> tuple: try: import torch diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index fc976cf897..d252fa0668 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -70,7 +70,6 @@ ensure_wire_safe_mask_format, ) from inference.core.entities.requests.model_selection import ( - ModelSelectionRequest, model_selection_cache_key, model_selection_kwargs, ) @@ -145,7 +144,6 @@ ) from inference.core.entities.responses.secure_gateway import SecureGatewayHealthResponse from inference.core.entities.responses.server_state import ( - ModelLoadResponse, ModelsDescriptions, ServerVersionInfo, ) @@ -1574,7 +1572,7 @@ def process_inference_request( service_secret (Optional[str]): The service secret. Returns: - InferenceResponse: The response containing the inference results. + Response: The HTTP response containing serialized inference results. """ api_key = api_key_fallback(api_key) if api_key is not None: @@ -1962,7 +1960,7 @@ def registry(): @app.post( "/model/add", - response_model=ModelLoadResponse, + response_model=ModelsDescriptions, summary="Load a model", description="Load the model with the given model ID", ) @@ -2015,10 +2013,8 @@ def model_add( models_descriptions=models_descriptions ) - return ModelLoadResponse( - **descriptions.model_dump(), - selected_model_id=cache_key if selectors else None, - ) + descriptions.selected_model_id = cache_key if selectors else None + return descriptions @app.post( "/model/remove", @@ -4872,7 +4868,9 @@ def legacy_infer_from_request( f"State of model registry: {self.model_manager.describe_models()}" ) try: - selection = ModelSelectionRequest( + selection = InferenceRequest( + id=str(uuid4()), + model_id=request_model_id, model_package_id=model_package_id, backend=backend, quantization=quantization, diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index 93c7c6504a..dafcdd9846 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -11,7 +11,6 @@ from inference.core.devices.utils import GLOBAL_INFERENCE_SERVER_ID from inference.core.entities.requests.inference import InferenceRequest from inference.core.entities.requests.model_selection import ( - ModelSelectionRequest, model_selection_kwargs, ) from inference.core.entities.responses.inference import InferenceResponse @@ -622,7 +621,7 @@ def infer_from_request_sync( async def model_infer(self, model_id: str, request: InferenceRequest, **kwargs): model = self._get_model_reference(model_id=model_id) - if isinstance(request, ModelSelectionRequest): + if isinstance(request, InferenceRequest): self._validate_loaded_model_selection( model, model_selection_kwargs(request) ) @@ -632,7 +631,7 @@ def model_infer_sync( self, model_id: str, request: InferenceRequest, **kwargs ) -> Union[List[InferenceResponse], InferenceResponse]: model = self._get_model_reference(model_id=model_id) - if isinstance(request, ModelSelectionRequest): + if isinstance(request, InferenceRequest): self._validate_loaded_model_selection( model, model_selection_kwargs(request) ) diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py index 2ad6bcf2f5..7a593f4bc9 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -14,6 +14,7 @@ def client(monkeypatch): monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) monkeypatch.setattr(http_api, "LAMBDA", False) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) + monkeypatch.setattr(http_api, "DEPTH_ESTIMATION_ENABLED", True) interface = http_api.HttpInterface(model_manager=MagicMock()) with TestClient(interface.app) as client: yield client @@ -51,6 +52,29 @@ def test_http_rejects_selection_when_flag_is_disabled(client, monkeypatch, path) assert "USE_INFERENCE_MODELS" in response.text +@pytest.mark.parametrize( + "path", ["/infer/depth-estimation", "/infer/depth-estimation/project/1"] +) +@pytest.mark.parametrize( + "selectors", + [{"backend": "trt"}, {"quantization": "fp16"}, {"model_package_id": "engine-1"}], +) +def test_depth_rejects_model_package_selection(client, monkeypatch, path, selectors): + from inference.core import env + + monkeypatch.setattr(env, "USE_INFERENCE_MODELS", True) + response = client.post( + path, + json={ + "model_id": "project/1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + **selectors, + }, + ) + assert response.status_code == 422 + assert "not supported for depth estimation" in response.text + + @pytest.fixture def package_client(monkeypatch): from types import SimpleNamespace From 99169bec9bd4c60548a4aaae3aa84539b622f38d Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 09:56:21 -0230 Subject: [PATCH 10/21] Document package selection through the generated OpenAPI schema --- docs/http-model-package-selection.md | 58 ------------------ inference/core/entities/requests/inference.py | 22 ++++++- inference/core/interfaces/http/http_api.py | 61 +++++++++++++++++-- inference_sdk/http/entities.py | 8 +++ .../http/test_model_selection_openapi.py | 55 +++++++++++++++++ 5 files changed, 138 insertions(+), 66 deletions(-) delete mode 100644 docs/http-model-package-selection.md create mode 100644 tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py diff --git a/docs/http-model-package-selection.md b/docs/http-model-package-selection.md deleted file mode 100644 index d41e0e3718..0000000000 --- a/docs/http-model-package-selection.md +++ /dev/null @@ -1,58 +0,0 @@ -# Select a model package over HTTP - -Model package selection requires a server with `USE_INFERENCE_MODELS=true`. -Explicit selection requires the `inference-models` release that exposes the resolved package descriptor and an updated dependency pin. -With an older library, ordinary inference works, but explicit selection fails because the server cannot verify the loaded package. -Before that release, a Docker build with the updated library source can validate the change. -A request can specify `model_package_id`, or `backend`, `quantization`, or both. -The server rejects a package ID combined with either other selector. -A successful request uses a package that satisfies all explicit selectors. -An unavailable, incompatible, or disabled package produces an error. - -| SDK configuration field | Selection | -| --- | --- | -| `model_package_id` | Exact package ID; cannot accompany either other field. | -| `backend` | Required backend, such as `onnx` or `trt`. | -| `quantization` | Required precision, such as `fp32` or `fp16`; can accompany `backend`. | - -Malformed or conflicting selectors return HTTP 422. -Unsupported selector values and unavailable backend/precision combinations return HTTP 400. -Unknown package IDs return HTTP 404 with `Model or model package not found.` - -Standard computer vision endpoints accept these fields in the v1 JSON body. -The legacy `/{project}/{version}` endpoint accepts them as query parameters. -`/model/add` and `/model/remove` also accept them in the JSON body. - -```python -from inference_sdk import InferenceConfiguration, InferenceHTTPClient - -client = InferenceHTTPClient(api_url="http://localhost:9001", api_key="YOUR_API_KEY") -with client.use_configuration(InferenceConfiguration(backend="trt", quantization="fp16")): - result = client.infer("image.jpg", model_id="project/1") - client.unload_model("project/1") - -with client.use_configuration(InferenceConfiguration(model_package_id="PACKAGE_ID")): - client.load_model("project/1") - result = client.infer("image.jpg", model_id="project/1") -``` - -The configuration applies to synchronous and asynchronous SDK calls. -For inference and model loading, the SDK requires a server acknowledgment for explicit selection. -It raises an error if an older server silently ignores the selectors. -HTTP success responses include `X-Roboflow-Model-Selection: applied` for selected inference and model loading. - -Each model ID, selector combination, and client credential identifies a separate cache entry. -Each entry counts toward the existing LRU limit. -The registry exposes its cache handle as `model_id`. -The server generates opaque handles with a private random key. -`/model/add` returns `selected_model_id` for explicit selections; the SDK uses this identifier to find the loaded model. -Handles can change after a server restart. They cannot be used for inference or model loading. -Two selector combinations can load separate instances of the same package. -Requests without selectors retain their existing automatic selection and cache entry. -Registration of a preferred package does not replace a resident instance. - -`unload_model()` removes the entry for the active configuration. -The SDK sends the public model ID, selectors, and credential so the server can find the current entry. -`/model/remove` accepts the original model ID with the same selectors and credential, or the registry handle without selectors. -An active request retains its model reference until inference completes. -Concurrent package loads still require sufficient device memory. diff --git a/inference/core/entities/requests/inference.py b/inference/core/entities/requests/inference.py index d96dc332a3..5241388a60 100644 --- a/inference/core/entities/requests/inference.py +++ b/inference/core/entities/requests/inference.py @@ -64,17 +64,33 @@ class InferenceRequest(BaseRequest): model_package_id: Optional[str] = Field( default=None, min_length=1, - description="Exact model package ID. Cannot be combined with backend or quantization.", + description=( + "Exact model package ID. Requires USE_INFERENCE_MODELS=true. " + "Cannot be combined with backend or quantization. A successful request " + "uses this package; an unknown package ID returns HTTP 404. " + "Without selectors, the server uses its automatic package selection." + ), ) backend: Optional[str] = Field( default=None, min_length=1, - description="Required package backend, such as trt or onnx.", + description=( + "Required package backend, such as trt or onnx. Requires " + "USE_INFERENCE_MODELS=true. Can be combined with quantization, but " + "not model_package_id. A successful request satisfies both selectors. " + "Unsupported values or unavailable combinations return HTTP 400." + ), ) quantization: Optional[str] = Field( default=None, min_length=1, - description="Required package quantization, such as fp16 or fp32.", + description=( + "Required package quantization, such as fp16 or fp32. Requires " + "USE_INFERENCE_MODELS=true. Can be combined with backend, but " + "not model_package_id. This describes the package, not the precision " + "of every runtime operation. Unsupported values or unavailable " + "combinations return HTTP 400." + ), ) @model_validator(mode="after") diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index d252fa0668..24d1be882c 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -432,6 +432,19 @@ async def dispatch(self, request, call_next): AUTH_CACHE_TTL_SECONDS = 3600 SHORT_AUTH_CACHE_TTL_SECONDS = 60 REQUEST_RECEIVED_LOG_MESSAGE = "Request received" +MODEL_SELECTION_RESPONSES: Dict[Union[int, str], Dict[str, Any]] = { + 200: { + "headers": { + MODEL_SELECTION_HEADER: { + "description": ( + "Present when explicit model package selectors were applied. " + "A successful response satisfies every explicit selector." + ), + "schema": {"type": "string", "enum": ["applied"]}, + } + } + } +} # Probe/health endpoints whose access-log lines are demoted to DEBUG. HEALTH_CHECK_LOG_PATHS = frozenset( { @@ -1962,7 +1975,16 @@ def registry(): "/model/add", response_model=ModelsDescriptions, summary="Load a model", - description="Load the model with the given model ID", + description=( + "Load the model with the given model ID. Optional package selectors " + "require USE_INFERENCE_MODELS=true. Each model ID, selector combination, " + "and credential identifies a separate entry in the existing LRU cache. " + "Explicit selections return selected_model_id, an opaque registry handle. " + "Use the public model ID and selectors for inference and model loading. " + "Requests without selectors keep their automatic selection and cache entry. " + "Registering a preferred package does not replace a loaded model." + ), + responses=MODEL_SELECTION_RESPONSES, ) @with_route_exceptions def model_add( @@ -2020,7 +2042,13 @@ def model_add( "/model/remove", response_model=ModelsDescriptions, summary="Remove a model", - description="Remove the model with the given model ID", + description=( + "Remove the model with the given model ID. To remove a selected entry, " + "supply its public model ID with the same selectors and credential used " + "to load it, or supply its registry handle without selectors. " + "Other entries remain loaded. Active inference requests retain their " + "model reference until completion." + ), ) @with_route_exceptions def model_remove(request: ClearModelRequest): @@ -2080,6 +2108,7 @@ def model_clear(): ], summary="Object detection infer", description="Run inference with the specified object detection model", + responses=MODEL_SELECTION_RESPONSES, response_model_exclude_none=True, ) @with_route_exceptions @@ -2115,6 +2144,7 @@ def infer_object_detection( ], summary="Instance segmentation infer", description="Run inference with the specified instance segmentation model", + responses=MODEL_SELECTION_RESPONSES, ) @with_route_exceptions @usage_collector("request") @@ -2149,6 +2179,7 @@ def infer_instance_segmentation( ], summary="Semantic segmentation infer", description="Run inference with the specified semantic segmentation model", + responses=MODEL_SELECTION_RESPONSES, ) @with_route_exceptions @usage_collector("request") @@ -2186,6 +2217,7 @@ def infer_semantic_segmentation( ], summary="Classification infer", description="Run inference with the specified classification model", + responses=MODEL_SELECTION_RESPONSES, ) @with_route_exceptions @usage_collector("request") @@ -2218,6 +2250,7 @@ def infer_classification( response_model=Union[KeypointsDetectionInferenceResponse, StubResponse], summary="Keypoints detection infer", description="Run inference with the specified keypoints detection model", + responses=MODEL_SELECTION_RESPONSES, ) @with_route_exceptions @usage_collector("request") @@ -4621,6 +4654,7 @@ def notebook_start(browserless: bool = False): # Legacy object detection inference path for backwards compatibility @app.get( "/{dataset_id}/{version_id:str}", + responses=MODEL_SELECTION_RESPONSES, # Order matters in this response model Union. It will use the first matching model. For example, Object Detection Inference Response is a subset of Instance segmentation inference response, so instance segmentation must come first in order for the matching logic to work. response_model=Union[ InstanceSegmentationInferenceResponse, @@ -4637,6 +4671,7 @@ def notebook_start(browserless: bool = False): ) @app.post( "/{dataset_id}/{version_id:str}", + responses=MODEL_SELECTION_RESPONSES, # Order matters in this response model Union. It will use the first matching model. For example, Object Detection Inference Response is a subset of Instance segmentation inference response, so instance segmentation must come first in order for the matching logic to work. response_model=Union[ InstanceSegmentationInferenceResponse, @@ -4670,9 +4705,25 @@ def legacy_infer_from_request( None, description="Roboflow API Key that will be passed to the model during initialization for artifact retrieval", ), - model_package_id: Optional[str] = Query(None, min_length=1), - backend: Optional[str] = Query(None, min_length=1), - quantization: Optional[str] = Query(None, min_length=1), + model_package_id: Optional[str] = Query( + None, + min_length=1, + description=InferenceRequest.model_fields[ + "model_package_id" + ].description, + ), + backend: Optional[str] = Query( + None, + min_length=1, + description=InferenceRequest.model_fields["backend"].description, + ), + quantization: Optional[str] = Query( + None, + min_length=1, + description=InferenceRequest.model_fields[ + "quantization" + ].description, + ), confidence: Confidence = Query( 0.4, description=( diff --git a/inference_sdk/http/entities.py b/inference_sdk/http/entities.py index bea1729260..3492d9590a 100644 --- a/inference_sdk/http/entities.py +++ b/inference_sdk/http/entities.py @@ -140,6 +140,14 @@ class InferenceConfiguration: max_detections: The maximum number of detections for the inference. iou_threshold: The intersection over union threshold for the inference. stroke_width: The stroke width for the inference. + model_package_id: Exact package ID. Cannot accompany backend or quantization. + backend: Required package backend, such as onnx or trt. + quantization: Required package quantization, such as fp32 or fp16. + + Package selection requires USE_INFERENCE_MODELS=true on the server. + A successful inference satisfies every explicit selector. The SDK rejects + responses without the server's selection acknowledgment. These settings + also apply to model loading and unloading, for synchronous and asynchronous calls. """ confidence_threshold: Optional[Confidence] = None diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py new file mode 100644 index 0000000000..ce7bd925d7 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py @@ -0,0 +1,55 @@ +from unittest.mock import MagicMock + +import pytest + +from inference.core.constants import MODEL_SELECTION_HEADER +from inference.core.interfaces.http import http_api + + +@pytest.fixture +def openapi_schema(monkeypatch): + monkeypatch.setattr(http_api, "InferenceInstrumentator", MagicMock()) + monkeypatch.setattr(http_api, "OFFLINE_MODE", True) + monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) + monkeypatch.setattr(http_api, "LAMBDA", False) + monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) + return http_api.HttpInterface(model_manager=MagicMock()).app.openapi() + + +def test_openapi_describes_selectors_in_body_and_legacy_query(openapi_schema): + schemas = openapi_schema["components"]["schemas"] + legacy = openapi_schema["paths"]["/{dataset_id}/{version_id}"]["post"] + parameters = {parameter["name"]: parameter for parameter in legacy["parameters"]} + + for selector in ("model_package_id", "backend", "quantization"): + description = schemas["ObjectDetectionInferenceRequest"]["properties"][ + selector + ]["description"] + assert "USE_INFERENCE_MODELS=true" in description + assert parameters[selector]["description"] == description + for request in ("AddModelRequest", "ClearModelRequest"): + assert ( + schemas[request]["properties"][selector]["description"] == description + ) + + +@pytest.mark.parametrize( + "path,method", + [ + ("/model/add", "post"), + ("/infer/object_detection", "post"), + ("/infer/instance_segmentation", "post"), + ("/infer/semantic_segmentation", "post"), + ("/infer/classification", "post"), + ("/infer/keypoints_detection", "post"), + ("/{dataset_id}/{version_id}", "get"), + ("/{dataset_id}/{version_id}", "post"), + ], +) +def test_openapi_documents_selection_acknowledgment(openapi_schema, path, method): + response = openapi_schema["paths"][path][method]["responses"]["200"] + assert response["headers"][MODEL_SELECTION_HEADER]["schema"] == { + "type": "string", + "enum": ["applied"], + } + assert "application/json" in response["content"] From 33ccdb7326ebe67e07b345546a9b3f48325fa41b Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 10:00:41 -0230 Subject: [PATCH 11/21] Align package cache fixtures with base model metadata --- .../core/managers/test_model_package_cache.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/inference/unit_tests/core/managers/test_model_package_cache.py b/tests/inference/unit_tests/core/managers/test_model_package_cache.py index 62352b290e..d8c5fbcbef 100644 --- a/tests/inference/unit_tests/core/managers/test_model_package_cache.py +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -19,6 +19,7 @@ from inference.core.managers.active_learning import ActiveLearningManager from inference.core.managers.base import ModelManager from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache +from inference.core.models import base as model_module from inference.core.models.base import Model @@ -30,11 +31,13 @@ class PackageModel(Model): def __init__(self, model_id, api_key, backend=None, **kwargs): self.model_id = model_id - self.resolved_model = SimpleNamespace( - model_id=model_id, - model_package_id=f"{backend or 'onnx'}-package", - backend=backend or "onnx", - quantization="fp16" if backend == "trt" else "fp32", + self._model = SimpleNamespace( + resolved_model=SimpleNamespace( + model_id=model_id, + model_package_id=f"{backend or 'onnx'}-package", + backend=backend or "onnx", + quantization="fp16" if backend == "trt" else "fp32", + ) ) def infer_from_request(self, request) -> Any: @@ -44,6 +47,11 @@ def clear_cache(self, delete_from_disk=True): pass +@pytest.fixture(autouse=True) +def enable_model_metadata(monkeypatch): + monkeypatch.setattr(model_module, "USE_INFERENCE_MODELS", True) + + def test_package_variants_coexist_without_changing_automatic_model(): registry = Mock() registry.get_model.return_value = PackageModel From e3e0451ce1e455e844a7ec32c22282470e2ffb19 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 10:41:17 -0230 Subject: [PATCH 12/21] Shorten model selection field descriptions --- inference/core/entities/requests/inference.py | 22 +++---------------- 1 file changed, 3 insertions(+), 19 deletions(-) diff --git a/inference/core/entities/requests/inference.py b/inference/core/entities/requests/inference.py index 5241388a60..7422a44566 100644 --- a/inference/core/entities/requests/inference.py +++ b/inference/core/entities/requests/inference.py @@ -64,33 +64,17 @@ class InferenceRequest(BaseRequest): model_package_id: Optional[str] = Field( default=None, min_length=1, - description=( - "Exact model package ID. Requires USE_INFERENCE_MODELS=true. " - "Cannot be combined with backend or quantization. A successful request " - "uses this package; an unknown package ID returns HTTP 404. " - "Without selectors, the server uses its automatic package selection." - ), + description="Exact model package ID.", ) backend: Optional[str] = Field( default=None, min_length=1, - description=( - "Required package backend, such as trt or onnx. Requires " - "USE_INFERENCE_MODELS=true. Can be combined with quantization, but " - "not model_package_id. A successful request satisfies both selectors. " - "Unsupported values or unavailable combinations return HTTP 400." - ), + description="Required package backend, such as trt or onnx.", ) quantization: Optional[str] = Field( default=None, min_length=1, - description=( - "Required package quantization, such as fp16 or fp32. Requires " - "USE_INFERENCE_MODELS=true. Can be combined with backend, but " - "not model_package_id. This describes the package, not the precision " - "of every runtime operation. Unsupported values or unavailable " - "combinations return HTTP 400." - ), + description="Required package quantization, such as fp16 or fp32.", ) @model_validator(mode="after") From 28e64c5b202c5faf68fa815890db189d07d1dd03 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 10:56:55 -0230 Subject: [PATCH 13/21] Support action recognition package selection and independent management requests --- .../entities/requests/action_recognition.py | 5 +- inference/core/entities/requests/inference.py | 21 ++--- .../core/entities/requests/model_selection.py | 14 ++++ .../core/entities/requests/server_state.py | 42 ++++++++-- inference/core/interfaces/http/http_api.py | 29 +++---- .../core/models/inference_models_adapters.py | 6 +- .../test_legacy_action_recognition_route.py | 11 ++- .../interfaces/http/test_model_selection.py | 83 +++++++++++++++++-- .../http/test_model_selection_openapi.py | 39 ++++++++- .../models/test_action_recognition_adapter.py | 30 +++++++ 10 files changed, 229 insertions(+), 51 deletions(-) diff --git a/inference/core/entities/requests/action_recognition.py b/inference/core/entities/requests/action_recognition.py index 63c6d4e88f..1570ec686d 100644 --- a/inference/core/entities/requests/action_recognition.py +++ b/inference/core/entities/requests/action_recognition.py @@ -2,7 +2,7 @@ from pydantic import BaseModel, Field -from inference.core.entities.requests.inference import BaseRequest +from inference.core.entities.requests.inference import InferenceRequest class InferenceRequestVideo(BaseModel): @@ -32,7 +32,7 @@ class InferenceRequestVideo(BaseModel): ) -class ActionRecognitionInferenceRequest(BaseRequest): +class ActionRecognitionInferenceRequest(InferenceRequest): """Request for action recognition over a video clip. Attributes: @@ -44,6 +44,7 @@ class ActionRecognitionInferenceRequest(BaseRequest): """ model_id: str = Field( + default=..., examples=["workspace/action-recognition-1"], description="The model to classify with", ) diff --git a/inference/core/entities/requests/inference.py b/inference/core/entities/requests/inference.py index 7422a44566..819d3120b4 100644 --- a/inference/core/entities/requests/inference.py +++ b/inference/core/entities/requests/inference.py @@ -11,9 +11,12 @@ ) from pydantic.json_schema import SkipJsonSchema -from inference.core import env, logger +from inference.core import logger from inference.core.entities.common import ApiKey, ModelID, ModelType -from inference.core.entities.requests.model_selection import model_selection_kwargs +from inference.core.entities.requests.model_selection import ( + model_selection_kwargs, + validate_model_selection, +) from inference_sdk.http.entities import Confidence @@ -77,19 +80,7 @@ class InferenceRequest(BaseRequest): description="Required package quantization, such as fp16 or fp32.", ) - @model_validator(mode="after") - def validate_model_selection(self): - if self.model_package_id is not None and ( - self.backend is not None or self.quantization is not None - ): - raise ValueError( - "model_package_id cannot be combined with backend or quantization." - ) - if model_selection_kwargs(self) and not env.USE_INFERENCE_MODELS: - raise ValueError( - "Model package selection requires USE_INFERENCE_MODELS=true." - ) - return self + validate_model_selection = model_validator(mode="after")(validate_model_selection) class InferenceRequestImage(BaseModel): diff --git a/inference/core/entities/requests/model_selection.py b/inference/core/entities/requests/model_selection.py index cd716e96eb..6182a35504 100644 --- a/inference/core/entities/requests/model_selection.py +++ b/inference/core/entities/requests/model_selection.py @@ -3,11 +3,25 @@ import secrets from typing import Optional +from inference.core import env from inference.models.aliases import resolve_roboflow_model_alias _MODEL_SELECTION_SECRET = secrets.token_bytes(32) +def validate_model_selection(request): + selectors = model_selection_kwargs(request) + if "model_package_id" in selectors and ( + "backend" in selectors or "quantization" in selectors + ): + raise ValueError( + "model_package_id cannot be combined with backend or quantization." + ) + if selectors and not env.USE_INFERENCE_MODELS: + raise ValueError("Model package selection requires USE_INFERENCE_MODELS=true.") + return request + + def model_selection_kwargs(request) -> dict: return { name: value diff --git a/inference/core/entities/requests/server_state.py b/inference/core/entities/requests/server_state.py index b38b2a4499..12387c033b 100644 --- a/inference/core/entities/requests/server_state.py +++ b/inference/core/entities/requests/server_state.py @@ -1,12 +1,12 @@ from typing import Optional -from pydantic import ConfigDict +from pydantic import BaseModel, ConfigDict, Field, model_validator from inference.core.entities.common import ApiKey, ModelID, ModelType -from inference.core.entities.requests.inference import InferenceRequest +from inference.core.entities.requests.model_selection import validate_model_selection -class AddModelRequest(InferenceRequest): +class AddModelRequest(BaseModel): """Request to add a model to the inference server. Attributes: @@ -19,9 +19,24 @@ class AddModelRequest(InferenceRequest): model_id: str = ModelID model_type: Optional[str] = ModelType api_key: Optional[str] = ApiKey - - -class ClearModelRequest(InferenceRequest): + model_package_id: Optional[str] = Field( + default=None, min_length=1, description="Exact model package ID." + ) + backend: Optional[str] = Field( + default=None, + min_length=1, + description="Required package backend, such as trt or onnx.", + ) + quantization: Optional[str] = Field( + default=None, + min_length=1, + description="Required package quantization, such as fp16 or fp32.", + ) + + validate_model_selection = model_validator(mode="after")(validate_model_selection) + + +class ClearModelRequest(BaseModel): """Request to clear a model from the inference server. Attributes: @@ -31,3 +46,18 @@ class ClearModelRequest(InferenceRequest): model_config = ConfigDict(protected_namespaces=()) model_id: str = ModelID api_key: Optional[str] = ApiKey + model_package_id: Optional[str] = Field( + default=None, min_length=1, description="Exact model package ID." + ) + backend: Optional[str] = Field( + default=None, + min_length=1, + description="Required package backend, such as trt or onnx.", + ) + quantization: Optional[str] = Field( + default=None, + min_length=1, + description="Required package quantization, such as fp16 or fp32.", + ) + + validate_model_selection = model_validator(mode="after")(validate_model_selection) diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index 24d1be882c..656067b9af 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -4425,6 +4425,7 @@ def depth_estimation_with_model_id( @app.post( "/infer/action_recognition", response_model=ActionRecognitionInferenceResponse, + responses=MODEL_SELECTION_RESPONSES, summary="Action Recognition", description=( "Classify the actions in a video clip. The model states " @@ -4464,23 +4465,16 @@ def infer_action_recognition( request (Request): The HTTP request. Returns: - ActionRecognitionInferenceResponse: The classified ranges - covering the clip. + Response: The HTTP response containing classified ranges. """ logger.debug("Reached /infer/action_recognition") - api_key = api_key_fallback(api_key) - if api_key is not None: - inference_request.api_key = api_key model_id = inference_request.model_id - self.model_manager.add_model( - model_id, - inference_request.api_key, + response = process_inference_request( + inference_request=inference_request, + api_key=api_key, countinference=countinference, service_secret=service_secret, ) - response = self.model_manager.infer_from_request_sync( - model_id, inference_request - ) if LAMBDA: actor = request.scope["aws.event"]["requestContext"][ "authorizer" @@ -4943,11 +4937,6 @@ def legacy_infer_from_request( ) task_type = self.model_manager.get_task_type(cache_key, api_key=api_key) - if selectors and task_type == "action-recognition": - raise HTTPException( - status_code=422, - detail="Model package selection is not supported for action recognition requests.", - ) if task_type == "action-recognition": # The payload is a clip, so none of the image-shaped # arguments below apply to it. The `image` query parameter @@ -4959,7 +4948,7 @@ def legacy_infer_from_request( # model_id, so the lookup asks for that. Under Lambda # request_model_id is the authorizer's endpoint and # names nothing the manager holds. - model_id, + cache_key if selectors else model_id, ActionRecognitionInferenceRequest( api_key=api_key, model_id=model_id, @@ -4969,10 +4958,14 @@ def legacy_infer_from_request( class_filter=_parse_legacy_class_filter( class_filter=class_filter ), + **selectors, ), ) logger.debug("Response ready.") - return orjson_response(inference_response) + response = orjson_response(inference_response) + if selectors: + response.headers[MODEL_SELECTION_HEADER] = "applied" + return response inference_request_type = ObjectDetectionInferenceRequest args: Dict[str, Any] = {} if task_type == "instance-segmentation": diff --git a/inference/core/models/inference_models_adapters.py b/inference/core/models/inference_models_adapters.py index 46ab6f0f3d..dbc0cace3e 100644 --- a/inference/core/models/inference_models_adapters.py +++ b/inference/core/models/inference_models_adapters.py @@ -2189,7 +2189,11 @@ def load_action_recognition_model( backend=_get_requested_inference_models_backends(kwargs), **kwargs, ) - return _as_action_recognition_model(model=loaded_model, model_id=model_id) + model = _as_action_recognition_model(model=loaded_model, model_id=model_id) + resolved_model = getattr(loaded_model, "resolved_model", None) + if resolved_model is not None: + setattr(model, "resolved_model", resolved_model) + return model def _as_action_recognition_model(model: Any, model_id: str) -> ActionRecognitionModel: diff --git a/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py b/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py index f0de06740d..add194c1a2 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py @@ -78,9 +78,13 @@ def _build_interface(monkeypatch, lambda_mode: bool): @pytest.mark.parametrize("lambda_mode", [False, True]) +@pytest.mark.parametrize("selectors", [{}, {"model_package_id": "video-package"}]) def test_legacy_action_recognition_infers_with_the_registered_identifier( - monkeypatch, lambda_mode: bool + monkeypatch, lambda_mode: bool, selectors ) -> None: + from inference.core import env + + monkeypatch.setattr(env, "USE_INFERENCE_MODELS", True) interface, model_manager = _build_interface(monkeypatch, lambda_mode=lambda_mode) app = _LambdaScope(interface.app) if lambda_mode else interface.app @@ -90,6 +94,7 @@ def test_legacy_action_recognition_infers_with_the_registered_identifier( params={ "api_key": "query-api-key", "image": "https://example.com/clip.mp4", + **selectors, }, ) @@ -100,7 +105,9 @@ def test_legacy_action_recognition_infers_with_the_registered_identifier( assert registered_under == PATH_MODEL_ID # The inference call has to name that same identifier. inferred_with = model_manager.infer_from_request_sync.call_args.args[0] - assert inferred_with == registered_under + assert inferred_with == add_model_call.kwargs.get( + "model_cache_key", registered_under + ) def test_lambda_request_model_id_really_does_differ(monkeypatch) -> None: diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py index 7a593f4bc9..089b642a94 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -15,19 +15,30 @@ def client(monkeypatch): monkeypatch.setattr(http_api, "LAMBDA", False) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) monkeypatch.setattr(http_api, "DEPTH_ESTIMATION_ENABLED", True) + monkeypatch.setattr(http_api, "ACTION_RECOGNITION_ENABLED", True) interface = http_api.HttpInterface(model_manager=MagicMock()) with TestClient(interface.app) as client: yield client @pytest.mark.parametrize("selector", [{"backend": "trt"}, {"quantization": "fp16"}]) -def test_http_rejects_package_id_combined_with_other_selectors(client, selector): - response = client.post( +@pytest.mark.parametrize( + "path", + [ "/infer/object_detection", + "/infer/action_recognition", + "/model/add", + "/model/remove", + ], +) +def test_http_rejects_package_id_combined_with_other_selectors(client, selector, path): + response = client.post( + path, json={ "model_id": "project/1", "model_package_id": "package-1", "image": {"type": "url", "value": "https://example.com/image.jpg"}, + "video": {"type": "url", "value": "https://example.com/clip.mp4"}, **selector, }, ) @@ -35,7 +46,15 @@ def test_http_rejects_package_id_combined_with_other_selectors(client, selector) assert "model_package_id" in response.text -@pytest.mark.parametrize("path", ["/infer/object_detection", "/model/add"]) +@pytest.mark.parametrize( + "path", + [ + "/infer/object_detection", + "/infer/action_recognition", + "/model/add", + "/model/remove", + ], +) def test_http_rejects_selection_when_flag_is_disabled(client, monkeypatch, path): from inference.core import env @@ -46,6 +65,7 @@ def test_http_rejects_selection_when_flag_is_disabled(client, monkeypatch, path) "model_id": "project/1", "backend": "trt", "image": {"type": "url", "value": "https://example.com/image.jpg"}, + "video": {"type": "url", "value": "https://example.com/clip.mp4"}, }, ) assert response.status_code == 422 @@ -76,7 +96,7 @@ def test_depth_rejects_model_package_selection(client, monkeypatch, path, select @pytest.fixture -def package_client(monkeypatch): +def package_client(monkeypatch, request): from types import SimpleNamespace from inference.core import env @@ -88,7 +108,7 @@ class PackageResponse(BaseModel): served_package: str class PackageModel: - task_type = "object-detection" + task_type = getattr(request, "param", "object-detection") batch_size = 1 img_size_h = 32 img_size_w = 32 @@ -113,6 +133,19 @@ def __init__( ) def infer_from_request(self, request): + if self.task_type == "action-recognition": + from inference.core.entities.responses.action_recognition import ( + ActionRecognitionInferenceResponse, + ) + from inference.core.entities.responses.inference import ResolvedModel + + return ActionRecognitionInferenceResponse( + timeline=[], + source_fps=30, + frame_count=30, + windows_classified=1, + resolved_model=ResolvedModel(**vars(self.resolved_model)), + ) return PackageResponse(served_package=self.resolved_model.model_package_id) def clear_cache(self, delete_from_disk=True): @@ -125,6 +158,7 @@ def clear_cache(self, delete_from_disk=True): monkeypatch.setattr(http_api, "OFFLINE_MODE", True) monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) monkeypatch.setattr(http_api, "LAMBDA", False) + monkeypatch.setattr(http_api, "ACTION_RECOGNITION_ENABLED", True) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) registry = MagicMock() monkeypatch.setattr(http_api, "ALLOW_ORIGINS", ["https://client.example"]) @@ -288,3 +322,42 @@ def test_lmm_rejects_model_package_selection(package_client): ) assert response.status_code == 422 assert "not supported for LMM" in response.text + + +@pytest.mark.parametrize("package_client", ["action-recognition"], indirect=True) +@pytest.mark.parametrize("legacy", [False, True]) +@pytest.mark.parametrize( + "selectors", + [{"backend": "trt", "quantization": "fp16"}, {"model_package_id": "engine-1"}], +) +def test_action_recognition_selected_package_coexists_with_default( + package_client, legacy, selectors +): + def infer(selection): + if legacy: + return package_client.post( + "/project/1", + params={"image": "https://example.com/clip.mp4", **selection}, + ) + return package_client.post( + "/infer/action_recognition", + json={ + "model_id": "project/1", + "video": {"type": "url", "value": "https://example.com/clip.mp4"}, + "disable_model_monitoring": True, + **selection, + }, + ) + + automatic = infer({}) + selected = infer(selectors) + again = infer({}) + assert automatic.status_code == selected.status_code == again.status_code == 200 + assert automatic.json()["resolved_model"]["model_package_id"] == "onnx-1" + assert again.json()["resolved_model"]["model_package_id"] == "onnx-1" + assert selected.json()["resolved_model"]["model_package_id"] == "engine-1" + assert selected.headers["X-Roboflow-Model-Selection"] == "applied" + assert "X-Roboflow-Model-Selection" not in automatic.headers + mismatch = infer({"backend": "trt", "quantization": "fp32"}) + assert mismatch.status_code == 400 + assert "X-Roboflow-Model-Selection" not in mismatch.headers diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py index ce7bd925d7..d61e3102af 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py @@ -12,6 +12,7 @@ def openapi_schema(monkeypatch): monkeypatch.setattr(http_api, "OFFLINE_MODE", True) monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) monkeypatch.setattr(http_api, "LAMBDA", False) + monkeypatch.setattr(http_api, "ACTION_RECOGNITION_ENABLED", True) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) return http_api.HttpInterface(model_manager=MagicMock()).app.openapi() @@ -25,14 +26,47 @@ def test_openapi_describes_selectors_in_body_and_legacy_query(openapi_schema): description = schemas["ObjectDetectionInferenceRequest"]["properties"][ selector ]["description"] - assert "USE_INFERENCE_MODELS=true" in description + assert description assert parameters[selector]["description"] == description - for request in ("AddModelRequest", "ClearModelRequest"): + for request in ( + "AddModelRequest", + "ClearModelRequest", + "ActionRecognitionInferenceRequest", + ): assert ( schemas[request]["properties"][selector]["description"] == description ) +@pytest.mark.parametrize( + "schema_name,expected_fields", + [ + ( + "AddModelRequest", + { + "model_id", + "model_type", + "api_key", + "model_package_id", + "backend", + "quantization", + }, + ), + ( + "ClearModelRequest", + {"model_id", "api_key", "model_package_id", "backend", "quantization"}, + ), + ], +) +def test_model_management_schemas_only_include_model_management_fields( + openapi_schema, schema_name, expected_fields +): + assert ( + set(openapi_schema["components"]["schemas"][schema_name]["properties"]) + == expected_fields + ) + + @pytest.mark.parametrize( "path,method", [ @@ -42,6 +76,7 @@ def test_openapi_describes_selectors_in_body_and_legacy_query(openapi_schema): ("/infer/semantic_segmentation", "post"), ("/infer/classification", "post"), ("/infer/keypoints_detection", "post"), + ("/infer/action_recognition", "post"), ("/{dataset_id}/{version_id}", "get"), ("/{dataset_id}/{version_id}", "post"), ], diff --git a/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py b/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py index 63b7da74c0..32a5ce7087 100644 --- a/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py +++ b/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py @@ -363,3 +363,33 @@ def test_load_action_recognition_model_passes_the_zero_shot_id_through() -> None auto_model.from_pretrained.call_args.kwargs["model_id_or_path"] == "nvidia/cosmos-3-edge-action-recognition" ) + + +@pytest.mark.parametrize( + "selectors", + [ + {"backend": "torch", "quantization": "fp16"}, + {"model_package_id": "video-package"}, + ], +) +def test_action_loader_preserves_package_selection_when_wrapping_a_reasoner(selectors): + from inference.core.models import inference_models_adapters as adapters + + resolved = SimpleNamespace( + model_id="workspace/model", + model_package_id="video-package", + backend="torch", + quantization="fp16", + ) + reasoner = SimpleNamespace(resolved_model=resolved) + wrapped = _FakeModel(responses=[]) + with patch.object( + adapters.AutoModel, "from_pretrained", return_value=reasoner + ) as load, patch.object( + adapters, "_as_action_recognition_model", return_value=wrapped + ): + model = adapters.load_action_recognition_model("workspace/model", **selectors) + + assert getattr(model, "resolved_model", None) == resolved + for name, value in selectors.items(): + assert load.call_args.kwargs[name] == value From cd57a84ccbbb9540b806b4a8d2d47b0fc797021c Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 11:01:11 -0230 Subject: [PATCH 14/21] Remove redundant model selection gates --- inference/core/interfaces/http/http_api.py | 2 +- inference/core/managers/base.py | 10 ++-------- tests/inference/unit_tests/core/managers/test_base.py | 5 +++-- 3 files changed, 6 insertions(+), 11 deletions(-) diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index 656067b9af..203f27970f 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -4948,7 +4948,7 @@ def legacy_infer_from_request( # model_id, so the lookup asks for that. Under Lambda # request_model_id is the authorizer's endpoint and # names nothing the manager holds. - cache_key if selectors else model_id, + cache_key, ActionRecognitionInferenceRequest( api_key=api_key, model_id=model_id, diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index dafcdd9846..0dafbd3f71 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -621,20 +621,14 @@ def infer_from_request_sync( async def model_infer(self, model_id: str, request: InferenceRequest, **kwargs): model = self._get_model_reference(model_id=model_id) - if isinstance(request, InferenceRequest): - self._validate_loaded_model_selection( - model, model_selection_kwargs(request) - ) + self._validate_loaded_model_selection(model, model_selection_kwargs(request)) return model.infer_from_request(request) def model_infer_sync( self, model_id: str, request: InferenceRequest, **kwargs ) -> Union[List[InferenceResponse], InferenceResponse]: model = self._get_model_reference(model_id=model_id) - if isinstance(request, InferenceRequest): - self._validate_loaded_model_selection( - model, model_selection_kwargs(request) - ) + self._validate_loaded_model_selection(model, model_selection_kwargs(request)) return model.infer_from_request(request) def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: diff --git a/tests/inference/unit_tests/core/managers/test_base.py b/tests/inference/unit_tests/core/managers/test_base.py index 920af621ba..2b79bf7afd 100644 --- a/tests/inference/unit_tests/core/managers/test_base.py +++ b/tests/inference/unit_tests/core/managers/test_base.py @@ -3,6 +3,7 @@ import pytest +from inference.core.entities.requests.inference import BaseRequest from inference.core.exceptions import InferenceModelNotFound from inference.core.managers import base as base_module from inference.core.managers.base import ModelManager @@ -127,7 +128,7 @@ async def test_infer_from_request_when_model_is_available() -> None: model_registry = MagicMock() model_manager = ModelManager(model_registry=model_registry) model_manager._models = {"some/1": MagicMock()} - request = MagicMock() + request = BaseRequest(id="request-id", disable_model_monitoring=True) # when result = await model_manager.infer_from_request(model_id="some/1", request=request) @@ -148,7 +149,7 @@ async def test_infer_from_request_when_model_is_available_but_exception_raised() error = ValueError() model_mock.infer_from_request.side_effect = error model_manager._models = {"some/1": model_mock} - request = MagicMock() + request = BaseRequest(id="request-id", disable_model_monitoring=True) # when with pytest.raises(ValueError) as thrown_exception: From 640dd9728c0feaa67d39bab943af9c6025acecc8 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 11:01:45 -0230 Subject: [PATCH 15/21] Document model package selection exceptions --- inference/core/exceptions.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/inference/core/exceptions.py b/inference/core/exceptions.py index c737ffa9d0..fe027669e5 100644 --- a/inference/core/exceptions.py +++ b/inference/core/exceptions.py @@ -39,11 +39,19 @@ class InferenceModelNotFound(Exception): class ModelPackageSelectionError(Exception): - pass + """Raised when the requested model package selection cannot be satisfied. + + Attributes: + message (str): Optional message describing the error. + """ class ModelPackageNotFoundError(Exception): - pass + """Raised when the requested model package is not found. + + Attributes: + message (str): Optional message describing the error. + """ class InvalidEnvironmentVariableError(Exception): From dd146f42bc84e800b61fe7205cc2324f361db399 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 11:11:12 -0230 Subject: [PATCH 16/21] Use confirmed handles for selected unloads and move cache identity policy --- .../core/entities/requests/model_selection.py | 23 ---------- inference/core/interfaces/http/http_api.py | 2 +- inference/core/managers/decorators/base.py | 9 ++-- .../managers/decorators/fixed_size_cache.py | 3 +- inference/core/managers/model_selection.py | 23 ++++++++++ inference_sdk/http/client.py | 29 ++++++++++-- .../interfaces/http/test_model_selection.py | 14 +++--- .../unit_tests/http/test_model_selection.py | 45 +++++++++++++------ 8 files changed, 96 insertions(+), 52 deletions(-) create mode 100644 inference/core/managers/model_selection.py diff --git a/inference/core/entities/requests/model_selection.py b/inference/core/entities/requests/model_selection.py index 6182a35504..59354913b6 100644 --- a/inference/core/entities/requests/model_selection.py +++ b/inference/core/entities/requests/model_selection.py @@ -1,12 +1,4 @@ -import hmac -import json -import secrets -from typing import Optional - from inference.core import env -from inference.models.aliases import resolve_roboflow_model_alias - -_MODEL_SELECTION_SECRET = secrets.token_bytes(32) def validate_model_selection(request): @@ -28,18 +20,3 @@ def model_selection_kwargs(request) -> dict: for name in ("model_package_id", "backend", "quantization") if (value := getattr(request, name, None)) is not None } - - -def model_selection_cache_key( - model_id: str, selectors: dict, api_key: Optional[str] = None -) -> str: - if not selectors: - return model_id - model_id = resolve_roboflow_model_alias(model_id) - payload = json.dumps( - {"model_id": model_id, "selectors": selectors, "api_key": api_key}, - sort_keys=True, - separators=(",", ":"), - ).encode() - digest = hmac.digest(_MODEL_SELECTION_SECRET, payload, "sha256").hex() - return f"{model_id}:package:{digest}" diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index 203f27970f..65392e70a1 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -70,7 +70,6 @@ ensure_wire_safe_mask_format, ) from inference.core.entities.requests.model_selection import ( - model_selection_cache_key, model_selection_kwargs, ) from inference.core.entities.requests.owlv2 import OwlV2InferenceRequest @@ -331,6 +330,7 @@ request_model_ids, request_workflow_id, ) +from inference.core.managers.model_selection import model_selection_cache_key from inference.core.managers.prometheus import InferenceInstrumentator from inference.core.roboflow_api import ( assume_identity_authorised_workspace_db_id, diff --git a/inference/core/managers/decorators/base.py b/inference/core/managers/decorators/base.py index 0774379799..d55da39f8d 100644 --- a/inference/core/managers/decorators/base.py +++ b/inference/core/managers/decorators/base.py @@ -6,7 +6,12 @@ from inference.core.entities.requests.inference import InferenceRequest from inference.core.entities.responses.inference import InferenceResponse from inference.core.env import API_KEY -from inference.core.managers.base import Model, ModelManager, model_load_options +from inference.core.managers.base import ( + Model, + ModelManager, + model_load_options, + validate_public_model_id, +) from inference.core.managers.model_load_collector import request_model_ids from inference.core.models.types import PreprocessReturnMetadata from inference.core.roboflow_api import ModelEndpointType @@ -75,8 +80,6 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ - from inference.core.managers.base import validate_public_model_id - validate_public_model_id(model_id, model_id_alias) cache_key = model_cache_key or model_id if cache_key in self: diff --git a/inference/core/managers/decorators/fixed_size_cache.py b/inference/core/managers/decorators/fixed_size_cache.py index 73dbd66876..3b21d970cf 100644 --- a/inference/core/managers/decorators/fixed_size_cache.py +++ b/inference/core/managers/decorators/fixed_size_cache.py @@ -23,6 +23,7 @@ ModelManager, acquire_with_timeout, model_load_options, + validate_public_model_id, ) from inference.core.managers.decorators.base import ModelManagerDecorator from inference.core.managers.entities import ModelDescription @@ -76,8 +77,6 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ - from inference.core.managers.base import validate_public_model_id - validate_public_model_id(model_id, model_id_alias) if MODELS_CACHE_AUTH_ENABLED and not OFFLINE_MODE: if not _check_if_api_key_has_access_to_model( diff --git a/inference/core/managers/model_selection.py b/inference/core/managers/model_selection.py new file mode 100644 index 0000000000..36b2d9d03b --- /dev/null +++ b/inference/core/managers/model_selection.py @@ -0,0 +1,23 @@ +import hmac +import json +import secrets +from typing import Optional + +from inference.models.aliases import resolve_roboflow_model_alias + +_MODEL_SELECTION_SECRET = secrets.token_bytes(32) + + +def model_selection_cache_key( + model_id: str, selectors: dict, api_key: Optional[str] = None +) -> str: + if not selectors: + return model_id + model_id = resolve_roboflow_model_alias(model_id) + payload = json.dumps( + {"model_id": model_id, "selectors": selectors, "api_key": api_key}, + sort_keys=True, + separators=(",", ":"), + ).encode() + digest = hmac.digest(_MODEL_SELECTION_SECRET, payload, "sha256").hex() + return f"{model_id}:package:{digest}" diff --git a/inference_sdk/http/client.py b/inference_sdk/http/client.py index 0b282e8220..2a4e2c9413 100644 --- a/inference_sdk/http/client.py +++ b/inference_sdk/http/client.py @@ -1158,6 +1158,9 @@ def __remember_model_selection( def unload_model(self, model_id: str) -> RegisteredModels: """Unload a model from the server. + With package selection configured, load the selected model through this + client before unloading it. + Args: model_id (str): The identifier of the model to unload. @@ -1166,17 +1169,20 @@ def unload_model(self, model_id: str) -> RegisteredModels: Raises: WrongClientModeError: If not in API v1 mode. + ModelNotInitializedError: If this client has no confirmed handle for the selected model. HTTPCallErrorError: If there is an error in the HTTP call. HTTPClientError: If there is an error with the server connection. """ self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) selectors = self.__inference_configuration.to_model_selection_parameters() + removal_model_id = self.__resolve_model_id_for_unload( + de_aliased_model_id, selectors + ) response = requests.post( f"{self.__api_url}/model/remove", json={ - "model_id": de_aliased_model_id, - **selectors, + "model_id": removal_model_id, **(self.__legacy_api_key_payload() if selectors else {}), }, headers=self.__headers_with_auth(DEFAULT_HEADERS), @@ -1198,12 +1204,14 @@ async def unload_model_async(self, model_id: str) -> RegisteredModels: self.__ensure_v1_client_mode() de_aliased_model_id = resolve_roboflow_model_alias(model_id=model_id) selectors = self.__inference_configuration.to_model_selection_parameters() + removal_model_id = self.__resolve_model_id_for_unload( + de_aliased_model_id, selectors + ) async with aiohttp.ClientSession() as session: async with session.post( f"{self.__api_url}/model/remove", json={ - "model_id": de_aliased_model_id, - **selectors, + "model_id": removal_model_id, **(self.__legacy_api_key_payload() if selectors else {}), }, headers=self.__headers_with_auth(DEFAULT_HEADERS), @@ -1220,6 +1228,19 @@ async def unload_model_async(self, model_id: str) -> RegisteredModels: self.__selected_model = None return RegisteredModels.from_dict(response_payload) + def __resolve_model_id_for_unload( + self, model_id: str, selectors: Dict[str, str] + ) -> str: + if not selectors: + return model_id + selection_key = (model_id, tuple(sorted(selectors.items()))) + try: + return self.__model_selection_ids[selection_key] + except KeyError as error: + raise ModelNotInitializedError( + "Load the selected model with this client before unloading it." + ) from error + @wrap_errors def unload_all_models(self) -> RegisteredModels: self.__ensure_v1_client_mode() diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py index 089b642a94..19b4682751 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -175,8 +175,9 @@ def clear_cache(self, delete_from_disk=True): "selectors", [{"backend": "trt", "quantization": "fp16"}, {"model_package_id": "engine-1"}], ) +@pytest.mark.parametrize("remove_by_handle", [False, True]) def test_http_selected_package_coexists_with_automatic_package( - package_client, selectors + package_client, selectors, remove_by_handle ): payload = { "model_id": "project/1", @@ -195,9 +196,12 @@ def test_http_selected_package_coexists_with_automatic_package( assert selected.json()["served_package"] == "engine-1" registered = package_client.get("/model/registry").json()["models"] assert len(registered) == 2 - removed = package_client.post( - "/model/remove", json={"model_id": "project/1", **selectors} - ) + removal = {"model_id": "project/1", **selectors} + if remove_by_handle: + loaded = package_client.post("/model/add", json=removal) + assert loaded.status_code == 200 + removal = {"model_id": loaded.json()["selected_model_id"]} + removed = package_client.post("/model/remove", json=removal) assert removed.status_code == 200 assert [model["model_id"] for model in removed.json()["models"]] == ["project/1"] @@ -263,7 +267,7 @@ def test_aliased_legacy_selection_shares_v1_entry_and_can_be_removed( def test_selected_model_handle_depends_on_server_secret(package_client, monkeypatch): - from inference.core.entities.requests import model_selection + from inference.core.managers import model_selection payload = {"model_id": "project/1", "backend": "trt", "api_key": "owner"} first = package_client.post("/model/add", json=payload) diff --git a/tests/inference_sdk/unit_tests/http/test_model_selection.py b/tests/inference_sdk/unit_tests/http/test_model_selection.py index 99b8811caf..fc463b3bef 100644 --- a/tests/inference_sdk/unit_tests/http/test_model_selection.py +++ b/tests/inference_sdk/unit_tests/http/test_model_selection.py @@ -6,7 +6,7 @@ from inference_sdk import InferenceConfiguration, InferenceHTTPClient from inference_sdk.http.entities import ApiKeyTransport -from inference_sdk.http.errors import InvalidParameterError +from inference_sdk.http.errors import InvalidParameterError, ModelNotInitializedError def register_response( @@ -138,17 +138,21 @@ async def test_async_v1_description_loads_selected_variant_without_loading_defau } -def test_unload_selected_variant_sends_public_id_and_selectors(requests_mock): - selectors: dict[str, Any] = {"model_package_id": "engine-1"} - requests_mock.post("http://server/model/remove", json={"models": []}) +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_selected_unload_requires_a_confirmed_handle_before_sending_request( + requests_mock, asynchronous +): client = InferenceHTTPClient(api_url="http://server", api_key="key") - client.configure(InferenceConfiguration(**selectors)) - client.unload_model("project/1") - assert requests_mock.last_request.json() == { - "model_id": "project/1", - "api_key": "key", - **selectors, - } + client.configure(InferenceConfiguration(model_package_id="engine-1")) + with aioresponses() as responses: + register_response( + requests_mock, responses, "post", "/model/remove", registry_payload() + ) + with pytest.raises(ModelNotInitializedError, match="Load the selected model"): + await call_client(client, "unload_model", asynchronous, "project/1") + assert not requests_mock.called + assert not responses.requests @pytest.mark.asyncio @@ -290,25 +294,38 @@ async def test_unload_forgets_selected_handle(requests_mock, asynchronous, unloa @pytest.mark.asyncio @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("transport", [ApiKeyTransport.LEGACY, ApiKeyTransport.HEADER]) -async def test_unload_selected_model_uses_public_alias_and_auth( +async def test_unload_selected_model_uses_confirmed_handle_and_auth( requests_mock, asynchronous, transport ): client = InferenceHTTPClient(api_url="http://server", api_key="key") client.configure(InferenceConfiguration(backend="trt", api_key_transport=transport)) with aioresponses() as responses: + register_response( + requests_mock, + responses, + "post", + "/model/add", + registry_payload("server-handle", selected_model_id="server-handle"), + {"X-Roboflow-Model-Selection": "applied"}, + ) + await call_client(client, "load_model", asynchronous, "yolov8n-640") register_response( requests_mock, responses, "post", "/model/remove", registry_payload() ) await call_client(client, "unload_model", asynchronous, "yolov8n-640") if asynchronous: - request = next(iter(responses.requests.values()))[0].kwargs + request = next( + calls[0].kwargs + for (_, url), calls in responses.requests.items() + if url.path == "/model/remove" + ) payload, headers = request["json"], request["headers"] else: payload, headers = ( requests_mock.last_request.json(), requests_mock.last_request.headers, ) - expected = {"model_id": "coco/3", "backend": "trt"} + expected = {"model_id": "server-handle"} if transport is ApiKeyTransport.LEGACY: expected["api_key"] = "key" else: From e1524525b9eb749364485001cecb4c764a2ac88d Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Thu, 17 Sep 2026 11:32:26 -0230 Subject: [PATCH 17/21] Check selector schema contracts instead of description text --- .../http/test_model_selection_openapi.py | 23 +++++++++++-------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py index d61e3102af..ea72ecc607 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py @@ -17,25 +17,30 @@ def openapi_schema(monkeypatch): return http_api.HttpInterface(model_manager=MagicMock()).app.openapi() -def test_openapi_describes_selectors_in_body_and_legacy_query(openapi_schema): +def test_openapi_exposes_optional_string_selectors_in_body_and_legacy_query( + openapi_schema, +): schemas = openapi_schema["components"]["schemas"] legacy = openapi_schema["paths"]["/{dataset_id}/{version_id}"]["post"] parameters = {parameter["name"]: parameter for parameter in legacy["parameters"]} for selector in ("model_package_id", "backend", "quantization"): - description = schemas["ObjectDetectionInferenceRequest"]["properties"][ - selector - ]["description"] - assert description - assert parameters[selector]["description"] == description + parameter = parameters[selector] + assert parameter["in"] == "query" + assert parameter["required"] is False + assert {"type": "string", "minLength": 1} in parameter["schema"]["anyOf"] + assert {"type": "null"} in parameter["schema"]["anyOf"] for request in ( + "ObjectDetectionInferenceRequest", "AddModelRequest", "ClearModelRequest", "ActionRecognitionInferenceRequest", ): - assert ( - schemas[request]["properties"][selector]["description"] == description - ) + schema = schemas[request] + field = schema["properties"][selector] + assert selector not in schema.get("required", []) + assert {"type": "string", "minLength": 1} in field["anyOf"] + assert {"type": "null"} in field["anyOf"] @pytest.mark.parametrize( From e256ff7ce5e11598fa7e41f56b914a63c9323316 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Fri, 18 Sep 2026 16:23:28 -0230 Subject: [PATCH 18/21] Match CI import formatting and document cache HMAC alert --- inference/core/interfaces/http/http_api.py | 4 +--- inference/core/managers/base.py | 4 +--- inference/core/managers/model_selection.py | 1 + 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index 65392e70a1..3b110bae49 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -69,9 +69,7 @@ SemanticSegmentationInferenceRequest, ensure_wire_safe_mask_format, ) -from inference.core.entities.requests.model_selection import ( - model_selection_kwargs, -) +from inference.core.entities.requests.model_selection import model_selection_kwargs from inference.core.entities.requests.owlv2 import OwlV2InferenceRequest from inference.core.entities.requests.perception_encoder import ( PerceptionEncoderCompareRequest, diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index 0dafbd3f71..3ae5281f07 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -10,9 +10,7 @@ from inference.core.cache.serializers import to_cachable_inference_item from inference.core.devices.utils import GLOBAL_INFERENCE_SERVER_ID from inference.core.entities.requests.inference import InferenceRequest -from inference.core.entities.requests.model_selection import ( - model_selection_kwargs, -) +from inference.core.entities.requests.model_selection import model_selection_kwargs from inference.core.entities.responses.inference import InferenceResponse from inference.core.env import ( DISABLE_INFERENCE_CACHE, diff --git a/inference/core/managers/model_selection.py b/inference/core/managers/model_selection.py index 36b2d9d03b..be42f08af8 100644 --- a/inference/core/managers/model_selection.py +++ b/inference/core/managers/model_selection.py @@ -19,5 +19,6 @@ def model_selection_cache_key( sort_keys=True, separators=(",", ":"), ).encode() + # codeql[py/weak-sensitive-data-hashing]: HMAC with a random server secret identifies cache entries; this is not password storage. digest = hmac.digest(_MODEL_SELECTION_SECRET, payload, "sha256").hex() return f"{model_id}:package:{digest}" From 8b390b285ded209375bd20505a95f4f3f15969e6 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Mon, 28 Sep 2026 17:15:18 -0400 Subject: [PATCH 19/21] Fix package selection validation, SDK video requests, and telemetry --- inference/core/managers/base.py | 47 +++-- inference/core/managers/decorators/base.py | 3 +- .../managers/decorators/fixed_size_cache.py | 3 +- inference/core/managers/model_selection.py | 4 + inference/core/models/base.py | 2 + .../core/models/inference_models_adapters.py | 21 +++ .../enterprise/parallel/dispatch_manager.py | 12 +- inference_models/inference_models/errors.py | 4 + .../models/auto_loaders/auto_negotiation.py | 28 ++- .../models/auto_loaders/core.py | 8 + .../auto_loaders/test_auto_negotiation.py | 62 ++++++- .../models/auto_loaders/test_core.py | 49 ++++++ inference_sdk/http/client.py | 17 +- inference_sdk/http/errors.py | 6 + inference_sdk/http/utils/model_selection.py | 7 +- .../interfaces/http/test_model_selection.py | 56 +++++- .../http/test_model_selection_errors.py | 64 ++++++- .../http/test_model_selection_openapi.py | 9 +- .../http/test_parallel_model_selection.py | 66 +++++++ .../core/managers/test_model_package_cache.py | 18 ++ .../models/test_action_recognition_adapter.py | 9 +- .../models/test_pinned_package_policy.py | 80 +++++++++ .../unit_tests/http/test_model_selection.py | 163 +++++++++++++++++- 23 files changed, 707 insertions(+), 31 deletions(-) create mode 100644 tests/inference/unit_tests/core/interfaces/http/test_parallel_model_selection.py create mode 100644 tests/inference/unit_tests/models/test_pinned_package_policy.py diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index 3ae5281f07..d7dd16f22b 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -6,6 +6,7 @@ import numpy as np from fastapi.encoders import jsonable_encoder +import inference_models.errors as inference_models_errors from inference.core.cache import model_monitoring as model_monitoring_cache_module from inference.core.cache.serializers import to_cachable_inference_item from inference.core.devices.utils import GLOBAL_INFERENCE_SERVER_ID @@ -37,6 +38,7 @@ model_load_info, request_model_ids, ) +from inference.core.managers.model_selection import public_model_id from inference.core.managers.pingback import PingbackInfo from inference.core.models.base import Model, PreprocessReturnMetadata from inference.core.registries.base import ModelRegistry @@ -192,7 +194,7 @@ def add_model( ) ids_collector = request_model_ids.get(None) if ids_collector is not None: - ids_collector.add(resolved_identifier) + ids_collector.add(lookup_identifier) model_lock = self._get_lock_for_a_model(model_id=resolved_identifier) with acquire_with_timeout(lock=model_lock) as acquired: if not acquired: @@ -207,7 +209,7 @@ def add_model( ) return try: - with start_span("model.load", {"model.id": resolved_identifier}): + with start_span("model.load", {"model.id": lookup_identifier}): logger.debug("ModelManager - model initialisation...") t_load_start = time.perf_counter() vram_before = _get_cuda_memory_allocated() @@ -217,6 +219,12 @@ def add_model( countinference=countinference, service_secret=service_secret, ) + if selectors and not getattr( + model_class, "supports_model_package_selection", False + ): + raise ModelPackageSelectionError( + "This model does not support explicit package selection." + ) extra_init_kwargs = {} extra_init_kwargs.update(selectors) @@ -266,7 +274,7 @@ def add_model( load_time = time.perf_counter() - t_load_start vram_delta = getattr(model, "_vram_bytes", None) set_span_attribute("model.load_time_seconds", load_time) - record_model_loaded(resolved_identifier, load_time) + record_model_loaded(lookup_identifier, load_time) logger.info( "Model loaded: model_id=%s, load_time=%.2fs, task_type=%s, vram_bytes=%s", resolved_identifier, @@ -278,11 +286,18 @@ def add_model( collector = model_load_info.get(None) if collector is not None: collector.record( - model_id=resolved_identifier, load_time=load_time + model_id=lookup_identifier, load_time=load_time ) except Exception as error: record_error(error) self._dispose_model_lock(model_id=resolved_identifier) + if selectors and isinstance( + error, + getattr(inference_models_errors, "ModelPackagePolicyError", ()), + ): + raise ModelPackageSelectionError( + "The requested model package is incompatible with this server's backend policy or runtime." + ) from error if selectors and isinstance(error, NoModelPackagesAvailableError): if model_package_id is not None: raise ModelPackageNotFoundError() from error @@ -428,14 +443,19 @@ async def infer_from_request( self.pingback.fallback_api_key = request.api_key with start_span( "model.infer", - {"model.id": model_id, "model.infer.caller": "infer_from_request"}, + { + "model.id": public_model_id(model_id), + "model.infer.caller": "infer_from_request", + }, ): try: t_infer_start = time.perf_counter() rtn_val = await self.model_infer( model_id=model_id, request=request, **kwargs ) - record_inference(model_id, time.perf_counter() - t_infer_start) + record_inference( + public_model_id(model_id), time.perf_counter() - t_infer_start + ) logger.debug( f"ModelManager - inference from request finished for model_id={model_id}." ) @@ -534,14 +554,19 @@ def infer_from_request_sync( self.pingback.fallback_api_key = request.api_key with start_span( "model.infer", - {"model.id": model_id, "model.infer.caller": "infer_from_request_sync"}, + { + "model.id": public_model_id(model_id), + "model.infer.caller": "infer_from_request_sync", + }, ): try: t_infer_start = time.perf_counter() rtn_val = self.model_infer_sync( model_id=model_id, request=request, **kwargs ) - record_inference(model_id, time.perf_counter() - t_infer_start) + record_inference( + public_model_id(model_id), time.perf_counter() - t_infer_start + ) logger.debug( f"ModelManager - inference from request finished for model_id={model_id}." ) @@ -633,7 +658,7 @@ def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: with start_span( "model.infer", { - "model.id": model_id, + "model.id": public_model_id(model_id), "model.infer.caller": "run_tensor_native_inference", }, ): @@ -641,7 +666,9 @@ def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: t_infer_start = time.perf_counter() model = self._get_model_reference(model_id=model_id) result = model.run_tensor_native_inference(**kwargs) - record_inference(model_id, time.perf_counter() - t_infer_start) + record_inference( + public_model_id(model_id), time.perf_counter() - t_infer_start + ) return result except Exception as error: record_error(error) diff --git a/inference/core/managers/decorators/base.py b/inference/core/managers/decorators/base.py index d55da39f8d..f9d0812ca8 100644 --- a/inference/core/managers/decorators/base.py +++ b/inference/core/managers/decorators/base.py @@ -13,6 +13,7 @@ validate_public_model_id, ) from inference.core.managers.model_load_collector import request_model_ids +from inference.core.managers.model_selection import public_model_id from inference.core.models.types import PreprocessReturnMetadata from inference.core.roboflow_api import ModelEndpointType @@ -96,7 +97,7 @@ def add_model( ) ids_collector = request_model_ids.get(None) if ids_collector is not None: - ids_collector.add(cache_key) + ids_collector.add(public_model_id(cache_key)) return self.model_manager.add_model( model_id, diff --git a/inference/core/managers/decorators/fixed_size_cache.py b/inference/core/managers/decorators/fixed_size_cache.py index 3b21d970cf..a8abd6335d 100644 --- a/inference/core/managers/decorators/fixed_size_cache.py +++ b/inference/core/managers/decorators/fixed_size_cache.py @@ -28,6 +28,7 @@ from inference.core.managers.decorators.base import ModelManagerDecorator from inference.core.managers.entities import ModelDescription from inference.core.managers.model_load_collector import request_model_ids +from inference.core.managers.model_selection import public_model_id from inference.core.registries.roboflow import ( ModelEndpointType, _check_if_api_key_has_access_to_model, @@ -95,7 +96,7 @@ def add_model( ) ids_collector = request_model_ids.get(None) if ids_collector is not None: - ids_collector.add(queue_id) + ids_collector.add(public_model_id(queue_id)) if queue_id in self: self.validate_model_selection( queue_id, diff --git a/inference/core/managers/model_selection.py b/inference/core/managers/model_selection.py index be42f08af8..ea840b9671 100644 --- a/inference/core/managers/model_selection.py +++ b/inference/core/managers/model_selection.py @@ -22,3 +22,7 @@ def model_selection_cache_key( # codeql[py/weak-sensitive-data-hashing]: HMAC with a random server secret identifies cache entries; this is not password storage. digest = hmac.digest(_MODEL_SELECTION_SECRET, payload, "sha256").hex() return f"{model_id}:package:{digest}" + + +def public_model_id(model_id: str) -> str: + return model_id.partition(":package:")[0] diff --git a/inference/core/models/base.py b/inference/core/models/base.py index 547da9f601..d7d2463894 100644 --- a/inference/core/models/base.py +++ b/inference/core/models/base.py @@ -110,6 +110,8 @@ class Model(BaseInference): clear_cache(): Clears any cache if necessary. """ + supports_model_package_selection = False + @property def resolved_model(self) -> Any: model = getattr(self, "_model", None) diff --git a/inference/core/models/inference_models_adapters.py b/inference/core/models/inference_models_adapters.py index dbc0cace3e..78b8c786e5 100644 --- a/inference/core/models/inference_models_adapters.py +++ b/inference/core/models/inference_models_adapters.py @@ -217,6 +217,15 @@ def _get_enabled_inference_models_backends() -> List[str]: def _get_requested_inference_models_backends(kwargs: dict) -> Union[str, List[str]]: from inference.core.exceptions import ModelPackageSelectionError + if kwargs.get("model_package_id") is not None: + if ( + "validate_model_package" + not in signature(AutoModel.from_pretrained).parameters + ): + raise ModelPackageSelectionError( + "This inference-models version cannot validate pinned package policy." + ) + kwargs["validate_model_package"] = True requested_backend = kwargs.pop("backend", None) enabled_backends = _get_enabled_inference_models_backends() if requested_backend is None: @@ -271,6 +280,8 @@ def _fixed_input_hw_from_backend(backend: Any) -> Optional[Tuple[int, int]]: class InferenceModelsObjectDetectionAdapter(Model): + supports_model_package_selection = True + def __init__(self, model_id: str, api_key: str = None, **kwargs): super().__init__() @@ -437,6 +448,8 @@ def draw_predictions( class InferenceModelsInstanceSegmentationAdapter(Model): + supports_model_package_selection = True + def __init__(self, model_id: str, api_key: str = None, **kwargs): super().__init__() @@ -1150,6 +1163,8 @@ def rle_masks2poly(masks: InstancesRLEMasks) -> List[np.ndarray]: class InferenceModelsKeyPointsDetectionAdapter(Model): + supports_model_package_selection = True + def __init__(self, model_id: str, api_key: str = None, **kwargs): super().__init__() @@ -1371,6 +1386,8 @@ def model_keypoints_to_response( class InferenceModelsClassificationAdapter(Model): + supports_model_package_selection = True + def __init__(self, model_id: str, api_key: str = None, **kwargs): super().__init__() @@ -1774,6 +1791,8 @@ def draw_predictions(inference_request, inference_response, class_names: List[st class InferenceModelsSemanticSegmentationAdapter(Model): + supports_model_package_selection = True + def __init__(self, model_id: str, api_key: str = None, **kwargs): super().__init__() @@ -2062,6 +2081,8 @@ class InferenceModelsActionRecognitionAdapter(Model): and the trailing remainder is dropped, which is how training validates. """ + supports_model_package_selection = True + def __init__(self, model_id: str, api_key: str = None, **kwargs): super().__init__() self.metrics = {"num_inferences": 0, "avg_inference_time": 0.0} diff --git a/inference/enterprise/parallel/dispatch_manager.py b/inference/enterprise/parallel/dispatch_manager.py index 0e6518094a..8cdcc131d2 100644 --- a/inference/enterprise/parallel/dispatch_manager.py +++ b/inference/enterprise/parallel/dispatch_manager.py @@ -11,6 +11,7 @@ ) from inference.core.entities.responses.inference import response_from_type from inference.core.env import NUM_PARALLEL_TASKS +from inference.core.exceptions import ModelPackageSelectionError from inference.core.managers.base import ModelManager from inference.core.registries.base import ModelRegistry from inference.core.registries.roboflow import get_model_type @@ -164,8 +165,17 @@ def add_model( endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, service_secret: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[str] = None, + quantization: Optional[str] = None, + model_cache_key: Optional[str] = None, ) -> None: - pass + if any( + value is not None for value in (model_package_id, backend, quantization) + ): + raise ModelPackageSelectionError( + "Model package selection is unavailable in parallel mode." + ) def __contains__(self, model_id: str) -> bool: return True diff --git a/inference_models/inference_models/errors.py b/inference_models/inference_models/errors.py index 778d6f2f67..5f4662a209 100644 --- a/inference_models/inference_models/errors.py +++ b/inference_models/inference_models/errors.py @@ -91,6 +91,10 @@ class ModelPackageNegotiationError(BaseInferenceModelsError): pass +class ModelPackagePolicyError(ModelPackageNegotiationError): + """The pinned model package does not satisfy backend or runtime policy.""" + + class UnknownBackendTypeError(ModelPackageNegotiationError): pass diff --git a/inference_models/inference_models/models/auto_loaders/auto_negotiation.py b/inference_models/inference_models/models/auto_loaders/auto_negotiation.py index a67114c4d0..bb68270f9f 100644 --- a/inference_models/inference_models/models/auto_loaders/auto_negotiation.py +++ b/inference_models/inference_models/models/auto_loaders/auto_negotiation.py @@ -11,6 +11,7 @@ AssumptionError, InvalidRequestedBatchSizeError, ModelPackageNegotiationError, + ModelPackagePolicyError, NoModelPackagesAvailableError, UnknownBackendTypeError, UnknownQuantizationError, @@ -74,6 +75,7 @@ def negotiate_model_packages( trt_engine_host_code_allowed: bool = True, nms_fusion_preferences: Optional[Union[bool, dict]] = None, verbose: bool = False, + validate_model_package: bool = False, ) -> List[ModelPackageMetadata]: verbose_info( "The following model packages were exposed by weights provider:", @@ -100,6 +102,30 @@ def negotiate_model_packages( f"source and cannot be loaded while `allow_untrusted_packages=False`.", help_url="https://inference-models.roboflow.com/errors/package-negotiation/#nomodelpackagesavailableerror", ) + if validate_model_package: + eligible_packages = [selected_package] + if requested_backends is not None: + eligible_packages, _ = filter_model_packages_by_requested_backend( + model_packages=eligible_packages, + requested_backends=requested_backends, + verbose=verbose, + ) + if eligible_packages: + eligible_packages, _ = ( + filter_model_packages_matching_runtime_environment( + model_packages=eligible_packages, + device=device, + onnx_execution_providers=onnx_execution_providers, + trt_engine_host_code_allowed=trt_engine_host_code_allowed, + verbose=verbose, + ) + ) + if not eligible_packages: + raise ModelPackagePolicyError( + message=f"Requested model package `{requested_model_package_id}` is not available " + f"under the allowed backends and runtime environment.", + help_url="https://inference-models.roboflow.com/errors/package-negotiation/#nomodelpackagesavailableerror", + ) return [selected_package] model_packages, discarded_packages = remove_packages_not_matching_implementation( model_architecture=model_architecture, @@ -479,7 +505,7 @@ def model_package_matches_batch_size_request( def filter_model_packages_matching_runtime_environment( model_packages: List[ModelPackageMetadata], - device: torch.device, + device: Optional[torch.device], onnx_execution_providers: Optional[List[Union[str, tuple]]], trt_engine_host_code_allowed: bool, verbose: bool = False, diff --git a/inference_models/inference_models/models/auto_loaders/core.py b/inference_models/inference_models/models/auto_loaders/core.py index b8a7010920..82547427a6 100644 --- a/inference_models/inference_models/models/auto_loaders/core.py +++ b/inference_models/inference_models/models/auto_loaders/core.py @@ -1530,6 +1530,7 @@ def from_pretrained( content_addressed_artifact_cache: Optional[ ContentAddressedArtifactCache ] = None, + validate_model_package: bool = False, **kwargs, ) -> AnyModel: """Load and initialize a computer vision model with automatic backend selection. @@ -1562,6 +1563,7 @@ def from_pretrained( requested backend/quantization. Use `AutoModel.describe_model()` to see available packages. + validate_model_package: Check a pinned package against backend and runtime policy. backend: Preferred inference backend(s). Can be: - Single backend: "torch", "onnx", "trt" (TensorRT), "hugging-face" - List of allowed backends: ["trt", "torch"] (the normal @@ -1853,6 +1855,11 @@ def register_file_created_for_model_package( "provider": weights_provider, "model_id": model_id_or_path, "requested_model_package_id": model_package_id, + **( + {"validate_model_package": True} + if validate_model_package + else {} + ), "requested_backends": _canonicalize_unordered_request_values( backend, case_insensitive=True, @@ -2043,6 +2050,7 @@ def attempt_cached_load(cache_hash: str) -> Optional[AnyModel]: task_type=model_metadata.task_type, model_packages=model_metadata.model_packages, requested_model_package_id=model_package_id, + validate_model_package=validate_model_package, requested_backends=backend, requested_batch_size=batch_size, requested_quantization=quantization, diff --git a/inference_models/tests/unit_tests/models/auto_loaders/test_auto_negotiation.py b/inference_models/tests/unit_tests/models/auto_loaders/test_auto_negotiation.py index 4c979b6dc2..807275f712 100644 --- a/inference_models/tests/unit_tests/models/auto_loaders/test_auto_negotiation.py +++ b/inference_models/tests/unit_tests/models/auto_loaders/test_auto_negotiation.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import Any, Tuple from unittest import mock from unittest.mock import MagicMock @@ -10,6 +10,7 @@ AmbiguousModelPackageResolutionError, InvalidRequestedBatchSizeError, ModelPackageNegotiationError, + ModelPackagePolicyError, NoModelPackagesAvailableError, UnknownBackendTypeError, UnknownQuantizationError, @@ -4362,3 +4363,62 @@ def test_filter_model_packages_based_on_model_features_when_package_not_should_b # then assert len(remaining_packages) == 1 assert len(discarded_packages) == 0 + + +@pytest.mark.parametrize("validate_model_package", [False, True]) +def test_pinned_package_respects_backend_policy_only_when_validation_enabled( + validate_model_package: bool, +) -> None: + package = ModelPackageMetadata( + package_id="pinned-trt", + backend=BackendType.TRT, + quantization=Quantization.FP16, + static_batch_size=1, + package_artefacts=[], + trusted_source=True, + ) + kwargs: dict[str, Any] = dict( + model_architecture="rfdetr", + task_type="object-detection", + model_packages=[package], + requested_model_package_id="pinned-trt", + requested_backends=["onnx"], + validate_model_package=validate_model_package, + ) + if validate_model_package: + with pytest.raises(ModelPackagePolicyError): + negotiate_model_packages(**kwargs) + else: + assert negotiate_model_packages(**kwargs) == [package] + + +@mock.patch.object( + auto_negotiation, "filter_model_packages_matching_runtime_environment" +) +@pytest.mark.parametrize("runtime_compatible", [False, True]) +def test_pinned_package_runtime_validation( + runtime_filter, runtime_compatible: bool +) -> None: + package = ModelPackageMetadata( + package_id="pinned-trt", + backend=BackendType.TRT, + quantization=Quantization.FP16, + static_batch_size=1, + package_artefacts=[], + trusted_source=True, + ) + runtime_filter.return_value = ([package] if runtime_compatible else [], []) + kwargs: dict[str, Any] = dict( + model_architecture="rfdetr", + task_type="object-detection", + model_packages=[package], + requested_model_package_id="pinned-trt", + requested_backends=["trt"], + validate_model_package=True, + ) + if runtime_compatible: + assert negotiate_model_packages(**kwargs) == [package] + else: + with pytest.raises(ModelPackagePolicyError): + negotiate_model_packages(**kwargs) + runtime_filter.assert_called_once() diff --git a/inference_models/tests/unit_tests/models/auto_loaders/test_core.py b/inference_models/tests/unit_tests/models/auto_loaders/test_core.py index b7eb9f934d..0acfdf5dfd 100644 --- a/inference_models/tests/unit_tests/models/auto_loaders/test_core.py +++ b/inference_models/tests/unit_tests/models/auto_loaders/test_core.py @@ -4205,3 +4205,52 @@ def test_matching_package_loader_forwards_artifact_cache() -> None: initialize.call_args.kwargs["content_addressed_artifact_cache"] is artifact_cache ) + + +def test_pinned_package_backend_policy_rejects_before_loading() -> None: + from inference_models.errors import ModelPackagePolicyError + + package = ModelPackageMetadata( + package_id="pinned-trt", + backend=BackendType.TRT, + quantization=Quantization.FP16, + static_batch_size=1, + package_artefacts=[], + trusted_source=True, + ) + metadata = ModelMetadata( + model_id="workspace/model/1", + model_architecture="yolov8", + task_type="object-detection", + model_packages=[package], + ) + with mock.patch.object( + core, "get_model_from_provider", return_value=metadata + ), mock.patch.object(core, "attempt_loading_matching_model_packages") as load: + with pytest.raises(ModelPackagePolicyError): + core.AutoModel.from_pretrained( + "workspace/model/1", + model_package_id="pinned-trt", + backend=["onnx"], + validate_model_package=True, + use_auto_resolution_cache=False, + ) + load.assert_not_called() + + +def test_pinned_validation_has_distinct_resolution_cache_identity() -> None: + with mock.patch.object( + core, "attempt_loading_model_with_auto_load_cache", return_value=MagicMock() + ) as cached_load: + for validate in [False, True]: + core.AutoModel.from_pretrained( + "workspace/model/1", + model_package_id="pinned-trt", + backend=["onnx"], + validate_model_package=validate, + auto_resolution_cache=MagicMock(), + ) + hashes = [ + call.kwargs["auto_negotiation_hash"] for call in cached_load.call_args_list + ] + assert hashes[0] != hashes[1] diff --git a/inference_sdk/http/client.py b/inference_sdk/http/client.py index 2a4e2c9413..f254cd1930 100644 --- a/inference_sdk/http/client.py +++ b/inference_sdk/http/client.py @@ -52,6 +52,7 @@ InvalidParameterError, ModelNotInitializedError, ModelNotSelectedError, + ModelSelectionNotConfirmedError, ModelTaskTypeNotSupportedError, RetryError, WrongClientModeError, @@ -1146,7 +1147,7 @@ def __remember_model_selection( registered_models = RegisteredModels.from_dict(response_payload) if selectors: if not registered_models.selected_model_id: - raise InvalidParameterError( + raise ModelSelectionNotConfirmedError( "The server did not return the selected model ID. Upgrade the inference server to use model selection." ) self.__model_selection_ids[(model_id, tuple(sorted(selectors.items())))] = ( @@ -2773,6 +2774,7 @@ def infer_on_video_from_api_v0( ) response = requests.post(url, params=params, data=data, headers=headers) api_key_safe_raise_for_status(response=response) + ensure_model_selection_applied(response.headers, params) return response.json() async def infer_on_video_from_api_v0_async( @@ -2789,6 +2791,7 @@ async def infer_on_video_from_api_v0_async( url, params=params, data=data, headers=headers ) as response: response.raise_for_status() + ensure_model_selection_applied(response.headers, params) return await response.json() def infer_on_video_from_api_v1( @@ -2807,6 +2810,7 @@ def infer_on_video_from_api_v1( url, json=payload, headers=self.__headers_with_auth(DEFAULT_HEADERS) ) api_key_safe_raise_for_status(response=response) + ensure_model_selection_applied(response.headers, payload) return response.json() async def infer_on_video_from_api_v1_async( @@ -2827,6 +2831,7 @@ async def infer_on_video_from_api_v1_async( url, json=payload, headers=self.__headers_with_auth(DEFAULT_HEADERS) ) as response: response.raise_for_status() + ensure_model_selection_applied(response.headers, payload) return await response.json() def __resolve_video_model_id(self, model_id: Optional[str]) -> str: @@ -2846,7 +2851,10 @@ def __build_v0_video_request( raise InvalidModelIdentifier( f"Invalid model id: {model_id}. Expected format: project_id/model_version_id." ) - params = self.__legacy_api_key_payload() + params = { + **self.__legacy_api_key_payload(), + **self.__inference_configuration.to_model_selection_parameters(), + } class_filter = self.__inference_configuration.class_filter if class_filter: params["class_filter"] = ",".join(class_filter) @@ -2872,7 +2880,10 @@ def __build_v1_video_request( f"for one image or infer_on_stream() to classify a video frame by frame." ) video_type, video = _resolve_video_payload(video_reference=video_reference) - payload = self.__initialise_payload() + payload = { + **self.__initialise_payload(), + **self.__inference_configuration.to_model_selection_parameters(), + } payload["model_id"] = model_id payload["video"] = {"type": video_type, "value": video} class_filter = self.__inference_configuration.class_filter diff --git a/inference_sdk/http/errors.py b/inference_sdk/http/errors.py index 3d94153b98..6e3a80c066 100644 --- a/inference_sdk/http/errors.py +++ b/inference_sdk/http/errors.py @@ -54,6 +54,12 @@ def __str__(self) -> str: return self.__repr__() +class ModelSelectionNotConfirmedError(HTTPClientError): + """Raised when the server does not acknowledge model package selection.""" + + pass + + class InvalidInputFormatError(HTTPClientError): """Error for invalid input format.""" diff --git a/inference_sdk/http/utils/model_selection.py b/inference_sdk/http/utils/model_selection.py index 1f0f0dd183..e0631d5915 100644 --- a/inference_sdk/http/utils/model_selection.py +++ b/inference_sdk/http/utils/model_selection.py @@ -1,6 +1,6 @@ from typing import Mapping -from inference_sdk.http.errors import InvalidParameterError +from inference_sdk.http.errors import ModelSelectionNotConfirmedError MODEL_SELECTION_HEADER = "X-Roboflow-Model-Selection" @@ -11,7 +11,8 @@ def ensure_model_selection_applied(headers: Mapping[str, str], selectors: dict) for name in ("model_package_id", "backend", "quantization") ): if headers.get(MODEL_SELECTION_HEADER) != "applied": - raise InvalidParameterError( + raise ModelSelectionNotConfirmedError( "The server did not confirm model package selection. Upgrade the server " - "and enable USE_INFERENCE_MODELS before requesting a specific package." + "and enable USE_INFERENCE_MODELS, or check whether a proxy removed the " + "X-Roboflow-Model-Selection response header." ) diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py index 19b4682751..af3ea97509 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from pydantic import BaseModel @@ -7,6 +7,13 @@ from inference.core.interfaces.http import http_api +@pytest.fixture(autouse=True) +def disable_usage_push(monkeypatch): + monkeypatch.setattr( + http_api.usage_collector, "async_push_usage_payloads", AsyncMock() + ) + + @pytest.fixture def client(monkeypatch): monkeypatch.setattr(http_api, "InferenceInstrumentator", MagicMock()) @@ -108,6 +115,7 @@ class PackageResponse(BaseModel): served_package: str class PackageModel: + supports_model_package_selection = True task_type = getattr(request, "param", "object-detection") batch_size = 1 img_size_h = 32 @@ -194,6 +202,7 @@ def test_http_selected_package_coexists_with_automatic_package( automatic.json()["served_package"] == again.json()["served_package"] == "onnx-1" ) assert selected.json()["served_package"] == "engine-1" + assert selected.headers["X-Model-Id"] == "project/1" registered = package_client.get("/model/registry").json()["models"] assert len(registered) == 2 removal = {"model_id": "project/1", **selectors} @@ -365,3 +374,48 @@ def infer(selection): mismatch = infer({"backend": "trt", "quantization": "fp32"}) assert mismatch.status_code == 400 assert "X-Roboflow-Model-Selection" not in mismatch.headers + + +@pytest.mark.parametrize("legacy", [False, True]) +def test_selected_package_telemetry_uses_public_model_id( + package_client, monkeypatch, legacy +): + from inference.core.managers import base + + spans = MagicMock() + loaded = MagicMock() + inferred = MagicMock() + monkeypatch.setattr(base, "start_span", spans) + monkeypatch.setattr(base, "record_model_loaded", loaded) + monkeypatch.setattr(base, "record_inference", inferred) + selectors = {"backend": "trt", "quantization": "fp16"} + if legacy: + response = package_client.post( + "/project/1", + params={"image": "https://example.com/image.jpg", **selectors}, + ) + else: + response = package_client.post( + "/infer/object_detection", + json={ + "model_id": "project/1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + "disable_model_monitoring": True, + **selectors, + }, + ) + assert response.status_code == 200 + assert response.headers["X-Model-Id"] == "project/1" + assert loaded.call_args.args[0] == "project/1" + assert inferred.call_args.args[0] == "project/1" + model_spans = [ + call.args[1]["model.id"] + for call in spans.call_args_list + if len(call.args) > 1 and "model.id" in call.args[1] + ] + assert model_spans == ["project/1", "project/1"] + selected = package_client.post( + "/model/add", json={"model_id": "project/1", **selectors} + ) + assert ":package:" in selected.json()["selected_model_id"] + assert selected.headers["X-Model-Id"] == "project/1" diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py index cd6cb9c525..070d896186 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py @@ -1,8 +1,9 @@ -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest from starlette.testclient import TestClient +import inference_models.errors as inference_models_errors from inference.core import env from inference.core.interfaces.http import http_api from inference.core.managers import base @@ -16,6 +17,13 @@ ) +@pytest.fixture(autouse=True) +def disable_usage_push(monkeypatch): + monkeypatch.setattr( + http_api.usage_collector, "async_push_usage_payloads", AsyncMock() + ) + + @pytest.fixture def failing_package_client(monkeypatch): monkeypatch.setattr(env, "USE_INFERENCE_MODELS", True) @@ -27,6 +35,7 @@ def failing_package_client(monkeypatch): monkeypatch.setattr(http_api, "LAMBDA", False) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) constructor = MagicMock() + constructor.supports_model_package_selection = True registry = MagicMock() registry.get_model.return_value = constructor manager = base.ModelManager(registry, content_addressed_artifact_cache=MagicMock()) @@ -91,3 +100,56 @@ def test_selector_errors_do_not_mask_server_failures( }, ) assert response.status_code == status + + +@pytest.mark.parametrize( + "selectors", + [{"backend": "onnx"}, {"quantization": "fp32"}, {"model_package_id": "package-1"}], +) +def test_unsupported_model_selection_is_rejected_before_construction( + failing_package_client, selectors +): + client, constructor = failing_package_client + constructor.supports_model_package_selection = False + + response = client.post( + "/model/add", json={"model_id": "clip/ViT-B-16", **selectors} + ) + + assert response.status_code == 400 + assert "does not support explicit package selection" in response.json()["message"] + constructor.assert_not_called() + + +@pytest.mark.parametrize( + "path", ["/model/add", "/infer/object_detection", "/project/1"] +) +def test_pinned_package_policy_rejection_returns_bad_request( + failing_package_client, path +): + error_type = getattr(inference_models_errors, "ModelPackagePolicyError", None) + if error_type is None: + pytest.skip( + "Installed inference-models does not expose pinned policy validation" + ) + client, constructor = failing_package_client + constructor.side_effect = error_type("Unavailable on this runtime") + if path == "/project/1": + response = client.post( + path, + params={ + "model_package_id": "package-1", + "image": "https://example.com/image.jpg", + }, + ) + else: + response = client.post( + path, + json={ + "model_id": "project/1", + "model_package_id": "package-1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + }, + ) + assert response.status_code == 400 + assert "backend policy or runtime" in response.json()["message"] diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py index ea72ecc607..3563cecb16 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py @@ -1,4 +1,4 @@ -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest @@ -6,6 +6,13 @@ from inference.core.interfaces.http import http_api +@pytest.fixture(autouse=True) +def disable_usage_push(monkeypatch): + monkeypatch.setattr( + http_api.usage_collector, "async_push_usage_payloads", AsyncMock() + ) + + @pytest.fixture def openapi_schema(monkeypatch): monkeypatch.setattr(http_api, "InferenceInstrumentator", MagicMock()) diff --git a/tests/inference/unit_tests/core/interfaces/http/test_parallel_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_parallel_model_selection.py new file mode 100644 index 0000000000..8c96b9a74a --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/http/test_parallel_model_selection.py @@ -0,0 +1,66 @@ +import importlib +import sys +from types import ModuleType +from unittest.mock import AsyncMock, MagicMock + +import pytest +from starlette.testclient import TestClient + +from inference.core import env +from inference.core.interfaces.http import http_api + + +@pytest.fixture +def parallel_client(monkeypatch): + tasks = ModuleType("inference.enterprise.parallel.tasks") + setattr(tasks, "preprocess", MagicMock()) + monkeypatch.setitem(sys.modules, tasks.__name__, tasks) + module_name = "inference.enterprise.parallel.dispatch_manager" + monkeypatch.delitem(sys.modules, module_name, raising=False) + module = importlib.import_module(module_name) + monkeypatch.setattr(env, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr(http_api, "InferenceInstrumentator", MagicMock()) + monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) + monkeypatch.setattr(http_api, "LAMBDA", False) + monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) + monkeypatch.setattr( + http_api.usage_collector, "async_push_usage_payloads", AsyncMock() + ) + checker = MagicMock() + manager = module.DispatchModelManager( + MagicMock(), checker, content_addressed_artifact_cache=MagicMock() + ) + try: + with TestClient(http_api.HttpInterface(manager).app) as client: + yield client, checker + finally: + sys.modules.pop(module_name, None) + + +@pytest.mark.parametrize( + "path", ["/model/add", "/infer/object_detection", "/project/1"] +) +@pytest.mark.parametrize( + "selectors", + [{"backend": "onnx"}, {"quantization": "fp32"}, {"model_package_id": "package-1"}], +) +def test_parallel_routes_reject_selection_before_dispatch( + parallel_client, path, selectors +): + client, checker = parallel_client + if path == "/project/1": + response = client.post( + path, params={"image": "https://example.com/image.jpg", **selectors} + ) + else: + response = client.post( + path, + json={ + "model_id": "project/1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + **selectors, + }, + ) + assert response.status_code == 400 + assert "parallel mode" in response.json()["message"] + checker.add_task.assert_not_called() diff --git a/tests/inference/unit_tests/core/managers/test_model_package_cache.py b/tests/inference/unit_tests/core/managers/test_model_package_cache.py index d8c5fbcbef..b9186ddbb9 100644 --- a/tests/inference/unit_tests/core/managers/test_model_package_cache.py +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -24,6 +24,7 @@ class PackageModel(Model): + supports_model_package_selection = True task_type = "object-detection" batch_size = 1 img_size_h = 32 @@ -47,6 +48,23 @@ def clear_cache(self, delete_from_disk=True): pass +def test_clip_selection_rejects_before_loading_packages(monkeypatch): + from inference.models.clip.clip_inference_models import InferenceModelsClipAdapter + from inference_models import AutoModel + + loader = Mock() + monkeypatch.setattr(AutoModel, "from_pretrained", loader) + monkeypatch.setattr(base_module, "USE_INFERENCE_MODELS", True) + registry = Mock() + registry.get_model.return_value = InferenceModelsClipAdapter + manager = ModelManager(registry, content_addressed_artifact_cache=Mock()) + + with pytest.raises(ModelPackageSelectionError, match="does not support"): + manager.add_model("clip/ViT-B-16", "key", backend="onnx") + + loader.assert_not_called() + + @pytest.fixture(autouse=True) def enable_model_metadata(monkeypatch): monkeypatch.setattr(model_module, "USE_INFERENCE_MODELS", True) diff --git a/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py b/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py index 32a5ce7087..f376d1f6ca 100644 --- a/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py +++ b/tests/inference/unit_tests/core/models/test_action_recognition_adapter.py @@ -383,8 +383,15 @@ def test_action_loader_preserves_package_selection_when_wrapping_a_reasoner(sele ) reasoner = SimpleNamespace(resolved_model=resolved) wrapped = _FakeModel(responses=[]) + + def from_pretrained(model_id_or_path, validate_model_package=False, **kwargs): + return reasoner + with patch.object( - adapters.AutoModel, "from_pretrained", return_value=reasoner + adapters.AutoModel, + "from_pretrained", + autospec=from_pretrained, + return_value=reasoner, ) as load, patch.object( adapters, "_as_action_recognition_model", return_value=wrapped ): diff --git a/tests/inference/unit_tests/models/test_pinned_package_policy.py b/tests/inference/unit_tests/models/test_pinned_package_policy.py new file mode 100644 index 0000000000..919cd9b1af --- /dev/null +++ b/tests/inference/unit_tests/models/test_pinned_package_policy.py @@ -0,0 +1,80 @@ +import importlib +from types import SimpleNamespace + +import pytest + +from inference.core.exceptions import ModelPackageSelectionError + + +@pytest.fixture +def adapters(monkeypatch): + module = importlib.import_module("inference.core.models.inference_models_adapters") + monkeypatch.setattr(module, "VALID_INFERENCE_MODELS_BACKENDS", {"onnx", "trt"}) + monkeypatch.setattr(module, "DISABLED_INFERENCE_MODELS_BACKENDS", set()) + return module + + +def test_pinned_package_requires_library_policy_validation(adapters, monkeypatch): + calls = [] + + def from_pretrained(model_id_or_path, validate_model_package=False, **kwargs): + calls.append((model_id_or_path, validate_model_package, kwargs)) + return SimpleNamespace(pre_process=lambda x: x, class_names=[]) + + monkeypatch.setattr( + adapters, "AutoModel", SimpleNamespace(from_pretrained=from_pretrained) + ) + + adapters.InferenceModelsObjectDetectionAdapter( + model_id="workspace/model/1", model_package_id="package-id" + ) + + assert len(calls) == 1 + model_id, validate_model_package, kwargs = calls[0] + assert model_id == "workspace/model/1" + assert validate_model_package is True + assert kwargs["backend"] == ["onnx", "trt"] + assert kwargs["model_package_id"] == "package-id" + + +def test_pinned_package_rejects_library_without_policy_validation( + adapters, monkeypatch +): + calls = [] + + def from_pretrained(model_id_or_path, **kwargs): + calls.append(kwargs) + return SimpleNamespace(pre_process=lambda x: x, class_names=[]) + + monkeypatch.setattr( + adapters, "AutoModel", SimpleNamespace(from_pretrained=from_pretrained) + ) + + with pytest.raises(ModelPackageSelectionError, match="cannot validate pinned"): + adapters.InferenceModelsObjectDetectionAdapter( + model_id="workspace/model/1", model_package_id="package-id" + ) + + assert calls == [] + + +@pytest.mark.parametrize("kwargs", [{}, {"backend": "onnx"}]) +def test_unpinned_selection_preserves_old_library_compatibility( + adapters, monkeypatch, kwargs +): + calls = [] + + def from_pretrained(model_id_or_path, **kwargs): + calls.append(kwargs) + return SimpleNamespace(pre_process=lambda x: x, class_names=[]) + + monkeypatch.setattr( + adapters, "AutoModel", SimpleNamespace(from_pretrained=from_pretrained) + ) + + adapters.InferenceModelsObjectDetectionAdapter( + model_id="workspace/model/1", **kwargs + ) + + assert len(calls) == 1 + assert "validate_model_package" not in calls[0] diff --git a/tests/inference_sdk/unit_tests/http/test_model_selection.py b/tests/inference_sdk/unit_tests/http/test_model_selection.py index fc463b3bef..a42f52a5a2 100644 --- a/tests/inference_sdk/unit_tests/http/test_model_selection.py +++ b/tests/inference_sdk/unit_tests/http/test_model_selection.py @@ -1,3 +1,4 @@ +import re from typing import Any import numpy as np @@ -6,7 +7,10 @@ from inference_sdk import InferenceConfiguration, InferenceHTTPClient from inference_sdk.http.entities import ApiKeyTransport -from inference_sdk.http.errors import InvalidParameterError, ModelNotInitializedError +from inference_sdk.http.errors import ( + ModelNotInitializedError, + ModelSelectionNotConfirmedError, +) def register_response( @@ -48,7 +52,7 @@ def test_v0_selection_requires_server_acknowledgment(requests_mock, acknowledged np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1" ) == {"predictions": []} else: - with pytest.raises(InvalidParameterError, match="did not confirm"): + with pytest.raises(ModelSelectionNotConfirmedError, match="did not confirm"): client.infer(np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1") assert requests_mock.last_request.qs["backend"] == ["trt"] assert requests_mock.last_request.qs["quantization"] == ["fp16"] @@ -57,8 +61,6 @@ def test_v0_selection_requires_server_acknowledgment(requests_mock, acknowledged @pytest.mark.asyncio @pytest.mark.parametrize("acknowledged", [True, False]) async def test_async_v0_selection_requires_server_acknowledgment(acknowledged): - import re - headers = {"X-Roboflow-Model-Selection": "applied"} if acknowledged else {} client = InferenceHTTPClient(api_url="http://server", api_key="key").select_api_v0() client.configure(InferenceConfiguration(model_package_id="engine-1")) @@ -73,7 +75,9 @@ async def test_async_v0_selection_requires_server_acknowledgment(acknowledged): np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1" ) == {"predictions": []} else: - with pytest.raises(InvalidParameterError, match="did not confirm"): + with pytest.raises( + ModelSelectionNotConfirmedError, match="did not confirm" + ): await client.infer_async( np.zeros((2, 2, 3), dtype=np.uint8), model_id="project/1" ) @@ -241,7 +245,7 @@ async def test_selected_load_rejects_missing_server_handle( registry_payload("project/1"), {"X-Roboflow-Model-Selection": "applied"}, ) - with pytest.raises(InvalidParameterError, match="selected model ID"): + with pytest.raises(ModelSelectionNotConfirmedError, match="selected model ID"): await call_client(client, method, asynchronous, "project/1") @@ -331,3 +335,150 @@ async def test_unload_selected_model_uses_confirmed_handle_and_auth( else: assert headers["Authorization"] == "Bearer key" assert payload == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("api_version", [0, 1]) +@pytest.mark.parametrize("acknowledged", [False, True]) +@pytest.mark.parametrize( + "selectors", + [{"model_package_id": "engine-1"}, {"backend": "trt", "quantization": "fp16"}], +) +async def test_video_selection_is_sent_and_requires_acknowledgment( + requests_mock, asynchronous, api_version, acknowledged, selectors +): + client = InferenceHTTPClient(api_url="http://server", api_key="key") + getattr(client, f"select_api_v{api_version}")() + client.configure(InferenceConfiguration(**selectors)) + headers = {"X-Roboflow-Model-Selection": "applied"} if acknowledged else {} + path = "/project/1" if api_version == 0 else "/infer/action_recognition" + with aioresponses() as responses: + if api_version == 1: + register_response( + requests_mock, + responses, + "get", + "/model/registry?api_key=key", + { + "models": [ + {"model_id": "server-handle", "task_type": "action-recognition"} + ] + }, + ) + register_response( + requests_mock, + responses, + "post", + "/model/add", + { + "models": [ + {"model_id": "server-handle", "task_type": "action-recognition"} + ], + "selected_model_id": "server-handle", + }, + {"X-Roboflow-Model-Selection": "applied"}, + ) + requests_mock.post( + f"http://server{path}", json={"timeline": []}, headers=headers + ) + responses.post( + re.compile(r"http://server" + path.split("?")[0] + r".*"), + payload={"timeline": []}, + headers=headers, + ) + if acknowledged: + assert await call_client( + client, + "infer_on_video", + asynchronous, + "https://example.com/clip.mp4", + model_id="project/1", + ) == {"timeline": []} + else: + with pytest.raises(ModelSelectionNotConfirmedError, match="proxy"): + await call_client( + client, + "infer_on_video", + asynchronous, + "https://example.com/clip.mp4", + model_id="project/1", + ) + if asynchronous: + (_, url), calls = list(responses.requests.items())[-1] + sent = calls[0].kwargs["json"] if api_version else dict(url.query) + else: + sent = ( + requests_mock.last_request.json() + if api_version + else { + key: values[0] + for key, values in requests_mock.last_request.qs.items() + } + ) + for name, value in selectors.items(): + assert sent[name] == value + if api_version: + assert sent["model_id"] == "project/1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("acknowledged", [False, True]) +@pytest.mark.parametrize( + "selectors", + [{"model_package_id": "engine-1"}, {"backend": "trt", "quantization": "fp16"}], +) +async def test_v1_inference_sends_selection_in_body_and_requires_acknowledgment( + requests_mock, asynchronous, acknowledged, selectors +): + client = InferenceHTTPClient(api_url="http://server", api_key="key") + client.configure(InferenceConfiguration(**selectors)) + with aioresponses() as responses: + register_response( + requests_mock, + responses, + "get", + "/model/registry?api_key=key", + registry_payload("server-handle"), + ) + register_response( + requests_mock, + responses, + "post", + "/model/add", + registry_payload("server-handle", selected_model_id="server-handle"), + {"X-Roboflow-Model-Selection": "applied"}, + ) + register_response( + requests_mock, + responses, + "post", + "/infer/object_detection", + {"predictions": []}, + {"X-Roboflow-Model-Selection": "applied"} if acknowledged else {}, + ) + image = np.zeros((2, 2, 3), dtype=np.uint8) + if acknowledged: + assert await call_client( + client, "infer", asynchronous, image, model_id="project/1" + ) == {"predictions": []} + else: + with pytest.raises( + ModelSelectionNotConfirmedError, match="did not confirm" + ): + await call_client( + client, "infer", asynchronous, image, model_id="project/1" + ) + if asynchronous: + sent = next( + calls[0].kwargs["json"] + for (_, url), calls in responses.requests.items() + if url.path == "/infer/object_detection" + ) + else: + sent = requests_mock.last_request.json() + assert sent["model_id"] == "project/1" + assert sent["image"]["type"] == "base64" + for name, value in selectors.items(): + assert sent[name] == value From edefb5008bbe4016d5a7177fd1492e5c68055f4b Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 29 Sep 2026 10:38:30 -0400 Subject: [PATCH 20/21] Share loaded models by resolved package identity --- inference/core/interfaces/http/http_api.py | 65 +--- inference/core/managers/active_learning.py | 17 +- inference/core/managers/base.py | 349 ++++++++++++++++-- inference/core/managers/decorators/base.py | 27 +- .../managers/decorators/fixed_size_cache.py | 23 +- inference/core/managers/model_selection.py | 14 +- .../core/models/inference_models_adapters.py | 33 ++ .../enterprise/parallel/dispatch_manager.py | 5 + .../models/auto_loaders/core.py | 118 ++++++ .../test_resolve_model_packages.py | 210 +++++++++++ .../core/interfaces/http/test_http_api.py | 13 +- .../test_legacy_action_recognition_route.py | 8 +- .../interfaces/http/test_model_selection.py | 227 +++++++++++- .../http/test_model_selection_errors.py | 15 +- .../test_workflows_model_manager_hook.py | 18 +- .../test_workflows_shared_package.py | 62 ++++ .../core/managers/test_active_learning.py | 44 +++ .../core/managers/test_model_package_cache.py | 3 + .../managers/test_shared_package_aliases.py | 97 +++++ 19 files changed, 1224 insertions(+), 124 deletions(-) create mode 100644 inference_models/tests/unit_tests/models/auto_loaders/test_resolve_model_packages.py create mode 100644 tests/inference/unit_tests/core/interfaces/test_workflows_shared_package.py create mode 100644 tests/inference/unit_tests/core/managers/test_active_learning.py create mode 100644 tests/inference/unit_tests/core/managers/test_shared_package_aliases.py diff --git a/inference/core/interfaces/http/http_api.py b/inference/core/interfaces/http/http_api.py index 32264a698f..94f5c2d977 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -335,7 +335,6 @@ request_model_ids, request_workflow_id, ) -from inference.core.managers.model_selection import model_selection_cache_key from inference.core.managers.prometheus import InferenceInstrumentator from inference.core.roboflow_api import ( assume_identity_authorised_workspace_db_id, @@ -1613,27 +1612,16 @@ def process_inference_request( status_code=422, detail="Model package selection is not supported for LMM requests.", ) - cache_key = model_selection_cache_key( - de_aliased_model_id, selectors, inference_request.api_key - ) - selection_args = ( - {**selectors, "model_cache_key": cache_key} if selectors else {} - ) - self.model_manager.add_model( + cache_key = self.model_manager.load_model( de_aliased_model_id, inference_request.api_key, model_id_alias=model_id_alias, countinference=countinference, service_secret=service_secret, - **selection_args, - ) - inference_model_id = ( - requested_model_id - if model_id_alias is not None - else de_aliased_model_id + **selectors, ) resp = self.model_manager.infer_from_request_sync( - cache_key if selectors else inference_model_id, + cache_key, inference_request, **kwargs, ) @@ -1982,11 +1970,11 @@ def registry(): summary="Load a model", description=( "Load the model with the given model ID. Optional package selectors " - "require USE_INFERENCE_MODELS=true. Each model ID, selector combination, " - "and credential identifies a separate entry in the existing LRU cache. " + "require USE_INFERENCE_MODELS=true. Requests resolving to the same model " + "package share one entry in the existing LRU cache after authorization. " "Explicit selections return selected_model_id, an opaque registry handle. " "Use the public model ID and selectors for inference and model loading. " - "Requests without selectors keep their automatic selection and cache entry. " + "Requests without selectors retain their loaded automatic selection. " "Registering a preferred package does not replace a loaded model." ), responses=MODEL_SELECTION_RESPONSES, @@ -2015,23 +2003,12 @@ def model_add( ) logger.info(f"Loading model: {de_aliased_model_id}") selectors = model_selection_kwargs(request) - cache_key = model_selection_cache_key( - de_aliased_model_id, selectors, request.api_key - ) - selection_args = ( - { - **selectors, - "model_cache_key": cache_key, - } - if selectors - else {} - ) - self.model_manager.add_model( + cache_key = self.model_manager.load_model( de_aliased_model_id, request.api_key, countinference=countinference, service_secret=service_secret, - **selection_args, + **selectors, ) if selectors: response.headers[MODEL_SELECTION_HEADER] = "applied" @@ -2040,7 +2017,11 @@ def model_add( models_descriptions=models_descriptions ) - descriptions.selected_model_id = cache_key if selectors else None + descriptions.selected_model_id = ( + self.model_manager.get_model_registry_key(cache_key) + if selectors + else None + ) return descriptions @app.post( @@ -2049,8 +2030,8 @@ def model_add( summary="Remove a model", description=( "Remove the model with the given model ID. To remove a selected entry, " - "supply its public model ID with the same selectors and credential used " - "to load it, or supply its registry handle without selectors. " + "supply its public model ID with selectors resolving to that package, " + "or supply its registry handle without selectors. " "Other entries remain loaded. Active inference requests retain their " "model reference until completion." ), @@ -2070,10 +2051,10 @@ def model_remove(request: ClearModelRequest): model_id=request.model_id ) self.model_manager.remove( - model_selection_cache_key( + self.model_manager.resolve_selection_cache_key( de_aliased_model_id, - model_selection_kwargs(request), api_key_override(request.api_key), + model_selection_kwargs(request), ) ) models_descriptions = self.model_manager.describe_models() @@ -4996,15 +4977,11 @@ def legacy_infer_from_request( except ValidationError as error: raise HTTPException(status_code=422, detail=str(error)) from error selectors = model_selection_kwargs(selection) - cache_key = model_selection_cache_key(model_id, selectors, api_key) - selection_args = ( - {**selectors, "model_cache_key": cache_key} if selectors else {} - ) - self.model_manager.add_model( + cache_key = self.model_manager.load_model( request_model_id, api_key, model_id_alias=model_id, - **selection_args, + **selectors, countinference=countinference, service_secret=service_secret, ) @@ -5017,10 +4994,6 @@ def legacy_infer_from_request( # base64 body grows the clip by a third and is held whole # in memory. inference_response = self.model_manager.infer_from_request_sync( - # add_model above registers under the alias, which is - # model_id, so the lookup asks for that. Under Lambda - # request_model_id is the authorizer's endpoint and - # names nothing the manager holds. cache_key, ActionRecognitionInferenceRequest( api_key=api_key, diff --git a/inference/core/managers/active_learning.py b/inference/core/managers/active_learning.py index c4737138a3..b4f45f5c59 100644 --- a/inference/core/managers/active_learning.py +++ b/inference/core/managers/active_learning.py @@ -4,12 +4,15 @@ from fastapi import BackgroundTasks from inference.core import logger +from inference.core.active_learning.configuration import ( + construct_cache_key_for_active_learning_config, +) from inference.core.active_learning.middlewares import ActiveLearningMiddleware from inference.core.cache.base import BaseCache from inference.core.entities.requests.inference import InferenceRequest from inference.core.entities.requests.model_selection import model_selection_kwargs from inference.core.entities.responses.inference import InferenceResponse -from inference.core.env import DISABLE_PREPROC_AUTO_ORIENT +from inference.core.env import DISABLE_PREPROC_AUTO_ORIENT, USE_INFERENCE_MODELS from inference.core.managers.base import ModelManager from inference.core.registries.base import ModelRegistry from inference.models.aliases import resolve_roboflow_model_alias @@ -85,7 +88,7 @@ def register( try: public_model_id = ( (request.model_id or model_id) - if model_selection_kwargs(request) + if USE_INFERENCE_MODELS or model_selection_kwargs(request) else model_id ) resolved_model_id = resolve_roboflow_model_alias(model_id=public_model_id) @@ -96,6 +99,12 @@ def register( or resolved_model_id.split("/")[0] ) middleware_key = f"{model_id}->{target_dataset}" + if USE_INFERENCE_MODELS: + middleware_key = construct_cache_key_for_active_learning_config( + api_key=request.api_key or "", + target_dataset=target_dataset, + model_id=resolved_model_id, + ) self.ensure_middleware_initialised( model_id=resolved_model_id, request=request, @@ -105,7 +114,9 @@ def register( self.register_datapoint( prediction=prediction, model_id=( - model_id if model_selection_kwargs(request) else resolved_model_id + model_id + if USE_INFERENCE_MODELS or model_selection_kwargs(request) + else resolved_model_id ), request=request, middleware_key=middleware_key, diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index d7dd16f22b..af2f946950 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -1,7 +1,8 @@ import time from contextlib import contextmanager -from threading import Lock +from threading import Lock, RLock from typing import TYPE_CHECKING, Any, Dict, Generator, List, Optional, Tuple, Union +from weakref import WeakValueDictionary import numpy as np from fastapi.encoders import jsonable_encoder @@ -38,7 +39,10 @@ model_load_info, request_model_ids, ) -from inference.core.managers.model_selection import public_model_id +from inference.core.managers.model_selection import ( + model_selection_cache_key, + public_model_id, +) from inference.core.managers.pingback import PingbackInfo from inference.core.models.base import Model, PreprocessReturnMetadata from inference.core.registries.base import ModelRegistry @@ -54,6 +58,7 @@ set_span_attribute, start_span, ) +from inference.models.aliases import resolve_roboflow_model_alias from inference.usage_tracking.model_types import bind_usage_model_descriptor from inference_models.errors import ( NoModelPackagesAvailableError, @@ -80,6 +85,8 @@ def __init__( ): self.model_registry = model_registry self._models: Dict[str, Model] = models if models is not None else {} + self._automatic_model_keys: Dict[str, str] = {} + self._selection_locks: WeakValueDictionary[str, Any] = WeakValueDictionary() self._model_request_aliases: Dict[str, set] = {} self._model_request_paths: Dict[str, set] = {} self.pingback = None @@ -130,10 +137,267 @@ def __workflows_bind__( step_error_handler=step_error_handler, ) + def model_selection_lock(self, model_id: str): + """Serialize package loads for one public model ID. + + Args: + model_id: Public model identifier. + + Returns: + Lock shared by concurrent selections. + """ + with self._state_lock: + return self._selection_locks.setdefault(model_id, RLock()) + + def get_model_cache_key(self, model_id: str) -> str: + """Resolve the automatic selection to its physical cache key. + + Args: + model_id: Public model identifier. + + Returns: + Identifier used for the cache or registry lookup. + """ + if model_id in self._models: + return model_id + + canonical_id = resolve_roboflow_model_alias(model_id) + default_key = canonical_id if canonical_id in self._models else model_id + return self._automatic_model_keys.get( + model_id, self._automatic_model_keys.get(canonical_id, default_key) + ) + + def get_model_registry_key(self, cache_key: str) -> str: + """Return the public default ID or selected registry handle. + + Args: + cache_key: Physical key of the loaded package. + + Returns: + Identifier used for the cache or registry lookup. + """ + return next( + ( + alias + for alias, key in list(self._automatic_model_keys.items()) + if key == cache_key + ), + cache_key, + ) + + def set_automatic_model_key(self, model_id: str, cache_key: str) -> None: + """Bind automatic requests to the selected cache entry. + + Args: + model_id: Public model identifier. + cache_key: Physical key of the loaded package. + """ + self._automatic_model_keys[model_id] = cache_key + + def resolve_model_packages( + self, model_id: str, api_key: Optional[str], selectors: dict, **kwargs + ): + """Authorize and resolve eligible packages without loading weights. + + Args: + model_id: Public model identifier. + api_key: Credentials for this resolution. + selectors: Explicit package constraints. + **kwargs: Provider request context. + + Returns: + Ranked package descriptors, or None with an older library. + """ + from inference.core.models.inference_models_adapters import ( + resolve_model_packages, + ) + + return resolve_model_packages(model_id, api_key, selectors, **kwargs) + + def supports_package_selection( + self, model_id: str, api_key: Optional[str], **kwargs + ) -> bool: + """Check whether the model adapter supports package selection. + + Args: + model_id: Public model identifier. + api_key: Credentials for registry lookup. + **kwargs: Provider request context. + + Returns: + Whether the adapter supports package selection. + """ + cached = self._models.get(self.get_model_cache_key(model_id)) + registry: Any = self.model_registry + model_class = ( + type(cached) + if cached is not None + else registry.get_model(model_id, api_key, **kwargs) + ) + return getattr(model_class, "supports_model_package_selection", False) is True + + def load_model( + self, + model_id: str, + api_key: Optional[str], + model_id_alias: Optional[str] = None, + countinference: Optional[bool] = None, + service_secret: Optional[str] = None, + **selectors, + ) -> str: + """Authorize and load one instance per resolved model package. + + Args: + model_id: Model identifier used to initialize the adapter. + api_key: Credentials for resolution and authorization. + model_id_alias: Public identifier recorded for this request. + countinference: Whether the provider must enforce credits. + service_secret: Internal service credential for provider requests. + **selectors: Optional package ID, backend, or quantization constraints. + + Returns: + Cache key for the loaded package. + + Raises: + ModelPackageSelectionError: The server cannot satisfy the selection. + """ + validate_public_model_id(model_id, model_id_alias) + public_id = model_id_alias or model_id + context: Dict[str, Any] = dict( + countinference=countinference, service_secret=service_secret + ) + existing_key = self.get_model_cache_key(model_id) + lookup_id = model_id if existing_key in self.models() else public_id + if not USE_INFERENCE_MODELS or not self.supports_package_selection( + lookup_id, api_key, **context + ): + cache_key = existing_key if existing_key in self.models() else public_id + self.add_model( + model_id, + api_key, + model_id_alias=model_id_alias, + model_cache_key=cache_key, + **context, + **selectors, + ) + return cache_key + public_id = resolve_roboflow_model_alias(public_id) + with self.model_selection_lock(public_id): + existing_key = self.get_model_cache_key(public_id) + existing = self.models().get(existing_key) + requested = dict(selectors) + if not selectors and existing is not None: + metadata = getattr(existing, "resolved_model", None) + if metadata is not None: + requested = {"model_package_id": metadata.model_package_id} + try: + resolved = self.resolve_model_packages( + public_id, api_key, requested, **context + ) + except Exception as error: + _raise_model_selection_error(error, requested) + raise + if resolved is None: + if selectors: + raise ModelPackageSelectionError( + "This inference-models version cannot resolve package selection." + ) + self.add_model( + model_id, + api_key, + model_id_alias=model_id_alias, + model_cache_key=public_id, + **context, + ) + return public_id + candidates = resolved + for candidate in candidates: + matching_key = self.find_model_package_key( + public_id, candidate.model_package_id + ) + cache_key = matching_key or model_selection_cache_key( + public_id, candidate.model_package_id + ) + try: + self.add_model( + model_id, + api_key, + model_id_alias=model_id_alias, + **context, + model_package_id=candidate.model_package_id, + model_cache_key=cache_key, + ) + except inference_models_errors.ModelLoadingError: + if candidate is candidates[-1]: + raise + continue + if not selectors: + self.set_automatic_model_key(public_id, cache_key) + return cache_key + raise ModelPackageSelectionError( + "No model package satisfies the requested selection on this server." + ) + + def find_model_package_key( + self, model_id: str, model_package_id: str + ) -> Optional[str]: + """Find an existing instance of a resolved package. + + Args: + model_id: Public model identifier. + model_package_id: Resolved package identifier. + + Returns: + Existing cache key, or None if the package is absent. + """ + return next( + ( + key + for key, model in list(self.models().items()) + if public_model_id(key) == model_id + and getattr( + getattr(model, "resolved_model", None), "model_package_id", None + ) + == model_package_id + ), + None, + ) + + def resolve_selection_cache_key( + self, model_id: str, api_key: Optional[str], selectors: dict + ) -> str: + """Resolve selectors to the cache entry used for removal. + + Args: + model_id: Public model identifier or registry handle. + api_key: Credentials for package resolution. + selectors: Optional package constraints. + + Returns: + Cache key for the matching package or automatic selection. + """ + if not selectors: + return self.get_model_cache_key(model_id) + validate_public_model_id(model_id) + try: + resolved = self.resolve_model_packages(model_id, api_key, selectors) + except Exception as error: + _raise_model_selection_error(error, selectors) + raise + if not resolved: + raise ModelPackageSelectionError( + "This inference-models version cannot resolve package selection." + ) + for candidate in resolved: + key = self.find_model_package_key(model_id, candidate.model_package_id) + if key is not None: + return key + return model_selection_cache_key(model_id, resolved[0].model_package_id) + def add_model( self, model_id: str, - api_key: str, + api_key: Optional[str], model_id_alias: Optional[str] = None, endpoint_type: ModelEndpointType = ModelEndpointType.ORT, countinference: Optional[bool] = None, @@ -150,6 +414,17 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ + if _should_resolve_model_package(model_cache_key, endpoint_type): + self.load_model( + model_id, + api_key, + model_id_alias=model_id_alias, + countinference=countinference, + service_secret=service_secret, + **model_load_options(model_package_id, backend, quantization, None), + ) + return None + validate_public_model_id(model_id, model_id_alias) selectors = { name: value @@ -194,7 +469,7 @@ def add_model( ) ids_collector = request_model_ids.get(None) if ids_collector is not None: - ids_collector.add(lookup_identifier) + ids_collector.add(public_model_id(resolved_identifier)) model_lock = self._get_lock_for_a_model(model_id=resolved_identifier) with acquire_with_timeout(lock=model_lock) as acquired: if not acquired: @@ -291,25 +566,7 @@ def add_model( except Exception as error: record_error(error) self._dispose_model_lock(model_id=resolved_identifier) - if selectors and isinstance( - error, - getattr(inference_models_errors, "ModelPackagePolicyError", ()), - ): - raise ModelPackageSelectionError( - "The requested model package is incompatible with this server's backend policy or runtime." - ) from error - if selectors and isinstance(error, NoModelPackagesAvailableError): - if model_package_id is not None: - raise ModelPackageNotFoundError() from error - raise ModelPackageSelectionError( - "No model package satisfies the requested selection on this server." - ) from error - if backend is not None and isinstance(error, UnknownBackendTypeError): - raise ModelPackageSelectionError("Unknown backend.") from error - if quantization is not None and isinstance( - error, UnknownQuantizationError - ): - raise ModelPackageSelectionError("Unknown quantization.") from error + _raise_model_selection_error(error, selectors) raise error def validate_model_selection( @@ -820,6 +1077,7 @@ def remove(self, model_id: str, delete_from_disk: bool = True) -> None: model_id (str): The identifier of the model. """ try: + model_id = self.get_model_cache_key(model_id) logger.debug(f"Removing model {model_id} from base model manager") model_lock = self._get_lock_for_a_model(model_id=model_id) with acquire_with_timeout(lock=model_lock) as acquired: @@ -834,6 +1092,11 @@ def remove(self, model_id: str, delete_from_disk: bool = True) -> None: task_type = getattr(model, "task_type", "unknown") model.clear_cache(delete_from_disk=delete_from_disk) del self._models[model_id] + self._automatic_model_keys = { + alias: key + for alias, key in list(self._automatic_model_keys.items()) + if key != model_id + } logger.info( "Model unloaded: model_id=%s, task_type=%s, vram_bytes=%s, remaining_models=%d", model_id, @@ -859,7 +1122,7 @@ def clear(self) -> None: def _get_model_reference(self, model_id: str) -> Model: try: - return self._models[model_id] + return self._models[self.get_model_cache_key(model_id)] except KeyError as error: raise InferenceModelNotFound( f"Model with id {model_id} not loaded." @@ -874,7 +1137,7 @@ def __contains__(self, model_id: str) -> bool: Returns: bool: Whether the model is in the manager. """ - return model_id in self._models + return self.get_model_cache_key(model_id) in self._models def __getitem__(self, key: str) -> Model: """Retrieve a model from the manager by key. @@ -914,7 +1177,7 @@ def models(self) -> Dict[str, Model]: def describe_models(self) -> List[ModelDescription]: return [ ModelDescription( - model_id=model_id, + model_id=self.get_model_registry_key(model_id), task_type=model.task_type, batch_size=getattr(model, "batch_size", None), input_width=getattr(model, "img_size_w", None), @@ -1019,3 +1282,37 @@ def try_releasing_cuda_memory() -> None: pass except Exception as error: logger.warning(f"Attempted to purge CUDA memory but failed with error: {error}") + + +def _raise_model_selection_error(error: Exception, selectors: dict) -> None: + if selectors and isinstance( + error, + getattr(inference_models_errors, "ModelPackagePolicyError", ()), + ): + raise ModelPackageSelectionError( + "The requested model package is incompatible with this server's backend policy or runtime." + ) from error + if selectors and isinstance(error, NoModelPackagesAvailableError): + if selectors.get("model_package_id") is not None: + raise ModelPackageNotFoundError() from error + raise ModelPackageSelectionError( + "No model package satisfies the requested selection on this server." + ) from error + if selectors.get("backend") is not None and isinstance( + error, UnknownBackendTypeError + ): + raise ModelPackageSelectionError("Unknown backend.") from error + if selectors.get("quantization") is not None and isinstance( + error, UnknownQuantizationError + ): + raise ModelPackageSelectionError("Unknown quantization.") from error + + +def _should_resolve_model_package( + model_cache_key: Optional[str], endpoint_type: ModelEndpointType +) -> bool: + return ( + USE_INFERENCE_MODELS + and model_cache_key is None + and endpoint_type == ModelEndpointType.ORT + ) diff --git a/inference/core/managers/decorators/base.py b/inference/core/managers/decorators/base.py index f9d0812ca8..d41ed27be4 100644 --- a/inference/core/managers/decorators/base.py +++ b/inference/core/managers/decorators/base.py @@ -9,6 +9,7 @@ from inference.core.managers.base import ( Model, ModelManager, + _should_resolve_model_package, model_load_options, validate_public_model_id, ) @@ -61,6 +62,28 @@ def pingback(self): def content_addressed_artifact_cache(self): return self.model_manager.content_addressed_artifact_cache + def model_selection_lock(self, model_id): + return self.model_manager.model_selection_lock(model_id) + + def get_model_cache_key(self, model_id): + return self.model_manager.get_model_cache_key(model_id) + + def get_model_registry_key(self, cache_key): + return self.model_manager.get_model_registry_key(cache_key) + + def set_automatic_model_key(self, model_id, cache_key): + return self.model_manager.set_automatic_model_key(model_id, cache_key) + + def resolve_model_packages(self, model_id, api_key, selectors, **kwargs): + return self.model_manager.resolve_model_packages( + model_id, api_key, selectors, **kwargs + ) + + def supports_package_selection(self, model_id, api_key, **kwargs): + return self.model_manager.supports_package_selection( + model_id, api_key, **kwargs + ) + def add_model( self, model_id: str, @@ -83,7 +106,9 @@ def add_model( """ validate_public_model_id(model_id, model_id_alias) cache_key = model_cache_key or model_id - if cache_key in self: + if cache_key in self and not _should_resolve_model_package( + model_cache_key, endpoint_type + ): self.validate_model_selection( cache_key, model_package_id=model_package_id, diff --git a/inference/core/managers/decorators/fixed_size_cache.py b/inference/core/managers/decorators/fixed_size_cache.py index a8abd6335d..bfe8e2097a 100644 --- a/inference/core/managers/decorators/fixed_size_cache.py +++ b/inference/core/managers/decorators/fixed_size_cache.py @@ -21,6 +21,7 @@ from inference.core.managers.base import ( Model, ModelManager, + _should_resolve_model_package, acquire_with_timeout, model_load_options, validate_public_model_id, @@ -78,6 +79,17 @@ def add_model( model (Model): The model instance. endpoint_type (ModelEndpointType, optional): The endpoint type to use for the model. """ + if _should_resolve_model_package(model_cache_key, endpoint_type): + self.load_model( + model_id, + api_key, + model_id_alias=model_id_alias, + countinference=countinference, + service_secret=service_secret, + **model_load_options(model_package_id, backend, quantization, None), + ) + return None + validate_public_model_id(model_id, model_id_alias) if MODELS_CACHE_AUTH_ENABLED and not OFFLINE_MODE: if not _check_if_api_key_has_access_to_model( @@ -91,8 +103,8 @@ def add_model( f"API key {api_key} does not have access to model {model_id}" ) - queue_id = model_cache_key or self._resolve_queue_id( - model_id=model_id, model_id_alias=model_id_alias + queue_id = model_cache_key or self.get_model_cache_key( + self._resolve_queue_id(model_id=model_id, model_id_alias=model_id_alias) ) ids_collector = request_model_ids.get(None) if ids_collector is not None: @@ -141,7 +153,10 @@ def add_model( skipped_pinned = [] while evicted_count < 3 and self._key_queue: to_remove_model_id = self._key_queue.popleft() - if to_remove_model_id in self._pinned_models: + if any( + self.get_model_cache_key(pinned_id) == to_remove_model_id + for pinned_id in self._pinned_models + ): skipped_pinned.append(to_remove_model_id) continue super().remove( @@ -207,6 +222,7 @@ def clear(self) -> None: self.remove(model_id) def remove(self, model_id: str, delete_from_disk: bool = True) -> Model: + model_id = self.get_model_cache_key(model_id) with acquire_with_timeout( lock=self._queue_lock, timeout=HOT_MODELS_QUEUE_LOCK_ACQUIRE_TIMEOUT ) as acquired: @@ -286,6 +302,7 @@ def _resolve_queue_id( return model_id if model_id_alias is None else model_id_alias def _refresh_model_position_in_a_queue(self, model_id: str) -> None: + model_id = self.get_model_cache_key(model_id) with acquire_with_timeout( lock=self._queue_lock, timeout=HOT_MODELS_QUEUE_LOCK_ACQUIRE_TIMEOUT ) as acquired: diff --git a/inference/core/managers/model_selection.py b/inference/core/managers/model_selection.py index ea840b9671..9715dcd2e6 100644 --- a/inference/core/managers/model_selection.py +++ b/inference/core/managers/model_selection.py @@ -1,25 +1,15 @@ import hmac import json import secrets -from typing import Optional from inference.models.aliases import resolve_roboflow_model_alias _MODEL_SELECTION_SECRET = secrets.token_bytes(32) -def model_selection_cache_key( - model_id: str, selectors: dict, api_key: Optional[str] = None -) -> str: - if not selectors: - return model_id +def model_selection_cache_key(model_id: str, model_package_id: str) -> str: model_id = resolve_roboflow_model_alias(model_id) - payload = json.dumps( - {"model_id": model_id, "selectors": selectors, "api_key": api_key}, - sort_keys=True, - separators=(",", ":"), - ).encode() - # codeql[py/weak-sensitive-data-hashing]: HMAC with a random server secret identifies cache entries; this is not password storage. + payload = json.dumps([model_id, model_package_id], separators=(",", ":")).encode() digest = hmac.digest(_MODEL_SELECTION_SECRET, payload, "sha256").hex() return f"{model_id}:package:{digest}" diff --git a/inference/core/models/inference_models_adapters.py b/inference/core/models/inference_models_adapters.py index 78b8c786e5..0135552079 100644 --- a/inference/core/models/inference_models_adapters.py +++ b/inference/core/models/inference_models_adapters.py @@ -237,6 +237,39 @@ def _get_requested_inference_models_backends(kwargs: dict) -> Union[str, List[st return requested_backend +def resolve_model_packages( + model_id: str, api_key: Optional[str], selectors: dict, **kwargs +): + """Resolve packages under this server's backend and trust policy. + + Args: + model_id: Public model identifier or alias. + api_key: Credentials for the current request. + selectors: Optional package ID, backend, or quantization constraints. + **kwargs: Provider context for credit checks and internal requests. + + Returns: + Ranked package descriptors, or None with an older library. + """ + if not callable(getattr(AutoModel, "resolve_model_packages", None)): + return None + selection = dict(selectors) + backend = _get_requested_inference_models_backends( + {"backend": selection.pop("backend", None)} + ) + return AutoModel.resolve_model_packages( + model_id=resolve_roboflow_model_alias(model_id), + api_key=api_key or API_KEY, + backend=backend, + allow_untrusted_packages=ALLOW_INFERENCE_MODELS_UNTRUSTED_PACKAGES, + weights_provider_extra_headers=get_extra_weights_provider_headers( + countinference=kwargs.get("countinference"), + service_secret=kwargs.get("service_secret"), + ), + **selection, + ) + + def _supports_independent_stage_execution(pre_process) -> bool: """Return whether preprocessing declares the composed-execution control.""" try: diff --git a/inference/enterprise/parallel/dispatch_manager.py b/inference/enterprise/parallel/dispatch_manager.py index 8cdcc131d2..178543b415 100644 --- a/inference/enterprise/parallel/dispatch_manager.py +++ b/inference/enterprise/parallel/dispatch_manager.py @@ -157,6 +157,11 @@ def model_infer_sync(self, model_id: str, request: InferenceRequest, **kwargs): return responses return responses[0] + def supports_package_selection( + self, model_id: str, api_key: Optional[str], **kwargs + ) -> bool: + return False + def add_model( self, model_id: str, diff --git a/inference_models/inference_models/models/auto_loaders/core.py b/inference_models/inference_models/models/auto_loaders/core.py index 82547427a6..a6223fc7ad 100644 --- a/inference_models/inference_models/models/auto_loaders/core.py +++ b/inference_models/inference_models/models/auto_loaders/core.py @@ -1489,6 +1489,124 @@ def verify_offline_model( check_hashes=check_hashes, ) + @classmethod + def resolve_model_packages( + cls, + model_id: str, + *, + weights_provider: str = "roboflow", + api_key: Optional[str] = None, + model_package_id: Optional[str] = None, + backend: Optional[ + Union[str, BackendType, List[Union[str, BackendType]]] + ] = None, + batch_size: Optional[Union[int, Tuple[int, int]]] = None, + quantization: Optional[ + Union[str, Quantization, List[Union[str, Quantization]]] + ] = None, + device: Union[torch.device, str] = DEFAULT_DEVICE, + onnx_execution_providers: Optional[List[Union[str, tuple]]] = None, + allow_untrusted_packages: bool = False, + trt_engine_host_code_allowed: bool = True, + nms_fusion_preferences: Optional[Union[bool, dict]] = None, + weights_provider_extra_query_params: Optional[List[Tuple[str, str]]] = None, + weights_provider_extra_headers: Optional[Dict[str, str]] = None, + verbose: bool = False, + ) -> List[ResolvedModelMetadata]: + """Resolve eligible remote packages without loading model weights. + + Each call retrieves metadata with the supplied credentials. Resolution + uses the same package ranking as ``from_pretrained``. Callers can try + candidates in order if initialization fails. Exact package IDs return + one candidate and must satisfy backend and runtime policy. + + Args: + model_id: Model identifier or alias exposed by the weights provider. + weights_provider: Registered metadata provider. + api_key: Credentials for this metadata request. + model_package_id: Exact package identifier to resolve. + backend: Allowed backends. + batch_size: Required batch size or supported range. + quantization: Allowed weight precisions. + device: Target device used for runtime compatibility checks. + onnx_execution_providers: Allowed ONNX execution providers. + allow_untrusted_packages: Whether untrusted packages are eligible. + trt_engine_host_code_allowed: Whether TRT engines may execute host code. + nms_fusion_preferences: Constraints on fused non-maximum suppression. + weights_provider_extra_query_params: Extra provider query parameters. + weights_provider_extra_headers: Extra provider request headers. + verbose: Whether to print package negotiation details. + + Returns: + Eligible package metadata in descending preference order. + + Raises: + UnauthorizedModelAccessError: The provider rejects the credentials. + ForbiddenModelAccessError: The provider denies model access. + NoModelPackagesAvailableError: No package satisfies the constraints. + ModelPackagePolicyError: The pinned package violates runtime policy. + InvalidParameterError: The device cannot be parsed. + """ + _validate_remote_model_id(model_id=model_id) + if OFFLINE_MODE and OFFLINE_MODE_WARM_UP: + raise ModelRetrievalError( + message="OFFLINE_MODE and OFFLINE_MODE_WARM_UP are mutually exclusive.", + help_url="https://inference-models.roboflow.com/errors/model-retrieval/#modelretrievalerror", + ) + if OFFLINE_MODE and weights_provider == "roboflow": + weights_provider = ROBOFLOW_OFFLINE_WEIGHTS_PROVIDER + + api_key = _resolve_effective_api_key(api_key=api_key, provider=weights_provider) + if isinstance(device, str): + try: + device = torch.device(device) + except RuntimeError as error: + raise InvalidParameterError( + message="Could not parse `device`; use a valid torch device such as 'cpu' or 'cuda:0'.", + help_url="https://inference-models.roboflow.com/errors/input-validation/#invalidparametererror", + ) from error + + metadata = get_model_from_provider( + provider=weights_provider, + model_id=model_id, + api_key=api_key, + weights_provider_extra_query_params=weights_provider_extra_query_params, + weights_provider_extra_headers=weights_provider_extra_headers, + ) + if not isinstance(metadata.model_id, str) or not metadata.model_id.strip(): + raise CorruptedModelPackageError( + message=f"Weights provider {weights_provider} returned an empty or invalid canonical model ID.", + help_url="https://inference-models.roboflow.com/errors/model-loading/#corruptedmodelpackageerror", + ) + + packages = negotiate_model_packages( + model_architecture=metadata.model_architecture, + task_type=metadata.task_type, + model_packages=metadata.model_packages, + requested_model_package_id=model_package_id, + validate_model_package=True, + requested_backends=backend, + requested_batch_size=batch_size, + requested_quantization=quantization, + device=device, + onnx_execution_providers=onnx_execution_providers, + allow_untrusted_packages=allow_untrusted_packages, + trt_engine_host_code_allowed=trt_engine_host_code_allowed, + nms_fusion_preferences=nms_fusion_preferences, + verbose=verbose, + ) + resolved_packages = [ + ResolvedModelMetadata( + model_id=metadata.model_id, + model_package_id=package.package_id, + backend=package.backend.value, + quantization=(package.quantization or Quantization.UNKNOWN).value, + ) + for package in packages + ] + + return resolved_packages + @classmethod def from_pretrained( cls, diff --git a/inference_models/tests/unit_tests/models/auto_loaders/test_resolve_model_packages.py b/inference_models/tests/unit_tests/models/auto_loaders/test_resolve_model_packages.py new file mode 100644 index 0000000000..46dbe775d5 --- /dev/null +++ b/inference_models/tests/unit_tests/models/auto_loaders/test_resolve_model_packages.py @@ -0,0 +1,210 @@ +from unittest import mock + +import pytest +from packaging.version import Version + +from inference_models import AutoModel +from inference_models.entities import ResolvedModelMetadata +from inference_models.models.auto_loaders import auto_negotiation, core +from inference_models.models.auto_loaders.entities import BackendType +from inference_models.runtime_introspection.core import RuntimeXRayResult +from inference_models.weights_providers.entities import ( + ModelMetadata, + ModelPackageMetadata, + ONNXPackageDetails, + Quantization, +) + + +@pytest.fixture(autouse=True) +def onnx_cpu_runtime(monkeypatch: pytest.MonkeyPatch) -> None: + runtime = RuntimeXRayResult( + gpu_available=False, + gpu_devices=[], + gpu_devices_cc=[], + driver_version=None, + cuda_version=None, + trt_version=None, + jetson_type=None, + l4t_version=None, + os_version="ubuntu-22.04", + torch_available=True, + torch_version=Version("2.7.0"), + torchvision_version=Version("0.22.0"), + onnxruntime_version=Version("1.21.0"), + available_onnx_execution_providers={"CPUExecutionProvider"}, + hf_transformers_available=False, + trt_python_package_available=False, + ) + monkeypatch.setattr(auto_negotiation, "x_ray_runtime_environment", lambda: runtime) + + +def test_resolve_model_packages_without_loading_weights() -> None: + metadata = ModelMetadata( + model_id="workspace/model/1", + model_architecture="yolov8", + task_type="object-detection", + model_packages=[ + ModelPackageMetadata( + package_id="onnx-fp32", + backend=BackendType.ONNX, + quantization=Quantization.FP32, + package_artefacts=[], + trusted_source=True, + onnx_package_details=ONNXPackageDetails(opset=17), + ) + ], + ) + with mock.patch.object( + core, "get_model_from_provider", return_value=metadata + ) as provider, mock.patch.object(AutoModel, "from_pretrained") as load: + result = AutoModel.resolve_model_packages( + "model-alias", + api_key="current-key", + backend="onnx", + quantization="fp32", + device="cpu", + onnx_execution_providers=["CPUExecutionProvider"], + ) + + assert result == [ + ResolvedModelMetadata( + model_id="workspace/model/1", + model_package_id="onnx-fp32", + backend="onnx", + quantization="fp32", + ) + ] + assert provider.call_args.kwargs["api_key"] == "current-key" + load.assert_not_called() + + +def test_resolve_model_packages_checks_each_callers_access() -> None: + from inference_models.errors import UnauthorizedModelAccessError + + package = ModelPackageMetadata( + package_id="onnx-fp32", + backend=BackendType.ONNX, + quantization=Quantization.FP32, + package_artefacts=[], + trusted_source=True, + onnx_package_details=ONNXPackageDetails(opset=17), + ) + metadata = ModelMetadata( + model_id="workspace/model/1", + model_architecture="yolov8", + task_type="object-detection", + model_packages=[package], + ) + denied = UnauthorizedModelAccessError(message="Access denied") + with mock.patch.object( + core, "get_model_from_provider", side_effect=[metadata, denied] + ) as provider: + AutoModel.resolve_model_packages( + "workspace/model/1", + api_key="authorized-key", + device="cpu", + onnx_execution_providers=["CPUExecutionProvider"], + ) + with pytest.raises(UnauthorizedModelAccessError): + AutoModel.resolve_model_packages( + "workspace/model/1", api_key="unauthorized-key" + ) + + assert [call.kwargs["api_key"] for call in provider.call_args_list] == [ + "authorized-key", + "unauthorized-key", + ] + + +@pytest.mark.parametrize("selectors", [{"backend": "onnx"}, {}]) +def test_resolve_model_packages_rejects_pinned_trt_on_cpu(selectors: dict) -> None: + from inference_models.errors import ModelPackagePolicyError + + metadata = ModelMetadata( + model_id="workspace/model/1", + model_architecture="yolov8", + task_type="object-detection", + model_packages=[ + ModelPackageMetadata( + package_id="trt-fp16", + backend=BackendType.TRT, + quantization=Quantization.FP16, + package_artefacts=[], + trusted_source=True, + ) + ], + ) + with mock.patch.object(core, "get_model_from_provider", return_value=metadata): + with pytest.raises(ModelPackagePolicyError): + AutoModel.resolve_model_packages( + "workspace/model/1", + model_package_id="trt-fp16", + device="cpu", + **selectors, + ) + + +@pytest.mark.parametrize( + "selectors, expected_packages", + [ + ({"backend": "onnx", "quantization": "fp32"}, ["onnx-fp32"]), + ({"quantization": "fp16"}, ["onnx-fp16"]), + ({"model_package_id": "onnx-fp16"}, ["onnx-fp16"]), + ({"quantization": ["fp32", "fp16"]}, ["onnx-fp16", "onnx-fp32"]), + ], +) +def test_resolve_model_packages_selects_matching_package( + selectors: dict, expected_packages: list +) -> None: + metadata = ModelMetadata( + model_id="workspace/model/1", + model_architecture="yolov8", + task_type="object-detection", + model_packages=[ + ModelPackageMetadata( + package_id=f"onnx-{quantization.value}", + backend=BackendType.ONNX, + quantization=quantization, + package_artefacts=[], + trusted_source=True, + onnx_package_details=ONNXPackageDetails(opset=17), + ) + for quantization in [Quantization.FP32, Quantization.FP16] + ], + ) + with mock.patch.object(core, "get_model_from_provider", return_value=metadata): + result = AutoModel.resolve_model_packages( + "workspace/model/1", + device="cpu", + onnx_execution_providers=["CPUExecutionProvider"], + **selectors, + ) + + assert [package.model_package_id for package in result] == expected_packages + + +def test_resolve_model_packages_uses_offline_provider_and_request_headers() -> None: + from inference_models.errors import UnauthorizedModelAccessError + + with mock.patch.object(core, "OFFLINE_MODE", True), mock.patch.object( + core, "OFFLINE_MODE_WARM_UP", False + ), mock.patch.object( + core, + "get_model_from_provider", + side_effect=UnauthorizedModelAccessError(message="Access denied"), + ) as provider: + with pytest.raises(UnauthorizedModelAccessError): + AutoModel.resolve_model_packages( + "workspace/model/1", + api_key="current-key", + weights_provider_extra_headers={"X-Request-Id": "request-1"}, + ) + + assert ( + provider.call_args.kwargs["provider"] == core.ROBOFLOW_OFFLINE_WEIGHTS_PROVIDER + ) + assert provider.call_args.kwargs["api_key"] == "current-key" + assert provider.call_args.kwargs["weights_provider_extra_headers"] == { + "X-Request-Id": "request-1" + } diff --git a/tests/inference/unit_tests/core/interfaces/http/test_http_api.py b/tests/inference/unit_tests/core/interfaces/http/test_http_api.py index 10ef575aa7..cfa9e47545 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_http_api.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_http_api.py @@ -386,6 +386,7 @@ def test_infer_lmm_with_model_id_uses_alias_registry_key(monkeypatch) -> None: ) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) model_manager = MagicMock() + model_manager.load_model.return_value = "florence-2-base" model_manager.pingback = None model_manager.num_errors = 0 model_manager.infer_from_request_sync.return_value = _DummyResponse() @@ -408,7 +409,7 @@ def test_infer_lmm_with_model_id_uses_alias_registry_key(monkeypatch) -> None: assert response.status_code == 200 assert response.json() == {"ok": True} - model_manager.add_model.assert_called_once_with( + model_manager.load_model.assert_called_once_with( "florence-pretrains/3", "query-api-key", model_id_alias="florence-2-base", @@ -1970,7 +1971,7 @@ def test_serverless_auth_middleware_query_key_wins_over_bearer_header( def test_route_level_api_key_reaches_model_manager( monkeypatch, location, expected_api_key ) -> None: - # The api_key handed to model_manager.add_model is the same value the + # The api_key handed to model_manager.load_model is the same value the # usage collector attributes billing to - this matrix is the guard for # the chokepoint fallbacks. interface, model_manager = _build_plain_interface(monkeypatch) @@ -1985,8 +1986,8 @@ def test_route_level_api_key_reaches_model_manager( ) assert response.status_code == 200, f"location={location}" - assert model_manager.add_model.call_count == 1 - assert model_manager.add_model.call_args.args[1] == expected_api_key + assert model_manager.load_model.call_count == 1 + assert model_manager.load_model.call_args.args[1] == expected_api_key def test_route_level_query_key_wins_over_bearer_header(monkeypatch) -> None: @@ -2001,7 +2002,7 @@ def test_route_level_query_key_wins_over_bearer_header(monkeypatch) -> None: ) assert response.status_code == 200 - assert model_manager.add_model.call_args.args[1] == "query-key" + assert model_manager.load_model.call_args.args[1] == "query-key" def test_route_level_bearer_header_wins_over_body_key(monkeypatch) -> None: @@ -2017,7 +2018,7 @@ def test_route_level_bearer_header_wins_over_body_key(monkeypatch) -> None: ) assert response.status_code == 200 - assert model_manager.add_model.call_args.args[1] == "header-key" + assert model_manager.load_model.call_args.args[1] == "header-key" def test_serverless_middleware_bearer_header_wins_over_body_key( diff --git a/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py b/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py index add194c1a2..f8d4932ed3 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_legacy_action_recognition_route.py @@ -69,6 +69,10 @@ def _build_interface(monkeypatch, lambda_mode: bool): # patching LAMBDA after import leaves the name unbound. monkeypatch.setattr(http_api, "trackUsage", MagicMock(), raising=False) model_manager = MagicMock() + model_manager.load_model.side_effect = ( + lambda model_id, api_key, model_id_alias=None, **kwargs: model_id_alias + or model_id + ) model_manager.pingback = None model_manager.num_errors = 0 model_manager.get_task_type.return_value = "action-recognition" @@ -100,7 +104,7 @@ def test_legacy_action_recognition_infers_with_the_registered_identifier( assert response.status_code == 200, response.text # add_model registers under model_id_alias when one is given. - add_model_call = model_manager.add_model.call_args + add_model_call = model_manager.load_model.call_args registered_under = add_model_call.kwargs["model_id_alias"] assert registered_under == PATH_MODEL_ID # The inference call has to name that same identifier. @@ -123,5 +127,5 @@ def test_lambda_request_model_id_really_does_differ(monkeypatch) -> None: }, ) - assert model_manager.add_model.call_args.args[0] == RESOLVED_REQUEST_MODEL_ID + assert model_manager.load_model.call_args.args[0] == RESOLVED_REQUEST_MODEL_ID assert RESOLVED_REQUEST_MODEL_ID != PATH_MODEL_ID diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py index af3ea97509..5e09bd7333 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -114,6 +114,42 @@ def package_client(monkeypatch, request): class PackageResponse(BaseModel): served_package: str + from inference_models import AutoModel + + def resolve_package( + model_id, + api_key=None, + model_package_id=None, + backend=None, + quantization=None, + **kwargs + ): + if api_key == "other-workspace" and model_package_id == "engine-1": + from inference_models.errors import UnauthorizedModelAccessError + + raise UnauthorizedModelAccessError( + "Package is not accessible to this workspace." + ) + selected_backend = ( + "trt" + if model_package_id == "engine-1" + or backend == "trt" + or quantization == "fp16" + else "onnx" + ) + return [ + SimpleNamespace( + model_id=model_id, + model_package_id="engine-1" if selected_backend == "trt" else "onnx-1", + backend=selected_backend, + quantization="fp16" if selected_backend == "trt" else "fp32", + ) + ] + + monkeypatch.setattr( + AutoModel, "resolve_model_packages", resolve_package, raising=False + ) + class PackageModel: supports_model_package_selection = True task_type = getattr(request, "param", "object-detection") @@ -275,23 +311,6 @@ def test_aliased_legacy_selection_shares_v1_entry_and_can_be_removed( assert removed.json()["models"] == [] -def test_selected_model_handle_depends_on_server_secret(package_client, monkeypatch): - from inference.core.managers import model_selection - - payload = {"model_id": "project/1", "backend": "trt", "api_key": "owner"} - first = package_client.post("/model/add", json=payload) - assert first.status_code == 200 - handle = first.json()["selected_model_id"] - assert ( - package_client.post("/model/add", json=payload).json()["selected_model_id"] - == handle - ) - monkeypatch.setattr(model_selection, "_MODEL_SELECTION_SECRET", b"different-server") - second = package_client.post("/model/add", json=payload) - assert second.status_code == 200 - assert second.json()["selected_model_id"] != handle - - def test_browser_can_read_package_selection_acknowledgment(package_client): response = package_client.post( "/model/add", @@ -419,3 +438,177 @@ def test_selected_package_telemetry_uses_public_model_id( ) assert ":package:" in selected.json()["selected_model_id"] assert selected.headers["X-Model-Id"] == "project/1" + + +def test_equivalent_selectors_and_credentials_share_one_package(package_client): + handles = [] + for selectors, key in [ + ({"backend": "trt"}, "owner"), + ({"quantization": "fp16"}, "second-authorized-key"), + ({"model_package_id": "engine-1"}, "owner"), + ]: + response = package_client.post( + "/model/add", json={"model_id": "project/1", "api_key": key, **selectors} + ) + assert response.status_code == 200, response.text + handles.append(response.json()["selected_model_id"]) + assert len(set(handles)) == 1 + assert len(package_client.get("/model/registry").json()["models"]) == 1 + + +@pytest.mark.parametrize("selected_first", [False, True]) +def test_default_and_explicit_requests_share_the_same_package( + package_client, selected_first +): + payload = {"model_id": "project/1", "api_key": "owner"} + requests = [{}, {"backend": "onnx"}] + if selected_first: + requests.reverse() + for selection in requests: + response = package_client.post("/model/add", json={**payload, **selection}) + assert response.status_code == 200, response.text + models = package_client.get("/model/registry").json()["models"] + assert len(models) == 1 + assert models[0]["model_id"] == "project/1" + response = package_client.post("/model/remove", json=payload) + assert response.status_code == 200 + assert response.json()["models"] == [] + + +def test_concurrent_equivalent_requests_load_one_package(package_client): + from concurrent.futures import ThreadPoolExecutor + + selections = [ + {"backend": "trt"}, + {"quantization": "fp16"}, + {"model_package_id": "engine-1"}, + ] + with ThreadPoolExecutor(max_workers=3) as executor: + responses = list( + executor.map( + lambda selection: package_client.post( + "/model/add", json={"model_id": "project/1", **selection} + ), + selections, + ) + ) + assert all(response.status_code == 200 for response in responses) + assert len({response.json()["selected_model_id"] for response in responses}) == 1 + assert len(package_client.get("/model/registry").json()["models"]) == 1 + + +@pytest.mark.parametrize("reuse_cached_fallback", [False, True]) +def test_selection_tries_next_eligible_package_after_load_failure( + package_client, monkeypatch, reuse_cached_fallback +): + from types import SimpleNamespace + + from inference.core.managers.base import ModelManager + from inference_models import AutoModel + from inference_models.errors import ModelPackageAlternativesExhaustedError + + monkeypatch.setattr( + AutoModel, + "resolve_model_packages", + lambda **kwargs: [ + SimpleNamespace(model_package_id="broken-onnx"), + SimpleNamespace(model_package_id="onnx-1"), + ], + ) + if reuse_cached_fallback: + monkeypatch.setattr( + AutoModel, + "resolve_model_packages", + lambda **kwargs: [SimpleNamespace(model_package_id="onnx-1")], + ) + assert ( + package_client.post( + "/model/add", json={"model_id": "project/1", "backend": "onnx"} + ).status_code + == 200 + ) + monkeypatch.setattr( + AutoModel, + "resolve_model_packages", + lambda **kwargs: [ + SimpleNamespace(model_package_id="broken-onnx"), + SimpleNamespace(model_package_id="onnx-1"), + ], + ) + original_add = ModelManager.add_model + attempted_packages = [] + + def add_model(self, *args, **kwargs): + package_id = kwargs.get("model_package_id") + attempted_packages.append(package_id) + if package_id == "broken-onnx": + raise ModelPackageAlternativesExhaustedError("Invalid weights") + return original_add(self, *args, **kwargs) + + monkeypatch.setattr(ModelManager, "add_model", add_model) + selection = {} if reuse_cached_fallback else {"backend": "onnx"} + response = package_client.post( + "/model/add", json={"model_id": "project/1", **selection} + ) + assert response.status_code == 200, response.text + assert attempted_packages == ( + ["broken-onnx"] if reuse_cached_fallback else ["broken-onnx", "onnx-1"] + ) + assert len(response.json()["models"]) == 1 + inferred = package_client.post( + "/infer/object_detection", + json={ + "model_id": "project/1", + **selection, + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + }, + ) + assert inferred.status_code == 200, inferred.text + assert inferred.json()["served_package"] == "onnx-1" + + +def test_eviction_removes_default_alias_and_allows_reload(package_client): + payload = {"model_id": "project/1"} + selected = package_client.post("/model/add", json={**payload, "backend": "onnx"}) + assert selected.status_code == 200 + assert package_client.post("/model/add", json=payload).status_code == 200 + for model_id in ["other/1", "third/1", "fourth/1"]: + assert ( + package_client.post("/model/add", json={"model_id": model_id}).status_code + == 200 + ) + assert all( + model["model_id"] != "project/1" + for model in package_client.get("/model/registry").json()["models"] + ) + reloaded = package_client.post("/model/add", json=payload) + assert reloaded.status_code == 200, reloaded.text + matching = [ + model for model in reloaded.json()["models"] if model["model_id"] == "project/1" + ] + assert len(matching) == 1 + inferred = package_client.post( + "/infer/object_detection", + json={ + **payload, + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + }, + ) + assert inferred.status_code == 200, inferred.text + assert inferred.json()["served_package"] == "onnx-1" + + +def test_older_library_keeps_default_loading_but_rejects_selection( + package_client, monkeypatch +): + from inference_models import AutoModel + + monkeypatch.delattr(AutoModel, "resolve_model_packages", raising=False) + automatic = package_client.post("/model/add", json={"model_id": "project/1"}) + selected = package_client.post( + "/model/add", json={"model_id": "project/1", "backend": "onnx"} + ) + assert automatic.status_code == 200 + assert selected.status_code == 400 + assert "cannot resolve package selection" in selected.json()["message"] + assert len(package_client.get("/model/registry").json()["models"]) == 1 diff --git a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py index 070d896186..d047c2a22a 100644 --- a/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py @@ -1,3 +1,4 @@ +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest @@ -34,6 +35,15 @@ def failing_package_client(monkeypatch): monkeypatch.setattr(http_api, "GCP_SERVERLESS", False) monkeypatch.setattr(http_api, "LAMBDA", False) monkeypatch.setattr(http_api, "DEDICATED_DEPLOYMENT_WORKSPACE_URL", None) + monkeypatch.setattr( + base.ModelManager, + "resolve_model_packages", + lambda self, model_id, api_key, selectors, **kwargs: [ + SimpleNamespace( + model_package_id=selectors.get("model_package_id", "package-1") + ) + ], + ) constructor = MagicMock() constructor.supports_model_package_selection = True registry = MagicMock() @@ -87,9 +97,12 @@ def test_missing_package_and_cache_handle_return_same_public_error( ], ) def test_selector_errors_do_not_mask_server_failures( - failing_package_client, path, selectors, error, status + failing_package_client, path, selectors, error, status, monkeypatch ): client, constructor = failing_package_client + monkeypatch.setattr( + base.ModelManager, "resolve_model_packages", MagicMock(side_effect=error) + ) constructor.side_effect = error response = client.post( path, diff --git a/tests/inference/unit_tests/core/interfaces/test_workflows_model_manager_hook.py b/tests/inference/unit_tests/core/interfaces/test_workflows_model_manager_hook.py index 0f742655c4..5cdbd6311a 100644 --- a/tests/inference/unit_tests/core/interfaces/test_workflows_model_manager_hook.py +++ b/tests/inference/unit_tests/core/interfaces/test_workflows_model_manager_hook.py @@ -456,13 +456,17 @@ def test_engine_run_reaches_raw_manager_through_wired_provider(wrapped): "model_id": MODEL_ID, } ) - # The decorator short-circuits to `record_request_metadata` when the - # model is already present; the raw manager takes `add_model`. Either - # way, the translated request must reach the inner manager. - if wrapped: - inner.record_request_metadata.assert_called_once() - else: - inner.add_model.assert_called_once_with(model_id=MODEL_ID, api_key=API_KEY) + # Package resolution must reach the manager even for a warm model so + # that it can authorize the current caller before sharing the instance. + assert isinstance(inner.add_model, MagicMock) + inner.add_model.assert_called_once() + load_call = inner.add_model.call_args + assert ( + load_call.args[0] if load_call.args else load_call.kwargs["model_id"] + ) == MODEL_ID + assert ( + load_call.args[1] if load_call.args else load_call.kwargs["api_key"] + ) == API_KEY call = inner.infer_from_request_sync.call_args request = call.kwargs.get("request") or call.args[1] assert isinstance(request, ObjectDetectionInferenceRequest) diff --git a/tests/inference/unit_tests/core/interfaces/test_workflows_shared_package.py b/tests/inference/unit_tests/core/interfaces/test_workflows_shared_package.py new file mode 100644 index 0000000000..2867c05b55 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/test_workflows_shared_package.py @@ -0,0 +1,62 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +from inference.core.entities.requests.inference import InferenceRequest +from inference.core.interfaces.workflows_models_provider import ( + ModelManagerModelsProvider, +) +from inference.core.managers import base as base_module +from inference.core.managers.base import ModelManager +from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache +from inference.core.models import base as model_module +from inference.core.models.base import Model +from inference_models import AutoModel + + +def test_workflow_reuses_package_loaded_with_http_selection(monkeypatch): + monkeypatch.setattr(base_module, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr(base_module, "DISABLE_INFERENCE_CACHE", True) + monkeypatch.setattr(model_module, "USE_INFERENCE_MODELS", True) + metadata = SimpleNamespace( + model_id="project/1", + model_package_id="onnx-package", + backend="onnx", + quantization="fp32", + ) + resolve = Mock(return_value=[metadata]) + monkeypatch.setattr(AutoModel, "resolve_model_packages", resolve) + loaded = [] + + class PackageModel(Model): + supports_model_package_selection = True + task_type = "object-detection" + + def __init__(self, **kwargs): + self._model = SimpleNamespace(resolved_model=metadata) + loaded.append(self) + + def infer_from_request(self, request): + return {"instance": id(self), "api_key": request.api_key} + + registry = Mock() + registry.get_model.return_value = PackageModel + manager = WithFixedSizeCache( + ModelManager(registry, content_addressed_artifact_cache=Mock()), max_size=3 + ) + selected_key = manager.load_model("project/1", "http-key", backend="onnx") + provider = ModelManagerModelsProvider(manager) + + assert provider.add_model("project/1", "workflow-key") is None + result = provider.infer_from_request_sync( + model_id="project/1", + request=InferenceRequest( + id="request-id", model_id="project/1", api_key="workflow-key" + ), + ) + + assert len(loaded) == len(manager.keys()) == 1 + assert result == {"instance": id(manager[selected_key]), "api_key": "workflow-key"} + assert [call.kwargs["api_key"] for call in resolve.call_args_list] == [ + "http-key", + "workflow-key", + ] diff --git a/tests/inference/unit_tests/core/managers/test_active_learning.py b/tests/inference/unit_tests/core/managers/test_active_learning.py new file mode 100644 index 0000000000..6008b8a526 --- /dev/null +++ b/tests/inference/unit_tests/core/managers/test_active_learning.py @@ -0,0 +1,44 @@ +from unittest.mock import Mock, patch + +import pytest + +from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest +from inference.core.managers import active_learning as active_learning_module +from inference.core.managers.active_learning import ActiveLearningManager +from inference.core.models.base import Model + + +@pytest.mark.parametrize("backend", [None, "onnx"]) +def test_shared_package_keeps_active_learning_credentials_separate( + monkeypatch, backend +): + monkeypatch.setattr(active_learning_module, "USE_INFERENCE_MODELS", True) + manager = ActiveLearningManager( + model_registry=Mock(), cache=Mock(), content_addressed_artifact_cache=Mock() + ) + model_id = "project/1:package:shared" + manager._models[model_id] = Mock(spec=Model, task_type="object-detection") + first_middleware, second_middleware = Mock(), Mock() + with patch( + "inference.core.managers.active_learning.ActiveLearningMiddleware.init", + side_effect=[first_middleware, second_middleware], + ) as initialize: + for api_key in ("first-key", "second-key", "first-key"): + request = ObjectDetectionInferenceRequest( + model_id="project/1", + backend=backend, + api_key=api_key, + image=[], + id="request-id", + ) + manager.register(prediction=Mock(), model_id=model_id, request=request) + + assert [call.kwargs["api_key"] for call in initialize.call_args_list] == [ + "first-key", + "second-key", + ] + assert all( + call.kwargs["model_id"] == "project/1" for call in initialize.call_args_list + ) + assert first_middleware.register_batch.call_count == 2 + assert second_middleware.register_batch.call_count == 1 diff --git a/tests/inference/unit_tests/core/managers/test_model_package_cache.py b/tests/inference/unit_tests/core/managers/test_model_package_cache.py index b9186ddbb9..b12d71790c 100644 --- a/tests/inference/unit_tests/core/managers/test_model_package_cache.py +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -68,6 +68,9 @@ def test_clip_selection_rejects_before_loading_packages(monkeypatch): @pytest.fixture(autouse=True) def enable_model_metadata(monkeypatch): monkeypatch.setattr(model_module, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr( + ModelManager, "resolve_model_packages", lambda *args, **kwargs: None + ) def test_package_variants_coexist_without_changing_automatic_model(): diff --git a/tests/inference/unit_tests/core/managers/test_shared_package_aliases.py b/tests/inference/unit_tests/core/managers/test_shared_package_aliases.py new file mode 100644 index 0000000000..bc4c9f8f1e --- /dev/null +++ b/tests/inference/unit_tests/core/managers/test_shared_package_aliases.py @@ -0,0 +1,97 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from inference.core.managers import base as base_module +from inference.core.managers.base import ModelManager +from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache +from inference.core.models import base as model_module +from inference.core.models.base import Model +from inference_models import AutoModel + + +@pytest.mark.parametrize("first_selectors", [{}, {"backend": "onnx"}]) +def test_aliases_share_one_package_and_retain_request_attribution( + monkeypatch, first_selectors +): + monkeypatch.setattr(base_module, "USE_INFERENCE_MODELS", True) + monkeypatch.setattr(model_module, "USE_INFERENCE_MODELS", True) + metadata = SimpleNamespace( + model_id="coco/41", + model_package_id="onnx-package", + backend="onnx", + quantization="fp32", + ) + resolution = Mock(return_value=[metadata]) + monkeypatch.setattr(AutoModel, "resolve_model_packages", resolution, raising=False) + loaded = [] + + class PackageModel(Model): + supports_model_package_selection = True + task_type = "object-detection" + + def __init__(self, **kwargs): + self._model = SimpleNamespace(resolved_model=metadata) + loaded.append(self) + + registry = Mock() + registry.get_model.return_value = PackageModel + manager = WithFixedSizeCache( + ModelManager(registry, content_addressed_artifact_cache=Mock()), max_size=3 + ) + first_key = manager.load_model( + "coco/41", "first-key", model_id_alias="yolo26n-640", **first_selectors + ) + second_key = manager.load_model( + "coco/41", "second-key", model_id_alias="yolov26n-640", backend="onnx" + ) + third_key = manager.load_model("coco/41", "third-key") + assert manager.add_model("yolo26n-640", "fourth-key") is None + + assert first_key == second_key == third_key + for alias in ("yolo26n-640", "yolov26n-640", "coco/41"): + assert alias in manager + assert manager[alias] is loaded[0] + assert manager.get_task_type(alias) == "object-detection" + assert len(loaded) == len(manager.keys()) == 1 + assert all( + call.kwargs["model_id"] == "coco/41" for call in resolution.call_args_list + ) + assert {"yolo26n-640", "yolov26n-640"}.issubset( + manager.describe_models()[0].request_aliases + ) + + manager.pin_model("yolo26n-640") + manager.max_size = 1 + manager.add_model("another-project/1", "key") + + assert manager[first_key] is loaded[0] + assert manager["coco/41"] is loaded[0] + assert len(manager.keys()) == 2 + + +def test_alias_lookup_without_package_resolution(monkeypatch): + monkeypatch.setattr(base_module, "USE_INFERENCE_MODELS", True) + monkeypatch.delattr(AutoModel, "resolve_model_packages", raising=False) + loaded = [] + + class LegacyPackageModel(Model): + supports_model_package_selection = True + task_type = "object-detection" + + def __init__(self, **kwargs): + loaded.append(self) + + registry = Mock() + registry.get_model.return_value = LegacyPackageModel + manager = WithFixedSizeCache( + ModelManager(registry, content_addressed_artifact_cache=Mock()), max_size=3 + ) + + manager.add_model("yolo26n-640", "key") + manager.add_model("coco/41", "key") + + assert len(loaded) == len(manager.keys()) == 1 + assert manager["yolo26n-640"] is manager["coco/41"] + assert manager.get_task_type("yolo26n-640") == "object-detection" From 91ee525251498bc8d4bc61107be251ce7a3c1c44 Mon Sep 17 00:00:00 2001 From: Lee Clement Date: Tue, 29 Sep 2026 13:00:03 -0400 Subject: [PATCH 21/21] Fix workflow tests for shared package caching --- ..._workflow_after_continue_if_with_crop_batch_lineage.py | 4 ++-- ...ow_continue_if_inside_inner_with_crop_batch_lineage.py | 8 ++++---- .../execution/test_workflow_with_sahi.py | 6 +++--- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py index af0eab0c1c..880134e347 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py @@ -294,7 +294,7 @@ def test_inlined_continue_if_echo_matches_inner_workflow( ) inner = echo_child_workflow() - with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + with mock.patch.object(model_manager, "add_model"), mock.patch.object( ModelManager, "infer_from_request_sync", new=infer_mock, @@ -413,7 +413,7 @@ def test_inlined_continue_if_echo_matches_inner_workflow_tensor_native( ) inner = echo_child_workflow() - with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + with mock.patch.object(model_manager, "add_model"), mock.patch.object( ModelManager, "get_class_names", new=mock.MagicMock(side_effect=_get_class_names_factory()), diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py index e9960fcfcc..170c49dd41 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py @@ -440,7 +440,7 @@ def test_inlined_continue_if_inside_inner_matches_flat_workflow( ) inner = _inner_continue_if_then_pick() - with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + with mock.patch.object(model_manager, "add_model"), mock.patch.object( ModelManager, "infer_from_request_sync", new=infer_mock, @@ -568,7 +568,7 @@ def test_inlined_continue_if_inside_inner_matches_flat_workflow_tensor_native( ) inner = _inner_continue_if_then_pick() - with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + with mock.patch.object(model_manager, "add_model"), mock.patch.object( ModelManager, "get_class_names", side_effect=_get_class_names_native ), mock.patch.object( ModelManager, @@ -610,7 +610,7 @@ def test_inlined_continue_if_last_in_inner_echo_on_outer_matches_flat_workflow( ) inner = _inner_continue_if_only_outer_echo_name_matches_parent() - with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + with mock.patch.object(model_manager, "add_model"), mock.patch.object( ModelManager, "infer_from_request_sync", new=infer_mock, @@ -666,7 +666,7 @@ def test_inlined_continue_if_last_in_inner_echo_on_outer_matches_flat_workflow_t ) inner = _inner_continue_if_only_outer_echo_name_matches_parent() - with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + with mock.patch.object(model_manager, "add_model"), mock.patch.object( ModelManager, "get_class_names", side_effect=_get_class_names_native ), mock.patch.object( ModelManager, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py b/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py index ff13787e89..9c86bed6e2 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py @@ -667,7 +667,7 @@ def test_sahi_workflow_provides_the_same_result_as_sahi_applied_directly( model_id="yolov8n-640", api_key=None, ) - model = raw_model_manager.models()["yolov8n-640"] + model = raw_model_manager["yolov8n-640"] def slicer_callback(image_slice: np.ndarray): inference_image = {"type": "numpy_object", "value": image_slice} @@ -761,7 +761,7 @@ def test_sahi_workflow_provides_the_same_result_as_sahi_applied_directly_tensor_ model_id="yolov8n-640", api_key=None, ) - model = raw_model_manager.models()["yolov8n-640"] + model = raw_model_manager["yolov8n-640"] def slicer_callback(image_slice: np.ndarray): inference_image = {"type": "numpy_object", "value": image_slice} @@ -873,7 +873,7 @@ def test_sahi_workflow_provides_the_same_result_as_sahi_applied_directly_with_te model_id="yolov8n-640", api_key=None, ) - model = raw_model_manager.models()["yolov8n-640"] + model = raw_model_manager["yolov8n-640"] def slicer_callback(image_slice: np.ndarray): inference_image = {"type": "numpy_object", "value": image_slice}