add generic predictor pool

Author: baltzell

reconstruction/ai/src/main/java/org/jlab/service/ai/DCClsComboEngine.java

@@ -12,7 +12,6 @@
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
-import ai.djl.ndarray.types.Shape;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ModelNotFoundException;
 import ai.djl.repository.zoo.ZooModel;
@@ -22,8 +21,6 @@
 import ai.djl.translate.TranslatorContext;
 import java.io.IOException;
 import java.nio.file.Paths;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.BlockingQueue;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 import org.jlab.clas.reco.ReconstructionEngine;
@@ -57,27 +54,8 @@ public class DCClsComboEngine extends ReconstructionEngine {
 
     final static int SUPERLAYERS = 6;    
 
-    // -------- Predictor Pool --------
-    public static class PredictorPool {
-        final BlockingQueue<Predictor<float[][], float[]>> pool;
-        public PredictorPool(int size, ZooModel<float[][], float[]> model) {
-            pool = new ArrayBlockingQueue<>(size);
-            for (int i=0; i<size; i++) {
-                try {
-                    pool.add(model.newPredictor());
-                } catch (Exception e) {
-                    Logger.getLogger(PredictorPool.class.getName()).log(Level.WARNING, "Failed to create predictor", e);
-                }
-            }
-        }
-        public Predictor<float[][], float[]> take() throws InterruptedException { return pool.take(); }
-        public void put(Predictor<float[][], float[]> p) throws InterruptedException { if (p!=null) pool.put(p); }
-        public void shutdownAll() { for (Predictor p: pool) { try { p.close(); } catch (Exception ignored) {} } }
-    }
-    
     public DCClsComboEngine() {
         super("DCClsComboEngine","tongtong","1.0");
-        
     }
 
     @Override

reconstruction/ai/src/main/java/org/jlab/service/ai/DCDenoiseEngine.java

@@ -17,8 +17,6 @@
 import ai.djl.translate.TranslateException;
 
 import java.io.IOException;
-import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.ArrayBlockingQueue;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
@@ -44,24 +42,6 @@ public class DCDenoiseEngine extends ReconstructionEngine {
     ZooModel<float[][][], float[][][]> model;
     PredictorPool predictors;
 
-    // -------- Predictor Pool --------
-    public static class PredictorPool {
-        final BlockingQueue<Predictor<float[][][], float[][][]>> pool;
-        public PredictorPool(int size, ZooModel<float[][][], float[][][]> model) {
-            pool = new ArrayBlockingQueue<>(size);
-            for (int i=0; i<size; i++) {
-                try {
-                    pool.add(model.newPredictor());
-                } catch (Exception e) {
-                    Logger.getLogger(PredictorPool.class.getName()).log(Level.WARNING, "Failed to create predictor", e);
-                }
-            }
-        }
-        public Predictor<float[][][], float[][][]> take() throws InterruptedException { return pool.take(); }
-        public void put(Predictor<float[][][], float[][][]> p) throws InterruptedException { if (p!=null) pool.put(p); }
-        public void shutdownAll() { for (Predictor p: pool) { try { p.close(); } catch (Exception ignored) {} } }
-    }
-
     public DCDenoiseEngine() {
         super("DenoiseEngine","lleztlab","1.0");
     }

reconstruction/ai/src/main/java/org/jlab/service/ai/PredictorPool.java

@@ -0,0 +1,39 @@
+package org.jlab.service.ai;
+
+import ai.djl.inference.Predictor;
+import ai.djl.repository.zoo.ZooModel;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.concurrent.BlockingQueue;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+public class PredictorPool <T,U> {
+    
+        final BlockingQueue<Predictor<T, U>> pool;
+
+        public PredictorPool(int size, ZooModel<T, U> model) {
+            pool = new ArrayBlockingQueue<>(size);
+            for (int i=0; i<size; i++) {
+                try {
+                    pool.add(model.newPredictor());
+                } catch (Exception e) {
+                    Logger.getLogger(PredictorPool.class.getName()).log(Level.WARNING, "Failed to create predictor", e);
+                }
+            }
+        }
+
+        public Predictor<T, U> take() throws InterruptedException {
+            return pool.take();
+        }
+
+        public void put(Predictor<T, U> p) throws InterruptedException {
+            if (p!=null) pool.put(p);
+        }
+
+        public void shutdownAll() {
+            for (Predictor p: pool) {
+                try { p.close(); }
+                catch (Exception ignored) {}
+            }
+        }
+}