Spaces:
Sleeping
Sleeping
import time | |
from typing import Dict, Optional | |
from fastapi import BackgroundTasks | |
from inference.core import logger | |
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.responses.inference import InferenceResponse | |
from inference.core.env import DISABLE_PREPROC_AUTO_ORIENT | |
from inference.core.managers.base import ModelManager | |
from inference.core.registries.base import ModelRegistry | |
ACTIVE_LEARNING_ELIGIBLE_PARAM = "active_learning_eligible" | |
DISABLE_ACTIVE_LEARNING_PARAM = "disable_active_learning" | |
BACKGROUND_TASKS_PARAM = "background_tasks" | |
class ActiveLearningManager(ModelManager): | |
def __init__( | |
self, | |
model_registry: ModelRegistry, | |
cache: BaseCache, | |
middlewares: Optional[Dict[str, ActiveLearningMiddleware]] = None, | |
): | |
super().__init__(model_registry=model_registry) | |
self._cache = cache | |
self._middlewares = middlewares if middlewares is not None else {} | |
async def infer_from_request( | |
self, model_id: str, request: InferenceRequest, **kwargs | |
) -> InferenceResponse: | |
prediction = await super().infer_from_request( | |
model_id=model_id, request=request, **kwargs | |
) | |
active_learning_eligible = kwargs.get(ACTIVE_LEARNING_ELIGIBLE_PARAM, False) | |
active_learning_disabled_for_request = getattr( | |
request, DISABLE_ACTIVE_LEARNING_PARAM, False | |
) | |
if ( | |
not active_learning_eligible | |
or active_learning_disabled_for_request | |
or request.api_key is None | |
): | |
return prediction | |
self.register(prediction=prediction, model_id=model_id, request=request) | |
return prediction | |
def register( | |
self, prediction: InferenceResponse, model_id: str, request: InferenceRequest | |
) -> None: | |
try: | |
self.ensure_middleware_initialised(model_id=model_id, request=request) | |
self.register_datapoint( | |
prediction=prediction, | |
model_id=model_id, | |
request=request, | |
) | |
except Exception as error: | |
# Error handling to be decided | |
logger.warning( | |
f"Error in datapoint registration for Active Learning. Details: {error}. " | |
f"Error is suppressed in favour of normal operations of API." | |
) | |
def ensure_middleware_initialised( | |
self, model_id: str, request: InferenceRequest | |
) -> None: | |
if model_id in self._middlewares: | |
return None | |
start = time.perf_counter() | |
logger.debug(f"Initialising AL middleware for {model_id}") | |
self._middlewares[model_id] = ActiveLearningMiddleware.init( | |
api_key=request.api_key, | |
model_id=model_id, | |
cache=self._cache, | |
) | |
end = time.perf_counter() | |
logger.debug(f"Middleware init latency: {(end - start) * 1000} ms") | |
def register_datapoint( | |
self, prediction: InferenceResponse, model_id: str, request: InferenceRequest | |
) -> None: | |
start = time.perf_counter() | |
inference_inputs = getattr(request, "image", None) | |
if inference_inputs is None: | |
logger.warning( | |
"Could not register datapoint, as inference input has no `image` field." | |
) | |
return None | |
if not issubclass(type(inference_inputs), list): | |
inference_inputs = [inference_inputs] | |
if not issubclass(type(prediction), list): | |
results_dicts = [prediction.dict(by_alias=True, exclude={"visualization"})] | |
else: | |
results_dicts = [ | |
e.dict(by_alias=True, exclude={"visualization"}) for e in prediction | |
] | |
prediction_type = self.get_task_type(model_id=model_id) | |
disable_preproc_auto_orient = ( | |
getattr(request, "disable_preproc_auto_orient", False) | |
or DISABLE_PREPROC_AUTO_ORIENT | |
) | |
self._middlewares[model_id].register_batch( | |
inference_inputs=inference_inputs, | |
predictions=results_dicts, | |
prediction_type=prediction_type, | |
disable_preproc_auto_orient=disable_preproc_auto_orient, | |
) | |
end = time.perf_counter() | |
logger.debug(f"Registration: {(end - start) * 1000} ms") | |
class BackgroundTaskActiveLearningManager(ActiveLearningManager): | |
async def infer_from_request( | |
self, model_id: str, request: InferenceRequest, **kwargs | |
) -> InferenceResponse: | |
active_learning_eligible = kwargs.get(ACTIVE_LEARNING_ELIGIBLE_PARAM, False) | |
active_learning_disabled_for_request = getattr( | |
request, DISABLE_ACTIVE_LEARNING_PARAM, False | |
) | |
kwargs[ACTIVE_LEARNING_ELIGIBLE_PARAM] = False # disabling AL in super-classes | |
prediction = await super().infer_from_request( | |
model_id=model_id, request=request, **kwargs | |
) | |
if ( | |
not active_learning_eligible | |
or active_learning_disabled_for_request | |
or request.api_key is None | |
): | |
return prediction | |
if BACKGROUND_TASKS_PARAM not in kwargs: | |
logger.warning( | |
"BackgroundTaskActiveLearningManager used against rules - `background_tasks` argument not " | |
"provided making Active Learning registration running sequentially." | |
) | |
self.register(prediction=prediction, model_id=model_id, request=request) | |
else: | |
background_tasks: BackgroundTasks = kwargs["background_tasks"] | |
background_tasks.add_task( | |
self.register, prediction=prediction, model_id=model_id, request=request | |
) | |
return prediction | |