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/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 f195e3adf0..819d3120b4 100644 --- a/inference/core/entities/requests/inference.py +++ b/inference/core/entities/requests/inference.py @@ -1,11 +1,22 @@ 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.entities.common import ApiKey, ModelID, ModelType +from inference.core.entities.requests.model_selection import ( + model_selection_kwargs, + validate_model_selection, +) from inference_sdk.http.entities import Confidence @@ -53,6 +64,23 @@ 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.", + ) + 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 InferenceRequestImage(BaseModel): @@ -134,6 +162,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 new file mode 100644 index 0000000000..59354913b6 --- /dev/null +++ b/inference/core/entities/requests/model_selection.py @@ -0,0 +1,22 @@ +from inference.core import env + + +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 + 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..12387c033b 100644 --- a/inference/core/entities/requests/server_state.py +++ b/inference/core/entities/requests/server_state.py @@ -1,8 +1,9 @@ from typing import Optional -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field, model_validator from inference.core.entities.common import ApiKey, ModelID, ModelType +from inference.core.entities.requests.model_selection import validate_model_selection class AddModelRequest(BaseModel): @@ -18,6 +19,21 @@ class AddModelRequest(BaseModel): model_id: str = ModelID model_type: Optional[str] = ModelType 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) class ClearModelRequest(BaseModel): @@ -29,3 +45,19 @@ class ClearModelRequest(BaseModel): 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/entities/responses/server_state.py b/inference/core/entities/responses/server_state.py index 8571061005..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.", ) diff --git a/inference/core/exceptions.py b/inference/core/exceptions.py index 3692942603..658a4b699c 100644 --- a/inference/core/exceptions.py +++ b/inference/core/exceptions.py @@ -48,6 +48,22 @@ class InferenceModelNotFound(Exception): """ +class ModelPackageSelectionError(Exception): + """Raised when the requested model package selection cannot be satisfied. + + Attributes: + message (str): Optional message describing the error. + """ + + +class ModelPackageNotFoundError(Exception): + """Raised when the requested model package is not found. + + Attributes: + message (str): Optional message describing the error. + """ + + class InvalidMaskDecodeArgument(Exception): """Raised when an invalid argument is provided for mask decoding. diff --git a/inference/core/interfaces/http/error_handlers.py b/inference/core/interfaces/http/error_handlers.py index 5faa2436a7..346a34a743 100644 --- a/inference/core/interfaces/http/error_handlers.py +++ b/inference/core/interfaces/http/error_handlers.py @@ -23,6 +23,8 @@ ModelArtefactError, ModelDeploymentNotSupportedError, ModelManagerLockAcquisitionError, + ModelPackageNotFoundError, + ModelPackageSelectionError, OnnxProviderNotAvailable, PayloadTooLargeError, PaymentRequiredError, @@ -499,6 +501,13 @@ 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: logger.exception("%s: %s", type(error).__name__, error) resp = JSONResponse( @@ -979,6 +988,13 @@ 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: 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 de4fd2e4d4..94f5c2d977 100644 --- a/inference/core/interfaces/http/http_api.py +++ b/inference/core/interfaces/http/http_api.py @@ -41,6 +41,7 @@ MODEL_ID_HEADER, MODEL_LOAD_DETAILS_HEADER, MODEL_LOAD_TIME_HEADER, + MODEL_SELECTION_HEADER, PROCESSING_TIME_HEADER, TRACE_ID_HEADER, WORKFLOW_ID_HEADER, @@ -71,6 +72,7 @@ SemanticSegmentationInferenceRequest, ensure_wire_safe_mask_format, ) +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, @@ -434,6 +436,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( { @@ -760,6 +775,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, @@ -1564,7 +1580,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,13 +1589,15 @@ 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: 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 ) @@ -1588,24 +1606,29 @@ def process_inference_request( if de_aliased_model_id != requested_model_id else None ) - self.model_manager.add_model( + 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 = 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, - ) - 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( - inference_model_id, + cache_key, 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, @@ -1945,11 +1968,21 @@ 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. 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 retain their loaded automatic selection. " + "Registering a preferred package does not replace a loaded model." + ), + responses=MODEL_SELECTION_RESPONSES, ) @with_route_exceptions def model_add( request: AddModelRequest, + response: Response, countinference: Optional[bool] = None, service_secret: Optional[str] = None, ): @@ -1969,22 +2002,39 @@ def model_add( model_id=request.model_id ) logger.info(f"Loading model: {de_aliased_model_id}") - self.model_manager.add_model( + selectors = model_selection_kwargs(request) + cache_key = self.model_manager.load_model( de_aliased_model_id, request.api_key, countinference=countinference, service_secret=service_secret, + **selectors, ) + 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 ) + descriptions.selected_model_id = ( + self.model_manager.get_model_registry_key(cache_key) + if selectors + else None + ) + return descriptions + @app.post( "/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 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." + ), ) @with_route_exceptions def model_remove(request: ClearModelRequest): @@ -2000,7 +2050,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( + self.model_manager.resolve_selection_cache_key( + de_aliased_model_id, + api_key_override(request.api_key), + model_selection_kwargs(request), + ) + ) models_descriptions = self.model_manager.describe_models() return ModelsDescriptions.from_models_descriptions( models_descriptions=models_descriptions @@ -2038,6 +2094,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 @@ -2073,6 +2130,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") @@ -2107,6 +2165,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") @@ -2144,6 +2203,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") @@ -2176,6 +2236,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") @@ -4418,6 +4479,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 " @@ -4457,23 +4519,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" @@ -4647,6 +4702,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, @@ -4663,6 +4719,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, @@ -4696,6 +4753,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, + 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=( @@ -4890,15 +4966,27 @@ def legacy_infer_from_request( logger.debug( f"State of model registry: {self.model_manager.describe_models()}" ) - self.model_manager.add_model( + try: + selection = InferenceRequest( + id=str(uuid4()), + model_id=request_model_id, + 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 = self.model_manager.load_model( request_model_id, api_key, model_id_alias=model_id, + **selectors, 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 task_type == "action-recognition": # The payload is a clip, so none of the image-shaped # arguments below apply to it. The `image` query parameter @@ -4906,11 +4994,7 @@ 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. - model_id, + cache_key, ActionRecognitionInferenceRequest( api_key=api_key, model_id=model_id, @@ -4920,12 +5004,16 @@ 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() + args: Dict[str, Any] = {} if task_type == "instance-segmentation": inference_request_type = InstanceSegmentationInferenceRequest args = { @@ -4964,10 +5052,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, @@ -4977,9 +5066,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/managers/active_learning.py b/inference/core/managers/active_learning.py index ba5c5b24c4..b4f45f5c59 100644 --- a/inference/core/managers/active_learning.py +++ b/inference/core/managers/active_learning.py @@ -4,11 +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 @@ -82,7 +86,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 USE_INFERENCE_MODELS or 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 = ( @@ -90,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, @@ -98,7 +113,11 @@ def register( ) self.register_datapoint( prediction=prediction, - model_id=resolved_model_id, + 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 c51c785c1b..af2f946950 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -1,15 +1,18 @@ 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 +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 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_INFERENCE_CACHE, @@ -25,6 +28,8 @@ from inference.core.exceptions import ( InferenceModelNotFound, ModelManagerLockAcquisitionError, + ModelPackageNotFoundError, + ModelPackageSelectionError, RoboflowAPINotAuthorizedError, ) from inference.core.logger import logger @@ -34,6 +39,10 @@ model_load_info, request_model_ids, ) +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 @@ -49,7 +58,13 @@ 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, + UnknownBackendTypeError, + UnknownQuantizationError, +) if TYPE_CHECKING: from inference_models.utils.content_addressed_artifact_cache import ( @@ -70,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 @@ -120,14 +137,275 @@ 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, 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 +414,37 @@ 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 + 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 +460,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, @@ -159,7 +469,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(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: @@ -168,23 +478,31 @@ 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." ) 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() model_class = self.model_registry.get_model( - resolved_identifier, + lookup_identifier, api_key, 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) if USE_INFERENCE_MODELS: extra_init_kwargs["torchscript_state_global_lock"] = ( self.torchscript_state_global_lock @@ -199,8 +517,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: @@ -230,7 +549,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, @@ -242,13 +561,67 @@ 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) + _raise_model_selection_error(error, selectors) 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 ): @@ -327,14 +700,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}." ) @@ -433,14 +811,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}." ) @@ -518,19 +901,21 @@ 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) + 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) + 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: with start_span( "model.infer", { - "model.id": model_id, + "model.id": public_model_id(model_id), "model.infer.caller": "run_tensor_native_inference", }, ): @@ -538,7 +923,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) @@ -690,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: @@ -704,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, @@ -729,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." @@ -744,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. @@ -784,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), @@ -851,6 +1244,31 @@ 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 ModelPackageNotFoundError() + + +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 @@ -864,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 d17384a400..d41ed27be4 100644 --- a/inference/core/managers/decorators/base.py +++ b/inference/core/managers/decorators/base.py @@ -6,8 +6,15 @@ 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, + _should_resolve_model_package, + model_load_options, + 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 @@ -55,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, @@ -63,6 +92,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 +104,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: + validate_public_model_id(model_id, model_id_alias) + cache_key = model_cache_key or model_id + 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, + 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(public_model_id(cache_key)) return self.model_manager.add_model( model_id, @@ -88,8 +131,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..bfe8e2097a 100644 --- a/inference/core/managers/decorators/fixed_size_cache.py +++ b/inference/core/managers/decorators/fixed_size_cache.py @@ -18,10 +18,18 @@ ModelManagerLockAcquisitionError, RoboflowAPINotAuthorizedError, ) -from inference.core.managers.base import Model, ModelManager, acquire_with_timeout +from inference.core.managers.base import ( + Model, + ModelManager, + _should_resolve_model_package, + 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 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, @@ -59,6 +67,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 +79,18 @@ 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( api_key=api_key, @@ -79,13 +103,19 @@ def add_model( f"API key {api_key} does not have access to model {model_id}" ) - queue_id = 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: - ids_collector.add(queue_id) + ids_collector.add(public_model_id(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." ) @@ -123,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( @@ -165,6 +198,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( @@ -186,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: @@ -265,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/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/model_selection.py b/inference/core/managers/model_selection.py new file mode 100644 index 0000000000..9715dcd2e6 --- /dev/null +++ b/inference/core/managers/model_selection.py @@ -0,0 +1,18 @@ +import hmac +import json +import secrets + +from inference.models.aliases import resolve_roboflow_model_alias + +_MODEL_SELECTION_SECRET = secrets.token_bytes(32) + + +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_package_id], separators=(",", ":")).encode() + 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/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/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 85814074f8..0135552079 100644 --- a/inference/core/models/inference_models_adapters.py +++ b/inference/core/models/inference_models_adapters.py @@ -214,6 +214,62 @@ 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: + 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 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: @@ -257,6 +313,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__() @@ -271,7 +329,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, @@ -423,6 +481,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__() @@ -437,7 +497,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, @@ -1136,6 +1196,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__() @@ -1150,7 +1212,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, @@ -1357,6 +1419,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__() @@ -1370,7 +1434,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, @@ -1760,6 +1824,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__() @@ -1774,7 +1840,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, @@ -2048,6 +2114,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} @@ -2172,10 +2240,14 @@ 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) + 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/inference/enterprise/parallel/dispatch_manager.py b/inference/enterprise/parallel/dispatch_manager.py index 0e6518094a..178543b415 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 @@ -156,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, @@ -164,8 +170,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..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, @@ -1530,6 +1648,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 +1681,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 +1973,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 +2168,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_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/inference_sdk/http/client.py b/inference_sdk/http/client.py index 343289c30a..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, @@ -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, @@ -293,6 +295,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 @@ -908,15 +913,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() - matching_model = filter_model_descriptions( - descriptions=registered_models.models, - model_id=de_aliased_model_id, + 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=de_aliased_model_id, + 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 @@ -944,17 +969,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() - matching_model = filter_model_descriptions( - descriptions=registered_models.models, - model_id=de_aliased_model_id, + 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=de_aliased_model_id, + 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 @@ -1028,19 +1073,28 @@ 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, + **selectors, **self.__legacy_api_key_payload(), }, headers=self.__headers_with_auth(DEFAULT_HEADERS), ) response.raise_for_status() + ensure_model_selection_applied( + response.headers, + 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( @@ -1062,8 +1116,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, + **selectors, **self.__legacy_api_key_payload(), } async with aiohttp.ClientSession() as session: @@ -1073,15 +1129,39 @@ async def load_model_async( headers=self.__headers_with_auth(DEFAULT_HEADERS), ) as response: response.raise_for_status() + ensure_model_selection_applied( + response.headers, + 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 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())))] = ( + registered_models.selected_model_id + ) + return registered_models @wrap_errors 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. @@ -1090,20 +1170,29 @@ 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, + "model_id": removal_model_id, + **(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 @@ -1115,16 +1204,24 @@ 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) + 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, + "model_id": removal_model_id, + **(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 @@ -1132,6 +1229,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() @@ -1142,6 +1252,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 @@ -1155,6 +1266,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 @@ -2662,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( @@ -2678,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( @@ -2696,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( @@ -2716,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: @@ -2735,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) @@ -2761,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/entities.py b/inference_sdk/http/entities.py index 0e82e17783..3492d9590a 100644 --- a/inference_sdk/http/entities.py +++ b/inference_sdk/http/entities.py @@ -75,6 +75,7 @@ class RegisteredModels(DataClassJsonMixin): """ models: List[ModelDescription] + selected_model_id: Optional[str] = None class HTTPClientMode(str, Enum): @@ -139,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 @@ -183,8 +192,17 @@ 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 ( + 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 +219,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 +300,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"), @@ -333,6 +361,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 +392,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"), 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/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..e0631d5915 --- /dev/null +++ b/inference_sdk/http/utils/model_selection.py @@ -0,0 +1,18 @@ +from typing import Mapping + +from inference_sdk.http.errors import ModelSelectionNotConfirmedError + +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 ModelSelectionNotConfirmedError( + "The server did not confirm model package selection. Upgrade the server " + "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_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 f0de06740d..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" @@ -78,9 +82,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,17 +98,20 @@ def test_legacy_action_recognition_infers_with_the_registered_identifier( params={ "api_key": "query-api-key", "image": "https://example.com/clip.mp4", + **selectors, }, ) 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. 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: @@ -116,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 new file mode 100644 index 0000000000..5e09bd7333 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection.py @@ -0,0 +1,614 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest +from pydantic import BaseModel +from starlette.testclient import TestClient + +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()) + 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) + 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"}]) +@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, + }, + ) + assert response.status_code == 422 + assert "model_package_id" in response.text + + +@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 + + 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"}, + "video": {"type": "url", "value": "https://example.com/clip.mp4"}, + }, + ) + assert response.status_code == 422 + 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, request): + 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 + + 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") + 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): + 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): + 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, "ACTION_RECOGNITION_ENABLED", True) + 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 + ) + 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"}], +) +@pytest.mark.parametrize("remove_by_handle", [False, True]) +def test_http_selected_package_coexists_with_automatic_package( + package_client, selectors, remove_by_handle +): + 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" + 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} + 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"] + + +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_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_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", + "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 + + +@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 + + +@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" + + +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 new file mode 100644 index 0000000000..d047c2a22a --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_errors.py @@ -0,0 +1,168 @@ +from types import SimpleNamespace +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 +from inference_models.errors import ( + ModelLoadingError, + ModelPackageNegotiationError, + NoModelPackagesAvailableError, + RuntimeIntrospectionError, + UnknownBackendTypeError, + UnknownQuantizationError, +) + + +@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) + 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) + 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() + 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, 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, + json={ + "model_id": "project/1", + "image": {"type": "url", "value": "https://example.com/image.jpg"}, + **selectors, + }, + ) + 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 new file mode 100644 index 0000000000..3563cecb16 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/http/test_model_selection_openapi.py @@ -0,0 +1,102 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from inference.core.constants import MODEL_SELECTION_HEADER +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()) + 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() + + +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"): + 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", + ): + 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( + "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", + [ + ("/model/add", "post"), + ("/infer/object_detection", "post"), + ("/infer/instance_segmentation", "post"), + ("/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"), + ], +) +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"] 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/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_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: 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..b12d71790c --- /dev/null +++ b/tests/inference/unit_tests/core/managers/test_model_package_cache.py @@ -0,0 +1,232 @@ +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 ( + 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 +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 + + +class PackageModel(Model): + supports_model_package_selection = True + 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._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: + return {"backend": self.resolved_model.backend} + + 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) + monkeypatch.setattr( + ModelManager, "resolve_model_packages", lambda *args, **kwargs: None + ) + + +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", 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 + ) + + +@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", InferenceRequest(id="request-id", model_id="project/1") + ) == {"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(ModelPackageNotFoundError): + 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( + 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 = ObjectDetectionInferenceRequest( + 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) -> Any: + 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", + InferenceRequest(id="request-id", model_id="project/1"), + ) + 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", InferenceRequest(id="request-id", model_id="project/1") + ) == {"backend": "trt"} + finally: + finish.set() + assert active.result(timeout=10) == {"backend": "onnx"} 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" 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..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 @@ -363,3 +363,40 @@ 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=[]) + + def from_pretrained(model_id_or_path, validate_model_package=False, **kwargs): + return reasoner + + with patch.object( + adapters.AutoModel, + "from_pretrained", + autospec=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 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_entities.py b/tests/inference_sdk/unit_tests/http/test_entities.py index 01f1e99186..9a73502dc7 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,29 @@ 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) + + +def test_positional_configuration_preserves_confidence_threshold() -> None: + configuration = InferenceConfiguration(0.5) + assert configuration.confidence_threshold == 0.5 + assert configuration.to_model_selection_parameters() == {} 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..a42f52a5a2 --- /dev/null +++ b/tests/inference_sdk/unit_tests/http/test_model_selection.py @@ -0,0 +1,484 @@ +import re +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 ApiKeyTransport +from inference_sdk.http.errors import ( + ModelNotInitializedError, + ModelSelectionNotConfirmedError, +) + + +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 {} + 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(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"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("acknowledged", [True, False]) +async def test_async_v0_selection_requires_server_acknowledgment(acknowledged): + 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( + ModelSelectionNotConfirmedError, 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 = "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"}], + "selected_model_id": cache_key, + }, + 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 = "another-server-handle" + 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"}], + "selected_model_id": cache_key, + }, + 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, + } + + +@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(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 +@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(ModelSelectionNotConfirmedError, 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_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( + 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": "server-handle"} + if transport is ApiKeyTransport.LEGACY: + expected["api_key"] = "key" + 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 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}