Make OnnxEvaluator closable

1 files changed, 12 insertions, 1 deletions

diff --git a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java
index 9961c24005c..66eb8caabd0 100644
--- a/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java
+++ b/model-integration/src/main/java/ai/vespa/modelintegration/evaluator/OnnxEvaluator.java
@@ -20,7 +20,7 @@ import java.util.Map;
  *
  * @author lesters
  */
-public class OnnxEvaluator {
+public class OnnxEvaluator implements AutoCloseable {
 
     private final OrtEnvironment environment;
     private final OrtSession session;
@@ -88,6 +88,17 @@ public class OnnxEvaluator {
         }
     }
 
+    @Override
+    public void close() throws IllegalStateException {
+        try {
+            session.close();
+        } catch (OrtException e) {
+            throw new IllegalStateException("Failed to close ONNX session", e);
+        } catch (IllegalStateException e) {
+            throw new IllegalStateException("Already closed", e);
+        }
+    }
+
     private static OrtSession createSession(String modelPath, OrtEnvironment environment, OnnxEvaluatorOptions options, boolean tryCuda) {
         if (options == null) {
             options = new OnnxEvaluatorOptions();