From 5794dd21f48c5f25c30b83e7a99eece51de73cfd Mon Sep 17 00:00:00 2001 From: Andrews Arokiam Date: Wed, 29 Nov 2023 13:26:03 +0530 Subject: [PATCH] Added missing `response_headers` variable in predict call. Signed-off-by: Andrews Arokiam --- python/kserve/test/test_server.py | 3 ++- python/lgbserver/lgbserver/test_model.py | 24 +++++++++---------- .../graph/success_200_isvc/model.py | 3 ++- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/python/kserve/test/test_server.py b/python/kserve/test/test_server.py index 65e4bc81eb4..97e12905337 100644 --- a/python/kserve/test/test_server.py +++ b/python/kserve/test/test_server.py @@ -333,6 +333,7 @@ def test_infer_parameters_v2(self, http_server_client): input_data = json.dumps(req.to_rest()).encode('utf-8') with patch.object(DummyModel, 'predict', new_callable=mock.Mock) as mock_predict: + response_headers = {} mock_predict.return_value = InferResponse(model_name="TestModel", response_id="123", parameters={ "test-str": "dummy", @@ -347,7 +348,7 @@ def test_infer_parameters_v2(self, http_server_client): "test-str": "dummy", "test-bool": True, "test-int": 100 - })]) + })]), response_headers resp = http_server_client.post('/v2/models/TestModel/infer', content=input_data) mock_predict.assert_called_with(req, mock.ANY) diff --git a/python/lgbserver/lgbserver/test_model.py b/python/lgbserver/lgbserver/test_model.py index 6e1c6b98fad..0368a043ef1 100644 --- a/python/lgbserver/lgbserver/test_model.py +++ b/python/lgbserver/lgbserver/test_model.py @@ -47,39 +47,37 @@ def test_model(): 'petal_width_(cm)': {0: 0.2}, 'sepal_length_(cm)': {0: 5.1}} response, response_headers = model.predict({"inputs": [request, request]}) - assert numpy.argmax(response["predictions"][0]) == 2 + assert numpy.argmax(response["predictions"][0]) == 0 - response, response_headers = model.predict( - {"instances": [request, request]}) - assert numpy.argmax(response["predictions"][0]) == 2 + response, response_headers = model.predict({"instances": [request, request]}) + assert numpy.argmax(response["predictions"][0]) == 0 request = [ {'sepal_width_(cm)': 3.5}, {'petal_length_(cm)': 1.4}, {'petal_width_(cm)': 0.2}, {'sepal_length_(cm)': 5.1} ] - response = model.predict({"inputs": [request, request]}) + response, response_headers = model.predict({"inputs": [request, request]}) assert numpy.argmax(response["predictions"][0]) == 0 - response = model.predict({"instances": [request, request]}) + response, response_headers = model.predict({"instances": [request, request]}) assert numpy.argmax(response["predictions"][0]) == 0 request = [ {'sepal_width_(cm)': 3.5}, {'petal_length_(cm)': 1.4}, {'petal_width_(cm)': 0.2} ] - response = model.predict({"inputs": [request, request]}) + response, response_headers = model.predict({"inputs": [request, request]}) assert numpy.argmax(response["predictions"][0]) == 0 - response = model.predict({"instances": [request, request]}) + response, response_headers = model.predict({"instances": [request, request]}) assert numpy.argmax(response["predictions"][0]) == 0 # test v2 handler infer_input = InferInput(name="input-0", shape=[2, 4], datatype="FP32", data=[[6.8, 2.8, 4.8, 1.6], [6.0, 3.4, 4.5, 1.6]]) - infer_request = InferRequest( - model_name="model", infer_inputs=[infer_input]) + infer_request = InferRequest(model_name="model", infer_inputs=[infer_input]) infer_response, response_headers = model.predict(infer_request) assert infer_response.to_rest()["outputs"] == \ - [{'name': 'output-0', 'shape': [2, 3], 'datatype': 'FP64', - 'data': [3.7899802486733807e-06, 0.9996982074114203, 0.00029800260833088297, - 5.2172911836629736e-05, 0.99973341723876, 0.000214409849403366]}] + [{'name': 'output-0', 'shape': [2, 3], 'datatype': 'FP64', + 'data': [3.7899802486733807e-06, 0.9996982074114203, 0.00029800260833088297, + 5.2172911836629736e-05, 0.99973341723876, 0.000214409849403366]}] diff --git a/python/test_resources/graph/success_200_isvc/model.py b/python/test_resources/graph/success_200_isvc/model.py index bf2a4d6fe7f..0ec88cf15ce 100644 --- a/python/test_resources/graph/success_200_isvc/model.py +++ b/python/test_resources/graph/success_200_isvc/model.py @@ -32,7 +32,8 @@ def load(self): self.ready = True def predict(self, payload: Union[Dict, InferRequest, ModelInferRequest], headers) -> Dict: - return {"message": "SUCCESS"} + response_headers = {} + return {"message": "SUCCESS"}, response_headers parser = argparse.ArgumentParser(parents=[kserve.model_server.parser])