[python] Add more OOM handling (#2515)

deepjavalibrary · Nov 4, 2024 · 481f184 · 481f184
1 parent 8e6d936
commit 481f184
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python_engine.py b/engines/python/setup/djl_python_engine.py
@@ -166,8 +166,10 @@ def run_server(self):
                         f"Invalid output type: {type(outputs)}")
             except Exception as e:
                 logging.exception("Failed invoke service.invoke_handler()")
-                if type(e).__name__ == "OutOfMemoryError" or type(
-                        e).__name__ == "MemoryError":
+                if (type(e).__name__ == "OutOfMemoryError"
+                        or type(e).__name__ == "MemoryError"
+                        or "No available memory for the cache blocks" in str(e)
+                        or "CUDA error: out of memory" in str(e)):
                     outputs = Output(code=507, message=str(e))
                 else:
                     outputs = Output().error(str(e))