Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,407 Bytes
df6c67d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
from collections import deque
from typing import List, Optional
from inference.core.entities.requests.inference import InferenceRequest
from inference.core.entities.responses.inference import InferenceResponse
from inference.core.managers.base import Model, ModelManager
from inference.core.managers.decorators.base import ModelManagerDecorator
from inference.core.managers.entities import ModelDescription
class WithFixedSizeCache(ModelManagerDecorator):
def __init__(self, model_manager: ModelManager, max_size: int = 8):
"""Cache decorator, models will be evicted based on the last utilization (`.infer` call). Internally, a [double-ended queue](https://docs.python.org/3/library/collections.html#collections.deque) is used to keep track of model utilization.
Args:
model_manager (ModelManager): Instance of a ModelManager.
max_size (int, optional): Max number of models at the same time. Defaults to 8.
"""
super().__init__(model_manager)
self.max_size = max_size
self._key_queue = deque(self.model_manager.keys())
def add_model(
self, model_id: str, api_key: str, model_id_alias: Optional[str] = None
):
"""Adds a model to the manager and evicts the least recently used if the cache is full.
Args:
model_id (str): The identifier of the model.
model (Model): The model instance.
"""
queue_id = self._resolve_queue_id(
model_id=model_id, model_id_alias=model_id_alias
)
if model_id in self:
self._key_queue.remove(queue_id)
self._key_queue.append(queue_id)
return
should_pop = len(self) == self.max_size
if should_pop:
to_remove_model_id = self._key_queue.popleft()
self.remove(to_remove_model_id)
self._key_queue.append(queue_id)
try:
return super().add_model(model_id, api_key, model_id_alias=model_id_alias)
except Exception as error:
self._key_queue.remove(model_id)
raise error
def clear(self) -> None:
"""Removes all models from the manager."""
for model_id in list(self.keys()):
self.remove(model_id)
def remove(self, model_id: str) -> Model:
try:
self._key_queue.remove(model_id)
except ValueError:
pass
return super().remove(model_id)
async def infer_from_request(
self, model_id: str, request: InferenceRequest, **kwargs
) -> InferenceResponse:
"""Processes a complete inference request and updates the cache.
Args:
model_id (str): The identifier of the model.
request (InferenceRequest): The request to process.
Returns:
InferenceResponse: The response from the inference.
"""
self._key_queue.remove(model_id)
self._key_queue.append(model_id)
return await super().infer_from_request(model_id, request, **kwargs)
def infer_only(self, model_id: str, request, img_in, img_dims, batch_size=None):
"""Performs only the inference part of a request and updates the cache.
Args:
model_id (str): The identifier of the model.
request: The request to process.
img_in: Input image.
img_dims: Image dimensions.
batch_size (int, optional): Batch size.
Returns:
Response from the inference-only operation.
"""
self._key_queue.remove(model_id)
self._key_queue.append(model_id)
return super().infer_only(model_id, request, img_in, img_dims, batch_size)
def preprocess(self, model_id: str, request):
"""Processes the preprocessing part of a request and updates the cache.
Args:
model_id (str): The identifier of the model.
request (InferenceRequest): The request to preprocess.
"""
self._key_queue.remove(model_id)
self._key_queue.append(model_id)
return super().preprocess(model_id, request)
def describe_models(self) -> List[ModelDescription]:
return self.model_manager.describe_models()
def _resolve_queue_id(
self, model_id: str, model_id_alias: Optional[str] = None
) -> str:
return model_id if model_id_alias is None else model_id_alias
|