handle crashes gracefully

bioimage-io · Oct 3, 2024 · 050116e · 050116e
1 parent e21a11f
commit 050116e
Showing 1 changed file with 12 additions and 3 deletions.
diff --git a/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/PytorchJavaCPPInterface.java b/src/main/java/io/bioimage/modelrunner/pytorch/javacpp/PytorchJavaCPPInterface.java
@@ -188,8 +188,11 @@ private void launchModelLoadOnProcess() throws IOException, InterruptedException
 			throw new RuntimeException();
 		else if (task.status == TaskStatus.FAILED)
 			throw new RuntimeException();
-		else if (task.status == TaskStatus.CRASHED)
+		else if (task.status == TaskStatus.CRASHED) {
+			this.runner.close();
+			runner = null;
 			throw new RuntimeException();
+		}
 	}
 
 	/**
@@ -301,8 +304,11 @@ void runInterprocessing(List<Tensor<T>> inputTensors, List<Tensor<R>> outputTens
 				throw new RuntimeException();
 			else if (task.status == TaskStatus.FAILED)
 				throw new RuntimeException();
-			else if (task.status == TaskStatus.CRASHED)
+			else if (task.status == TaskStatus.CRASHED) {
+				this.runner.close();
+				runner = null;
 				throw new RuntimeException();
+			}
 			for (int i = 0; i < outputTensors.size(); i ++) {
 	        	String name = (String) Types.decode(encOuts.get(i)).get(MEM_NAME_KEY);
 	        	SharedMemoryArray shm = shmaOutputList.stream()
@@ -364,8 +370,11 @@ public void closeModel() {
 				throw new RuntimeException();
 			else if (task.status == TaskStatus.FAILED)
 				throw new RuntimeException();
-			else if (task.status == TaskStatus.CRASHED)
+			else if (task.status == TaskStatus.CRASHED) {
+				this.runner.close();
+				runner = null;
 				throw new RuntimeException();
+			}
 			this.runner.close();
 			this.runner = null;
 			return;