diff --git a/datamint/api/endpoints/inference_api.py b/datamint/api/endpoints/inference_api.py index 26a215a5..977a16a5 100644 --- a/datamint/api/endpoints/inference_api.py +++ b/datamint/api/endpoints/inference_api.py @@ -10,7 +10,12 @@ from ..entity_base_api import EntityBaseApi, ApiConfig from datamint.entities.inferencejob import InferenceJob -from datamint.exceptions import JobTimeoutError, ItemNotFoundError +from datamint.exceptions import ( + JobTimeoutError, + ItemNotFoundError, + ValidationError, + ModelNotDeployedError, +) from datamint.mlflow.flavors.model_parser import parse_model_reference if TYPE_CHECKING: @@ -81,6 +86,27 @@ def _build_common_payload( payload["params"] = params return payload + def _submit_prediction( + self, + add_path: str, + payload: dict[str, Any], + *, + model_name: str, + model_version: int | None, + model_alias: str | None, + ) -> InferenceJob: + """POST a prediction/submission request and get the resulting job status. """ + try: + response = self._make_request('POST', f'/{self.endpoint_base}{add_path}', json=payload) + except (ValidationError, ItemNotFoundError) as e: + if 'no deployed image found' in str(e).lower(): + raise ModelNotDeployedError( + model_name, model_version=model_version, model_alias=model_alias + ) from e + raise + data = response.json() + return self.get_status(data['job_id']) + # ------------------------------------------------------------------ # Generic inference # ------------------------------------------------------------------ @@ -128,9 +154,9 @@ def submit( if file_paths is not None: payload["file_paths"] = file_paths - response = self._make_request('POST', f'/{self.endpoint_base}', json=payload) - data = response.json() - return self.get_status(data['job_id']) + return self._submit_prediction( + '', payload, model_name=model_name, model_version=model_version, model_alias=model_alias + ) # ------------------------------------------------------------------ # Status / cancel @@ -316,9 +342,9 @@ def predict_image( save_results=save_results, params=params, ) - response = self._make_request('POST', f'/{self.endpoint_base}/predict-image', json=payload) - data = response.json() - return self.get_status(data['job_id']) + return self._submit_prediction( + '/predict-image', payload, model_name=model_name, model_version=model_version, model_alias=model_alias + ) def predict_frame( self, @@ -357,9 +383,9 @@ def predict_frame( params=params, ) payload["frame_index"] = frame_index - response = self._make_request('POST', f'/{self.endpoint_base}/predict-frame', json=payload) - data = response.json() - return self.get_status(data['job_id']) + return self._submit_prediction( + '/predict-frame', payload, model_name=model_name, model_version=model_version, model_alias=model_alias + ) def predict_slice( self, @@ -401,9 +427,9 @@ def predict_slice( ) payload["slice_index"] = slice_index payload["axis"] = axis - response = self._make_request('POST', f'/{self.endpoint_base}/predict-slice', json=payload) - data = response.json() - return self.get_status(data['job_id']) + return self._submit_prediction( + '/predict-slice', payload, model_name=model_name, model_version=model_version, model_alias=model_alias + ) def predict_volume( self, @@ -439,8 +465,8 @@ def predict_volume( save_results=save_results, params=params, ) - response = self._make_request('POST', f'/{self.endpoint_base}/predict-volume', json=payload) - data = response.json() - return self.get_status(data['job_id']) + return self._submit_prediction( + '/predict-volume', payload, model_name=model_name, model_version=model_version, model_alias=model_alias + ) predict = submit # Alias for generic prediction endpoint diff --git a/datamint/exceptions.py b/datamint/exceptions.py index 55b81863..8482d776 100644 --- a/datamint/exceptions.py +++ b/datamint/exceptions.py @@ -118,6 +118,47 @@ def __str__(self) -> str: return super().__str__() +# --------------------------------------------------------------------------- +# Model deployment / inference +# --------------------------------------------------------------------------- + +class ModelNotDeployedError(DatamintException): + """Raised when trying to run inference on a model with no deployed image (HTTP 404).""" + + def __init__( + self, + model_name: str, + model_version: int | None = None, + model_alias: str | None = None, + ): + self.model_name = model_name + self.model_version = model_version + self.model_alias = model_alias + super().__init__(str(self)) + + def __str__(self) -> str: + if self.model_version is not None: + ref = f"{self.model_name}:{self.model_version}" + deploy_kwarg = f"model_version={self.model_version}" + elif self.model_alias is not None: + ref = f"{self.model_name}:{self.model_alias}" + deploy_kwarg = f"model_alias='{self.model_alias}'" + else: + ref = f"{self.model_name}:champion" + deploy_kwarg = None + + deploy_call = f"api.deploy_model.start('{self.model_name}'" + if deploy_kwarg: + deploy_call += f", {deploy_kwarg}" + deploy_call += ")" + + return ( + f"Model '{ref}' is not deployed, so it cannot run inference yet. " + f"Deploy it first with {deploy_call}, wait for the job to finish " + f"(api.deploy_model.wait(job)), then retry." + ) + + # --------------------------------------------------------------------------- # Async job timeouts # --------------------------------------------------------------------------- diff --git a/tests/test_inference_api.py b/tests/test_inference_api.py index 6ad3cdf9..1b15a240 100644 --- a/tests/test_inference_api.py +++ b/tests/test_inference_api.py @@ -3,6 +3,49 @@ from datamint.api.base_api import ApiConfig from datamint.api.endpoints.inference_api import InferenceApi +from datamint.exceptions import ItemNotFoundError, ModelNotDeployedError + + +def test_predict_image_raises_model_not_deployed_error( + api_config: ApiConfig, + make_client, + decoded_path, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + path = decoded_path(request) + if request.method == "POST" and path == "/datamint/api/v1/model-inference/predict-image": + return httpx.Response( + 404, + json={"detail": "No deployed image found for model 'my_model:champion'. Please deploy the model first."}, + ) + raise AssertionError(f"Unexpected request: {request.method} {request.url}") + + with make_client(handler) as client: + inference_api = InferenceApi(api_config, client=client) + + with pytest.raises(ModelNotDeployedError) as excinfo: + inference_api.predict_image("my_model", resource_id="11111111-1111-1111-1111-111111111111") + + err = excinfo.value + assert err.model_name == "my_model" + assert err.model_version is None + assert err.model_alias is None + assert "not deployed" in str(err) + assert "api.deploy_model.start('my_model')" in str(err) + + +def test_predict_image_unrelated_404_is_not_translated( + api_config: ApiConfig, + make_client, +) -> None: + def handler(request: httpx.Request) -> httpx.Response: + return httpx.Response(404, json={"detail": "Model 'my_model' not found"}) + + with make_client(handler) as client: + inference_api = InferenceApi(api_config, client=client) + + with pytest.raises(ItemNotFoundError): + inference_api.predict_image("my_model", resource_id="11111111-1111-1111-1111-111111111111") def test_inference_api_get_status_and_stream_status_job_id_deprecated_alias(