[python] Handles invalid retrun type case (#790)

deepjavalibrary · Jun 1, 2023 · 663f389 · 663f389
1 parent 8d8cae5
commit 663f389
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 0 deletions.
diff --git a/engines/python/setup/djl_python_engine.py b/engines/python/setup/djl_python_engine.py
@@ -116,6 +116,8 @@ def run_server(self):
                 outputs = self.service.invoke_handler(function_name, inputs)
                 if outputs is None:
                     outputs = Output(code=204, message="No content")
+                elif not isinstance(outputs, Output):
+                    outputs = Output().error(f"Invalid output type: {type(outputs)}")
             except Exception as e:
                 logging.exception("Failed invoke service.invoke_handler()")
                 if type(e).__name__ == "OutOfMemoryError" or type(

diff --git a/engines/python/src/test/java/ai/djl/python/engine/PyEngineTest.java b/engines/python/src/test/java/ai/djl/python/engine/PyEngineTest.java
@@ -384,6 +384,11 @@ public void testModelException() throws TranslateException, IOException, ModelEx
             String error = json.getAsJsonObject().get("error").getAsString();
             Assert.assertEquals(error, "model error");
 
+            input = new Input();
+            input.add("typeerror", "type error");
+            output = predictor.predict(input);
+            Assert.assertEquals(output.getCode(), 424);
+
             // Test empty input
             input = new Input();
             input.add("exception", "");

diff --git a/engines/python/src/test/resources/echo/model.py b/engines/python/src/test/resources/echo/model.py
@@ -35,6 +35,8 @@ def handle(inputs: Input):
     if inputs.contains_key("exception"):
         ex = inputs.get_as_string("exception")
         raise ValueError(ex)
+    elif inputs.contains_key("typeerror"):
+        return "invalid_type"
     elif inputs.contains_key("exit"):
         sys.exit()